@tduncan / Nvidia-Optimization-Stack-VLLM / commits / bf93e1f92c

[Refactor] Remove dead tests utils (#58803)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

Wentao Ye committed Sep 26, 2026 at 09:16 UTC bf93e1f92c40673454c938b277a0a99d49662657
16 files changed +4 -1337
tests/conftest.py
-8
@@ -317,14 +317,6 @@ def example_system_message() -> str:
317 return f.read()
318
319
320 -class DecoderPromptType(Enum):
321 - """For encoder/decoder models only."""
322 -
323 - CUSTOM = 1
324 - NONE = 2
325 - EMPTY_STR = 3
326 -
327 -
320 @pytest.fixture
321 def example_long_prompts() -> list[str]:
322 return [prompt for filename in _LONG_PROMPTS for prompt in _read_prompts(filename)]
tests/entrypoints/serve/dev/rlhf/conftest.py
-86
@@ -15,7 +15,6 @@ import subprocess
15 import sys
16 import threading
17 import time
18 -from collections.abc import Callable
18 from contextlib import contextmanager
19 from dataclasses import dataclass, field
20 from typing import Any
@@ -138,36 +137,6 @@ def server(
137 proc.kill()
138
139
141 -# ---------------------------------------------------------------------------
142 -# Polling helper (200-lie workaround)
143 -# ---------------------------------------------------------------------------
144 -
145 -
146 -def poll_until(
147 - predicate: Callable[[], bool],
148 - timeout: float = 10.0,
149 - interval: float = 0.5,
150 -) -> bool:
151 - """Poll predicate() until it returns True or timeout expires.
152 -
153 - Workaround for the vLLM sleep/wake "200-lie" — the HTTP endpoints may
154 - return 200 before the underlying operation is complete, so callers that
155 - need to verify state *after* an operation can use this helper instead of
156 - assuming the 200 means completion.
157 -
158 - Returns True if predicate became true within timeout, False otherwise.
159 - """
160 - deadline = time.time() + timeout
161 - while time.time() < deadline:
162 - try:
163 - if predicate():
164 - return True
165 - except Exception:
166 - pass
167 - time.sleep(interval)
168 - return False
169 -
170 -
140 # ---------------------------------------------------------------------------
141 # HTTP helpers — generation
142 # ---------------------------------------------------------------------------
@@ -191,27 +160,6 @@ def gen(url, prompt="The capital of France is", max_tokens=8, timeout=30):
160 return None
161
162
194 -def gen_with_logprobs(
195 - url, prompt="The capital of France is", max_tokens=8, logprobs=5, timeout=30
196 -):
197 - """Fire a /v1/completions request with logprobs; return JSON or None."""
198 - try:
199 - r = requests.post(
200 - f"{url}/v1/completions",
201 - json={
202 - "model": "m",
203 - "prompt": prompt,
204 - "max_tokens": max_tokens,
205 - "temperature": 0,
206 - "logprobs": logprobs,
207 - },
208 - timeout=timeout,
209 - )
210 - return r.json()
211 - except Exception:
212 - return None
213 -
214 -
163 def ok(resp) -> bool:
164 """True iff resp is a successful completion (has choices, no error key)."""
165 return (
@@ -399,37 +347,3 @@ def get_world_size(url, include_dp=True):
347 params={"include_dp": include_dp},
348 timeout=5,
349 )
402 -
403 -
404 -# ---------------------------------------------------------------------------
405 -# GPU / metrics helpers
406 -# ---------------------------------------------------------------------------
407 -
408 -
409 -def gpu_free_bytes(device: int = 0) -> int:
410 - """Read GPU free bytes via subprocess to avoid import-time torch init."""
411 - out = subprocess.check_output(
412 - [
413 - sys.executable,
414 - "-c",
415 - f"import torch; f,_=torch.accelerator.get_memory_info({device}); print(f)",
416 - ],
417 - timeout=10,
418 - )
419 - return int(out.strip())
420 -
421 -
422 -def sleep_metrics(url):
423 - """Return (awake, weights_offloaded, discard_all) from /metrics."""
424 - try:
425 - from prometheus_client.parser import text_string_to_metric_families
426 - except ImportError:
427 - return None, None, None
428 -
429 - r = requests.get(f"{url}/metrics", timeout=5)
430 - vals: dict = {}
431 - for family in text_string_to_metric_families(r.text):
432 - if family.name == "vllm:engine_sleep_state":
433 - for s in family.samples:
434 - vals[s.labels.get("sleep_state", "")] = s.value
435 - return vals.get("awake"), vals.get("weights_offloaded"), vals.get("discard_all")
tests/kernels/attention/test_attention.py
-40
@@ -309,46 +309,6 @@ def test_paged_attention(
309 torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
310
311
312 -def ref_multi_query_kv_attention(
313 - cu_seq_lens: list[int],
314 - query: torch.Tensor,
315 - key: torch.Tensor,
316 - value: torch.Tensor,
317 - scale: float,
318 - alibi_bias: list[torch.Tensor] | None,
319 - dtype: torch.dtype,
320 -) -> torch.Tensor:
321 - num_seqs = len(cu_seq_lens) - 1
322 - ref_outputs: list[torch.Tensor] = []
323 - if alibi_bias:
324 - assert len(alibi_bias) == num_seqs
325 - for i in range(num_seqs):
326 - start_idx = cu_seq_lens[i]
327 - end_idx = cu_seq_lens[i + 1]
328 - seq_len = end_idx - start_idx
329 -
330 - # Create attention mask. ALiBi already includes a tril causal mask.
331 - if alibi_bias:
332 - attn_mask = alibi_bias[i]
333 - else:
334 - attn_mask = torch.triu(
335 - torch.ones(seq_len, seq_len, dtype=dtype), diagonal=1
336 - )
337 - attn_mask = attn_mask * torch.finfo(dtype).min
338 - attn_mask = attn_mask.to(dtype=dtype)
339 -
340 - ref_output = ref_masked_attention(
341 - query[start_idx:end_idx],
342 - key[start_idx:end_idx],
343 - value[start_idx:end_idx],
344 - scale,
345 - attn_mask=attn_mask,
346 - )
347 - ref_outputs.append(ref_output)
348 -
349 - return torch.cat(ref_outputs, dim=0)
350 -
351 -
312 @pytest.mark.parametrize("attention_cls", [Attention, MMEncoderAttention])
313 def test_num_heads_not_divisible_by_num_kv_heads(attention_cls: type) -> None:
314 head_size = 64
tests/kernels/mamba/test_causal_conv1d.py
-31
@@ -120,37 +120,6 @@ def causal_conv1d_update_ref(
120 return (out if activation is None else F.silu(out)).to(dtype=dtype_in)
121
122
123 -@pytest.mark.parametrize("itype", [torch.bfloat16, torch.float])
124 -@pytest.mark.parametrize("silu_activation", [True])
125 -@pytest.mark.parametrize("has_bias", [True])
126 -def causal_conv1d_opcheck_fn(
127 - x: torch.Tensor,
128 - weight: torch.Tensor,
129 - bias: torch.Tensor | None = None,
130 - cu_seq_len: torch.Tensor | None = None,
131 - cache_indices: torch.Tensor | None = None,
132 - has_initial_state: torch.Tensor | None = None,
133 - conv_states: torch.Tensor | None = None,
134 - activation: str | None = "silu",
135 - null_block_id: int = NULL_BLOCK_ID,
136 -):
137 - """x: (batch, dim, seqlen)
138 - weight: (dim, width)
139 - bias: (dim,)
140 - seq_idx: (batch, seqlen)
141 - initial_states: (batch, dim, width - 1)
142 - final_states_out: (batch, dim, width - 1), to be written to
143 - activation: either None or "silu" or "swish"
144 -
145 - out: (batch, dim, seqlen)
146 - """
147 - if activation not in [None, "silu", "swish"]:
148 - raise NotImplementedError("activation must be None, silu, or swish")
149 - if x.stride(-1) != 1:
150 - x = x.contiguous()
151 - bias = bias.contiguous() if bias is not None else None
152 -
153 -
123 @pytest.mark.parametrize("itype", [torch.bfloat16])
124 @pytest.mark.parametrize("silu_activation", [False, True])
125 @pytest.mark.parametrize("has_bias", [False, True])
tests/kernels/moe/modular_kernel_tools/mk_objects.py
-18
@@ -448,24 +448,6 @@ if cutlass_fp4_supported() or has_flashinfer_cutlass_fused_moe():
448 ]
449
450
451 -def _slice(rank: int, num_local_experts: int, t: torch.Tensor) -> torch.Tensor:
452 - s = rank * num_local_experts
453 - e = s + num_local_experts
454 - return t[s:e]
455 -
456 -
457 -def make_cutlass_strides(
458 - e: int,
459 - n: int,
460 - k: int,
461 -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
462 - ab_strides1 = torch.full((e,), k, device="cuda", dtype=torch.int64)
463 - ab_strides2 = torch.full((e,), n, device="cuda", dtype=torch.int64)
464 - c_strides1 = torch.full((e,), 2 * n, device="cuda", dtype=torch.int64)
465 - c_strides2 = torch.full((e,), k, device="cuda", dtype=torch.int64)
466 - return ab_strides1, ab_strides2, c_strides1, c_strides2
467 -
468 -
451 def make_fused_experts(
452 fused_experts_type: mk.FusedMoEExpertsModular,
453 moe: FusedMoEConfig,
tests/kernels/moe/test_moe_layer.py
-51
@@ -275,11 +275,6 @@ def maybe_roundup_layer_hidden_size(
275 return hidden_size
276
277
278 -def rank_chunk(num: int, r: int, w: int) -> int:
279 - rem = num % w
280 - return (num // w) + (1 if r < rem else 0)
281 -
282 -
278 def chunk_by_rank(
279 t: torch.Tensor,
280 r: int,
@@ -294,19 +289,6 @@ def chunk_by_rank(
289 return t
290
291
297 -def maybe_chunk_by_rank(
298 - t: torch.Tensor | None,
299 - r: int,
300 - w: int,
301 - dim: int = 0,
302 - device: torch.device | None = None,
303 -) -> torch.Tensor | None:
304 - if t is not None:
305 - return chunk_by_rank(t, r, w, dim, device)
306 - else:
307 - return t
308 -
309 -
292 def tp_chunk_gate_up(
293 w: torch.Tensor,
294 tp_rank: int,
@@ -676,26 +658,6 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
658 return True, None
659
660
679 -def chunk_scales_by_rank(
680 - t: torch.Tensor | None,
681 - r: int,
682 - w: int,
683 - device: torch.device | None = None,
684 -) -> torch.Tensor | None:
685 - if t is not None and t.numel() > 1:
686 - # Calculate start index by summing chunk sizes for all previous ranks
687 - # start = sum(rank_chunk(t.shape[0], i, w) for i in range(r))
688 - # chunk = rank_chunk(t.shape[0], r, w)
689 - # t = t[start:(start + chunk)]
690 - chunk = rank_chunk(t.shape[0], r, w)
691 - t = t[(r * chunk) : max(t.shape[0], (r + 1) * chunk)]
692 -
693 - if t is not None and device is not None:
694 - t = t.to(device)
695 -
696 - return t
697 -
698 -
661 def chunk_scales(
662 t: torch.Tensor | None,
663 start: int,
@@ -771,19 +733,6 @@ def _quantize_fp8_halves(
733 )
734
735
774 -def quantization_to_quant_dtype(
775 - quantization: str | None,
776 -) -> torch.dtype | str | None:
777 - if quantization is None:
778 - return None
779 - elif quantization in ["fp8", "fp8_blocked", "modelopt_fp8"]:
780 - return fp8_dtype
781 - elif quantization in ["modelopt_fp4"]:
782 - return "nvfp4"
783 - else:
784 - raise NotImplementedError(f"Unsupported quantization: {quantization}")
785 -
786 -
736 def make_quant_config(
737 quantization: str | None,
738 w1: torch.Tensor,
tests/kernels/moe/utils.py
-29
@@ -552,16 +552,6 @@ class TestMLP(torch.nn.Module):
552 return x
553
554
555 -def make_naive_shared_experts(
556 - N: int,
557 - K: int,
558 - in_dtype: torch.dtype = torch.bfloat16,
559 -) -> torch.nn.Module:
560 - w1 = torch.randn((K, N * 2), device=DEVICE, dtype=in_dtype) / 15
561 - w2 = torch.randn((N, K), device=DEVICE, dtype=in_dtype) / 15
562 - return TestMLP(w1, w2, out_dtype=in_dtype)
563 -
564 -
555 class RealMLP(torch.nn.Module):
556 def __init__(
557 self,
@@ -668,25 +658,6 @@ def modular_triton_fused_moe(
658 )
659
660
671 -def make_shared_experts(
672 - N: int,
673 - K: int,
674 - in_dtype: torch.dtype = torch.bfloat16,
675 - quant_dtype: torch.dtype | str | None = None,
676 -) -> torch.nn.Module:
677 - (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights(
678 - 1,
679 - N,
680 - K,
681 - in_dtype=in_dtype,
682 - quant_dtype=quant_dtype,
683 - )
684 -
685 - return make_shared_experts_with_weights(
686 - N, K, in_dtype, w1, w2, w1_s=w1_s, w2_s=w2_s, quant_dtype=quant_dtype
687 - )
688 -
689 -
661 def check_accuracy(a, b, atol, rtol, percent):
662 """Allow a mismatch percentage of 1 - percent."""
663 if torch.any(torch.isnan(a)):
tests/kernels/quant_utils.py
-19
@@ -253,25 +253,6 @@ def dequant(
253 return t.to(out_dtype)
254
255
256 -def batched_dequant(
257 - t: torch.Tensor,
258 - scale: torch.Tensor | None,
259 - block_shape: list[int] | None,
260 - per_act_token_quant: bool,
261 - out_dtype: torch.dtype | None = torch.float32,
262 -) -> torch.Tensor:
263 - if scale is not None:
264 - assert t.shape[0] == scale.shape[0]
265 - out = torch.empty_like(t, dtype=out_dtype)
266 - for e in range(t.shape[0]):
267 - out[e] = dequant(
268 - t[e], scale[e], block_shape, per_act_token_quant, out_dtype
269 - )
270 - return out
271 -
272 - return t.to(out_dtype)
273 -
274 -
256 def native_batched_masked_quant_matmul(
257 A: torch.Tensor,
258 B: torch.Tensor,
tests/kernels/utils.py
+1 -721
@@ -2,11 +2,8 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Kernel test utils."""
4
5 -import itertools
6 -import random
5 from collections.abc import Sequence
8 -from numbers import Number
9 -from typing import Any, NamedTuple
6 +from typing import Any
7 from unittest.mock import patch
8
9 import torch
@@ -17,8 +14,6 @@ from vllm.model_executor.custom_op import op_registry
14 from vllm.model_executor.layers.activation import SiluAndMul
15 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
16 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
20 -from vllm.utils.torch_utils import make_tensor_with_pad
21 -from vllm.v1.attention.backend import AttentionType
17
18 # For now, disable "test_aot_dispatch_dynamic" since there are some
19 # bugs related to this test in PyTorch 2.4.
@@ -93,721 +88,6 @@ def _assert_deterministic(
88 )
89
90
96 -class QKVInputs(NamedTuple):
97 - """Data structure for representing unpacked attention inputs,
98 - query/key/values and their sequence lengths.
99 -
100 - Attributes:
101 - * {query,key,value}: unpacked (batch_size x padded_seq_len x
102 - num_heads x head_size) attention inputs
103 - * q_seq_lens: query sequence lengths list
104 - * kv_seq_lens: shared key/value sequence lengths list
105 -
106 - """
107 -
108 - query: torch.Tensor
109 - key: torch.Tensor
110 - value: torch.Tensor
111 - q_seq_lens: list[int]
112 - kv_seq_lens: list[int]
113 -
114 -
115 -class QKVO(NamedTuple):
116 - """Data structure for representing unpacked attention inputs,
117 - alongside unpacked known-correct attention output
118 -
119 - Attributes:
120 - * qkv: unpacked (batch_size x padded_seq_len x
121 - num_heads x head_size) attention inputs
122 - * ideal_output: unpacked (batch_size x padded_seq_len x
123 - num_heads x head_size) known-correct attention output
124 -
125 - """
126 -
127 - qkv: QKVInputs
128 - ideal_output: torch.Tensor
129 -
130 -
131 -class PackedQKVInputs(NamedTuple):
132 - """Data structure for representing packed attention inputs.
133 -
134 - Attributes:
135 - * {query,key,value}: packed (number_of_tokens x num_heads
136 - x head_size) attention inputs
137 - * q_start_loc_list: list of query start locations within packed tensor
138 - * kv_start_loc_list: shared list of key/value start locations within
139 - packed tensor
140 - * q_seq_lens: query sequence lengths list
141 - * kv_seq_lens: shared key/value sequence lengths list
142 -
143 - """
144 -
145 - query: torch.Tensor
146 - key: torch.Tensor
147 - value: torch.Tensor
148 - q_start_loc_list: list[int] | None
149 - kv_start_loc_list: list[int] | None
150 - q_seq_lens: list[int] | None
151 - kv_seq_lens: list[int] | None
152 -
153 -
154 -class PackedQKVO(NamedTuple):
155 - """Data structure for representing packed attention inputs,
156 - alongside packed known-correct attention output
157 -
158 - Attributes:
159 - * packed_qkv: packed (number_of_tokens x num_heads
160 - x head_size) attention inputs
161 - * ideal_output: packed (number_of_tokens x num_heads
162 - x head_size) known-correct attention output
163 -
164 - """
165 -
166 - packed_qkv: PackedQKVInputs | None
167 - ideal_output: torch.Tensor
168 -
169 -
170 -class KVMemoryMap(NamedTuple):
171 - """Data structure for encapsulating KV cache memory mapping.
172 -
173 - Attributes:
174 - * block_tables: KV cache block tables
175 - * slot_mapping: mapping of sequence offset to physical address
176 -
177 - """
178 -
179 - block_tables: torch.Tensor
180 - slot_mapping: torch.Tensor
181 -
182 -
183 -class PhaseTestParameters(NamedTuple):
184 - """Data structure for encapsulating the test parameters
185 - for a given test "phase" (prefill or decode phase) and attention
186 - scenario (encoder, decoder-self, encoder/decoder-cross)
187 -
188 - Attributes:
189 - * packed_qkvo: packed (number_of_tokens x num_heads
190 - x head_size) attention inputs & known-correct
191 - output
192 - * kv_mmap: KV cache memory mapping, specific to this test phase &
193 - attention scenario
194 -
195 - """
196 -
197 - packed_qkvo: PackedQKVO
198 - kv_mmap: KVMemoryMap | None
199 -
200 -
201 -def maybe_make_int_tensor(
202 - _list: list[int] | None,
203 - device: torch.device | str,
204 -) -> torch.Tensor:
205 - """Convert Python int list to a 1D int torch.Tensor on `device`.
206 -
207 - Returns:
208 - If _list is not None: 1D int torch.Tensor on `device`, `None` otherwise.
209 -
210 - """
211 - return (
212 - None if _list is None else torch.tensor(_list, dtype=torch.int, device=device)
213 - )
214 -
215 -
216 -def maybe_make_long_tensor(
217 - _list: list[int] | None,
218 - device: torch.device | str,
219 -) -> torch.Tensor:
220 - """Convert Python int list to a 1D long torch.Tensor on `device`.
221 -
222 - Returns:
223 - If _list is not None: 1D long torch.Tensor on `device`, `None` otherwise.
224 -
225 - """
226 - return (
227 - None if _list is None else torch.tensor(_list, dtype=torch.long, device=device)
228 - )
229 -
230 -
231 -def maybe_max(_list: list | None) -> Number | None:
232 - """Compute the maximum of a list, if it exists.
233 -
234 - Returns:
235 - If _list is not None: max(_list), `None` otherwise.
236 -
237 - """
238 - return None if _list is None else max(_list)
239 -
240 -
241 -def make_causal_mask(
242 - q_max_seq_len: int,
243 - kv_max_seq_len: int,
244 -) -> torch.Tensor:
245 - """Create a q_max_seq_len x kv_max_seq_len causal mask.
246 -
247 - Arguments:
248 - q_max_seq_len: query max seq len
249 - kv_max_seq_len: key/value max seq len
250 -
251 - Returns:
252 - 2D tensor, q_max_seq_len x kv_max_seq_len
253 -
254 - """
255 - # Create a matrix where entry (i, j) is True if i >= j
256 - mask = torch.triu(torch.ones(q_max_seq_len, kv_max_seq_len), diagonal=1)
257 - # Replace True with float('-inf') and False with 0
258 - mask = mask.masked_fill(mask == 1, float("-inf")).masked_fill(mask == 0, 0.0)
259 - return mask
260 -
261 -
262 -def ref_masked_attention(
263 - query: torch.Tensor,
264 - key: torch.Tensor,
265 - value: torch.Tensor,
266 - scale: float,
267 - custom_mask: torch.Tensor | None = None,
268 - q_seq_lens: list | None = None,
269 - kv_seq_lens: list | None = None,
270 -) -> torch.Tensor:
271 - """Reference ("golden") masked attention. Supports two types of masking:
272 -
273 - * Basic attention mask, utilizing {q,kv}_seq_lens args to mask out
274 - padding elements
275 - * Custom attention mask, which can force an arbitrary mask tensor, i.e.
276 - causal
277 -
278 - Arguments:
279 - query: batch_size x q_padded_seq_len x num_heads x head_size
280 - key: batch_size x kv_padded_seq_len x num_heads x head_size
281 - value: batch_size x kv_padded_seq_len x num_heads x head_size
282 - scale: Attention scale factor
283 - custom_mask: custom attention mask; good place to inject a causal attention mask
284 - q_seq_lens: list of unpadded query seq_lens for each batch index
285 - kv_seq_lens: list of unpadded key/value seq_lens for each batch index
286 -
287 - Returns:
288 - Attention result, batch_size x q_padded_seq_len x num_heads x head_size
289 -
290 - """
291 - assert q_seq_lens is not None
292 - assert kv_seq_lens is not None
293 -
294 - batch_size = query.shape[0]
295 - assert len(q_seq_lens) == batch_size
296 - assert len(kv_seq_lens) == batch_size
297 -
298 - attn_weights = scale * torch.einsum("bqhd,bkhd->bhqk", query, key).float()
299 -
300 - # Basic attention mask, derived from seq lens
301 - if (q_seq_lens is not None) or (kv_seq_lens is not None):
302 - attn_mask = torch.zeros_like(attn_weights)
303 - if q_seq_lens is not None:
304 - for bdx, plen in enumerate(q_seq_lens):
305 - attn_mask[bdx, :, plen:, :] = -torch.inf
306 - if kv_seq_lens is not None:
307 - for bdx, plen in enumerate(kv_seq_lens):
308 - attn_mask[bdx, :, :, plen:] = -torch.inf
309 -
310 - attn_weights = attn_weights + attn_mask.float()
311 -
312 - # Custom attention mask
313 - if custom_mask is not None:
314 - attn_weights = attn_weights + custom_mask.float()
315 -
316 - attn_weights = torch.softmax(attn_weights, dim=-1).to(value.dtype)
317 - out = torch.einsum("bhqk,bkhd->bqhd", attn_weights, value)
318 - return out
319 -
320 -
321 -def make_qkv(
322 - batch_size: int,
323 - max_q_seq_len: int,
324 - max_kv_seq_len: int | None,
325 - num_heads: int,
326 - head_size: int,
327 - device: torch.device | str,
328 - force_kv_seq_lens: list[int] | None = None,
329 - attn_type: AttentionType = AttentionType.ENCODER_DECODER,
330 - force_max_len: bool = False,
331 -) -> tuple[QKVInputs, QKVInputs, QKVInputs]:
332 - """Construct QKV test tensors for self- and cross-attention.
333 -
334 - Generates three query/key/value triplets:
335 -
336 - * "Baseline" query/key/value (for input to reference attention function)
337 - * "Prefill" query/key/value (last sequence offset zero'd out, for use as
338 - input to prefill kernel)
339 - * "Decode" query/key/value (only the last sequence offset from baseline,
340 - for use as input to decode kernel)
341 -
342 - Each Q/K/V triplet is associated with a list of q seqlens and a list of k/v
343 - seqlens
344 -
345 - Arguments:
346 - batch_size: number of sequences in the batch
347 - max_q_seq_len: max query seq len
348 - max_kv_seq_len: max key/value seq len
349 - num_heads: number of attention heads
350 - head_size: head dimension
351 - is_encoder_decoder_attn: if True, query seqlen may differ from
352 - key/value seqlen (as is often the case for cross-attention);
353 - o/w, query/key/value seqlens match at each batch index
354 - (max_kv_seq_len is unused)
355 - force_kv_seq_lens: if not None, overrides kv sequence lengths
356 - attn_type: encoder, decoder self, or enc/dec cross attention
357 - force_max_len: if True, all query seqlens are max_q_seq_len; o/w query
358 - seqlens are random in [2,max_q_seq_lens]. Same for key/value seqlens
359 - and max_kv_seq_len, unless forced by is_encoder_decoder_attn=False
360 - device: CPU or CUDA device
361 -
362 - Returns:
363 - Overall QKVInputs structure (containing full unpacked Q/K/V tensors)
364 - Prefill QKVInputs structure (containing all but the last sequence offset)
365 - Decode QKVInputs structure (containing all only the last sequence offset)
366 -
367 - """
368 - if force_max_len:
369 - q_seq_lens = [max_q_seq_len for _ in range(batch_size)]
370 - else:
371 - q_seq_lens = [random.randint(2, max_q_seq_len) for _ in range(batch_size)]
372 - kv_seq_lens = None
373 - if force_kv_seq_lens is not None:
374 - kv_seq_lens = force_kv_seq_lens
375 - elif attn_type != AttentionType.ENCODER_DECODER:
376 - # K,V seq lens match Q for self-attention
377 - kv_seq_lens = q_seq_lens
378 - else:
379 - # K,V seq lens are distinct from Q seq lens & random
380 - assert max_kv_seq_len is not None
381 - if force_max_len:
382 - kv_seq_lens = [max_kv_seq_len] * batch_size
383 - else:
384 - kv_seq_lens = [random.randint(2, max_kv_seq_len) for _ in range(batch_size)]
385 -
386 - query = torch.rand((batch_size, max_q_seq_len, num_heads, head_size)).to(device)
387 - key = torch.rand((batch_size, max_kv_seq_len, num_heads, head_size)).to(device)
388 - value = torch.rand((batch_size, max_kv_seq_len, num_heads, head_size)).to(device)
389 -
390 - prefill_query = torch.zeros((batch_size, max_q_seq_len, num_heads, head_size)).to(
391 - device
392 - )
393 - prefill_key = torch.zeros((batch_size, max_kv_seq_len, num_heads, head_size)).to(
394 - device
395 - )
396 - prefill_value = torch.zeros((batch_size, max_kv_seq_len, num_heads, head_size)).to(
397 - device
398 - )
399 -
400 - decode_query = torch.zeros((batch_size, 1, num_heads, head_size)).to(device)
401 - decode_key = torch.zeros((batch_size, 1, num_heads, head_size)).to(device)
402 - decode_value = torch.zeros((batch_size, 1, num_heads, head_size)).to(device)
403 -
404 - for bdx, (q_seq_len, kv_seq_len) in enumerate(zip(q_seq_lens, kv_seq_lens)):
405 - query[bdx, q_seq_len:, :, :] = 0
406 - key[bdx, kv_seq_len:, :, :] = 0
407 - value[bdx, kv_seq_len:, :, :] = 0
408 -
409 - prefill_query[bdx, 0 : (q_seq_len - 1), :, :] = query[
410 - bdx, 0 : (q_seq_len - 1), :, :
411 - ]
412 - prefill_key[bdx, 0 : (kv_seq_len - 1), :, :] = key[
413 - bdx, 0 : (kv_seq_len - 1), :, :
414 - ]
415 - prefill_value[bdx, 0 : (kv_seq_len - 1), :, :] = value[
416 - bdx, 0 : (kv_seq_len - 1), :, :
417 - ]
418 -
419 - decode_query[bdx, :, :, :] = query[bdx, (q_seq_len - 1) : q_seq_len, :, :]
420 - decode_key[bdx, :, :, :] = key[bdx, (kv_seq_len - 1) : kv_seq_len, :, :]
421 - decode_value[bdx, :, :, :] = value[bdx, (kv_seq_len - 1) : kv_seq_len, :, :]
422 -
423 - prefill_q_seq_lens = [plen - 1 for plen in q_seq_lens]
424 - prefill_kv_seq_lens = [plen - 1 for plen in kv_seq_lens]
425 -
426 - decode_q_seq_lens = [1 for _ in q_seq_lens]
427 - decode_kv_seq_lens = [1 for _ in kv_seq_lens]
428 -
429 - return (
430 - QKVInputs(
431 - query, # Overall QKV inputs
432 - key,
433 - value,
434 - q_seq_lens,
435 - kv_seq_lens,
436 - ),
437 - QKVInputs(
438 - prefill_query, # Prefill subset of QKV sequences
439 - prefill_key,
440 - prefill_value,
441 - prefill_q_seq_lens,
442 - prefill_kv_seq_lens,
443 - ),
444 - QKVInputs(
445 - decode_query, # Decode subset of KV sequences
446 - decode_key,
447 - decode_value,
448 - decode_q_seq_lens,
449 - decode_kv_seq_lens,
450 - ),
451 - )
452 -
453 -
454 -def pack_tensor(
455 - unpacked_tensor: torch.Tensor, seq_lens: list[int], device: torch.device | str
456 -) -> tuple[torch.Tensor, list[int]]:
457 - """Pack a batch_size x padded_seq_len x num_heads x head_size tensor into an
458 - unpadded number_of_tokens x num_heads x head_size tensor, where
459 - number_of_tokens = sum(seq_lens)
460 -
461 - Arguments:
462 - unpacked_tensor: batch_size x padded_seq_len x num_heads x head_size
463 - seq_lens: list of token counts for each seq
464 - device: CPU or CUDA device
465 -
466 - Returns:
467 - packed_tensor: number_of_tokens x num_heads x head_size
468 - start_loc_list: start idx of each batch elt in packed_tensor; [0] +
469 - list(itertools.accumulate(seq_lens))
470 -
471 - """
472 - num_tok = sum(seq_lens)
473 - num_heads = unpacked_tensor.shape[-2]
474 - head_size = unpacked_tensor.shape[-1]
475 - start_loc_list = [0] + list(itertools.accumulate(seq_lens))
476 - packed_tensor = torch.zeros((num_tok, num_heads, head_size), device=device)
477 -
478 - for bdx, (seq_len, start_loc) in enumerate(zip(seq_lens, start_loc_list)):
479 - packed_tensor[start_loc : (start_loc + seq_len), :, :] = unpacked_tensor[
480 - bdx, :seq_len, :, :
481 - ]
482 -
483 - return packed_tensor, start_loc_list
484 -
485 -
486 -def pack_qkv(qkv: QKVInputs, device: torch.device | str) -> PackedQKVInputs:
487 - """Individually pack each of Q, K and V, each with dimensions batch_size x
488 - padded_seq_len x num_heads x head_size, into respective number_of_tokens x
489 - num_heads x head_size tensors.
490 -
491 - For Q, number_of_tokens = sum(q_seq_lens).
492 -
493 - For K and V, number_of_tokens = sum(kv_seq_lens)
494 -
495 - Arguments:
496 - qkv: Unpacked (batch_size x padded_seq_len x num_heads x head_size)
497 - attention inputs
498 - device: CPU or CUDA device
499 -
500 - Returns:
501 - Packed (number_of_tokens x num_heads x head_size) QKV inputs
502 - derived from unpacked inputs
503 -
504 - """
505 - if qkv.query is None:
506 - packed_query = None
507 - q_start_loc_list = None
508 - else:
509 - packed_query, q_start_loc_list = pack_tensor(
510 - qkv.query, qkv.q_seq_lens, device=device
511 - )
512 - packed_key, kv_start_loc_list = pack_tensor(qkv.key, qkv.kv_seq_lens, device=device)
513 - packed_value, _ = pack_tensor(qkv.value, qkv.kv_seq_lens, device=device)
514 - return PackedQKVInputs(
515 - packed_query,
516 - packed_key,
517 - packed_value,
518 - q_start_loc_list,
519 - kv_start_loc_list,
520 - (None if q_start_loc_list is None else qkv.q_seq_lens),
521 - qkv.kv_seq_lens,
522 - )
523 -
524 -
525 -def _make_metadata_tensors(
526 - seq_lens: list[int] | None,
527 - context_lens: list[int] | None,
528 - encoder_seq_lens: list[int] | None,
529 - device: torch.device | str,
530 -) -> tuple[
531 - torch.Tensor,
532 - torch.Tensor,
533 - Any,
534 - Any,
535 - torch.Tensor | None,
536 - torch.Tensor,
537 - torch.Tensor,
538 - int | None,
539 -]:
540 - """Build scalar & tensor values required to build attention metadata structure.
541 -
542 - Arguments:
543 - seq_lens: list of token-counts for each decoder input seq
544 - context_lens: list of context length values for each seq
545 - encoder_seq_lens: list of token-counts for each encoder input seq
546 - device: CPU or CUDA device
547 -
548 - Returns:
549 - seq_lens_tensor: decoder seq_lens list, as tensor
550 - context_lens_tensor: context_lens list, as tensor
551 - max_context_len: max(context_lens)
552 - max_seq_len: max(seq_lens)
553 - seq_start_loc: start idx of each sequence
554 - encoder_seq_lens_tensor: encoder seq_lens list, as tensor
555 - encoder_seq_start_loc: start idx of each encoder sequence
556 - max_encoder_seq_len: encoder seq_lens list, as tensor
557 -
558 - """
559 - seq_lens_tensor = maybe_make_int_tensor(seq_lens, device)
560 - context_lens_tensor = maybe_make_int_tensor(context_lens, device)
561 - max_context_len = maybe_max(context_lens)
562 - max_seq_len = maybe_max(seq_lens)
563 -
564 - encoder_seq_lens_tensor = maybe_make_int_tensor(encoder_seq_lens, device)
565 - max_encoder_seq_len = None if encoder_seq_lens is None else max(encoder_seq_lens)
566 -
567 - seq_start_loc = None
568 -
569 - if seq_lens_tensor is not None:
570 - seq_start_loc = torch.zeros(
571 - seq_lens_tensor.shape[0] + 1,
572 - dtype=torch.int32,
573 - device=seq_lens_tensor.device,
574 - )
575 - torch.cumsum(
576 - seq_lens_tensor, dim=0, dtype=seq_start_loc.dtype, out=seq_start_loc[1:]
577 - )
578 -
579 - encoder_seq_start_loc = torch.zeros(
580 - encoder_seq_lens_tensor.shape[0] + 1,
581 - dtype=torch.int32,
582 - device=encoder_seq_lens_tensor.device,
583 - )
584 - torch.cumsum(
585 - encoder_seq_lens_tensor,
586 - dim=0,
587 - dtype=encoder_seq_start_loc.dtype,
588 - out=encoder_seq_start_loc[1:],
589 - )
590 -
591 - return (
592 - seq_lens_tensor,
593 - context_lens_tensor,
594 - max_context_len,
595 - max_seq_len,
596 - seq_start_loc,
597 - encoder_seq_lens_tensor,
598 - encoder_seq_start_loc,
599 - max_encoder_seq_len,
600 - )
601 -
602 -
603 -def make_kv_cache(
604 - num_blocks: int,
605 - num_heads: int,
606 - head_size: int,
607 - block_size: int,
608 - device: torch.device | str,
609 - backend: str,
610 - default_val: float = 0.0,
611 -) -> torch.Tensor:
612 - """Create a fake KV cache.
613 -
614 - Arguments:
615 - num_blocks: number of blocks in the KV cache
616 - num_heads: number of attention heads
617 - head_size: head dimension
618 - block_size: number of offsets within a block
619 - device: CPU or CUDA device
620 - backend: attention backend name; only 'FLASH_ATTN' is supported
621 - default_val: initialization value for KV cache elements
622 -
623 - Returns:
624 - kv_cache: 2 x num_blocks x block_size x num_heads x head_size
625 - for backend 'FLASH_ATTN'
626 -
627 - """
628 - if backend != "FLASH_ATTN":
629 - raise ValueError(f"Unknown backend value: '{backend}'. Expected 'FLASH_ATTN'.")
630 - kv_cache = torch.rand((2, num_blocks, block_size, num_heads, head_size)).to(device)
631 - if default_val is not None:
632 - kv_cache[:, :, :] = default_val
633 - return kv_cache
634 -
635 -
636 -def _num_tokens_to_min_blocks(num_tokens: int, block_size: int) -> int:
637 - """Compute the minimum number of blocks required to hold num_tokens tokens,
638 - given block_size
639 - """
640 - return (num_tokens + block_size - 1) // block_size
641 -
642 -
643 -def make_empty_slot_mapping_tensor(device: torch.device | str):
644 - return maybe_make_long_tensor([], device)
645 -
646 -
647 -def make_empty_block_tables_tensor(device: torch.device | str):
648 - return torch.tensor([], device=device)
649 -
650 -
651 -def split_slot_mapping(
652 - slot_mapping_list: torch.Tensor,
653 - seq_lens: list[int],
654 - device: torch.device | str,
655 -):
656 - """Split a slot mapping into valid prefill- and decode-phase slot mappings.
657 -
658 - Context:
659 - * Your goal is to test (1) prefill of N prompts, with prompt-lengths
660 - {K_i \\forall i \\in [0,N)}, followed by (2) decoding of a single token
661 - for all N prompts (N tokens total); the resultant sequence lengths
662 - after decode would be {K_i + 1 for i \\in [0,N)}
663 - * The test you want to do requires (1) having the prefill slot mapping
664 - for all tokens present during prefill, the number of which is
665 - M = \\sum_i{K_i}, and (2) having the decode slot mapping for all N
666 - decoded tokens
667 -
668 - This function consumes a single 1D slot mapping, which is the
669 - concatenation of N slot mappings each of length K_i + 1 (corresponding
670 - to the sequence lengths after decode), with a total length of
671 - P = \\sum_i{K_i + 1} = M + N
672 -
673 - The prefill-phase slot mapping results from excising the (K_i + 1)-th entry
674 - from each of the N subsequences in the slot mapping (i.e. omitting the
675 - decoded token's mapping.)
676 -
677 - The N excised entries are appended to obtain the decode-phase slot mapping
678 -
679 - Arguments:
680 - slot_mapping_list: Length-P 1D slot mapping (as list) reflecting all N
681 - post-decode sequences
682 - seq_lens: list of N post-decode sequence lengths (K_i + 1 in the
683 - description above)
684 - device: cuda, cpu, etc.
685 -
686 - Returns:
687 - prefill_slot_mapping: Length-M 1D slot mapping (as Tensor)
688 - reflecting all N prefill prompts
689 - decode_slot_mapping: Length-N 1D slot mapping (as Tensor) reflecting
690 - all N decoded tokens
691 -
692 - """
693 - prefill_slot_mapping = []
694 - decode_slot_mapping = []
695 -
696 - base_idx = 0
697 - for seq_len in seq_lens:
698 - prefill_slot_mapping.extend(
699 - slot_mapping_list[base_idx : (base_idx + seq_len - 1)]
700 - )
701 - decode_slot_mapping.append(slot_mapping_list[base_idx + seq_len - 1])
702 - base_idx += seq_len
703 -
704 - return (
705 - maybe_make_long_tensor(prefill_slot_mapping, device),
706 - maybe_make_long_tensor(decode_slot_mapping, device),
707 - )
708 -
709 -
710 -def make_block_tables_slot_mapping(
711 - block_size: int,
712 - seq_lens: list[int],
713 - device: torch.device | str,
714 - block_base_addr: int = 0,
715 -) -> tuple[torch.Tensor, list[int], int]:
716 - """Construct fake block tables & slot mappings.
717 -
718 - For a sequence with num_tokens tokens the minimum number
719 - of required KV cache blocks is
720 -
721 - num_blocks = (num_tokens + block_size - 1) // block_size
722 -
723 - Then the minimum KV cache size in blocks is
724 -
725 - total_cache_blocks = sum(num_blocks for all seqs)
726 -
727 - Then, the blocktable mapping counts downward from
728 -
729 - block_base_addr + total_cache_blocks
730 -
731 - to
732 -
733 - block_base_addr
734 -
735 -
736 - The constructed block-tables and slot-mapping are sized to the
737 - lengths of the sequences in their entirety (as reflected by seq_lens),
738 - i.e. the total of prefill prompt tokens + decoded tokens.
739 -
740 - Arguments:
741 - block_size: number of offsets per block
742 - seq_lens: list of token-counts for each sequence
743 - block_base_addr: the block table base address
744 - device: CPU or CUDA device
745 -
746 - Return:
747 - * block_tables_tensor: block table for sequence
748 - * slot_mapping_list: slot mapping for sequence
749 - * max_block_idx: the highest block address within this block table
750 -
751 - """
752 - # Provision minimum number of KV cache blocks
753 - num_blocks_list = [
754 - _num_tokens_to_min_blocks(num_tokens, block_size) for num_tokens in seq_lens
755 - ]
756 - max_block_table_len = max(num_blocks_list)
757 - block_table_pad_tokens = 10
758 -
759 - block_tables = []
760 - slot_mapping_list = []
761 - # Compute uppermost address of block table
762 - total_cache_blocks = sum(num_blocks_list)
763 - block_base_idx = block_base_addr + total_cache_blocks
764 - max_block_idx = block_base_idx
765 - for sdx, num_tokens in enumerate(seq_lens):
766 - num_blocks = num_blocks_list[sdx]
767 - block_table = list(range(block_base_idx, block_base_idx - num_blocks, -1))
768 - for idx in range(num_tokens):
769 - mapping_value = (idx % block_size) + block_table[
770 - idx // block_size
771 - ] * block_size
772 - slot_mapping_list.append(mapping_value)
773 -
774 - block_base_idx -= num_blocks
775 - block_tables.append(block_table)
776 -
777 - block_tables_tensor = make_tensor_with_pad(
778 - block_tables,
779 - max_len=max_block_table_len + block_table_pad_tokens,
780 - pad=0,
781 - dtype=torch.int,
782 - device=device,
783 - )
784 -
785 - return (block_tables_tensor, slot_mapping_list, max_block_idx)
786 -
787 -
788 -def assert_actual_matches_ideal(
789 - test_params: PhaseTestParameters, output_under_test: torch.Tensor, backend: str
790 -) -> None:
791 - """Assert that observed output matches the ideal output
792 - contained in the test parameters data structure.
793 -
794 - Arguments:
795 - test_params: Test parameters including packed ideal output
796 - output_under_test: actually observed output value
797 - backend: attention backend name; only 'FLASH_ATTN' is supported
798 -
799 - """
800 - ideal_output = test_params.packed_qkvo.ideal_output
801 - if backend != "FLASH_ATTN":
802 - raise ValueError(f"Unknown backend value: '{backend}'. Expected 'FLASH_ATTN'.")
803 - # For FlashAttention override the accuracy thresholds to non default
804 - # values since we notice a higher difference between the ideal and
805 - # actual output.
806 - torch.testing.assert_close(
807 - ideal_output, output_under_test.view_as(ideal_output), atol=0.01, rtol=0.016
808 - )
809 -
810 -
91 # Copied/modified from torch._refs.__init__.py
92 def fp8_allclose(
93 a: TensorLikeType,
tests/models/multimodal/generation/test_pixtral.py
+2 -62
@@ -6,15 +6,8 @@ from typing import TYPE_CHECKING, Any
6
7 import pytest
8 import torch
9 -from mistral_common.multimodal import download_image
10 -from mistral_common.protocol.instruct.chunk import ImageURLChunk
11 -from mistral_common.protocol.instruct.request import ChatCompletionRequest
12 -from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
13 -from mistral_common.tokens.tokenizers.multimodal import image_from_chunk
14 -from transformers import AutoProcessor
15 -
16 -from vllm import SamplingParams, TextPrompt, TokensPrompt
17 -from vllm.inputs import MultiModalDataBuiltins
9 +
10 +from vllm import SamplingParams
11 from vllm.logprobs import Logprob, SampleLogprobs
12 from vllm.model_executor.models.pixtral import _make_packed_sequence_metadata
13 from vllm.platforms import current_platform
@@ -56,59 +49,6 @@ def _create_msg_format(urls: list[str]) -> list[dict[str, Any]]:
49 ]
50
51
59 -def _create_msg_format_hf(urls: list[str]) -> list[dict[str, Any]]:
60 - return [
61 - {
62 - "role": "user",
63 - "content": [
64 - {
65 - "type": "text",
66 - "content": PROMPT,
67 - },
68 - *({"type": "image", "image": download_image(url)} for url in urls),
69 - ],
70 - }
71 - ]
72 -
73 -
74 -def _create_engine_inputs(urls: list[str]) -> TokensPrompt:
75 - msg = _create_msg_format(urls)
76 -
77 - tokenizer = MistralTokenizer.from_model("pixtral")
78 -
79 - request = ChatCompletionRequest(messages=msg) # type: ignore[type-var]
80 - tokenized = tokenizer.encode_chat_completion(request)
81 -
82 - engine_inputs = TokensPrompt(prompt_token_ids=tokenized.tokens)
83 -
84 - images = []
85 - for chunk in request.messages[0].content:
86 - if isinstance(chunk, ImageURLChunk):
87 - images.append(image_from_chunk(chunk))
88 -
89 - mm_data = MultiModalDataBuiltins(image=images)
90 - engine_inputs["multi_modal_data"] = mm_data
91 -
92 - return engine_inputs
93 -
94 -
95 -def _create_engine_inputs_hf(urls: list[str]) -> TextPrompt:
96 - msg = _create_msg_format_hf(urls)
97 -
98 - tokenizer = AutoProcessor.from_pretrained("mistral-community/pixtral-12b")
99 - prompt = tokenizer.apply_chat_template(msg)
100 -
101 - images = []
102 - for chunk in msg[0]["content"]:
103 - if chunk["type"] == "image":
104 - images.append(chunk["image"])
105 -
106 - mm_data = MultiModalDataBuiltins(image=images)
107 - engine_inputs = TextPrompt(prompt=prompt, multi_modal_data=mm_data)
108 -
109 - return engine_inputs
110 -
111 -
52 SAMPLING_PARAMS = SamplingParams(max_tokens=512, temperature=0.0, logprobs=5)
53 LIMIT_MM_PER_PROMPT = dict(image=4)
54
tests/models/multimodal/generation/vlm_utils/model_utils.py
+1 -44
@@ -6,7 +6,6 @@ typically specific to a small subset of models.
6 """
7
8 import types
9 -from pathlib import PosixPath
9
10 import numpy as np
11 import numpy.typing as npt
@@ -29,7 +28,7 @@ from vllm.logprobs import SampleLogprobs
28 from vllm.platforms import current_platform
29 from vllm.utils.collection_utils import is_list_of
30
32 -from .....conftest import HfRunner, ImageAsset, ImageTestAssets
31 +from .....conftest import HfRunner, ImageTestAssets
32 from .types import RunnerOutput
33
34
@@ -48,17 +47,6 @@ def blip2_vllm_to_hf_output(vllm_output: RunnerOutput, model: str) -> RunnerOutp
47 return hf_output_ids, hf_output_str, out_logprobs
48
49
51 -def qwen_vllm_to_hf_output(
52 - vllm_output: RunnerOutput, model: str
53 -) -> tuple[list[int], str, SampleLogprobs | None]:
54 - """Sanitize vllm output [qwen models] to be comparable with hf output."""
55 - output_ids, output_str, out_logprobs = vllm_output
56 -
57 - hf_output_str = output_str + "<|endoftext|>"
58 -
59 - return output_ids, hf_output_str, out_logprobs
60 -
61 -
50 def qwen2_vllm_to_hf_output(
51 vllm_output: RunnerOutput, model: str
52 ) -> tuple[list[int], str, SampleLogprobs | None]:
@@ -239,37 +227,6 @@ def get_llava_embeddings(image_assets: ImageTestAssets):
227 return [asset.image_embeds for asset in image_assets]
228
229
242 -####### Prompt path encoders for models that need models on disk
243 -def qwen_prompt_path_encoder(
244 - tmp_path: PosixPath, prompt: str, assets: list[ImageAsset] | ImageTestAssets
245 -) -> str:
246 - """Given a temporary dir path, export one or more image assets into the
247 - tempdir & replace its contents with the local path to the string so that
248 - the HF version of Qwen-VL can resolve the path and load the image in its
249 - forward() call.
250 -
251 - Args:
252 - tmp_path: Tempdir for test under consideration.
253 - prompt: Prompt with image placeholders.
254 - assets: list of image assets whose len equals the num placeholders.
255 -
256 - """
257 - # Ensure that the number of placeholders matches the number of assets;
258 - # If this is not true, the test is probably written incorrectly.
259 - assert prompt.count("<img></img>") == len(assets)
260 -
261 - # Replace the placeholders with local paths to the exported assets
262 - for asset in assets:
263 - image_tmp_path = tmp_path / f"{asset.name}.jpg"
264 - asset.pil_image.save(image_tmp_path)
265 - prompt = prompt.replace(
266 - "<img></img>",
267 - f"<img>{image_tmp_path}</img>",
268 - 1,
269 - )
270 - return prompt
271 -
272 -
230 ####### Model-specific HuggingFace runner patchers
231 def deepseekvl2_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
232 """Patches and returns an instance of the HfRunner to use for GLM4."""
tests/parser/engine/trace_builder.py
-20
@@ -1176,23 +1176,3 @@ def build_samples(model: str) -> tuple[Sample, ...]:
1176 builder = _BUILDERS[model]
1177 scenarios = _KIMI_K2_SCENARIOS if model == "kimi_k2" else SCENARIOS
1178 return tuple(s for s in (builder(sc) for sc in scenarios) if s is not None)
1179 -
1180 -
1181 -def build_sample(model: str, scenario: Scenario) -> Sample | None:
1182 - """Build a single sample for one model + scenario."""
1183 - return _BUILDERS[model](scenario)
1184 -
1185 -
1186 -def build_scaling_sample(
1187 - model: str, token_count: int, validate: bool = False
1188 -) -> Sample:
1189 - """Build a sample with approximately *token_count* tokens."""
1190 - sentence = "The quick brown fox jumps over the lazy dog. "
1191 - text = sentence * (token_count // 10 + 1)
1192 - scenario = Scenario(
1193 - id=f"scaling-{token_count}",
1194 - description=f"Scaling test with ~{token_count} tokens",
1195 - reasoning=text,
1196 - tool_calls=[_READ_TOOL],
1197 - )
1198 - return _BUILDERS[model](scenario, validate=validate)
tests/quantization/test_configs.py
-9
@@ -5,20 +5,11 @@
5 Run `pytest tests/quantization/test_configs.py --forked`.
6 """
7
8 -from dataclasses import dataclass
9 -
8 import pytest
9
10 from vllm.config import ModelConfig
11 from vllm.platforms import current_platform
12
15 -
16 -@dataclass
17 -class ModelPair:
18 - model_marlin: str
19 - model_gptq: str
20 -
21 -
13 # Model Id // Quantization Arg // Expected Type
14 MODEL_ARG_EXPTYPES = [
15 # AUTOGPTQ
tests/utils.py
-161
@@ -1,7 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -import asyncio
4 import atexit
5 import contextlib
6 import copy
@@ -23,7 +22,6 @@ from contextlib import ExitStack, contextmanager
22 from multiprocessing import Process, get_context
23 from pathlib import Path
24 from typing import TYPE_CHECKING, Any, Literal, cast
26 -from unittest.mock import patch
25
26 import anthropic
27 import cloudpickle
@@ -33,12 +31,9 @@ import pytest
31 import requests
32 import torch
33 import torch.nn.functional as F
36 -from huggingface_hub.constants import HF_HUB_OFFLINE
37 -from openai.types.completion import Completion
34 from typing_extensions import ParamSpec
35
36 import vllm.envs as envs
41 -from tests.models.utils import TextTextLogprobs
37 from vllm.distributed import (
38 ensure_model_parallel_initialized,
39 init_distributed_environment,
@@ -52,7 +47,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
47 from vllm.model_executor.model_loader import get_model_loader
48 from vllm.platforms import current_platform
49 from vllm.tokenizers import get_tokenizer
55 -from vllm.transformers_utils.repo_utils import hf_api
50 from vllm.utils.argparse_utils import FlexibleArgumentParser
51 from vllm.utils.mem_constants import GB_bytes
52 from vllm.utils.network_utils import get_open_port
@@ -69,24 +63,6 @@ logger = init_logger(__name__)
63 FP8_DTYPE = current_platform.fp8_dtype()
64
65
72 -def prewarm_hf_cache(assets: list[tuple[str, str]]) -> None:
73 - """Pre-populate the HF cache for (repo_id, filename) pairs that upstream
74 - trust_remote_code modules would otherwise fetch from third-party CDNs
75 - (often unreachable from US-based CI)."""
76 - if HF_HUB_OFFLINE:
77 - return
78 - for repo_id, filename in assets:
79 - try:
80 - hf_api().hf_hub_download(repo_id=repo_id, filename=filename)
81 - except Exception as e:
82 - logger.warning(
83 - "Failed to prefetch %s/%s: %r. Tests depending on this asset may fail.",
84 - repo_id,
85 - filename,
86 - e,
87 - )
88 -
89 -
66 if current_platform.is_cuda():
67 from vllm.third_party.pynvml import (
68 nvmlDeviceGetHandleByIndex,
@@ -1556,17 +1532,6 @@ def assert_rocm_custom_allreduce_backend_state_on_worker(
1532 )
1533
1534
1559 -@contextmanager
1560 -def error_on_warning(category: type[Warning] = Warning):
1561 - """Within the scope of this context manager, tests will fail if any warning
1562 - of the given category is emitted.
1563 - """
1564 - with warnings.catch_warnings():
1565 - warnings.filterwarnings("error", category=category)
1566 -
1567 - yield
1568 -
1569 -
1535 def get_physical_device_indices(devices: list[int]):
1536 visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
1537 if visible_devices is None:
@@ -2157,88 +2122,6 @@ def multi_gpu_only(*, num_gpus: int = 2):
2122 return wrapper
2123
2124
2160 -async def completions_with_server_args(
2161 - prompts: list[str],
2162 - model_name: str,
2163 - server_cli_args: list[str],
2164 - num_logprobs: int | None,
2165 - max_wait_seconds: int = 240,
2166 - max_tokens: int | list = 5,
2167 -) -> list[Completion]:
2168 - """Construct a remote OpenAI server, obtain an async client to the
2169 - server & invoke the completions API to obtain completions.
2170 -
2171 - Args:
2172 - prompts: test prompts
2173 - model_name: model to spin up on the vLLM server
2174 - server_cli_args: CLI args for starting the server
2175 - num_logprobs: Number of logprobs to report (or `None`)
2176 - max_wait_seconds: timeout interval for bringing up server.
2177 - Default: 240sec
2178 - max_tokens: max_tokens value for each of the given input prompts.
2179 - if only one max_token value is given, the same value is used
2180 - for all the prompts.
2181 -
2182 - Returns:
2183 - OpenAI Completion instance
2184 -
2185 - """
2186 - if isinstance(max_tokens, int):
2187 - max_tokens = [max_tokens] * len(prompts)
2188 -
2189 - assert len(max_tokens) == len(prompts)
2190 -
2191 - outputs = None
2192 - with RemoteOpenAIServer(
2193 - model_name, server_cli_args, max_wait_seconds=max_wait_seconds
2194 - ) as server:
2195 - client = server.get_async_client()
2196 - outputs = [
2197 - client.completions.create(
2198 - model=model_name,
2199 - prompt=[p],
2200 - temperature=0,
2201 - stream=False,
2202 - max_tokens=max_tok,
2203 - logprobs=num_logprobs,
2204 - )
2205 - for p, max_tok in zip(prompts, max_tokens)
2206 - ]
2207 - outputs = await asyncio.gather(*outputs)
2208 -
2209 - assert outputs is not None, "Completion API call failed."
2210 -
2211 - return outputs
2212 -
2213 -
2214 -def get_client_text_generations(completions: list[Completion]) -> list[str]:
2215 - """Extract generated tokens from the output of a
2216 - request made to an Open-AI-protocol completions endpoint.
2217 - """
2218 - assert all([len(x.choices) == 1 for x in completions])
2219 - return [x.choices[0].text for x in completions]
2220 -
2221 -
2222 -def get_client_text_logprob_generations(
2223 - completions: list[Completion],
2224 -) -> list[TextTextLogprobs]:
2225 - """Operates on the output of a request made to an Open-AI-protocol
2226 - completions endpoint; obtains top-rank logprobs for each token in
2227 - each {class}`SequenceGroup`
2228 - """
2229 - text_generations = get_client_text_generations(completions)
2230 - text = "".join(text_generations)
2231 - return [
2232 - (
2233 - text_generations,
2234 - text,
2235 - (None if x.logprobs is None else x.logprobs.top_logprobs),
2236 - )
2237 - for completion in completions
2238 - for x in completion.choices
2239 - ]
2240 -
2241 -
2125 def has_module_attribute(module_name, attribute_name):
2126 """Helper function to check if a module has a specific attribute."""
2127 try:
@@ -2267,50 +2150,6 @@ def get_attn_backend_list_based_on_platform() -> list[str]:
2150 raise ValueError("Unsupported platform")
2151
2152
2270 -@contextmanager
2271 -def override_cutlass_fp8_supported(value: bool):
2272 - with patch(
2273 - "vllm.model_executor.layers.quantization.utils.w8a8_utils.cutlass_fp8_supported",
2274 - return_value=value,
2275 - ):
2276 - yield
2277 -
2278 -
2279 -def disable_aiter_plain_rmsnorm(monkeypatch) -> None:
2280 - """Patch dispatch_rocm_rmsnorm_func so the plain (non-fused) rms_norm path
2281 - always uses the native float32 kernel for the duration of a test.
2282 -
2283 - The fused path (rms_norm2d_with_add, selected when with_fused_add=True) is
2284 - left on AITER -- only the plain path is redirected to native.
2285 -
2286 - AITER's plain rms_norm accumulates variance in bfloat16 (~1 ULP/call),
2287 - which drifts the KV cache over many decode steps. This drift is irrelevant
2288 - for a trained model (rank-1/rank-2 gap ~1-3 nats >> 1 ULP), but breaks
2289 - logprob comparison tests with randomly-initialised models like
2290 - TitanML/tiny-mixtral whose rank-1/rank-2 gap is only O(1/sqrt(V)) ~0.006
2291 - nats -- smaller than the accumulated per-step error.
2292 - """
2293 - import torch
2294 -
2295 - import vllm.model_executor.layers.layernorm as _ln_mod
2296 - from vllm.model_executor.layers.layernorm import rms_norm as _native
2297 -
2298 - _orig = _ln_mod.dispatch_rocm_rmsnorm_func
2299 -
2300 - def _native_plain(
2301 - with_fused_add: bool, dtype: torch.dtype, use_aiter: bool = False
2302 - ):
2303 - if (
2304 - use_aiter
2305 - and not with_fused_add
2306 - and dtype in (torch.float16, torch.bfloat16)
2307 - ):
2308 - return _native
2309 - return _orig(with_fused_add, dtype, use_aiter)
2310 -
2311 - monkeypatch.setattr(_ln_mod, "dispatch_rocm_rmsnorm_func", _native_plain)
2312 -
2313 -
2153 def prep_prompts(batch_size: int, ln_range: tuple[int, int] = (800, 1100)):
2154 """Generate prompts which a bunch of assignments,
2155 then asking for the value of one of them.
tests/v1/attention/utils.py
-21
@@ -262,27 +262,6 @@ def create_vllm_config(
262 )
263
264
265 -def create_dummy_kv_cache(
266 - block_size: int,
267 - num_kv_heads: int,
268 - head_size: int,
269 - dtype: torch.dtype,
270 - device: torch.device,
271 - num_blocks: int = 100,
272 -) -> torch.Tensor:
273 - """Create a dummy KV cache tensor for testing."""
274 - kv_cache = torch.randn(
275 - num_blocks,
276 - 2, # K and V
277 - block_size,
278 - num_kv_heads,
279 - head_size,
280 - dtype=dtype,
281 - device=device,
282 - )
283 - return kv_cache
284 -
285 -
265 @dataclass
266 class BackendConfig:
267 name: str
tests/v1/engine/utils.py
-17
@@ -196,23 +196,6 @@ def _create_random_top_token_test_matrix(
196 return matrix, prompt_token_ranks
197
198
199 -def decode_token(
200 - tok_id: int,
201 - tokenizer: PythonBackend,
202 -) -> str:
203 - """Reproduce the process of detokenizing a token for testing purposes.
204 -
205 - Args:
206 - tok_id: token id to detokenize
207 - tokenizer: tokenizer to use for detokenization
208 -
209 - Returns:
210 - string representation of token
211 -
212 - """
213 - return tokenizer.convert_ids_to_tokens(tok_id)
214 -
215 -
199 def generate_dummy_sample_logprobs(
200 sampled_tokens_list: list,
201 num_logprobs: int,