[Refactor] Remove dead tests utils (#58803)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
Wentao Ye
@@ -317,14 +317,6 @@ def example_system_message() -> str:
|
||||
return f.read()
|
||||
|
||||
|
||||
class DecoderPromptType(Enum):
|
||||
"""For encoder/decoder models only."""
|
||||
|
||||
CUSTOM = 1
|
||||
NONE = 2
|
||||
EMPTY_STR = 3
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def example_long_prompts() -> list[str]:
|
||||
return [prompt for filename in _LONG_PROMPTS for prompt in _read_prompts(filename)]
|
||||
|
||||
@@ -15,7 +15,6 @@ import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
@@ -138,36 +137,6 @@ def server(
|
||||
proc.kill()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Polling helper (200-lie workaround)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def poll_until(
|
||||
predicate: Callable[[], bool],
|
||||
timeout: float = 10.0,
|
||||
interval: float = 0.5,
|
||||
) -> bool:
|
||||
"""Poll predicate() until it returns True or timeout expires.
|
||||
|
||||
Workaround for the vLLM sleep/wake "200-lie" — the HTTP endpoints may
|
||||
return 200 before the underlying operation is complete, so callers that
|
||||
need to verify state *after* an operation can use this helper instead of
|
||||
assuming the 200 means completion.
|
||||
|
||||
Returns True if predicate became true within timeout, False otherwise.
|
||||
"""
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
if predicate():
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(interval)
|
||||
return False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HTTP helpers — generation
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -191,27 +160,6 @@ def gen(url, prompt="The capital of France is", max_tokens=8, timeout=30):
|
||||
return None
|
||||
|
||||
|
||||
def gen_with_logprobs(
|
||||
url, prompt="The capital of France is", max_tokens=8, logprobs=5, timeout=30
|
||||
):
|
||||
"""Fire a /v1/completions request with logprobs; return JSON or None."""
|
||||
try:
|
||||
r = requests.post(
|
||||
f"{url}/v1/completions",
|
||||
json={
|
||||
"model": "m",
|
||||
"prompt": prompt,
|
||||
"max_tokens": max_tokens,
|
||||
"temperature": 0,
|
||||
"logprobs": logprobs,
|
||||
},
|
||||
timeout=timeout,
|
||||
)
|
||||
return r.json()
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def ok(resp) -> bool:
|
||||
"""True iff resp is a successful completion (has choices, no error key)."""
|
||||
return (
|
||||
@@ -399,37 +347,3 @@ def get_world_size(url, include_dp=True):
|
||||
params={"include_dp": include_dp},
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GPU / metrics helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def gpu_free_bytes(device: int = 0) -> int:
|
||||
"""Read GPU free bytes via subprocess to avoid import-time torch init."""
|
||||
out = subprocess.check_output(
|
||||
[
|
||||
sys.executable,
|
||||
"-c",
|
||||
f"import torch; f,_=torch.accelerator.get_memory_info({device}); print(f)",
|
||||
],
|
||||
timeout=10,
|
||||
)
|
||||
return int(out.strip())
|
||||
|
||||
|
||||
def sleep_metrics(url):
|
||||
"""Return (awake, weights_offloaded, discard_all) from /metrics."""
|
||||
try:
|
||||
from prometheus_client.parser import text_string_to_metric_families
|
||||
except ImportError:
|
||||
return None, None, None
|
||||
|
||||
r = requests.get(f"{url}/metrics", timeout=5)
|
||||
vals: dict = {}
|
||||
for family in text_string_to_metric_families(r.text):
|
||||
if family.name == "vllm:engine_sleep_state":
|
||||
for s in family.samples:
|
||||
vals[s.labels.get("sleep_state", "")] = s.value
|
||||
return vals.get("awake"), vals.get("weights_offloaded"), vals.get("discard_all")
|
||||
|
||||
@@ -309,46 +309,6 @@ def test_paged_attention(
|
||||
torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
def ref_multi_query_kv_attention(
|
||||
cu_seq_lens: list[int],
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
scale: float,
|
||||
alibi_bias: list[torch.Tensor] | None,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
num_seqs = len(cu_seq_lens) - 1
|
||||
ref_outputs: list[torch.Tensor] = []
|
||||
if alibi_bias:
|
||||
assert len(alibi_bias) == num_seqs
|
||||
for i in range(num_seqs):
|
||||
start_idx = cu_seq_lens[i]
|
||||
end_idx = cu_seq_lens[i + 1]
|
||||
seq_len = end_idx - start_idx
|
||||
|
||||
# Create attention mask. ALiBi already includes a tril causal mask.
|
||||
if alibi_bias:
|
||||
attn_mask = alibi_bias[i]
|
||||
else:
|
||||
attn_mask = torch.triu(
|
||||
torch.ones(seq_len, seq_len, dtype=dtype), diagonal=1
|
||||
)
|
||||
attn_mask = attn_mask * torch.finfo(dtype).min
|
||||
attn_mask = attn_mask.to(dtype=dtype)
|
||||
|
||||
ref_output = ref_masked_attention(
|
||||
query[start_idx:end_idx],
|
||||
key[start_idx:end_idx],
|
||||
value[start_idx:end_idx],
|
||||
scale,
|
||||
attn_mask=attn_mask,
|
||||
)
|
||||
ref_outputs.append(ref_output)
|
||||
|
||||
return torch.cat(ref_outputs, dim=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("attention_cls", [Attention, MMEncoderAttention])
|
||||
def test_num_heads_not_divisible_by_num_kv_heads(attention_cls: type) -> None:
|
||||
head_size = 64
|
||||
|
||||
@@ -120,37 +120,6 @@ def causal_conv1d_update_ref(
|
||||
return (out if activation is None else F.silu(out)).to(dtype=dtype_in)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("itype", [torch.bfloat16, torch.float])
|
||||
@pytest.mark.parametrize("silu_activation", [True])
|
||||
@pytest.mark.parametrize("has_bias", [True])
|
||||
def causal_conv1d_opcheck_fn(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
cu_seq_len: torch.Tensor | None = None,
|
||||
cache_indices: torch.Tensor | None = None,
|
||||
has_initial_state: torch.Tensor | None = None,
|
||||
conv_states: torch.Tensor | None = None,
|
||||
activation: str | None = "silu",
|
||||
null_block_id: int = NULL_BLOCK_ID,
|
||||
):
|
||||
"""x: (batch, dim, seqlen)
|
||||
weight: (dim, width)
|
||||
bias: (dim,)
|
||||
seq_idx: (batch, seqlen)
|
||||
initial_states: (batch, dim, width - 1)
|
||||
final_states_out: (batch, dim, width - 1), to be written to
|
||||
activation: either None or "silu" or "swish"
|
||||
|
||||
out: (batch, dim, seqlen)
|
||||
"""
|
||||
if activation not in [None, "silu", "swish"]:
|
||||
raise NotImplementedError("activation must be None, silu, or swish")
|
||||
if x.stride(-1) != 1:
|
||||
x = x.contiguous()
|
||||
bias = bias.contiguous() if bias is not None else None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("itype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("silu_activation", [False, True])
|
||||
@pytest.mark.parametrize("has_bias", [False, True])
|
||||
|
||||
@@ -448,24 +448,6 @@ if cutlass_fp4_supported() or has_flashinfer_cutlass_fused_moe():
|
||||
]
|
||||
|
||||
|
||||
def _slice(rank: int, num_local_experts: int, t: torch.Tensor) -> torch.Tensor:
|
||||
s = rank * num_local_experts
|
||||
e = s + num_local_experts
|
||||
return t[s:e]
|
||||
|
||||
|
||||
def make_cutlass_strides(
|
||||
e: int,
|
||||
n: int,
|
||||
k: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
ab_strides1 = torch.full((e,), k, device="cuda", dtype=torch.int64)
|
||||
ab_strides2 = torch.full((e,), n, device="cuda", dtype=torch.int64)
|
||||
c_strides1 = torch.full((e,), 2 * n, device="cuda", dtype=torch.int64)
|
||||
c_strides2 = torch.full((e,), k, device="cuda", dtype=torch.int64)
|
||||
return ab_strides1, ab_strides2, c_strides1, c_strides2
|
||||
|
||||
|
||||
def make_fused_experts(
|
||||
fused_experts_type: mk.FusedMoEExpertsModular,
|
||||
moe: FusedMoEConfig,
|
||||
|
||||
@@ -275,11 +275,6 @@ def maybe_roundup_layer_hidden_size(
|
||||
return hidden_size
|
||||
|
||||
|
||||
def rank_chunk(num: int, r: int, w: int) -> int:
|
||||
rem = num % w
|
||||
return (num // w) + (1 if r < rem else 0)
|
||||
|
||||
|
||||
def chunk_by_rank(
|
||||
t: torch.Tensor,
|
||||
r: int,
|
||||
@@ -294,19 +289,6 @@ def chunk_by_rank(
|
||||
return t
|
||||
|
||||
|
||||
def maybe_chunk_by_rank(
|
||||
t: torch.Tensor | None,
|
||||
r: int,
|
||||
w: int,
|
||||
dim: int = 0,
|
||||
device: torch.device | None = None,
|
||||
) -> torch.Tensor | None:
|
||||
if t is not None:
|
||||
return chunk_by_rank(t, r, w, dim, device)
|
||||
else:
|
||||
return t
|
||||
|
||||
|
||||
def tp_chunk_gate_up(
|
||||
w: torch.Tensor,
|
||||
tp_rank: int,
|
||||
@@ -676,26 +658,6 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
|
||||
return True, None
|
||||
|
||||
|
||||
def chunk_scales_by_rank(
|
||||
t: torch.Tensor | None,
|
||||
r: int,
|
||||
w: int,
|
||||
device: torch.device | None = None,
|
||||
) -> torch.Tensor | None:
|
||||
if t is not None and t.numel() > 1:
|
||||
# Calculate start index by summing chunk sizes for all previous ranks
|
||||
# start = sum(rank_chunk(t.shape[0], i, w) for i in range(r))
|
||||
# chunk = rank_chunk(t.shape[0], r, w)
|
||||
# t = t[start:(start + chunk)]
|
||||
chunk = rank_chunk(t.shape[0], r, w)
|
||||
t = t[(r * chunk) : max(t.shape[0], (r + 1) * chunk)]
|
||||
|
||||
if t is not None and device is not None:
|
||||
t = t.to(device)
|
||||
|
||||
return t
|
||||
|
||||
|
||||
def chunk_scales(
|
||||
t: torch.Tensor | None,
|
||||
start: int,
|
||||
@@ -771,19 +733,6 @@ def _quantize_fp8_halves(
|
||||
)
|
||||
|
||||
|
||||
def quantization_to_quant_dtype(
|
||||
quantization: str | None,
|
||||
) -> torch.dtype | str | None:
|
||||
if quantization is None:
|
||||
return None
|
||||
elif quantization in ["fp8", "fp8_blocked", "modelopt_fp8"]:
|
||||
return fp8_dtype
|
||||
elif quantization in ["modelopt_fp4"]:
|
||||
return "nvfp4"
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported quantization: {quantization}")
|
||||
|
||||
|
||||
def make_quant_config(
|
||||
quantization: str | None,
|
||||
w1: torch.Tensor,
|
||||
|
||||
@@ -552,16 +552,6 @@ class TestMLP(torch.nn.Module):
|
||||
return x
|
||||
|
||||
|
||||
def make_naive_shared_experts(
|
||||
N: int,
|
||||
K: int,
|
||||
in_dtype: torch.dtype = torch.bfloat16,
|
||||
) -> torch.nn.Module:
|
||||
w1 = torch.randn((K, N * 2), device=DEVICE, dtype=in_dtype) / 15
|
||||
w2 = torch.randn((N, K), device=DEVICE, dtype=in_dtype) / 15
|
||||
return TestMLP(w1, w2, out_dtype=in_dtype)
|
||||
|
||||
|
||||
class RealMLP(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
@@ -668,25 +658,6 @@ def modular_triton_fused_moe(
|
||||
)
|
||||
|
||||
|
||||
def make_shared_experts(
|
||||
N: int,
|
||||
K: int,
|
||||
in_dtype: torch.dtype = torch.bfloat16,
|
||||
quant_dtype: torch.dtype | str | None = None,
|
||||
) -> torch.nn.Module:
|
||||
(_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights(
|
||||
1,
|
||||
N,
|
||||
K,
|
||||
in_dtype=in_dtype,
|
||||
quant_dtype=quant_dtype,
|
||||
)
|
||||
|
||||
return make_shared_experts_with_weights(
|
||||
N, K, in_dtype, w1, w2, w1_s=w1_s, w2_s=w2_s, quant_dtype=quant_dtype
|
||||
)
|
||||
|
||||
|
||||
def check_accuracy(a, b, atol, rtol, percent):
|
||||
"""Allow a mismatch percentage of 1 - percent."""
|
||||
if torch.any(torch.isnan(a)):
|
||||
|
||||
@@ -253,25 +253,6 @@ def dequant(
|
||||
return t.to(out_dtype)
|
||||
|
||||
|
||||
def batched_dequant(
|
||||
t: torch.Tensor,
|
||||
scale: torch.Tensor | None,
|
||||
block_shape: list[int] | None,
|
||||
per_act_token_quant: bool,
|
||||
out_dtype: torch.dtype | None = torch.float32,
|
||||
) -> torch.Tensor:
|
||||
if scale is not None:
|
||||
assert t.shape[0] == scale.shape[0]
|
||||
out = torch.empty_like(t, dtype=out_dtype)
|
||||
for e in range(t.shape[0]):
|
||||
out[e] = dequant(
|
||||
t[e], scale[e], block_shape, per_act_token_quant, out_dtype
|
||||
)
|
||||
return out
|
||||
|
||||
return t.to(out_dtype)
|
||||
|
||||
|
||||
def native_batched_masked_quant_matmul(
|
||||
A: torch.Tensor,
|
||||
B: torch.Tensor,
|
||||
|
||||
+1
-721
@@ -2,11 +2,8 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Kernel test utils."""
|
||||
|
||||
import itertools
|
||||
import random
|
||||
from collections.abc import Sequence
|
||||
from numbers import Number
|
||||
from typing import Any, NamedTuple
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
@@ -17,8 +14,6 @@ from vllm.model_executor.custom_op import op_registry
|
||||
from vllm.model_executor.layers.activation import SiluAndMul
|
||||
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
||||
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
|
||||
from vllm.utils.torch_utils import make_tensor_with_pad
|
||||
from vllm.v1.attention.backend import AttentionType
|
||||
|
||||
# For now, disable "test_aot_dispatch_dynamic" since there are some
|
||||
# bugs related to this test in PyTorch 2.4.
|
||||
@@ -93,721 +88,6 @@ def _assert_deterministic(
|
||||
)
|
||||
|
||||
|
||||
class QKVInputs(NamedTuple):
|
||||
"""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
|
||||
key: torch.Tensor
|
||||
value: torch.Tensor
|
||||
q_seq_lens: list[int]
|
||||
kv_seq_lens: list[int]
|
||||
|
||||
|
||||
class QKVO(NamedTuple):
|
||||
"""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
|
||||
ideal_output: torch.Tensor
|
||||
|
||||
|
||||
class PackedQKVInputs(NamedTuple):
|
||||
"""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
|
||||
* kv_start_loc_list: shared list of key/value start locations within
|
||||
packed tensor
|
||||
* q_seq_lens: query sequence lengths list
|
||||
* kv_seq_lens: shared key/value sequence lengths list
|
||||
|
||||
"""
|
||||
|
||||
query: torch.Tensor
|
||||
key: torch.Tensor
|
||||
value: torch.Tensor
|
||||
q_start_loc_list: list[int] | None
|
||||
kv_start_loc_list: list[int] | None
|
||||
q_seq_lens: list[int] | None
|
||||
kv_seq_lens: list[int] | None
|
||||
|
||||
|
||||
class PackedQKVO(NamedTuple):
|
||||
"""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
|
||||
ideal_output: torch.Tensor
|
||||
|
||||
|
||||
class KVMemoryMap(NamedTuple):
|
||||
"""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
|
||||
slot_mapping: torch.Tensor
|
||||
|
||||
|
||||
class PhaseTestParameters(NamedTuple):
|
||||
"""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
|
||||
kv_mmap: KVMemoryMap | None
|
||||
|
||||
|
||||
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`.
|
||||
|
||||
Returns:
|
||||
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)
|
||||
)
|
||||
|
||||
|
||||
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`.
|
||||
|
||||
Returns:
|
||||
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)
|
||||
)
|
||||
|
||||
|
||||
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.
|
||||
|
||||
"""
|
||||
return None if _list is None else max(_list)
|
||||
|
||||
|
||||
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.
|
||||
|
||||
Arguments:
|
||||
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
|
||||
|
||||
"""
|
||||
# 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
|
||||
mask = mask.masked_fill(mask == 1, float("-inf")).masked_fill(mask == 0, 0.0)
|
||||
return mask
|
||||
|
||||
|
||||
def ref_masked_attention(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
scale: float,
|
||||
custom_mask: torch.Tensor | None = None,
|
||||
q_seq_lens: list | None = None,
|
||||
kv_seq_lens: list | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Reference ("golden") masked attention. Supports two types of masking:
|
||||
|
||||
* Basic attention mask, utilizing {q,kv}_seq_lens args to mask out
|
||||
padding elements
|
||||
* Custom attention mask, which can force an arbitrary mask tensor, i.e.
|
||||
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
|
||||
|
||||
Returns:
|
||||
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
|
||||
|
||||
batch_size = query.shape[0]
|
||||
assert len(q_seq_lens) == batch_size
|
||||
assert len(kv_seq_lens) == batch_size
|
||||
|
||||
attn_weights = scale * torch.einsum("bqhd,bkhd->bhqk", query, key).float()
|
||||
|
||||
# Basic attention mask, derived from seq lens
|
||||
if (q_seq_lens is not None) or (kv_seq_lens is not None):
|
||||
attn_mask = torch.zeros_like(attn_weights)
|
||||
if q_seq_lens is not None:
|
||||
for bdx, plen in enumerate(q_seq_lens):
|
||||
attn_mask[bdx, :, plen:, :] = -torch.inf
|
||||
if kv_seq_lens is not None:
|
||||
for bdx, plen in enumerate(kv_seq_lens):
|
||||
attn_mask[bdx, :, :, plen:] = -torch.inf
|
||||
|
||||
attn_weights = attn_weights + attn_mask.float()
|
||||
|
||||
# Custom attention mask
|
||||
if custom_mask is not None:
|
||||
attn_weights = attn_weights + custom_mask.float()
|
||||
|
||||
attn_weights = torch.softmax(attn_weights, dim=-1).to(value.dtype)
|
||||
out = torch.einsum("bhqk,bkhd->bqhd", attn_weights, value)
|
||||
return out
|
||||
|
||||
|
||||
def make_qkv(
|
||||
batch_size: int,
|
||||
max_q_seq_len: int,
|
||||
max_kv_seq_len: int | None,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
device: torch.device | str,
|
||||
force_kv_seq_lens: list[int] | None = None,
|
||||
attn_type: AttentionType = AttentionType.ENCODER_DECODER,
|
||||
force_max_len: bool = False,
|
||||
) -> tuple[QKVInputs, QKVInputs, QKVInputs]:
|
||||
"""Construct QKV test tensors for self- and cross-attention.
|
||||
|
||||
Generates three query/key/value triplets:
|
||||
|
||||
* "Baseline" query/key/value (for input to reference attention function)
|
||||
* "Prefill" query/key/value (last sequence offset zero'd out, for use as
|
||||
input to prefill kernel)
|
||||
* "Decode" query/key/value (only the last sequence offset from baseline,
|
||||
for use as input to decode kernel)
|
||||
|
||||
Each Q/K/V triplet is associated with a list of q seqlens and a list of k/v
|
||||
seqlens
|
||||
|
||||
Arguments:
|
||||
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)
|
||||
|
||||
"""
|
||||
if force_max_len:
|
||||
q_seq_lens = [max_q_seq_len for _ in range(batch_size)]
|
||||
else:
|
||||
q_seq_lens = [random.randint(2, max_q_seq_len) for _ in range(batch_size)]
|
||||
kv_seq_lens = None
|
||||
if force_kv_seq_lens is not None:
|
||||
kv_seq_lens = force_kv_seq_lens
|
||||
elif attn_type != AttentionType.ENCODER_DECODER:
|
||||
# K,V seq lens match Q for self-attention
|
||||
kv_seq_lens = q_seq_lens
|
||||
else:
|
||||
# K,V seq lens are distinct from Q seq lens & random
|
||||
assert max_kv_seq_len is not None
|
||||
if force_max_len:
|
||||
kv_seq_lens = [max_kv_seq_len] * batch_size
|
||||
else:
|
||||
kv_seq_lens = [random.randint(2, max_kv_seq_len) for _ in range(batch_size)]
|
||||
|
||||
query = torch.rand((batch_size, max_q_seq_len, num_heads, head_size)).to(device)
|
||||
key = torch.rand((batch_size, max_kv_seq_len, num_heads, head_size)).to(device)
|
||||
value = torch.rand((batch_size, max_kv_seq_len, num_heads, head_size)).to(device)
|
||||
|
||||
prefill_query = torch.zeros((batch_size, max_q_seq_len, num_heads, head_size)).to(
|
||||
device
|
||||
)
|
||||
prefill_key = torch.zeros((batch_size, max_kv_seq_len, num_heads, head_size)).to(
|
||||
device
|
||||
)
|
||||
prefill_value = torch.zeros((batch_size, max_kv_seq_len, num_heads, head_size)).to(
|
||||
device
|
||||
)
|
||||
|
||||
decode_query = torch.zeros((batch_size, 1, num_heads, head_size)).to(device)
|
||||
decode_key = torch.zeros((batch_size, 1, num_heads, head_size)).to(device)
|
||||
decode_value = torch.zeros((batch_size, 1, num_heads, head_size)).to(device)
|
||||
|
||||
for bdx, (q_seq_len, kv_seq_len) in enumerate(zip(q_seq_lens, kv_seq_lens)):
|
||||
query[bdx, q_seq_len:, :, :] = 0
|
||||
key[bdx, kv_seq_len:, :, :] = 0
|
||||
value[bdx, kv_seq_len:, :, :] = 0
|
||||
|
||||
prefill_query[bdx, 0 : (q_seq_len - 1), :, :] = query[
|
||||
bdx, 0 : (q_seq_len - 1), :, :
|
||||
]
|
||||
prefill_key[bdx, 0 : (kv_seq_len - 1), :, :] = key[
|
||||
bdx, 0 : (kv_seq_len - 1), :, :
|
||||
]
|
||||
prefill_value[bdx, 0 : (kv_seq_len - 1), :, :] = value[
|
||||
bdx, 0 : (kv_seq_len - 1), :, :
|
||||
]
|
||||
|
||||
decode_query[bdx, :, :, :] = query[bdx, (q_seq_len - 1) : q_seq_len, :, :]
|
||||
decode_key[bdx, :, :, :] = key[bdx, (kv_seq_len - 1) : kv_seq_len, :, :]
|
||||
decode_value[bdx, :, :, :] = value[bdx, (kv_seq_len - 1) : kv_seq_len, :, :]
|
||||
|
||||
prefill_q_seq_lens = [plen - 1 for plen in q_seq_lens]
|
||||
prefill_kv_seq_lens = [plen - 1 for plen in kv_seq_lens]
|
||||
|
||||
decode_q_seq_lens = [1 for _ in q_seq_lens]
|
||||
decode_kv_seq_lens = [1 for _ in kv_seq_lens]
|
||||
|
||||
return (
|
||||
QKVInputs(
|
||||
query, # Overall QKV inputs
|
||||
key,
|
||||
value,
|
||||
q_seq_lens,
|
||||
kv_seq_lens,
|
||||
),
|
||||
QKVInputs(
|
||||
prefill_query, # Prefill subset of QKV sequences
|
||||
prefill_key,
|
||||
prefill_value,
|
||||
prefill_q_seq_lens,
|
||||
prefill_kv_seq_lens,
|
||||
),
|
||||
QKVInputs(
|
||||
decode_query, # Decode subset of KV sequences
|
||||
decode_key,
|
||||
decode_value,
|
||||
decode_q_seq_lens,
|
||||
decode_kv_seq_lens,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
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
|
||||
|
||||
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))
|
||||
|
||||
"""
|
||||
num_tok = sum(seq_lens)
|
||||
num_heads = unpacked_tensor.shape[-2]
|
||||
head_size = unpacked_tensor.shape[-1]
|
||||
start_loc_list = [0] + list(itertools.accumulate(seq_lens))
|
||||
packed_tensor = torch.zeros((num_tok, num_heads, head_size), device=device)
|
||||
|
||||
for bdx, (seq_len, start_loc) in enumerate(zip(seq_lens, start_loc_list)):
|
||||
packed_tensor[start_loc : (start_loc + seq_len), :, :] = unpacked_tensor[
|
||||
bdx, :seq_len, :, :
|
||||
]
|
||||
|
||||
return packed_tensor, start_loc_list
|
||||
|
||||
|
||||
def pack_qkv(qkv: QKVInputs, device: torch.device | str) -> PackedQKVInputs:
|
||||
"""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.
|
||||
|
||||
For Q, number_of_tokens = sum(q_seq_lens).
|
||||
|
||||
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
|
||||
|
||||
Returns:
|
||||
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
|
||||
else:
|
||||
packed_query, q_start_loc_list = pack_tensor(
|
||||
qkv.query, qkv.q_seq_lens, device=device
|
||||
)
|
||||
packed_key, kv_start_loc_list = pack_tensor(qkv.key, qkv.kv_seq_lens, device=device)
|
||||
packed_value, _ = pack_tensor(qkv.value, qkv.kv_seq_lens, device=device)
|
||||
return PackedQKVInputs(
|
||||
packed_query,
|
||||
packed_key,
|
||||
packed_value,
|
||||
q_start_loc_list,
|
||||
kv_start_loc_list,
|
||||
(None if q_start_loc_list is None else qkv.q_seq_lens),
|
||||
qkv.kv_seq_lens,
|
||||
)
|
||||
|
||||
|
||||
def _make_metadata_tensors(
|
||||
seq_lens: list[int] | None,
|
||||
context_lens: list[int] | None,
|
||||
encoder_seq_lens: list[int] | None,
|
||||
device: torch.device | str,
|
||||
) -> tuple[
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
Any,
|
||||
Any,
|
||||
torch.Tensor | None,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
int | None,
|
||||
]:
|
||||
"""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
|
||||
|
||||
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 = maybe_make_int_tensor(seq_lens, device)
|
||||
context_lens_tensor = maybe_make_int_tensor(context_lens, device)
|
||||
max_context_len = maybe_max(context_lens)
|
||||
max_seq_len = maybe_max(seq_lens)
|
||||
|
||||
encoder_seq_lens_tensor = maybe_make_int_tensor(encoder_seq_lens, device)
|
||||
max_encoder_seq_len = None if encoder_seq_lens is None else max(encoder_seq_lens)
|
||||
|
||||
seq_start_loc = None
|
||||
|
||||
if seq_lens_tensor is not None:
|
||||
seq_start_loc = torch.zeros(
|
||||
seq_lens_tensor.shape[0] + 1,
|
||||
dtype=torch.int32,
|
||||
device=seq_lens_tensor.device,
|
||||
)
|
||||
torch.cumsum(
|
||||
seq_lens_tensor, dim=0, dtype=seq_start_loc.dtype, out=seq_start_loc[1:]
|
||||
)
|
||||
|
||||
encoder_seq_start_loc = torch.zeros(
|
||||
encoder_seq_lens_tensor.shape[0] + 1,
|
||||
dtype=torch.int32,
|
||||
device=encoder_seq_lens_tensor.device,
|
||||
)
|
||||
torch.cumsum(
|
||||
encoder_seq_lens_tensor,
|
||||
dim=0,
|
||||
dtype=encoder_seq_start_loc.dtype,
|
||||
out=encoder_seq_start_loc[1:],
|
||||
)
|
||||
|
||||
return (
|
||||
seq_lens_tensor,
|
||||
context_lens_tensor,
|
||||
max_context_len,
|
||||
max_seq_len,
|
||||
seq_start_loc,
|
||||
encoder_seq_lens_tensor,
|
||||
encoder_seq_start_loc,
|
||||
max_encoder_seq_len,
|
||||
)
|
||||
|
||||
|
||||
def make_kv_cache(
|
||||
num_blocks: int,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
block_size: int,
|
||||
device: torch.device | str,
|
||||
backend: str,
|
||||
default_val: float = 0.0,
|
||||
) -> torch.Tensor:
|
||||
"""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
|
||||
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'
|
||||
|
||||
"""
|
||||
if backend != "FLASH_ATTN":
|
||||
raise ValueError(f"Unknown backend value: '{backend}'. Expected 'FLASH_ATTN'.")
|
||||
kv_cache = torch.rand((2, num_blocks, block_size, num_heads, head_size)).to(device)
|
||||
if default_val is not None:
|
||||
kv_cache[:, :, :] = default_val
|
||||
return 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,
|
||||
given block_size
|
||||
"""
|
||||
return (num_tokens + block_size - 1) // block_size
|
||||
|
||||
|
||||
def make_empty_slot_mapping_tensor(device: torch.device | str):
|
||||
return maybe_make_long_tensor([], device)
|
||||
|
||||
|
||||
def make_empty_block_tables_tensor(device: torch.device | str):
|
||||
return torch.tensor([], device=device)
|
||||
|
||||
|
||||
def split_slot_mapping(
|
||||
slot_mapping_list: torch.Tensor,
|
||||
seq_lens: list[int],
|
||||
device: torch.device | str,
|
||||
):
|
||||
"""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
|
||||
{K_i \\forall i \\in [0,N)}, followed by (2) decoding of a single token
|
||||
for all N prompts (N tokens total); the resultant sequence lengths
|
||||
after decode would be {K_i + 1 for i \\in [0,N)}
|
||||
* The test you want to do requires (1) having the prefill slot mapping
|
||||
for all tokens present during prefill, the number of which is
|
||||
M = \\sum_i{K_i}, and (2) having the decode slot mapping for all N
|
||||
decoded tokens
|
||||
|
||||
This function consumes a single 1D slot mapping, which is the
|
||||
concatenation of N slot mappings each of length K_i + 1 (corresponding
|
||||
to the sequence lengths after decode), with a total length of
|
||||
P = \\sum_i{K_i + 1} = M + N
|
||||
|
||||
The prefill-phase slot mapping results from excising the (K_i + 1)-th entry
|
||||
from each of the N subsequences in the slot mapping (i.e. omitting the
|
||||
decoded token's 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.
|
||||
|
||||
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 = []
|
||||
decode_slot_mapping = []
|
||||
|
||||
base_idx = 0
|
||||
for seq_len in seq_lens:
|
||||
prefill_slot_mapping.extend(
|
||||
slot_mapping_list[base_idx : (base_idx + seq_len - 1)]
|
||||
)
|
||||
decode_slot_mapping.append(slot_mapping_list[base_idx + seq_len - 1])
|
||||
base_idx += seq_len
|
||||
|
||||
return (
|
||||
maybe_make_long_tensor(prefill_slot_mapping, device),
|
||||
maybe_make_long_tensor(decode_slot_mapping, device),
|
||||
)
|
||||
|
||||
|
||||
def make_block_tables_slot_mapping(
|
||||
block_size: int,
|
||||
seq_lens: list[int],
|
||||
device: torch.device | str,
|
||||
block_base_addr: int = 0,
|
||||
) -> tuple[torch.Tensor, list[int], int]:
|
||||
"""Construct fake block tables & slot mappings.
|
||||
|
||||
For a sequence with num_tokens tokens the minimum number
|
||||
of required KV cache blocks is
|
||||
|
||||
num_blocks = (num_tokens + block_size - 1) // block_size
|
||||
|
||||
Then the minimum KV cache size in blocks is
|
||||
|
||||
total_cache_blocks = sum(num_blocks for all seqs)
|
||||
|
||||
Then, the blocktable mapping counts downward from
|
||||
|
||||
block_base_addr + total_cache_blocks
|
||||
|
||||
to
|
||||
|
||||
block_base_addr
|
||||
|
||||
|
||||
The constructed block-tables and slot-mapping are sized to the
|
||||
lengths of the sequences in their entirety (as reflected by seq_lens),
|
||||
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
|
||||
|
||||
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
|
||||
]
|
||||
max_block_table_len = max(num_blocks_list)
|
||||
block_table_pad_tokens = 10
|
||||
|
||||
block_tables = []
|
||||
slot_mapping_list = []
|
||||
# Compute uppermost address of block table
|
||||
total_cache_blocks = sum(num_blocks_list)
|
||||
block_base_idx = block_base_addr + total_cache_blocks
|
||||
max_block_idx = block_base_idx
|
||||
for sdx, num_tokens in enumerate(seq_lens):
|
||||
num_blocks = num_blocks_list[sdx]
|
||||
block_table = list(range(block_base_idx, block_base_idx - num_blocks, -1))
|
||||
for idx in range(num_tokens):
|
||||
mapping_value = (idx % block_size) + block_table[
|
||||
idx // block_size
|
||||
] * block_size
|
||||
slot_mapping_list.append(mapping_value)
|
||||
|
||||
block_base_idx -= num_blocks
|
||||
block_tables.append(block_table)
|
||||
|
||||
block_tables_tensor = make_tensor_with_pad(
|
||||
block_tables,
|
||||
max_len=max_block_table_len + block_table_pad_tokens,
|
||||
pad=0,
|
||||
dtype=torch.int,
|
||||
device=device,
|
||||
)
|
||||
|
||||
return (block_tables_tensor, slot_mapping_list, max_block_idx)
|
||||
|
||||
|
||||
def assert_actual_matches_ideal(
|
||||
test_params: PhaseTestParameters, output_under_test: torch.Tensor, backend: str
|
||||
) -> None:
|
||||
"""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
|
||||
|
||||
"""
|
||||
ideal_output = test_params.packed_qkvo.ideal_output
|
||||
if backend != "FLASH_ATTN":
|
||||
raise ValueError(f"Unknown backend value: '{backend}'. Expected 'FLASH_ATTN'.")
|
||||
# For FlashAttention override the accuracy thresholds to non default
|
||||
# values since we notice a higher difference between the ideal and
|
||||
# actual output.
|
||||
torch.testing.assert_close(
|
||||
ideal_output, output_under_test.view_as(ideal_output), atol=0.01, rtol=0.016
|
||||
)
|
||||
|
||||
|
||||
# Copied/modified from torch._refs.__init__.py
|
||||
def fp8_allclose(
|
||||
a: TensorLikeType,
|
||||
|
||||
@@ -6,15 +6,8 @@ from typing import TYPE_CHECKING, Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from mistral_common.multimodal import download_image
|
||||
from mistral_common.protocol.instruct.chunk import ImageURLChunk
|
||||
from mistral_common.protocol.instruct.request import ChatCompletionRequest
|
||||
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
|
||||
from mistral_common.tokens.tokenizers.multimodal import image_from_chunk
|
||||
from transformers import AutoProcessor
|
||||
|
||||
from vllm import SamplingParams, TextPrompt, TokensPrompt
|
||||
from vllm.inputs import MultiModalDataBuiltins
|
||||
from vllm import SamplingParams
|
||||
from vllm.logprobs import Logprob, SampleLogprobs
|
||||
from vllm.model_executor.models.pixtral import _make_packed_sequence_metadata
|
||||
from vllm.platforms import current_platform
|
||||
@@ -56,59 +49,6 @@ def _create_msg_format(urls: list[str]) -> list[dict[str, Any]]:
|
||||
]
|
||||
|
||||
|
||||
def _create_msg_format_hf(urls: list[str]) -> list[dict[str, Any]]:
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"content": PROMPT,
|
||||
},
|
||||
*({"type": "image", "image": download_image(url)} for url in urls),
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _create_engine_inputs(urls: list[str]) -> TokensPrompt:
|
||||
msg = _create_msg_format(urls)
|
||||
|
||||
tokenizer = MistralTokenizer.from_model("pixtral")
|
||||
|
||||
request = ChatCompletionRequest(messages=msg) # type: ignore[type-var]
|
||||
tokenized = tokenizer.encode_chat_completion(request)
|
||||
|
||||
engine_inputs = TokensPrompt(prompt_token_ids=tokenized.tokens)
|
||||
|
||||
images = []
|
||||
for chunk in request.messages[0].content:
|
||||
if isinstance(chunk, ImageURLChunk):
|
||||
images.append(image_from_chunk(chunk))
|
||||
|
||||
mm_data = MultiModalDataBuiltins(image=images)
|
||||
engine_inputs["multi_modal_data"] = mm_data
|
||||
|
||||
return engine_inputs
|
||||
|
||||
|
||||
def _create_engine_inputs_hf(urls: list[str]) -> TextPrompt:
|
||||
msg = _create_msg_format_hf(urls)
|
||||
|
||||
tokenizer = AutoProcessor.from_pretrained("mistral-community/pixtral-12b")
|
||||
prompt = tokenizer.apply_chat_template(msg)
|
||||
|
||||
images = []
|
||||
for chunk in msg[0]["content"]:
|
||||
if chunk["type"] == "image":
|
||||
images.append(chunk["image"])
|
||||
|
||||
mm_data = MultiModalDataBuiltins(image=images)
|
||||
engine_inputs = TextPrompt(prompt=prompt, multi_modal_data=mm_data)
|
||||
|
||||
return engine_inputs
|
||||
|
||||
|
||||
SAMPLING_PARAMS = SamplingParams(max_tokens=512, temperature=0.0, logprobs=5)
|
||||
LIMIT_MM_PER_PROMPT = dict(image=4)
|
||||
|
||||
|
||||
@@ -6,7 +6,6 @@ typically specific to a small subset of models.
|
||||
"""
|
||||
|
||||
import types
|
||||
from pathlib import PosixPath
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
@@ -29,7 +28,7 @@ from vllm.logprobs import SampleLogprobs
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.collection_utils import is_list_of
|
||||
|
||||
from .....conftest import HfRunner, ImageAsset, ImageTestAssets
|
||||
from .....conftest import HfRunner, ImageTestAssets
|
||||
from .types import RunnerOutput
|
||||
|
||||
|
||||
@@ -48,17 +47,6 @@ def blip2_vllm_to_hf_output(vllm_output: RunnerOutput, model: str) -> RunnerOutp
|
||||
return hf_output_ids, hf_output_str, out_logprobs
|
||||
|
||||
|
||||
def qwen_vllm_to_hf_output(
|
||||
vllm_output: RunnerOutput, model: str
|
||||
) -> tuple[list[int], str, SampleLogprobs | None]:
|
||||
"""Sanitize vllm output [qwen models] to be comparable with hf output."""
|
||||
output_ids, output_str, out_logprobs = vllm_output
|
||||
|
||||
hf_output_str = output_str + "<|endoftext|>"
|
||||
|
||||
return output_ids, hf_output_str, out_logprobs
|
||||
|
||||
|
||||
def qwen2_vllm_to_hf_output(
|
||||
vllm_output: RunnerOutput, model: str
|
||||
) -> tuple[list[int], str, SampleLogprobs | None]:
|
||||
@@ -239,37 +227,6 @@ def get_llava_embeddings(image_assets: ImageTestAssets):
|
||||
return [asset.image_embeds for asset in image_assets]
|
||||
|
||||
|
||||
####### Prompt path encoders for models that need models on disk
|
||||
def qwen_prompt_path_encoder(
|
||||
tmp_path: PosixPath, prompt: str, assets: list[ImageAsset] | ImageTestAssets
|
||||
) -> str:
|
||||
"""Given a temporary dir path, export one or more image assets into the
|
||||
tempdir & replace its contents with the local path to the string so that
|
||||
the HF version of Qwen-VL can resolve the path and load the image in its
|
||||
forward() call.
|
||||
|
||||
Args:
|
||||
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.
|
||||
assert prompt.count("<img></img>") == len(assets)
|
||||
|
||||
# Replace the placeholders with local paths to the exported assets
|
||||
for asset in assets:
|
||||
image_tmp_path = tmp_path / f"{asset.name}.jpg"
|
||||
asset.pil_image.save(image_tmp_path)
|
||||
prompt = prompt.replace(
|
||||
"<img></img>",
|
||||
f"<img>{image_tmp_path}</img>",
|
||||
1,
|
||||
)
|
||||
return prompt
|
||||
|
||||
|
||||
####### Model-specific HuggingFace runner patchers
|
||||
def deepseekvl2_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
|
||||
"""Patches and returns an instance of the HfRunner to use for GLM4."""
|
||||
|
||||
@@ -1176,23 +1176,3 @@ def build_samples(model: str) -> tuple[Sample, ...]:
|
||||
builder = _BUILDERS[model]
|
||||
scenarios = _KIMI_K2_SCENARIOS if model == "kimi_k2" else SCENARIOS
|
||||
return tuple(s for s in (builder(sc) for sc in scenarios) if s is not None)
|
||||
|
||||
|
||||
def build_sample(model: str, scenario: Scenario) -> Sample | None:
|
||||
"""Build a single sample for one model + scenario."""
|
||||
return _BUILDERS[model](scenario)
|
||||
|
||||
|
||||
def build_scaling_sample(
|
||||
model: str, token_count: int, validate: bool = False
|
||||
) -> Sample:
|
||||
"""Build a sample with approximately *token_count* tokens."""
|
||||
sentence = "The quick brown fox jumps over the lazy dog. "
|
||||
text = sentence * (token_count // 10 + 1)
|
||||
scenario = Scenario(
|
||||
id=f"scaling-{token_count}",
|
||||
description=f"Scaling test with ~{token_count} tokens",
|
||||
reasoning=text,
|
||||
tool_calls=[_READ_TOOL],
|
||||
)
|
||||
return _BUILDERS[model](scenario, validate=validate)
|
||||
|
||||
@@ -5,20 +5,11 @@
|
||||
Run `pytest tests/quantization/test_configs.py --forked`.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.config import ModelConfig
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelPair:
|
||||
model_marlin: str
|
||||
model_gptq: str
|
||||
|
||||
|
||||
# Model Id // Quantization Arg // Expected Type
|
||||
MODEL_ARG_EXPTYPES = [
|
||||
# AUTOGPTQ
|
||||
|
||||
-161
@@ -1,7 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
import atexit
|
||||
import contextlib
|
||||
import copy
|
||||
@@ -23,7 +22,6 @@ from contextlib import ExitStack, contextmanager
|
||||
from multiprocessing import Process, get_context
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Literal, cast
|
||||
from unittest.mock import patch
|
||||
|
||||
import anthropic
|
||||
import cloudpickle
|
||||
@@ -33,12 +31,9 @@ import pytest
|
||||
import requests
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from huggingface_hub.constants import HF_HUB_OFFLINE
|
||||
from openai.types.completion import Completion
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
import vllm.envs as envs
|
||||
from tests.models.utils import TextTextLogprobs
|
||||
from vllm.distributed import (
|
||||
ensure_model_parallel_initialized,
|
||||
init_distributed_environment,
|
||||
@@ -52,7 +47,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
from vllm.model_executor.model_loader import get_model_loader
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.tokenizers import get_tokenizer
|
||||
from vllm.transformers_utils.repo_utils import hf_api
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.utils.mem_constants import GB_bytes
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
@@ -69,24 +63,6 @@ logger = init_logger(__name__)
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
|
||||
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)."""
|
||||
if HF_HUB_OFFLINE:
|
||||
return
|
||||
for repo_id, filename in assets:
|
||||
try:
|
||||
hf_api().hf_hub_download(repo_id=repo_id, filename=filename)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Failed to prefetch %s/%s: %r. Tests depending on this asset may fail.",
|
||||
repo_id,
|
||||
filename,
|
||||
e,
|
||||
)
|
||||
|
||||
|
||||
if current_platform.is_cuda():
|
||||
from vllm.third_party.pynvml import (
|
||||
nvmlDeviceGetHandleByIndex,
|
||||
@@ -1556,17 +1532,6 @@ 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
|
||||
of the given category is emitted.
|
||||
"""
|
||||
with warnings.catch_warnings():
|
||||
warnings.filterwarnings("error", category=category)
|
||||
|
||||
yield
|
||||
|
||||
|
||||
def get_physical_device_indices(devices: list[int]):
|
||||
visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||||
if visible_devices is None:
|
||||
@@ -2157,88 +2122,6 @@ def multi_gpu_only(*, num_gpus: int = 2):
|
||||
return wrapper
|
||||
|
||||
|
||||
async def completions_with_server_args(
|
||||
prompts: list[str],
|
||||
model_name: str,
|
||||
server_cli_args: list[str],
|
||||
num_logprobs: int | None,
|
||||
max_wait_seconds: int = 240,
|
||||
max_tokens: int | list = 5,
|
||||
) -> list[Completion]:
|
||||
"""Construct a remote OpenAI server, obtain an async client to the
|
||||
server & invoke the completions API to obtain completions.
|
||||
|
||||
Args:
|
||||
prompts: test prompts
|
||||
model_name: model to spin up on the vLLM server
|
||||
server_cli_args: CLI args for starting the server
|
||||
num_logprobs: Number of logprobs to report (or `None`)
|
||||
max_wait_seconds: timeout interval for bringing up server.
|
||||
Default: 240sec
|
||||
max_tokens: max_tokens value for each of the given input prompts.
|
||||
if only one max_token value is given, the same value is used
|
||||
for all the prompts.
|
||||
|
||||
Returns:
|
||||
OpenAI Completion instance
|
||||
|
||||
"""
|
||||
if isinstance(max_tokens, int):
|
||||
max_tokens = [max_tokens] * len(prompts)
|
||||
|
||||
assert len(max_tokens) == len(prompts)
|
||||
|
||||
outputs = None
|
||||
with RemoteOpenAIServer(
|
||||
model_name, server_cli_args, max_wait_seconds=max_wait_seconds
|
||||
) as server:
|
||||
client = server.get_async_client()
|
||||
outputs = [
|
||||
client.completions.create(
|
||||
model=model_name,
|
||||
prompt=[p],
|
||||
temperature=0,
|
||||
stream=False,
|
||||
max_tokens=max_tok,
|
||||
logprobs=num_logprobs,
|
||||
)
|
||||
for p, max_tok in zip(prompts, max_tokens)
|
||||
]
|
||||
outputs = await asyncio.gather(*outputs)
|
||||
|
||||
assert outputs is not None, "Completion API call failed."
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
def get_client_text_generations(completions: list[Completion]) -> list[str]:
|
||||
"""Extract generated tokens from the output of a
|
||||
request made to an Open-AI-protocol completions endpoint.
|
||||
"""
|
||||
assert all([len(x.choices) == 1 for x in completions])
|
||||
return [x.choices[0].text for x in completions]
|
||||
|
||||
|
||||
def get_client_text_logprob_generations(
|
||||
completions: list[Completion],
|
||||
) -> list[TextTextLogprobs]:
|
||||
"""Operates on the output of a request made to an Open-AI-protocol
|
||||
completions endpoint; obtains top-rank logprobs for each token in
|
||||
each {class}`SequenceGroup`
|
||||
"""
|
||||
text_generations = get_client_text_generations(completions)
|
||||
text = "".join(text_generations)
|
||||
return [
|
||||
(
|
||||
text_generations,
|
||||
text,
|
||||
(None if x.logprobs is None else x.logprobs.top_logprobs),
|
||||
)
|
||||
for completion in completions
|
||||
for x in completion.choices
|
||||
]
|
||||
|
||||
|
||||
def has_module_attribute(module_name, attribute_name):
|
||||
"""Helper function to check if a module has a specific attribute."""
|
||||
try:
|
||||
@@ -2267,50 +2150,6 @@ def get_attn_backend_list_based_on_platform() -> list[str]:
|
||||
raise ValueError("Unsupported platform")
|
||||
|
||||
|
||||
@contextmanager
|
||||
def override_cutlass_fp8_supported(value: bool):
|
||||
with patch(
|
||||
"vllm.model_executor.layers.quantization.utils.w8a8_utils.cutlass_fp8_supported",
|
||||
return_value=value,
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
def disable_aiter_plain_rmsnorm(monkeypatch) -> None:
|
||||
"""Patch dispatch_rocm_rmsnorm_func so the plain (non-fused) rms_norm path
|
||||
always uses the native float32 kernel for the duration of a test.
|
||||
|
||||
The fused path (rms_norm2d_with_add, selected when with_fused_add=True) is
|
||||
left on AITER -- only the plain path is redirected to native.
|
||||
|
||||
AITER's plain rms_norm accumulates variance in bfloat16 (~1 ULP/call),
|
||||
which drifts the KV cache over many decode steps. This drift is irrelevant
|
||||
for a trained model (rank-1/rank-2 gap ~1-3 nats >> 1 ULP), but breaks
|
||||
logprob comparison tests with randomly-initialised models like
|
||||
TitanML/tiny-mixtral whose rank-1/rank-2 gap is only O(1/sqrt(V)) ~0.006
|
||||
nats -- smaller than the accumulated per-step error.
|
||||
"""
|
||||
import torch
|
||||
|
||||
import vllm.model_executor.layers.layernorm as _ln_mod
|
||||
from vllm.model_executor.layers.layernorm import rms_norm as _native
|
||||
|
||||
_orig = _ln_mod.dispatch_rocm_rmsnorm_func
|
||||
|
||||
def _native_plain(
|
||||
with_fused_add: bool, dtype: torch.dtype, use_aiter: bool = False
|
||||
):
|
||||
if (
|
||||
use_aiter
|
||||
and not with_fused_add
|
||||
and dtype in (torch.float16, torch.bfloat16)
|
||||
):
|
||||
return _native
|
||||
return _orig(with_fused_add, dtype, use_aiter)
|
||||
|
||||
monkeypatch.setattr(_ln_mod, "dispatch_rocm_rmsnorm_func", _native_plain)
|
||||
|
||||
|
||||
def prep_prompts(batch_size: int, ln_range: tuple[int, int] = (800, 1100)):
|
||||
"""Generate prompts which a bunch of assignments,
|
||||
then asking for the value of one of them.
|
||||
|
||||
@@ -262,27 +262,6 @@ def create_vllm_config(
|
||||
)
|
||||
|
||||
|
||||
def create_dummy_kv_cache(
|
||||
block_size: int,
|
||||
num_kv_heads: int,
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
num_blocks: int = 100,
|
||||
) -> torch.Tensor:
|
||||
"""Create a dummy KV cache tensor for testing."""
|
||||
kv_cache = torch.randn(
|
||||
num_blocks,
|
||||
2, # K and V
|
||||
block_size,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
return kv_cache
|
||||
|
||||
|
||||
@dataclass
|
||||
class BackendConfig:
|
||||
name: str
|
||||
|
||||
@@ -196,23 +196,6 @@ def _create_random_top_token_test_matrix(
|
||||
return matrix, prompt_token_ranks
|
||||
|
||||
|
||||
def decode_token(
|
||||
tok_id: int,
|
||||
tokenizer: PythonBackend,
|
||||
) -> str:
|
||||
"""Reproduce the process of detokenizing a token for testing purposes.
|
||||
|
||||
Args:
|
||||
tok_id: token id to detokenize
|
||||
tokenizer: tokenizer to use for detokenization
|
||||
|
||||
Returns:
|
||||
string representation of token
|
||||
|
||||
"""
|
||||
return tokenizer.convert_ids_to_tokens(tok_id)
|
||||
|
||||
|
||||
def generate_dummy_sample_logprobs(
|
||||
sampled_tokens_list: list,
|
||||
num_logprobs: int,
|
||||
|
||||
Reference in New Issue
Block a user