Add vLLM v0.18.1 source tree with KV transfer abort fix
third_party/vllm/ now tracked in git for direct patch management.
Based on vLLM v0.18.1 release with one patch applied:
vllm/v1/core/sched/scheduler.py:
Replace fatal assert with graceful skip when KV transfer callback
arrives for an already-aborted request during PD disaggregated serving.
Future vLLM modifications should be made directly in third_party/vllm/
and committed normally. The patches/ directory is kept as documentation
of what changed from upstream.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
0
third_party/vllm/tests/v1/__init__.py
vendored
Normal file
0
third_party/vllm/tests/v1/__init__.py
vendored
Normal file
739
third_party/vllm/tests/v1/attention/test_attention_backends.py
vendored
Normal file
739
third_party/vllm/tests/v1/attention/test_attention_backends.py
vendored
Normal file
@@ -0,0 +1,739 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for v1 attention backends without GPUModelRunner dependency."""
|
||||
|
||||
from functools import partial
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch.nn.attention.flex_attention import create_block_mask, flex_attention
|
||||
|
||||
from tests.v1.attention.utils import (
|
||||
BatchSpec,
|
||||
create_common_attn_metadata,
|
||||
create_standard_kv_cache_spec,
|
||||
create_vllm_config,
|
||||
try_backend_includes_kv_cache_update,
|
||||
try_get_attention_backend,
|
||||
)
|
||||
from vllm.config import ModelConfig
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.math_utils import cdiv
|
||||
from vllm.utils.torch_utils import (
|
||||
STR_DTYPE_TO_TORCH_DTYPE,
|
||||
is_torch_equal_or_newer,
|
||||
set_random_seed,
|
||||
)
|
||||
from vllm.v1.attention.backend import AttentionType, CommonAttentionMetadata
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
from vllm.v1.attention.backends.utils import (
|
||||
set_kv_cache_layout,
|
||||
)
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec
|
||||
|
||||
BACKENDS_TO_TEST = [
|
||||
AttentionBackendEnum.FLASH_ATTN,
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
AttentionBackendEnum.FLEX_ATTENTION,
|
||||
AttentionBackendEnum.TRITON_ATTN,
|
||||
AttentionBackendEnum.TREE_ATTN,
|
||||
"FLEX_ATTENTION_SLOW",
|
||||
]
|
||||
|
||||
# Remove flashinfer from the list if it's not available
|
||||
try:
|
||||
import flashinfer # noqa: F401
|
||||
except ImportError:
|
||||
BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASHINFER)
|
||||
|
||||
|
||||
def _convert_dtype_to_torch(dtype):
|
||||
"""Convert ModelDType to torch.dtype."""
|
||||
if isinstance(dtype, str):
|
||||
if dtype == "auto":
|
||||
return torch.float16 # Default dtype for testing
|
||||
elif dtype in STR_DTYPE_TO_TORCH_DTYPE:
|
||||
return STR_DTYPE_TO_TORCH_DTYPE[dtype]
|
||||
else:
|
||||
raise ValueError(f"Unknown dtype: {dtype}")
|
||||
elif isinstance(dtype, torch.dtype):
|
||||
return dtype
|
||||
else:
|
||||
raise ValueError(f"Unknown dtype: {dtype}")
|
||||
|
||||
|
||||
# Define common batch configurations
|
||||
BATCH_SPECS = {
|
||||
"small_decode": BatchSpec(seq_lens=[32, 40], query_lens=[1, 1]),
|
||||
"small_prefill": BatchSpec(seq_lens=[32, 40], query_lens=[8, 8]),
|
||||
"mixed_small": BatchSpec(seq_lens=[32, 40, 48, 56], query_lens=[1, 1, 5, 5]),
|
||||
"medium_decode": BatchSpec(
|
||||
seq_lens=[128, 256, 512, 1024, 128, 256, 512, 1024],
|
||||
query_lens=[1, 1, 1, 1, 1, 1, 1, 1],
|
||||
),
|
||||
"medium_prefill": BatchSpec(
|
||||
seq_lens=[256, 512, 1024, 2048], query_lens=[16, 16, 16, 16]
|
||||
),
|
||||
"mixed_medium": BatchSpec(
|
||||
seq_lens=[512, 1024, 2048, 512, 1024, 2048], query_lens=[1, 1, 1, 7, 7, 7]
|
||||
),
|
||||
"large_decode": BatchSpec(seq_lens=[2048] * 32, query_lens=[1] * 32),
|
||||
"large_prefill": BatchSpec(seq_lens=[4096] * 8, query_lens=[32] * 8),
|
||||
"mixed_large": BatchSpec(
|
||||
seq_lens=[1024, 2048, 4096, 1024, 2048, 4096], query_lens=[1, 1, 1, 32, 32, 32]
|
||||
),
|
||||
"single_decode": BatchSpec(seq_lens=[1024], query_lens=[1]),
|
||||
"single_prefill": BatchSpec(seq_lens=[1024], query_lens=[64]),
|
||||
# encoder-only
|
||||
"small_encoder_prefill": BatchSpec(
|
||||
seq_lens=[32, 64, 128, 256], query_lens=[32, 64, 128, 256]
|
||||
),
|
||||
"medium_encoder_prefill": BatchSpec(
|
||||
seq_lens=[256, 512, 1024, 2048], query_lens=[256, 512, 1024, 2048]
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def create_and_prepopulate_kv_cache(
|
||||
k_contexts: list[torch.Tensor],
|
||||
v_contexts: list[torch.Tensor],
|
||||
block_size: int,
|
||||
num_kv_heads: int,
|
||||
head_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
num_blocks: int,
|
||||
common_attn_metadata: CommonAttentionMetadata,
|
||||
randomize_blocks: bool = True,
|
||||
) -> torch.Tensor:
|
||||
"""Create and prepopulate a KV cache with context data.
|
||||
|
||||
Args:
|
||||
k_contexts: List of key context tensors for each sequence
|
||||
v_contexts: List of value context tensors for each sequence
|
||||
seq_lens: List of sequence lengths
|
||||
block_size: Size of each block
|
||||
num_kv_heads: Number of KV heads
|
||||
head_size: Size of each head
|
||||
dtype: Data type for the cache
|
||||
device: Device to create the cache on
|
||||
num_blocks: Total number of blocks in the cache
|
||||
block_table: Block table tensor to populate
|
||||
randomize_blocks: Whether to randomly permute blocks
|
||||
or use sequential order
|
||||
|
||||
Returns:
|
||||
Tuple of (kv_cache, updated_block_table)
|
||||
"""
|
||||
batch_size = len(k_contexts)
|
||||
seq_lens = common_attn_metadata.seq_lens.cpu()
|
||||
query_lens = (
|
||||
common_attn_metadata.query_start_loc_cpu[1:]
|
||||
- common_attn_metadata.query_start_loc_cpu[:-1]
|
||||
)
|
||||
context_lens = seq_lens - query_lens
|
||||
block_table = common_attn_metadata.block_table_tensor
|
||||
slot_mapping = common_attn_metadata.slot_mapping
|
||||
|
||||
# Create KV cache
|
||||
kv_cache = torch.zeros(
|
||||
2, num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
|
||||
)
|
||||
kv_cache_flat = kv_cache.view(2, -1, num_kv_heads, head_size)
|
||||
|
||||
# Populate the cache with the context tokens
|
||||
# Start from block_id=1 since block_id=0 is considered the null block
|
||||
start_block_idx = 1
|
||||
for i in range(batch_size):
|
||||
k_context, v_context = k_contexts[i], v_contexts[i]
|
||||
start = start_block_idx * block_size
|
||||
end = start + k_context.shape[0]
|
||||
kv_cache_flat[0, start:end, ...] = k_context
|
||||
kv_cache_flat[1, start:end, ...] = v_context
|
||||
|
||||
# Stay block aligned and allocate enough blocks for the new tokens
|
||||
start_block_idx += cdiv(int(seq_lens[i]), block_size)
|
||||
|
||||
blocks_end = start_block_idx
|
||||
|
||||
# Permute the context blocks (excluding block 0 which is null)
|
||||
if randomize_blocks:
|
||||
# Random permutation starting from block 1
|
||||
perm = torch.randperm(blocks_end - 1) + 1
|
||||
else:
|
||||
# Sequential order starting from block 1
|
||||
perm = torch.arange(1, blocks_end)
|
||||
|
||||
inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
|
||||
# Add 1 to account for starting from block 1
|
||||
inv_perm[1:] = torch.argsort(perm) + 1
|
||||
kv_cache[:, 1:blocks_end, ...] = kv_cache[:, perm, ...]
|
||||
|
||||
# Construct the right block table
|
||||
# Start from block_id=1 since block_id=0 is considered the null block
|
||||
start_block_idx = 1
|
||||
for i in range(batch_size):
|
||||
num_blocks_for_seq = cdiv(int(seq_lens[i]), block_size)
|
||||
start = start_block_idx
|
||||
end = start + num_blocks_for_seq
|
||||
block_table[i, :num_blocks_for_seq] = inv_perm[start:end]
|
||||
start_block_idx += num_blocks_for_seq
|
||||
|
||||
# Create a realistic slot mapping that corresponds to the block table
|
||||
for i in range(batch_size):
|
||||
token_offsets = torch.arange(int(query_lens[i])) + int(context_lens[i])
|
||||
block_indices = token_offsets // block_size
|
||||
token_inter_block_offsets = token_offsets % block_size
|
||||
start = common_attn_metadata.query_start_loc_cpu[i]
|
||||
end = common_attn_metadata.query_start_loc_cpu[i + 1]
|
||||
slot_mapping[start:end] = block_table[
|
||||
i, block_indices
|
||||
] * block_size + token_inter_block_offsets.to(device)
|
||||
|
||||
return kv_cache
|
||||
|
||||
|
||||
class MockAttentionLayer:
|
||||
"""A mock attention layer for testing."""
|
||||
|
||||
def __init__(self, device: torch.device):
|
||||
self._q_scale = torch.tensor(1.0, device=device)
|
||||
self._k_scale = torch.tensor(1.0, device=device)
|
||||
self._v_scale = torch.tensor(1.0, device=device)
|
||||
# Add float versions for flashinfer
|
||||
self._q_scale_float = 1.0
|
||||
self._k_scale_float = 1.0
|
||||
self._v_scale_float = 1.0
|
||||
|
||||
|
||||
def run_attention_backend(
|
||||
backend: AttentionBackendEnum,
|
||||
kv_cache_spec: FullAttentionSpec,
|
||||
layer_names: list[str],
|
||||
vllm_config,
|
||||
device: torch.device,
|
||||
common_attn_metadata: CommonAttentionMetadata,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
attn_type: AttentionType = AttentionType.DECODER,
|
||||
sliding_window: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Run attention computation using the specified backend's AttentionImpl."""
|
||||
|
||||
# Handle special case for FLEX_ATTENTION_SLOW
|
||||
actual_backend = backend
|
||||
|
||||
use_direct_block_mask = is_torch_equal_or_newer("2.9.0.dev0")
|
||||
if backend == "FLEX_ATTENTION_SLOW":
|
||||
actual_backend = AttentionBackendEnum.FLEX_ATTENTION
|
||||
use_direct_block_mask = False
|
||||
|
||||
builder_cls, impl_cls = try_get_attention_backend(actual_backend)
|
||||
|
||||
# Mock flashinfer's get_per_layer_parameters if needed
|
||||
if actual_backend == AttentionBackendEnum.FLASHINFER:
|
||||
import unittest.mock
|
||||
|
||||
from vllm.v1.attention.backends.utils import PerLayerParameters
|
||||
|
||||
def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
|
||||
# Return mock parameters for a single layer
|
||||
head_size = vllm_config.model_config.get_head_size()
|
||||
return {
|
||||
layer_name: PerLayerParameters(
|
||||
window_left=-1, # No sliding window
|
||||
logits_soft_cap=0.0, # No soft cap
|
||||
sm_scale=1.0 / (head_size**0.5), # Standard scale
|
||||
)
|
||||
for layer_name in layer_names
|
||||
}
|
||||
|
||||
with unittest.mock.patch(
|
||||
"vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
|
||||
mock_get_per_layer_parameters,
|
||||
):
|
||||
builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
|
||||
attn_metadata = builder.build(
|
||||
common_prefix_len=0,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
)
|
||||
else:
|
||||
# Build metadata
|
||||
builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
|
||||
if actual_backend == AttentionBackendEnum.FLEX_ATTENTION:
|
||||
builder.direct_build = use_direct_block_mask
|
||||
attn_metadata = builder.build(
|
||||
common_prefix_len=0,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
)
|
||||
|
||||
# Instantiate implementation
|
||||
num_heads = vllm_config.model_config.get_num_attention_heads(
|
||||
vllm_config.parallel_config
|
||||
)
|
||||
num_kv_heads = vllm_config.model_config.get_num_kv_heads(
|
||||
vllm_config.parallel_config
|
||||
)
|
||||
head_size = vllm_config.model_config.get_head_size()
|
||||
scale = 1.0 / (head_size**0.5)
|
||||
impl = impl_cls(
|
||||
num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
scale=scale,
|
||||
num_kv_heads=num_kv_heads,
|
||||
alibi_slopes=None,
|
||||
sliding_window=sliding_window,
|
||||
attn_type=attn_type,
|
||||
kv_cache_dtype="auto",
|
||||
)
|
||||
|
||||
# Create mock layer and output buffer
|
||||
mock_layer = MockAttentionLayer(device)
|
||||
output = torch.empty_like(query)
|
||||
|
||||
# Run forward pass
|
||||
# NOTE: The query, key, and value are already shaped correctly
|
||||
# in the calling test function.
|
||||
if not try_backend_includes_kv_cache_update(actual_backend):
|
||||
impl.do_kv_cache_update(
|
||||
mock_layer, key, value, kv_cache, attn_metadata.slot_mapping
|
||||
)
|
||||
output = impl.forward(
|
||||
mock_layer, query, key, value, kv_cache, attn_metadata, output=output
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def _test_backend_correctness(
|
||||
batch_spec: BatchSpec,
|
||||
model: str,
|
||||
backend_to_test: list[AttentionBackendEnum | str],
|
||||
mask_mod,
|
||||
*,
|
||||
attn_type: AttentionType = AttentionType.DECODER,
|
||||
block_size: int = 16,
|
||||
atol: float = 1e-2,
|
||||
rtol: float = 1e-2,
|
||||
tensor_parallel_size: int = 1,
|
||||
):
|
||||
"""
|
||||
Test that all backends produce similar outputs to a reference implementation
|
||||
using torch.nn.functional.scaled_dot_product_attention.
|
||||
|
||||
This test works by:
|
||||
1. Generating a batch of sequences with specified context and query lengths.
|
||||
2. Computing a ground-truth attention output using torch.sdpa on
|
||||
contiguous Q, K, and V tensors.
|
||||
3. Simulating vLLM's paged KV cache: It takes the context portion of the
|
||||
K/V tensors and manually places them into a paged buffer according to
|
||||
the test's (randomly generated) block table.
|
||||
4. Running each vLLM attention backend with the new queries and the
|
||||
simulated paged KV cache.
|
||||
5. Comparing the vLLM backend's output to the ground-truth SDPA output.
|
||||
|
||||
Note: When tensor_parallel_size > 1, we simulate the head partitioning
|
||||
by overriding the model config to use fewer heads, without requiring
|
||||
multiple GPUs. This tests that backends work correctly with different
|
||||
head counts.
|
||||
"""
|
||||
set_random_seed(42)
|
||||
|
||||
hf_config_override = None
|
||||
if tensor_parallel_size > 1:
|
||||
from vllm.config import ModelConfig
|
||||
|
||||
temp_config = ModelConfig(model=model, max_model_len=1)
|
||||
original_num_heads = temp_config.hf_text_config.num_attention_heads
|
||||
original_num_kv_heads = getattr(
|
||||
temp_config.hf_text_config, "num_key_value_heads", None
|
||||
)
|
||||
hf_config_override = {
|
||||
"num_attention_heads": original_num_heads // tensor_parallel_size,
|
||||
}
|
||||
if original_num_kv_heads is not None:
|
||||
hf_config_override["num_key_value_heads"] = max(
|
||||
1, original_num_kv_heads // tensor_parallel_size
|
||||
)
|
||||
|
||||
vllm_config = create_vllm_config(
|
||||
model_name=model,
|
||||
tensor_parallel_size=1, # Always use TP=1 to avoid multi-GPU requirements
|
||||
max_model_len=max(batch_spec.seq_lens),
|
||||
block_size=block_size,
|
||||
num_gpu_blocks=8192,
|
||||
hf_config_override=hf_config_override,
|
||||
)
|
||||
device = torch.device("cuda:0")
|
||||
|
||||
kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
|
||||
|
||||
# 1. Setup
|
||||
batch_size = batch_spec.batch_size
|
||||
seq_lens = batch_spec.seq_lens
|
||||
query_lens = batch_spec.query_lens
|
||||
num_q_heads = vllm_config.model_config.get_num_attention_heads(
|
||||
vllm_config.parallel_config
|
||||
)
|
||||
num_kv_heads = vllm_config.model_config.get_num_kv_heads(
|
||||
vllm_config.parallel_config
|
||||
)
|
||||
head_size = vllm_config.model_config.get_head_size()
|
||||
sliding_window = vllm_config.model_config.get_sliding_window()
|
||||
dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype)
|
||||
block_size = vllm_config.cache_config.block_size
|
||||
scale = 1.0 / (head_size**0.5)
|
||||
|
||||
# 2. Generate data and compute SDPA reference output
|
||||
all_q_vllm, all_k_vllm, all_v_vllm = [], [], []
|
||||
all_sdpa_outputs = []
|
||||
k_contexts, v_contexts = [], []
|
||||
|
||||
for i in range(batch_size):
|
||||
s_len = seq_lens[i]
|
||||
q_len = query_lens[i]
|
||||
context_len = s_len - q_len
|
||||
|
||||
# Generate Q, K, V for the whole sequence to be used in SDPA
|
||||
q = torch.randn(q_len, num_q_heads, head_size, dtype=dtype, device=device)
|
||||
k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
|
||||
v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
|
||||
|
||||
# SDPA expects (N, H, L, D), so unsqueeze batch and permute
|
||||
q_sdpa_in = q.unsqueeze(0).transpose(1, 2)
|
||||
k_sdpa_in = k_full.unsqueeze(0).transpose(1, 2)
|
||||
v_sdpa_in = v_full.unsqueeze(0).transpose(1, 2)
|
||||
|
||||
if num_q_heads != num_kv_heads:
|
||||
assert num_q_heads % num_kv_heads == 0, (
|
||||
f"num_q_heads ({num_q_heads}) must be divisible by "
|
||||
f"num_kv_heads ({num_kv_heads})"
|
||||
)
|
||||
repeats = num_q_heads // num_kv_heads
|
||||
k_sdpa_in = k_sdpa_in.repeat_interleave(repeats, dim=1)
|
||||
v_sdpa_in = v_sdpa_in.repeat_interleave(repeats, dim=1)
|
||||
|
||||
# Create causal mask: query token i attends to positions 0 to
|
||||
# (context_len + i)
|
||||
kv_len = s_len
|
||||
|
||||
final_mask_mod = partial(mask_mod, context_len=context_len)
|
||||
block_mask = create_block_mask(
|
||||
final_mask_mod, B=None, H=None, Q_LEN=q_len, KV_LEN=kv_len, device=device
|
||||
)
|
||||
sdpa_out_i = flex_attention(
|
||||
q_sdpa_in,
|
||||
k_sdpa_in,
|
||||
v_sdpa_in,
|
||||
block_mask=block_mask,
|
||||
scale=scale,
|
||||
enable_gqa=True,
|
||||
)
|
||||
|
||||
all_sdpa_outputs.append(sdpa_out_i.transpose(1, 2).squeeze(0))
|
||||
|
||||
# Inputs for vLLM backends are just the new tokens
|
||||
all_q_vllm.append(q)
|
||||
all_k_vllm.append(k_full[context_len:])
|
||||
all_v_vllm.append(v_full[context_len:])
|
||||
|
||||
# Contextual K/V data used to populate the paged cache
|
||||
k_contexts.append(k_full[:context_len])
|
||||
v_contexts.append(v_full[:context_len])
|
||||
|
||||
query_vllm = torch.cat(all_q_vllm, dim=0)
|
||||
key_vllm = torch.cat(all_k_vllm, dim=0)
|
||||
value_vllm = torch.cat(all_v_vllm, dim=0)
|
||||
sdpa_output = torch.cat(all_sdpa_outputs, dim=0)
|
||||
|
||||
common_attn_metadata = create_common_attn_metadata(
|
||||
batch_spec, vllm_config.cache_config.block_size, device
|
||||
)
|
||||
if attn_type == AttentionType.ENCODER_ONLY:
|
||||
# For encoder-only, all tokens are prefill tokens
|
||||
common_attn_metadata.causal = False
|
||||
|
||||
# 3. Simulate Paged KV Cache and a realistic slot_mapping
|
||||
kv_cache = create_and_prepopulate_kv_cache(
|
||||
k_contexts=k_contexts,
|
||||
v_contexts=v_contexts,
|
||||
block_size=block_size,
|
||||
num_kv_heads=num_kv_heads,
|
||||
head_size=head_size,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
num_blocks=vllm_config.cache_config.num_gpu_blocks or 1000,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
randomize_blocks=True,
|
||||
)
|
||||
|
||||
# 4. Run vLLM backends and compare
|
||||
# Note: flex_attention has known Triton kernel compatibility issues
|
||||
# with test infrastructures
|
||||
for backend_name in backend_to_test:
|
||||
# FlashAttentionm + FlexAttention:
|
||||
# [2, num_blocks, block_size, num_kv_heads, head_size]
|
||||
# FlashInfer + Triton:
|
||||
# [num_blocks, 2, block_size, num_kv_heads, head_size]
|
||||
# Select the appropriate KV cache format for each backend
|
||||
kv_cache_for_backend = kv_cache
|
||||
reset_kv_cache_layout = False
|
||||
if backend_name in (
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
AttentionBackendEnum.TRITON_ATTN,
|
||||
):
|
||||
kv_cache_for_backend = kv_cache.transpose(0, 1)
|
||||
|
||||
if backend_name == AttentionBackendEnum.FLASHINFER:
|
||||
# For FlashInfer default to HND layout and
|
||||
kv_cache_for_backend = (
|
||||
kv_cache_for_backend.transpose(2, 3).contiguous().transpose(2, 3)
|
||||
)
|
||||
set_kv_cache_layout("HND")
|
||||
reset_kv_cache_layout = True
|
||||
elif backend_name == AttentionBackendEnum.TRITON_ATTN:
|
||||
kv_cache_for_backend = kv_cache_for_backend.contiguous()
|
||||
|
||||
try:
|
||||
backend_output = run_attention_backend(
|
||||
backend_name,
|
||||
kv_cache_spec,
|
||||
["placeholder"],
|
||||
vllm_config,
|
||||
device,
|
||||
common_attn_metadata,
|
||||
query_vllm,
|
||||
key_vllm,
|
||||
value_vllm,
|
||||
kv_cache_for_backend,
|
||||
sliding_window=sliding_window,
|
||||
attn_type=attn_type,
|
||||
)
|
||||
finally:
|
||||
if reset_kv_cache_layout:
|
||||
set_kv_cache_layout(None)
|
||||
|
||||
# Check shape and dtype consistency
|
||||
assert backend_output.shape == sdpa_output.shape, (
|
||||
f"[{backend_name}] shape {backend_output.shape} != "
|
||||
f"SDPA shape {sdpa_output.shape}"
|
||||
)
|
||||
assert backend_output.dtype == sdpa_output.dtype, (
|
||||
f"[{backend_name}] dtype {backend_output.dtype} != "
|
||||
f"SDPA dtype {sdpa_output.dtype}"
|
||||
)
|
||||
|
||||
assert torch.isfinite(backend_output).all(), (
|
||||
f"[{backend_name}] produced non-finite values"
|
||||
)
|
||||
|
||||
# Check numerical similarity
|
||||
def error_msg(msg: str, backend_name: str):
|
||||
return f"[{backend_name}] output differs from SDPA baseline. {msg}"
|
||||
|
||||
torch.testing.assert_close(
|
||||
backend_output,
|
||||
sdpa_output,
|
||||
rtol=rtol,
|
||||
atol=atol,
|
||||
msg=partial(error_msg, backend_name=backend_name),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_spec_name",
|
||||
[
|
||||
"small_decode",
|
||||
"small_prefill",
|
||||
"mixed_small",
|
||||
"medium_decode",
|
||||
"medium_prefill",
|
||||
"mixed_medium",
|
||||
"large_decode",
|
||||
"large_prefill",
|
||||
"single_decode",
|
||||
"single_prefill",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])
|
||||
@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
|
||||
def test_causal_backend_correctness(
|
||||
default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int
|
||||
):
|
||||
"""Test backend's correctness with causal attention."""
|
||||
|
||||
def causal_mask_mod(
|
||||
b: torch.Tensor,
|
||||
h: torch.Tensor,
|
||||
q_idx: torch.Tensor,
|
||||
kv_idx: torch.Tensor,
|
||||
*,
|
||||
context_len: int,
|
||||
):
|
||||
return (q_idx + context_len) >= kv_idx
|
||||
|
||||
batch_spec = BATCH_SPECS[batch_spec_name]
|
||||
LARGE_BLOCK_BACKENDS = (
|
||||
[AttentionBackendEnum.FLEX_ATTENTION]
|
||||
if is_torch_equal_or_newer("2.9.0.dev0")
|
||||
else []
|
||||
)
|
||||
|
||||
if current_platform.is_rocm():
|
||||
SMALL_BLOCK_BACKENDS = [
|
||||
x
|
||||
for x in BACKENDS_TO_TEST
|
||||
if (
|
||||
x not in LARGE_BLOCK_BACKENDS
|
||||
and x is not AttentionBackendEnum.FLASH_ATTN
|
||||
)
|
||||
]
|
||||
else:
|
||||
SMALL_BLOCK_BACKENDS = [
|
||||
x for x in BACKENDS_TO_TEST if x not in LARGE_BLOCK_BACKENDS
|
||||
]
|
||||
|
||||
_test_backend_correctness(
|
||||
batch_spec,
|
||||
model,
|
||||
SMALL_BLOCK_BACKENDS,
|
||||
causal_mask_mod,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
)
|
||||
|
||||
# Fast FlexAttention needs to run with block_size=128
|
||||
if LARGE_BLOCK_BACKENDS:
|
||||
_test_backend_correctness(
|
||||
batch_spec,
|
||||
model,
|
||||
LARGE_BLOCK_BACKENDS,
|
||||
causal_mask_mod,
|
||||
block_size=128,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
)
|
||||
|
||||
|
||||
if current_platform.is_rocm():
|
||||
# FLASH_ATTN is not supported on ROCm
|
||||
SLIDING_WINDOW_BACKENDS_TO_TEST = [
|
||||
AttentionBackendEnum.FLEX_ATTENTION,
|
||||
AttentionBackendEnum.TRITON_ATTN,
|
||||
"FLEX_ATTENTION_SLOW",
|
||||
]
|
||||
else:
|
||||
SLIDING_WINDOW_BACKENDS_TO_TEST = [
|
||||
AttentionBackendEnum.FLASH_ATTN,
|
||||
AttentionBackendEnum.FLEX_ATTENTION,
|
||||
AttentionBackendEnum.TRITON_ATTN,
|
||||
"FLEX_ATTENTION_SLOW",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_spec_name",
|
||||
[
|
||||
"small_decode",
|
||||
"small_prefill",
|
||||
"mixed_medium",
|
||||
"large_decode",
|
||||
"large_prefill",
|
||||
"mixed_large",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("model", ["microsoft/Phi-tiny-MoE-instruct"])
|
||||
@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
|
||||
def test_sliding_window_backend_correctness(
|
||||
batch_spec_name: str, model: str, tensor_parallel_size: int
|
||||
):
|
||||
"""Test backend's correctness with sliding window attention."""
|
||||
|
||||
def sliding_window_mask_mod(
|
||||
b: torch.Tensor,
|
||||
h: torch.Tensor,
|
||||
q_idx: torch.Tensor,
|
||||
kv_idx: torch.Tensor,
|
||||
*,
|
||||
context_len: int,
|
||||
sliding_window: int,
|
||||
):
|
||||
causal_mask = q_idx + context_len >= kv_idx
|
||||
window_mask = q_idx + context_len - kv_idx < sliding_window
|
||||
return causal_mask & window_mask
|
||||
|
||||
batch_spec = BATCH_SPECS[batch_spec_name]
|
||||
model_config = ModelConfig(model=model, max_model_len=max(batch_spec.seq_lens))
|
||||
sliding_window = model_config.get_sliding_window()
|
||||
sliding_window_mask_mod_fn = partial(
|
||||
sliding_window_mask_mod, sliding_window=sliding_window
|
||||
)
|
||||
|
||||
LARGE_BLOCK_BACKENDS = (
|
||||
[AttentionBackendEnum.FLEX_ATTENTION]
|
||||
if is_torch_equal_or_newer("2.9.0.dev0")
|
||||
else []
|
||||
)
|
||||
SMALL_BLOCK_BACKENDS = [
|
||||
x for x in SLIDING_WINDOW_BACKENDS_TO_TEST if x not in LARGE_BLOCK_BACKENDS
|
||||
]
|
||||
_test_backend_correctness(
|
||||
batch_spec,
|
||||
model,
|
||||
SMALL_BLOCK_BACKENDS,
|
||||
sliding_window_mask_mod_fn,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
)
|
||||
|
||||
# Fast FlexAttention needs to run with block_size=128
|
||||
if LARGE_BLOCK_BACKENDS:
|
||||
_test_backend_correctness(
|
||||
batch_spec,
|
||||
model,
|
||||
LARGE_BLOCK_BACKENDS,
|
||||
sliding_window_mask_mod_fn,
|
||||
block_size=128,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_spec_name",
|
||||
[
|
||||
"small_encoder_prefill",
|
||||
"medium_encoder_prefill",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("model", ["google/embeddinggemma-300m"])
|
||||
@pytest.mark.parametrize("tensor_parallel_size", [1, 2])
|
||||
def test_sliding_window_encoder_backend_correctness(
|
||||
batch_spec_name: str, model: str, tensor_parallel_size: int
|
||||
):
|
||||
"""Test backend's correctness with sliding window attention."""
|
||||
|
||||
def bidi_sliding_window_mask_mod(
|
||||
b: torch.Tensor,
|
||||
h: torch.Tensor,
|
||||
q_idx: torch.Tensor,
|
||||
kv_idx: torch.Tensor,
|
||||
*,
|
||||
context_len: int,
|
||||
sliding_window: int,
|
||||
):
|
||||
return torch.abs(q_idx + context_len - kv_idx) < sliding_window
|
||||
|
||||
batch_spec = BATCH_SPECS[batch_spec_name]
|
||||
model_config = ModelConfig(model=model, max_model_len=max(batch_spec.seq_lens))
|
||||
sliding_window = model_config.get_sliding_window()
|
||||
sliding_window_mask_mod_fn = partial(
|
||||
bidi_sliding_window_mask_mod, sliding_window=sliding_window
|
||||
)
|
||||
|
||||
_test_backend_correctness(
|
||||
batch_spec,
|
||||
model,
|
||||
SLIDING_WINDOW_BACKENDS_TO_TEST,
|
||||
sliding_window_mask_mod_fn,
|
||||
attn_type=AttentionType.ENCODER_ONLY,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
)
|
||||
116
third_party/vllm/tests/v1/attention/test_attention_backends_selection.py
vendored
Normal file
116
third_party/vllm/tests/v1/attention/test_attention_backends_selection.py
vendored
Normal file
@@ -0,0 +1,116 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for mamba attention backend selectors."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.model_executor.layers.mamba.mamba_mixer import MambaMixer
|
||||
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
|
||||
from vllm.model_executor.layers.mamba.short_conv import ShortConv
|
||||
from vllm.model_executor.models.minimax_text_01 import MiniMaxText01LinearAttention
|
||||
from vllm.v1.attention.backends.linear_attn import LinearAttentionBackend
|
||||
from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionBackend
|
||||
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionBackend
|
||||
from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionBackend
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"layer_class, init_kwargs, expected_backend, expected_mamba_type",
|
||||
[
|
||||
(
|
||||
MambaMixer,
|
||||
dict(
|
||||
hidden_size=128,
|
||||
ssm_state_size=16,
|
||||
conv_kernel_size=4,
|
||||
intermediate_size=256,
|
||||
time_step_rank=8,
|
||||
use_conv_bias=True,
|
||||
use_bias=False,
|
||||
use_rms_norm=True,
|
||||
),
|
||||
Mamba1AttentionBackend,
|
||||
"mamba1",
|
||||
),
|
||||
(
|
||||
MambaMixer2,
|
||||
dict(
|
||||
hidden_size=128,
|
||||
ssm_state_size=16,
|
||||
conv_kernel_size=4,
|
||||
intermediate_size=256,
|
||||
use_conv_bias=True,
|
||||
use_bias=False,
|
||||
n_groups=1,
|
||||
num_heads=8,
|
||||
head_dim=32,
|
||||
),
|
||||
Mamba2AttentionBackend,
|
||||
"mamba2",
|
||||
),
|
||||
(
|
||||
MiniMaxText01LinearAttention,
|
||||
dict(
|
||||
hidden_size=128,
|
||||
hidden_inner_size=256,
|
||||
num_heads=8,
|
||||
head_dim=32,
|
||||
max_position=2048,
|
||||
block_size=64,
|
||||
num_hidden_layer=12,
|
||||
layer_idx=0,
|
||||
linear_layer_idx=0,
|
||||
),
|
||||
LinearAttentionBackend,
|
||||
"linear_attention",
|
||||
),
|
||||
(
|
||||
ShortConv,
|
||||
dict(
|
||||
config=SimpleNamespace(conv_L_cache=32, conv_bias=True),
|
||||
dim=128,
|
||||
layer_idx=0,
|
||||
),
|
||||
ShortConvAttentionBackend,
|
||||
"short_conv",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_mamba_layers_get_attn_backend(
|
||||
default_vllm_config,
|
||||
dist_init,
|
||||
layer_class,
|
||||
init_kwargs,
|
||||
expected_backend,
|
||||
expected_mamba_type,
|
||||
):
|
||||
"""Test that Mamba-like layers return the correct attention backend."""
|
||||
layer = layer_class(**init_kwargs)
|
||||
|
||||
backend_class = layer.get_attn_backend()
|
||||
assert backend_class is expected_backend
|
||||
assert layer.mamba_type == expected_mamba_type
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"layer_class,expected_backend,expected_mamba_type",
|
||||
[
|
||||
(MambaMixer, Mamba1AttentionBackend, "mamba1"),
|
||||
(MambaMixer2, Mamba2AttentionBackend, "mamba2"),
|
||||
(MiniMaxText01LinearAttention, LinearAttentionBackend, "linear_attention"),
|
||||
(ShortConv, ShortConvAttentionBackend, "short_conv"),
|
||||
],
|
||||
)
|
||||
def test_mamba_layers_have_unified_interface(
|
||||
layer_class, expected_backend, expected_mamba_type
|
||||
):
|
||||
"""Test that all Mamba layers have the unified get_attn_backend
|
||||
interface."""
|
||||
assert hasattr(layer_class, "get_attn_backend"), (
|
||||
f"{layer_class.__name__} should have get_attn_backend method"
|
||||
)
|
||||
assert hasattr(layer_class, "mamba_type"), (
|
||||
f"{layer_class.__name__} should have mamba_type property"
|
||||
)
|
||||
381
third_party/vllm/tests/v1/attention/test_attention_splitting.py
vendored
Normal file
381
third_party/vllm/tests/v1/attention/test_attention_splitting.py
vendored
Normal file
@@ -0,0 +1,381 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.v1.attention.test_attention_backends import BATCH_SPECS
|
||||
from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
|
||||
from vllm.v1.attention.backends.utils import (
|
||||
split_decodes_and_prefills,
|
||||
)
|
||||
from vllm.v1.worker.ubatch_utils import (
|
||||
UBatchSlice,
|
||||
_make_metadata_with_slice,
|
||||
maybe_create_ubatch_slices,
|
||||
slice_query_start_locs,
|
||||
split_attn_metadata,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_query_start_loc():
|
||||
"""Sample query_start_loc tensor for testing"""
|
||||
return torch.tensor([0, 5, 12, 20, 35, 50])
|
||||
|
||||
|
||||
def test_basic_slice_middle(sample_query_start_loc):
|
||||
"""Test slicing from middle of tensor"""
|
||||
req_slice = slice(1, 3) # slice from index 1 to 3
|
||||
result = slice_query_start_locs(sample_query_start_loc, req_slice)
|
||||
|
||||
expected = torch.tensor([0, 7, 15])
|
||||
assert torch.equal(result, expected)
|
||||
|
||||
|
||||
def test_slice_from_beginning(sample_query_start_loc):
|
||||
"""Test slicing from the beginning of tensor"""
|
||||
req_slice = slice(0, 2) # slice from index 0 to 2
|
||||
result = slice_query_start_locs(sample_query_start_loc, req_slice)
|
||||
|
||||
expected = torch.tensor([0, 5, 12])
|
||||
assert torch.equal(result, expected)
|
||||
|
||||
|
||||
def test_slice_to_end(sample_query_start_loc):
|
||||
"""Test slicing to the end of tensor"""
|
||||
req_slice = slice(3, 5) # slice from index 3 to 5 (last index)
|
||||
result = slice_query_start_locs(sample_query_start_loc, req_slice)
|
||||
|
||||
expected = torch.tensor([0, 15, 30])
|
||||
assert torch.equal(result, expected)
|
||||
|
||||
|
||||
def test_single_element_slice(sample_query_start_loc):
|
||||
"""Test slice that results in single element"""
|
||||
req_slice = slice(2, 3) # slice from index 2 to 3
|
||||
result = slice_query_start_locs(sample_query_start_loc, req_slice)
|
||||
|
||||
expected = torch.tensor([0, 8])
|
||||
assert torch.equal(result, expected)
|
||||
|
||||
|
||||
def test_full_tensor_slice(sample_query_start_loc):
|
||||
"""Test slicing the entire tensor"""
|
||||
req_slice = slice(0, 5) # slice entire tensor
|
||||
result = slice_query_start_locs(sample_query_start_loc, req_slice)
|
||||
|
||||
expected = torch.tensor([0, 5, 12, 20, 35, 50])
|
||||
assert torch.equal(result, expected)
|
||||
|
||||
|
||||
def test_slice_bounds_edge_cases(sample_query_start_loc):
|
||||
# Test slice that goes exactly to the last element
|
||||
req_slice = slice(4, 5) # Last index
|
||||
result = slice_query_start_locs(sample_query_start_loc, req_slice)
|
||||
|
||||
expected = torch.tensor([0, 15])
|
||||
assert torch.equal(result, expected)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def small_decode_metadata():
|
||||
"""Create metadata for small decode batch"""
|
||||
batch_spec = BATCH_SPECS["small_decode"]
|
||||
device = torch.device("cpu")
|
||||
return create_common_attn_metadata(batch_spec, block_size=16, device=device)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def large_decode_metadata():
|
||||
"""Create metadata for small decode batch"""
|
||||
batch_spec = BATCH_SPECS["large_decode"]
|
||||
device = torch.device("cpu")
|
||||
return create_common_attn_metadata(batch_spec, block_size=16, device=device)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mixed_small_metadata():
|
||||
"""Create metadata for mixed small batch"""
|
||||
batch_spec = BATCH_SPECS["mixed_small"]
|
||||
device = torch.device("cpu")
|
||||
return create_common_attn_metadata(batch_spec, block_size=16, device=device)
|
||||
|
||||
|
||||
# Tests for _make_metadata_with_slice
|
||||
def test_make_metadata_with_slice_decode_batch(small_decode_metadata):
|
||||
"""Test slicing decode batch metadata"""
|
||||
# Split first request only
|
||||
ubatch_slice = UBatchSlice(slice(0, 1), slice(0, 1))
|
||||
|
||||
result = _make_metadata_with_slice(ubatch_slice, small_decode_metadata)
|
||||
|
||||
# Check sliced results
|
||||
assert result.num_reqs == 1 # slice(0, 1) gives 1 requests
|
||||
assert result.num_actual_tokens == 1 # slice(0, 1) gives 1 token
|
||||
assert result.max_query_len == 1
|
||||
assert torch.equal(result.query_start_loc, torch.tensor([0, 1]))
|
||||
assert torch.equal(result.seq_lens, torch.tensor([32]))
|
||||
|
||||
|
||||
def test_make_metadata_with_slice_mixed_batch(mixed_small_metadata):
|
||||
"""Test slicing mixed batch metadata"""
|
||||
ubatch_slice = UBatchSlice(slice(1, 3), slice(1, 7)) # Requests 1-3, tokens 1-7
|
||||
|
||||
result = _make_metadata_with_slice(ubatch_slice, mixed_small_metadata)
|
||||
|
||||
assert result.num_reqs == 2 # slice(1, 3) gives 2 requests
|
||||
assert result.num_actual_tokens == 6 # slice(1, 7) gives 6 tokens
|
||||
assert result.max_query_len == 5
|
||||
assert torch.equal(result.query_start_loc, torch.tensor([0, 1, 6]))
|
||||
assert torch.equal(result.seq_lens, torch.tensor([40, 48]))
|
||||
|
||||
|
||||
def test_split_attn_metadata_decode_batch(large_decode_metadata):
|
||||
"""Test splitting decode batch into two equal parts"""
|
||||
num_tokens = large_decode_metadata.num_reqs
|
||||
mid_point = num_tokens // 2
|
||||
ubatch_slices = [
|
||||
UBatchSlice(slice(0, mid_point), slice(0, mid_point)),
|
||||
UBatchSlice(slice(mid_point, num_tokens), slice(mid_point, num_tokens)),
|
||||
]
|
||||
|
||||
results = split_attn_metadata(ubatch_slices, large_decode_metadata)
|
||||
|
||||
assert len(results) == 2
|
||||
|
||||
# Check first split
|
||||
assert results[0].num_reqs == mid_point
|
||||
assert results[0].num_actual_tokens == mid_point
|
||||
assert torch.equal(results[0].seq_lens, torch.tensor([2048] * mid_point))
|
||||
|
||||
# Check second split
|
||||
assert results[1].num_reqs == mid_point
|
||||
assert results[1].num_actual_tokens == mid_point
|
||||
assert torch.equal(results[1].seq_lens, torch.tensor([2048] * mid_point))
|
||||
|
||||
|
||||
def apply_split_decodes_and_prefills(
|
||||
query_lens: list[int],
|
||||
decode_threshold: int,
|
||||
require_uniform: bool,
|
||||
padded_num_tokens: int | None = None,
|
||||
):
|
||||
"""Helper function to apply split_decodes_and_prefills and return
|
||||
the results."""
|
||||
device = torch.device("cpu")
|
||||
seq_lens = [10 * (i + 1) for i in range(len(query_lens))]
|
||||
common_metadata = create_common_attn_metadata(
|
||||
BatchSpec(seq_lens=seq_lens, query_lens=query_lens),
|
||||
block_size=16,
|
||||
device=device,
|
||||
)
|
||||
|
||||
if padded_num_tokens is not None:
|
||||
common_metadata.num_actual_tokens = padded_num_tokens
|
||||
|
||||
return split_decodes_and_prefills(
|
||||
common_metadata,
|
||||
decode_threshold=decode_threshold,
|
||||
require_uniform=require_uniform,
|
||||
)
|
||||
|
||||
|
||||
def test_split_decodes_and_prefills_nonuniform_all_ones():
|
||||
query_lens = [1, 1, 1]
|
||||
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
|
||||
apply_split_decodes_and_prefills(query_lens, 1, False)
|
||||
)
|
||||
assert num_decodes == 3
|
||||
assert num_prefills == 0
|
||||
assert num_decode_tokens == 3
|
||||
assert num_prefill_tokens == 0
|
||||
|
||||
|
||||
def test_split_decodes_and_prefills_nonuniform_all_short_decodes():
|
||||
query_lens = [1, 2, 1, 3, 2, 1, 2]
|
||||
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
|
||||
apply_split_decodes_and_prefills(query_lens, 3, False)
|
||||
)
|
||||
assert num_decodes == 7
|
||||
assert num_prefills == 0
|
||||
assert num_decode_tokens == sum(query_lens)
|
||||
assert num_prefill_tokens == 0
|
||||
|
||||
|
||||
def test_split_decodes_and_prefills_nonuniform_all_prefills():
|
||||
query_lens = [4, 5, 6, 7]
|
||||
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
|
||||
apply_split_decodes_and_prefills(query_lens, 3, False)
|
||||
)
|
||||
assert num_decodes == 0
|
||||
assert num_prefills == 4
|
||||
assert num_decode_tokens == 0
|
||||
assert num_prefill_tokens == sum(query_lens)
|
||||
|
||||
|
||||
def test_split_decodes_and_prefills_nonuniform_mixed_batch():
|
||||
query_lens = [2, 1, 3, 4, 5, 6, 7, 8]
|
||||
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
|
||||
apply_split_decodes_and_prefills(query_lens, 4, False)
|
||||
)
|
||||
assert num_decodes == 4 # 2, 1, 3, 4 are all <= 4
|
||||
assert num_prefills == 4 # 5, 6, 7, 8 are all > 4
|
||||
assert num_decode_tokens == 10 # 2 + 1 + 3 + 4
|
||||
assert num_prefill_tokens == 26 # 5 + 6 + 7 + 8
|
||||
|
||||
|
||||
def test_split_decodes_and_prefills_uniform_all_ones():
|
||||
query_lens = [1, 1, 1]
|
||||
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
|
||||
apply_split_decodes_and_prefills(query_lens, 1, True)
|
||||
)
|
||||
assert num_decodes == 3
|
||||
assert num_prefills == 0
|
||||
assert num_decode_tokens == 3
|
||||
assert num_prefill_tokens == 0
|
||||
|
||||
|
||||
def test_split_decodes_and_prefills_uniform_all_short_decodes():
|
||||
query_lens = [2, 2, 1, 3, 2, 1, 2]
|
||||
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
|
||||
apply_split_decodes_and_prefills(query_lens, 3, True)
|
||||
)
|
||||
assert num_decodes == 2
|
||||
assert num_prefills == 5
|
||||
assert num_decode_tokens == 4
|
||||
assert num_prefill_tokens == (1 + 3 + 2 + 1 + 2)
|
||||
|
||||
|
||||
def test_split_decodes_and_prefills_uniform_all_prefills():
|
||||
query_lens = [4, 5, 6, 7]
|
||||
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
|
||||
apply_split_decodes_and_prefills(query_lens, 3, True)
|
||||
)
|
||||
assert num_decodes == 0
|
||||
assert num_prefills == 4
|
||||
assert num_decode_tokens == 0
|
||||
assert num_prefill_tokens == sum(query_lens)
|
||||
|
||||
|
||||
def test_split_decodes_and_prefills_uniform_mixed_batch_all_uniform_decodes():
|
||||
query_lens = [2, 2, 2, 4, 5, 6, 7, 8]
|
||||
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
|
||||
apply_split_decodes_and_prefills(query_lens, 4, True)
|
||||
)
|
||||
assert num_decodes == 3 # 2, 2, 2 are all <= 4 and uniform
|
||||
assert num_prefills == 5 # 4, 5, 6, 7, 8 are all > 4
|
||||
assert num_decode_tokens == 6 # 2 + 2 + 2
|
||||
assert num_prefill_tokens == 30 # 4 + 5 + 6 + 7 + 8
|
||||
|
||||
|
||||
def test_split_decodes_and_prefills_uniform_mixed_batch_non_uniform_decodes():
|
||||
query_lens = [2, 1, 2, 4, 5, 6, 7, 8]
|
||||
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
|
||||
apply_split_decodes_and_prefills(query_lens, 4, True)
|
||||
)
|
||||
assert num_decodes == 1 # only the first 2 is taken as decode
|
||||
assert num_prefills == 7 # 1, 2, 4, 5, 6, 7, 8 are all > 4 or non-uniform
|
||||
assert num_decode_tokens == 2 # only the first 2
|
||||
assert num_prefill_tokens == (sum(query_lens) - 2) # rest of the tokens
|
||||
|
||||
|
||||
def test_split_decodes_and_prefills_uniform_padded_batch_all_same():
|
||||
"""uniform batch where all query lengths are identical with 0 length padded reqs."""
|
||||
# All query lengths are 2, with decode_threshold=3 (so 2 <= 3)
|
||||
# This triggers the padded uniform path at line 891
|
||||
query_lens = [2, 2, 2, 0]
|
||||
padded_num_tokens = 8
|
||||
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
|
||||
apply_split_decodes_and_prefills(query_lens, 3, True, padded_num_tokens)
|
||||
)
|
||||
# With uniform batch, all requests are treated as decodes
|
||||
assert num_decodes == 4
|
||||
assert num_prefills == 0
|
||||
assert num_decode_tokens == padded_num_tokens
|
||||
assert num_prefill_tokens == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"seq_lens,query_lens,split_point,expected_first_reqs,expected_second_reqs",
|
||||
[
|
||||
# Split in the middle of request 1
|
||||
([32, 40], [8, 8], 12, 2, 1),
|
||||
# Split inside the first request
|
||||
([32, 40], [8, 8], 4, 1, 2),
|
||||
],
|
||||
)
|
||||
def test_prefill_split_across_ubatches(
|
||||
seq_lens, query_lens, split_point, expected_first_reqs, expected_second_reqs
|
||||
):
|
||||
"""Test splitting a prefill across ubatches"""
|
||||
import numpy as np
|
||||
|
||||
device = torch.device("cpu")
|
||||
batch_spec = BatchSpec(seq_lens=seq_lens, query_lens=query_lens)
|
||||
common = create_common_attn_metadata(batch_spec, block_size=16, device=device)
|
||||
|
||||
num_scheduled_tokens = np.array(query_lens, dtype=np.int32)
|
||||
qsl_np = common.query_start_loc_cpu.numpy()
|
||||
num_tokens = common.num_actual_tokens
|
||||
|
||||
ubatch_slices, _ = maybe_create_ubatch_slices(
|
||||
True,
|
||||
num_scheduled_tokens,
|
||||
num_tokens,
|
||||
batch_spec.batch_size,
|
||||
split_point=split_point,
|
||||
num_ubatches=2,
|
||||
)
|
||||
assert ubatch_slices is not None and len(ubatch_slices) == 2
|
||||
|
||||
first_meta = _make_metadata_with_slice(ubatch_slices[0], common)
|
||||
second_meta = _make_metadata_with_slice(ubatch_slices[1], common)
|
||||
|
||||
# Token counts match the split
|
||||
assert first_meta.num_actual_tokens == split_point
|
||||
assert second_meta.num_actual_tokens == num_tokens - split_point
|
||||
|
||||
# Number of requests per ubatch
|
||||
assert first_meta.num_reqs == expected_first_reqs
|
||||
assert second_meta.num_reqs == expected_second_reqs
|
||||
|
||||
# Identify which request is split and how many tokens are in the first chunk
|
||||
split_req_idx = int(np.searchsorted(qsl_np, split_point, side="right") - 1)
|
||||
tokens_in_first_chunk = split_point - int(qsl_np[split_req_idx])
|
||||
orig_q_lens = common.query_start_loc_cpu[1:] - common.query_start_loc_cpu[:-1]
|
||||
|
||||
# Check query length continuity: first-chunk + second-chunk == original qlen
|
||||
# First ubatch last request query length
|
||||
qlen_first_last = int(
|
||||
first_meta.query_start_loc_cpu[-1] - first_meta.query_start_loc_cpu[-2]
|
||||
)
|
||||
# Second ubatch first request query length
|
||||
qlen_second_first = int(
|
||||
second_meta.query_start_loc_cpu[1] - second_meta.query_start_loc_cpu[0]
|
||||
)
|
||||
assert qlen_first_last == tokens_in_first_chunk
|
||||
assert qlen_first_last + qlen_second_first == int(orig_q_lens[split_req_idx])
|
||||
|
||||
# Check seq_lens adjustments
|
||||
# Context lengths per original request
|
||||
context_lens = [s - q for s, q in zip(seq_lens, query_lens)]
|
||||
|
||||
# First ubatch: last request's seq_len should be
|
||||
# context + tokens_in_first_chunk
|
||||
expected_seqlen = context_lens[split_req_idx] + tokens_in_first_chunk
|
||||
assert int(first_meta.seq_lens[-1]) == expected_seqlen
|
||||
|
||||
# For full preceding requests in first ubatch, seq_lens should match
|
||||
# originals
|
||||
for i in range(first_meta.num_reqs - 1):
|
||||
assert int(first_meta.seq_lens[i]) == seq_lens[i]
|
||||
|
||||
# Second ubatch: first request (continuation) seq_len should be full
|
||||
# original
|
||||
assert int(second_meta.seq_lens[0]) == seq_lens[split_req_idx]
|
||||
# Any following full requests in second ubatch should match originals
|
||||
for j in range(1, second_meta.num_reqs):
|
||||
# Map to original request index
|
||||
orig_idx = split_req_idx + j
|
||||
assert int(second_meta.seq_lens[j]) == seq_lens[orig_idx]
|
||||
147
third_party/vllm/tests/v1/attention/test_batch_reordering.py
vendored
Normal file
147
third_party/vllm/tests/v1/attention/test_batch_reordering.py
vendored
Normal file
@@ -0,0 +1,147 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from vllm.v1.attention.backends.utils import reorder_batch_to_split_decodes_and_prefills
|
||||
|
||||
|
||||
class MockInputBatch:
|
||||
def __init__(self, req_ids, num_computed_tokens_cpu):
|
||||
self.req_ids = req_ids
|
||||
self.num_computed_tokens_cpu = num_computed_tokens_cpu
|
||||
|
||||
def swap_states(self, i, j):
|
||||
self.req_ids[i], self.req_ids[j] = self.req_ids[j], self.req_ids[i]
|
||||
self.num_computed_tokens_cpu[i], self.num_computed_tokens_cpu[j] = (
|
||||
self.num_computed_tokens_cpu[j],
|
||||
self.num_computed_tokens_cpu[i],
|
||||
)
|
||||
|
||||
|
||||
class MockSchedulerOutput:
|
||||
def __init__(self, num_scheduled_tokens):
|
||||
self.num_scheduled_tokens = num_scheduled_tokens
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReorderTestCase:
|
||||
requests: list[tuple[int, int]] # (num_scheduled_tokens, num_computed_tokens)
|
||||
expected_order: list[int]
|
||||
expected_modified: bool
|
||||
decode_threshold: int = 1
|
||||
|
||||
|
||||
# Test cases for batch reordering
|
||||
REORDER_TEST_CASES = {
|
||||
"all_decodes": ReorderTestCase(
|
||||
requests=[(1, 10), (1, 20), (1, 30)],
|
||||
expected_order=[0, 1, 2],
|
||||
expected_modified=False,
|
||||
),
|
||||
"all_prefills": ReorderTestCase(
|
||||
requests=[(100, 100), (200, 200), (300, 300)],
|
||||
expected_order=[0, 1, 2],
|
||||
expected_modified=False,
|
||||
),
|
||||
"mixed_interleaved": ReorderTestCase(
|
||||
requests=[(100, 100), (1, 10), (200, 200), (1, 20)],
|
||||
expected_order=[3, 1, 2, 0], # Only swap 0↔3, keep 1 and 2 in place
|
||||
expected_modified=True,
|
||||
),
|
||||
"already_ordered": ReorderTestCase(
|
||||
requests=[(1, 10), (1, 20), (100, 100), (200, 0)],
|
||||
expected_order=[0, 1, 2, 3],
|
||||
expected_modified=False,
|
||||
),
|
||||
"single_request": ReorderTestCase(
|
||||
requests=[(1, 10)],
|
||||
expected_order=[0],
|
||||
expected_modified=False,
|
||||
),
|
||||
"higher_threshold": ReorderTestCase(
|
||||
requests=[(2, 10), (3, 20), (5, 30), (6, 40)],
|
||||
expected_order=[0, 1, 2, 3],
|
||||
expected_modified=False,
|
||||
decode_threshold=4,
|
||||
),
|
||||
"decodes_at_end": ReorderTestCase(
|
||||
requests=[(100, 100), (200, 200), (1, 10), (1, 20)],
|
||||
expected_order=[2, 3, 0, 1],
|
||||
expected_modified=True,
|
||||
),
|
||||
"decode_extend_prefill": ReorderTestCase(
|
||||
requests=[(100, 0), (10, 50), (1, 10)],
|
||||
expected_order=[2, 1, 0],
|
||||
expected_modified=True,
|
||||
),
|
||||
"extend_prefill_only": ReorderTestCase(
|
||||
requests=[(100, 0), (10, 50), (200, 0), (20, 75)],
|
||||
expected_order=[3, 1, 2, 0], # Only swap 0↔3, keep 1 and 2 in place
|
||||
expected_modified=True,
|
||||
),
|
||||
"complicated_mixed_interleaved": ReorderTestCase(
|
||||
requests=[
|
||||
(1, 20),
|
||||
(1, 50),
|
||||
(374, 0),
|
||||
(300, 20),
|
||||
(1, 20),
|
||||
(256, 0),
|
||||
(1, 5),
|
||||
(27, 0),
|
||||
(1, 4),
|
||||
],
|
||||
expected_order=[0, 1, 6, 8, 4, 3, 2, 7, 5],
|
||||
expected_modified=True,
|
||||
),
|
||||
"new_request_single_token_prefill": ReorderTestCase(
|
||||
requests=[
|
||||
(100, 0),
|
||||
(1, 0), # New request with only 1 token (STILL prefill)
|
||||
(50, 100),
|
||||
(1, 10),
|
||||
],
|
||||
# Only index 3 is a true decode (has num_computed_tokens > 0)
|
||||
expected_order=[3, 2, 0, 1],
|
||||
expected_modified=True,
|
||||
),
|
||||
"multiple_new_requests_single_token_prefill": ReorderTestCase(
|
||||
requests=[
|
||||
(1, 0), # New prefill (1 token, no computed)
|
||||
(1, 0), # New prefill (1 token, no computed)
|
||||
(1, 50),
|
||||
(200, 0),
|
||||
],
|
||||
expected_order=[2, 1, 0, 3],
|
||||
expected_modified=True,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"test_case", REORDER_TEST_CASES.values(), ids=REORDER_TEST_CASES.keys()
|
||||
)
|
||||
def test_reorder_batch_to_split_decodes_and_prefills(test_case: ReorderTestCase):
|
||||
req_ids = [f"r{i}" for i in range(len(test_case.requests))]
|
||||
num_computed_tokens = np.array([r[1] for r in test_case.requests], dtype=np.int32)
|
||||
num_scheduled_tokens = {f"r{i}": r[0] for i, r in enumerate(test_case.requests)}
|
||||
|
||||
input_batch = MockInputBatch(req_ids, num_computed_tokens)
|
||||
scheduler_output = MockSchedulerOutput(num_scheduled_tokens)
|
||||
|
||||
modified = reorder_batch_to_split_decodes_and_prefills(
|
||||
input_batch, scheduler_output, decode_threshold=test_case.decode_threshold
|
||||
)
|
||||
|
||||
expected_req_ids = [f"r{i}" for i in test_case.expected_order]
|
||||
|
||||
assert modified == test_case.expected_modified, (
|
||||
f"Expected modified={test_case.expected_modified}, got {modified}"
|
||||
)
|
||||
assert input_batch.req_ids == expected_req_ids, (
|
||||
f"Expected order {expected_req_ids}, got {input_batch.req_ids}"
|
||||
)
|
||||
201
third_party/vllm/tests/v1/attention/test_chunked_local_attention.py
vendored
Normal file
201
third_party/vllm/tests/v1/attention/test_chunked_local_attention.py
vendored
Normal file
@@ -0,0 +1,201 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from dataclasses import dataclass
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
|
||||
from vllm.v1.attention.backends.utils import make_local_attention_virtual_batches
|
||||
|
||||
|
||||
@dataclass
|
||||
class LocalAttentionTestData:
|
||||
# Input parameters
|
||||
batch_spec: BatchSpec
|
||||
attn_chunk_size: int
|
||||
block_size: int
|
||||
# Expected return values
|
||||
expected_q_seqlens: list[int]
|
||||
expected_k_seqlens: list[int]
|
||||
expected_local_block_table: list[list[int]]
|
||||
|
||||
|
||||
test_data_list = [
|
||||
# Same as example in docstring of make_local_attention_virtual_batches
|
||||
# except block table has 9 columns instead of 10
|
||||
LocalAttentionTestData(
|
||||
batch_spec=BatchSpec(
|
||||
query_lens=[4, 10, 5],
|
||||
seq_lens=[6, 17, 9],
|
||||
),
|
||||
attn_chunk_size=4,
|
||||
block_size=2,
|
||||
expected_q_seqlens=[2, 2, 1, 4, 4, 1, 4, 1],
|
||||
expected_k_seqlens=[4, 2, 4, 4, 4, 1, 4, 1],
|
||||
# 2 pages per local branch
|
||||
# (chunk size 4 // block size 2)
|
||||
expected_local_block_table=[
|
||||
[0, 1], # local-batch 0, (batch 0, starting from k[0])
|
||||
[2, 3], # local-batch 1, (batch 0, starting from k[4])
|
||||
[11, 12], # local-batch 2, (batch 1, starting from k[4])
|
||||
[13, 14], # local-batch 3, (batch 1, starting from k[8])
|
||||
[15, 16], # local-batch 4, (batch 1, starting from k[12])
|
||||
[17, 17], # local-batch 5, (batch 1, starting from k[16])
|
||||
[20, 21], # local-batch 6, (batch 2, starting from k[4])
|
||||
[22, 23], # local-batch 7, (batch 2, starting from k[8])
|
||||
],
|
||||
),
|
||||
# Case where block indices are not clipped to block table ncols-1
|
||||
# because tokens_in_last_block == attn_chunk_size
|
||||
LocalAttentionTestData(
|
||||
batch_spec=BatchSpec(
|
||||
query_lens=[8],
|
||||
seq_lens=[12],
|
||||
),
|
||||
attn_chunk_size=4,
|
||||
block_size=2,
|
||||
expected_q_seqlens=[4, 4],
|
||||
expected_k_seqlens=[4, 4],
|
||||
expected_local_block_table=[
|
||||
[2, 3],
|
||||
[4, 5],
|
||||
],
|
||||
),
|
||||
# Case where all kv_seq positions are involved in attn
|
||||
LocalAttentionTestData(
|
||||
batch_spec=BatchSpec(
|
||||
query_lens=[7],
|
||||
# 10 - 7 = 3 previously computed tokens
|
||||
seq_lens=[10],
|
||||
),
|
||||
attn_chunk_size=4,
|
||||
block_size=2,
|
||||
expected_q_seqlens=[1, 4, 2],
|
||||
expected_k_seqlens=[4, 4, 2],
|
||||
expected_local_block_table=[
|
||||
[0, 1],
|
||||
[2, 3],
|
||||
[4, 4],
|
||||
],
|
||||
),
|
||||
# Case where attn_chunk_size > kv_seq_len
|
||||
# so no extra mini virtual batches are created
|
||||
LocalAttentionTestData(
|
||||
batch_spec=BatchSpec(
|
||||
query_lens=[4],
|
||||
seq_lens=[6],
|
||||
),
|
||||
# Larger than kv_seq_len
|
||||
attn_chunk_size=10,
|
||||
block_size=2,
|
||||
# No change to q_seqlens and k_seqlens
|
||||
expected_q_seqlens=[4],
|
||||
expected_k_seqlens=[6],
|
||||
# In this case, we only need a block-table like:
|
||||
# block_table = [ [0, 1, 2] ] # 1 batch, 3 pages
|
||||
# But we need to pad it to 5 pages per local batch
|
||||
# because currently the pages_per_local_batch
|
||||
# is calculated as (attn_chunk_size // block_size)
|
||||
expected_local_block_table=[
|
||||
[0, 1, 2, 2, 2],
|
||||
],
|
||||
),
|
||||
# Block size equal to chunk size
|
||||
# Expect single page per batch in local batch table
|
||||
LocalAttentionTestData(
|
||||
batch_spec=BatchSpec(
|
||||
query_lens=[6, 6],
|
||||
seq_lens=[8, 8],
|
||||
),
|
||||
attn_chunk_size=4,
|
||||
block_size=4,
|
||||
expected_q_seqlens=[2, 4, 2, 4],
|
||||
expected_k_seqlens=[4, 4, 4, 4],
|
||||
# Initial block table = [
|
||||
# [0, 1], < batch 0
|
||||
# [2, 3], < batch 1
|
||||
# ]
|
||||
expected_local_block_table=[
|
||||
[0], # local-batch 0, (batch 0, starting from k[0])
|
||||
[1], # local-batch 1, (batch 0, starting from k[4])
|
||||
[2], # local-batch 1, (batch 0, starting from k[0])
|
||||
[3], # local-batch 1, (batch 0, starting from k[4])
|
||||
],
|
||||
),
|
||||
# Case where query falls in the second attention chunk
|
||||
# k_toks > 0 1 2 3 4
|
||||
# q_toks v _____________
|
||||
# 0 | 1
|
||||
# 1 | 1 1
|
||||
# 2 | 1 1 1
|
||||
# 3 | 1 1 1 1
|
||||
# 4 | 1
|
||||
# where tokens 0,1,2,3 have been pre-computed
|
||||
LocalAttentionTestData(
|
||||
batch_spec=BatchSpec(
|
||||
query_lens=[1],
|
||||
seq_lens=[5],
|
||||
),
|
||||
attn_chunk_size=4,
|
||||
block_size=2,
|
||||
expected_q_seqlens=[1],
|
||||
expected_k_seqlens=[1],
|
||||
expected_local_block_table=[
|
||||
[2, 2],
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("test_data", test_data_list)
|
||||
def test_local_attention_virtual_batches(test_data: LocalAttentionTestData):
|
||||
device = torch.device("cuda:0")
|
||||
batch_spec = test_data.batch_spec
|
||||
attn_chunk_size = test_data.attn_chunk_size
|
||||
block_size = test_data.block_size
|
||||
expected_q_seqlens = test_data.expected_q_seqlens
|
||||
expected_k_seqlens = test_data.expected_k_seqlens
|
||||
expected_local_block_table = test_data.expected_local_block_table
|
||||
|
||||
# Create common attention metadata
|
||||
common_attn_metadata = create_common_attn_metadata(
|
||||
batch_spec,
|
||||
block_size,
|
||||
device,
|
||||
# Use torch.arange instead of torch.randint so we can assert on
|
||||
# block table tensor values. The block table will have shape
|
||||
# (num_batches, cdiv(max_seq_len, block_size)) and the values will be
|
||||
# arranged from 0 to cdiv(max_seq_len, block_size)-1
|
||||
arange_block_indices=True,
|
||||
)
|
||||
|
||||
# Call the function
|
||||
result, _ = make_local_attention_virtual_batches(
|
||||
attn_chunk_size, common_attn_metadata, block_size
|
||||
)
|
||||
|
||||
# Convert to numpy for easier comparison
|
||||
actual_q_seqlens = np.diff(result.query_start_loc_cpu.numpy())
|
||||
actual_k_seqlens = result.seq_lens_cpu.numpy()
|
||||
|
||||
# Check that all query lengths are less than or equal to attn_chunk_size
|
||||
assert all(q_len <= attn_chunk_size for q_len in actual_q_seqlens)
|
||||
# Check that all key lengths are less than or equal to attn_chunk_size
|
||||
assert all(k_len <= attn_chunk_size for k_len in actual_k_seqlens)
|
||||
# Check that the total number of query tokens is preserved
|
||||
assert sum(actual_q_seqlens) == sum(batch_spec.query_lens)
|
||||
|
||||
# Verify results
|
||||
np.testing.assert_array_equal(actual_q_seqlens, expected_q_seqlens)
|
||||
np.testing.assert_array_equal(actual_k_seqlens, expected_k_seqlens)
|
||||
|
||||
expected_block_table_tensor = torch.tensor(
|
||||
expected_local_block_table, dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
print(f"Expected block table:\n{expected_block_table_tensor}")
|
||||
print(f"Actual block table:\n{result.block_table_tensor}")
|
||||
|
||||
torch.testing.assert_close(result.block_table_tensor, expected_block_table_tensor)
|
||||
191
third_party/vllm/tests/v1/attention/test_gdn_metadata_builder.py
vendored
Normal file
191
third_party/vllm/tests/v1/attention/test_gdn_metadata_builder.py
vendored
Normal file
@@ -0,0 +1,191 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for GDNAttentionMetadataBuilder.build() — specifically the
|
||||
reclassification of non-spec decodes as prefills when spec decodes exist.
|
||||
Covers the fix for https://github.com/vllm-project/vllm/issues/34845.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.v1.attention.utils import (
|
||||
BatchSpec,
|
||||
create_common_attn_metadata,
|
||||
create_vllm_config,
|
||||
)
|
||||
from vllm.config import SpeculativeConfig
|
||||
from vllm.v1.attention.backends.gdn_attn import (
|
||||
GDNAttentionMetadata,
|
||||
GDNAttentionMetadataBuilder,
|
||||
)
|
||||
from vllm.v1.kv_cache_interface import MambaSpec
|
||||
|
||||
BLOCK_SIZE = 16
|
||||
DEVICE = torch.device("cpu")
|
||||
|
||||
|
||||
@dataclass
|
||||
class GDNBuildTestCase:
|
||||
"""Specification for a GDN metadata builder classification test."""
|
||||
|
||||
seq_lens: list[int]
|
||||
query_lens: list[int]
|
||||
num_decode_draft_tokens: list[int] | None # None = no spec config
|
||||
num_speculative_tokens: int
|
||||
expected_num_decodes: int
|
||||
expected_num_prefills: int
|
||||
expected_num_prefill_tokens: int
|
||||
expected_num_spec_decodes: int
|
||||
|
||||
|
||||
GDN_BUILD_TEST_CASES = {
|
||||
# The original #34845 crash: non-spec query_len=1 + spec decode
|
||||
"mixed_decode_and_spec_decode": GDNBuildTestCase(
|
||||
seq_lens=[65, 20],
|
||||
query_lens=[1, 3],
|
||||
num_decode_draft_tokens=[-1, 2],
|
||||
num_speculative_tokens=2,
|
||||
expected_num_decodes=0,
|
||||
expected_num_prefills=1,
|
||||
expected_num_prefill_tokens=1,
|
||||
expected_num_spec_decodes=1,
|
||||
),
|
||||
# All requests are spec decodes — no reclassification needed
|
||||
"pure_spec_decode": GDNBuildTestCase(
|
||||
seq_lens=[50, 30],
|
||||
query_lens=[3, 3],
|
||||
num_decode_draft_tokens=[2, 2],
|
||||
num_speculative_tokens=2,
|
||||
expected_num_decodes=0,
|
||||
expected_num_prefills=0,
|
||||
expected_num_prefill_tokens=0,
|
||||
expected_num_spec_decodes=2,
|
||||
),
|
||||
# No speculative config at all — standard decode path
|
||||
"pure_regular_decode": GDNBuildTestCase(
|
||||
seq_lens=[40, 30, 20],
|
||||
query_lens=[1, 1, 1],
|
||||
num_decode_draft_tokens=None,
|
||||
num_speculative_tokens=0,
|
||||
expected_num_decodes=3,
|
||||
expected_num_prefills=0,
|
||||
expected_num_prefill_tokens=0,
|
||||
expected_num_spec_decodes=0,
|
||||
),
|
||||
# Multi-token prefill alongside spec decode — no decode to reclassify
|
||||
"spec_decode_with_real_prefill": GDNBuildTestCase(
|
||||
seq_lens=[100, 20],
|
||||
query_lens=[50, 3],
|
||||
num_decode_draft_tokens=[-1, 2],
|
||||
num_speculative_tokens=2,
|
||||
expected_num_decodes=0,
|
||||
expected_num_prefills=1,
|
||||
expected_num_prefill_tokens=50,
|
||||
expected_num_spec_decodes=1,
|
||||
),
|
||||
# All three types in one batch — decode gets reclassified
|
||||
"prefill_decode_and_spec_decode": GDNBuildTestCase(
|
||||
seq_lens=[100, 65, 20],
|
||||
query_lens=[50, 1, 3],
|
||||
num_decode_draft_tokens=[-1, -1, 2],
|
||||
num_speculative_tokens=2,
|
||||
expected_num_decodes=0,
|
||||
expected_num_prefills=2,
|
||||
expected_num_prefill_tokens=51,
|
||||
expected_num_spec_decodes=1,
|
||||
),
|
||||
# Multiple non-spec query_len=1 requests all reclassified
|
||||
"multiple_decodes_reclassified": GDNBuildTestCase(
|
||||
seq_lens=[40, 50, 60, 20],
|
||||
query_lens=[1, 1, 1, 3],
|
||||
num_decode_draft_tokens=[-1, -1, -1, 2],
|
||||
num_speculative_tokens=2,
|
||||
expected_num_decodes=0,
|
||||
expected_num_prefills=3,
|
||||
expected_num_prefill_tokens=3,
|
||||
expected_num_spec_decodes=1,
|
||||
),
|
||||
# Zero-length padded sequence excluded from counts
|
||||
"zero_length_padding_with_spec": GDNBuildTestCase(
|
||||
seq_lens=[16, 65, 20],
|
||||
query_lens=[0, 1, 3],
|
||||
num_decode_draft_tokens=[-1, -1, 2],
|
||||
num_speculative_tokens=2,
|
||||
expected_num_decodes=0,
|
||||
expected_num_prefills=1,
|
||||
expected_num_prefill_tokens=1,
|
||||
expected_num_spec_decodes=1,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _create_gdn_builder(
|
||||
num_speculative_tokens: int = 0,
|
||||
) -> GDNAttentionMetadataBuilder:
|
||||
"""Create a GDNAttentionMetadataBuilder with minimal config."""
|
||||
vllm_config = create_vllm_config(block_size=BLOCK_SIZE)
|
||||
if num_speculative_tokens > 0:
|
||||
vllm_config.speculative_config = SpeculativeConfig(
|
||||
method="ngram",
|
||||
num_speculative_tokens=num_speculative_tokens,
|
||||
)
|
||||
mamba_spec = MambaSpec(
|
||||
block_size=BLOCK_SIZE,
|
||||
shapes=((16, 64),),
|
||||
dtypes=(torch.float16,),
|
||||
)
|
||||
return GDNAttentionMetadataBuilder(
|
||||
kv_cache_spec=mamba_spec,
|
||||
layer_names=["layer.0"],
|
||||
vllm_config=vllm_config,
|
||||
device=DEVICE,
|
||||
)
|
||||
|
||||
|
||||
def _build(
|
||||
builder: GDNAttentionMetadataBuilder,
|
||||
batch_spec: BatchSpec,
|
||||
num_decode_draft_tokens: list[int] | None = None,
|
||||
) -> GDNAttentionMetadata:
|
||||
"""Build GDN attention metadata, optionally with spec-decode kwargs."""
|
||||
common = create_common_attn_metadata(batch_spec, BLOCK_SIZE, DEVICE)
|
||||
kwargs: dict = {}
|
||||
if num_decode_draft_tokens is not None:
|
||||
kwargs["num_decode_draft_tokens_cpu"] = torch.tensor(
|
||||
num_decode_draft_tokens, dtype=torch.int32
|
||||
)
|
||||
kwargs["num_accepted_tokens"] = torch.ones(
|
||||
batch_spec.batch_size, dtype=torch.int32, device=DEVICE
|
||||
)
|
||||
return builder.build(common_prefix_len=0, common_attn_metadata=common, **kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"test_case", GDN_BUILD_TEST_CASES.values(), ids=GDN_BUILD_TEST_CASES.keys()
|
||||
)
|
||||
def test_gdn_build_classification(test_case: GDNBuildTestCase):
|
||||
"""Test that GDN metadata builder classifies requests correctly."""
|
||||
builder = _create_gdn_builder(test_case.num_speculative_tokens)
|
||||
batch = BatchSpec(seq_lens=test_case.seq_lens, query_lens=test_case.query_lens)
|
||||
meta = _build(builder, batch, test_case.num_decode_draft_tokens)
|
||||
|
||||
assert meta.num_decodes == test_case.expected_num_decodes
|
||||
assert meta.num_prefills == test_case.expected_num_prefills
|
||||
assert meta.num_prefill_tokens == test_case.expected_num_prefill_tokens
|
||||
assert meta.num_spec_decodes == test_case.expected_num_spec_decodes
|
||||
|
||||
|
||||
def test_has_initial_state_after_reclassification():
|
||||
"""After reclassification, num_prefills > 0 so the prefill kernel path
|
||||
should compute has_initial_state. For the reclassified request with
|
||||
context_lens > 0, the corresponding entry must be True."""
|
||||
builder = _create_gdn_builder(num_speculative_tokens=2)
|
||||
batch = BatchSpec(seq_lens=[65, 20], query_lens=[1, 3])
|
||||
meta = _build(builder, batch, num_decode_draft_tokens=[-1, 2])
|
||||
|
||||
assert meta.num_prefills > 0, "reclassification should produce prefills"
|
||||
assert meta.has_initial_state is not None
|
||||
# req0 has context_lens = 65 - 1 = 64 > 0, so has_initial_state[0] = True
|
||||
assert meta.has_initial_state[0].item() is True
|
||||
151
third_party/vllm/tests/v1/attention/test_mamba_update_block_table.py
vendored
Normal file
151
third_party/vllm/tests/v1/attention/test_mamba_update_block_table.py
vendored
Normal file
@@ -0,0 +1,151 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Regression test for https://github.com/vllm-project/vllm/issues/34865
|
||||
|
||||
When multiple KV cache groups share the same MambaSpec (as in Nemotron
|
||||
hybrid models), the metadata caching optimization reuses metadata from
|
||||
an earlier group via update_block_table(). In 'all' mode with CUDA graphs,
|
||||
update_block_table() must copy block_idx_last_scheduled_token and
|
||||
block_idx_last_computed_token to the *current* builder's persistent
|
||||
buffers, otherwise CUDA graph replay reads stale values from uninitialized
|
||||
buffers.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.config.compilation import CUDAGraphMode
|
||||
from vllm.v1.attention.backends.mamba_attn import (
|
||||
BaseMambaAttentionMetadata,
|
||||
BaseMambaAttentionMetadataBuilder,
|
||||
)
|
||||
from vllm.v1.kv_cache_interface import MambaSpec
|
||||
|
||||
|
||||
class _ConcreteMambaBuilder(
|
||||
BaseMambaAttentionMetadataBuilder[BaseMambaAttentionMetadata]
|
||||
):
|
||||
"""Minimal concrete subclass for testing (base class is ABC)."""
|
||||
|
||||
metadata_cls = BaseMambaAttentionMetadata
|
||||
|
||||
|
||||
def _make_vllm_config(block_size, max_model_len, max_num_seqs):
|
||||
"""Create a minimal mock VllmConfig with only the fields the builder
|
||||
accesses, avoiding any model download / HF config inspection."""
|
||||
return SimpleNamespace(
|
||||
cache_config=SimpleNamespace(mamba_cache_mode="all"),
|
||||
compilation_config=SimpleNamespace(
|
||||
cudagraph_mode=CUDAGraphMode.FULL,
|
||||
max_cudagraph_capture_size=None,
|
||||
),
|
||||
speculative_config=None,
|
||||
num_speculative_tokens=0,
|
||||
parallel_config=SimpleNamespace(decode_context_parallel_size=1),
|
||||
scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs),
|
||||
model_config=SimpleNamespace(max_model_len=max_model_len),
|
||||
)
|
||||
|
||||
|
||||
def test_update_block_table_copies_block_idx_to_persistent_buffers():
|
||||
"""update_block_table() must write block_idx tensors to the current
|
||||
builder's persistent buffers, not leave them pointing to a different
|
||||
builder's buffers."""
|
||||
|
||||
block_size = 16
|
||||
max_model_len = 256
|
||||
num_reqs = 4
|
||||
device = torch.device("cpu")
|
||||
|
||||
vllm_config = _make_vllm_config(block_size, max_model_len, num_reqs)
|
||||
|
||||
spec = MambaSpec(
|
||||
block_size=block_size,
|
||||
shapes=((1,), (1,)),
|
||||
dtypes=(torch.float32,),
|
||||
mamba_cache_mode="all",
|
||||
)
|
||||
|
||||
# Two builders simulating two KV cache groups with the same MambaSpec.
|
||||
builder_a = _ConcreteMambaBuilder(spec, ["layer0"], vllm_config, device)
|
||||
builder_b = _ConcreteMambaBuilder(spec, ["layer1"], vllm_config, device)
|
||||
|
||||
# Sanity: each builder has its own persistent buffer.
|
||||
assert (
|
||||
builder_a.block_idx_last_scheduled_token.data_ptr()
|
||||
!= builder_b.block_idx_last_scheduled_token.data_ptr()
|
||||
)
|
||||
|
||||
# Construct decode-only metadata as if builder_a.build() produced it.
|
||||
max_blocks = max_model_len // block_size
|
||||
seq_lens = torch.full((num_reqs,), 64, dtype=torch.int32, device=device)
|
||||
block_idx_vals = (seq_lens - 1) // block_size # [3, 3, 3, 3]
|
||||
|
||||
builder_a.block_idx_last_scheduled_token[:num_reqs].copy_(block_idx_vals)
|
||||
builder_a.block_idx_last_computed_token[:num_reqs].copy_(block_idx_vals)
|
||||
|
||||
metadata_a = BaseMambaAttentionMetadata(
|
||||
num_prefills=0,
|
||||
num_prefill_tokens=0,
|
||||
num_decodes=num_reqs,
|
||||
num_decode_tokens=num_reqs,
|
||||
num_reqs=num_reqs,
|
||||
has_initial_states_p=None,
|
||||
query_start_loc_p=None,
|
||||
num_computed_tokens_p=None,
|
||||
state_indices_tensor_p=None,
|
||||
query_start_loc_d=None,
|
||||
num_accepted_tokens=None,
|
||||
state_indices_tensor_d=builder_a.state_indices_tensor_d[:num_reqs],
|
||||
block_idx_last_scheduled_token=(
|
||||
builder_a.block_idx_last_scheduled_token[:num_reqs]
|
||||
),
|
||||
block_idx_first_scheduled_token_p=None,
|
||||
block_idx_last_computed_token=(
|
||||
builder_a.block_idx_last_computed_token[:num_reqs]
|
||||
),
|
||||
seq_lens=seq_lens,
|
||||
)
|
||||
|
||||
# Call update_block_table on builder_b (simulates the metadata caching
|
||||
# optimization reusing metadata from builder_a's group).
|
||||
blk_table = torch.randint(
|
||||
0, 100, (num_reqs, max_blocks), dtype=torch.int32, device=device
|
||||
)
|
||||
slot_mapping = torch.zeros(num_reqs, dtype=torch.int64, device=device)
|
||||
|
||||
metadata_b = builder_b.update_block_table(metadata_a, blk_table, slot_mapping)
|
||||
|
||||
# block_idx tensors must live in builder_b's persistent buffers.
|
||||
def shares_storage(tensor, buffer):
|
||||
return (
|
||||
tensor.untyped_storage().data_ptr() == buffer.untyped_storage().data_ptr()
|
||||
)
|
||||
|
||||
assert shares_storage(
|
||||
metadata_b.block_idx_last_scheduled_token,
|
||||
builder_b.block_idx_last_scheduled_token,
|
||||
), "block_idx_last_scheduled_token not in builder_b's persistent buffer"
|
||||
|
||||
assert shares_storage(
|
||||
metadata_b.block_idx_last_computed_token,
|
||||
builder_b.block_idx_last_computed_token,
|
||||
), "block_idx_last_computed_token not in builder_b's persistent buffer"
|
||||
|
||||
# Must NOT point to builder_a's buffers.
|
||||
assert not shares_storage(
|
||||
metadata_b.block_idx_last_scheduled_token,
|
||||
builder_a.block_idx_last_scheduled_token,
|
||||
), "block_idx_last_scheduled_token still points to builder_a's buffer"
|
||||
|
||||
# Values must be correct (copied from metadata_a).
|
||||
torch.testing.assert_close(
|
||||
metadata_b.block_idx_last_scheduled_token,
|
||||
block_idx_vals,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
metadata_b.block_idx_last_computed_token,
|
||||
block_idx_vals,
|
||||
)
|
||||
1126
third_party/vllm/tests/v1/attention/test_mla_backends.py
vendored
Normal file
1126
third_party/vllm/tests/v1/attention/test_mla_backends.py
vendored
Normal file
File diff suppressed because it is too large
Load Diff
372
third_party/vllm/tests/v1/attention/test_rocm_attention_backends_selection.py
vendored
Normal file
372
third_party/vllm/tests/v1/attention/test_rocm_attention_backends_selection.py
vendored
Normal file
@@ -0,0 +1,372 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for attention backend selectors."""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
from vllm.v1.attention.selector import AttentionSelectorConfig
|
||||
|
||||
# ROCm-specific attention backend selection tests
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not current_platform.is_rocm(), reason="ROCm-specific tests"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_vllm_config():
|
||||
"""Create a mock VllmConfig for testing."""
|
||||
config = MagicMock()
|
||||
config.model_config.dtype = torch.float16
|
||||
config.model_config.hf_config.architectures = ["LlamaForCausalLM"]
|
||||
config.cache_config.block_size = 16
|
||||
return config
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_on_gfx9():
|
||||
"""Mock gfx9 arch detection to return True."""
|
||||
with patch("vllm.platforms.rocm.on_gfx9", return_value=True):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_on_mi3xx():
|
||||
"""Mock mi3xx arch detection to return True."""
|
||||
with patch("vllm.platforms.rocm.on_mi3xx", return_value=True):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"env_vars, selected_backend, expected_backend_path",
|
||||
[
|
||||
# Test Case: Explicit FLEX_ATTENTION backend
|
||||
(
|
||||
{},
|
||||
"FLEX_ATTENTION",
|
||||
AttentionBackendEnum.FLEX_ATTENTION.get_path(),
|
||||
),
|
||||
# Test Case 1: Default (no env vars, no explicit backend)
|
||||
(
|
||||
{},
|
||||
None,
|
||||
AttentionBackendEnum.TRITON_ATTN.get_path(),
|
||||
),
|
||||
# Test Case 2: Explicit TRITON_ATTN backend
|
||||
(
|
||||
{},
|
||||
"TRITON_ATTN",
|
||||
AttentionBackendEnum.TRITON_ATTN.get_path(),
|
||||
),
|
||||
# Test Case 3: Explicit ROCM_ATTN backend
|
||||
(
|
||||
{},
|
||||
"ROCM_ATTN",
|
||||
AttentionBackendEnum.ROCM_ATTN.get_path(),
|
||||
),
|
||||
# Test Case 4: Explicit ROCM_AITER_FA backend
|
||||
(
|
||||
{},
|
||||
"ROCM_AITER_FA",
|
||||
AttentionBackendEnum.ROCM_AITER_FA.get_path(),
|
||||
),
|
||||
# Test Case 5: Explicit ROCM_AITER_UNIFIED_ATTN backend
|
||||
(
|
||||
{},
|
||||
"ROCM_AITER_UNIFIED_ATTN",
|
||||
AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path(),
|
||||
),
|
||||
# Test Case 6: VLLM_ROCM_USE_AITER=1
|
||||
# (defaults to AITER FA when MHA not explicitly disabled)
|
||||
(
|
||||
{"VLLM_ROCM_USE_AITER": "1"},
|
||||
None,
|
||||
AttentionBackendEnum.ROCM_AITER_FA.get_path(),
|
||||
),
|
||||
# Test Case 7: VLLM_ROCM_USE_AITER=1 + VLLM_ROCM_USE_AITER_MHA=1
|
||||
(
|
||||
{"VLLM_ROCM_USE_AITER": "1", "VLLM_ROCM_USE_AITER_MHA": "1"},
|
||||
None,
|
||||
AttentionBackendEnum.ROCM_AITER_FA.get_path(),
|
||||
),
|
||||
# Test Case 8: VLLM_ROCM_USE_AITER=1 + VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1
|
||||
(
|
||||
{
|
||||
"VLLM_ROCM_USE_AITER": "1",
|
||||
"VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION": "1",
|
||||
},
|
||||
None,
|
||||
AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path(),
|
||||
),
|
||||
# Test Case 9: VLLM_ROCM_USE_AITER=1 + explicit TRITON_ATTN
|
||||
(
|
||||
{"VLLM_ROCM_USE_AITER": "1"},
|
||||
"TRITON_ATTN",
|
||||
AttentionBackendEnum.TRITON_ATTN.get_path(),
|
||||
),
|
||||
# Test Case 10: VLLM_ROCM_USE_AITER=1 + VLLM_ROCM_USE_AITER_MHA=0
|
||||
# (explicitly disabled)
|
||||
(
|
||||
{"VLLM_ROCM_USE_AITER": "1", "VLLM_ROCM_USE_AITER_MHA": "0"},
|
||||
None,
|
||||
AttentionBackendEnum.TRITON_ATTN.get_path(),
|
||||
),
|
||||
# Test Case 11: VLLM_ROCM_USE_AITER=1 + explicit ROCM_ATTN
|
||||
(
|
||||
{"VLLM_ROCM_USE_AITER": "1"},
|
||||
"ROCM_ATTN",
|
||||
AttentionBackendEnum.ROCM_ATTN.get_path(),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_standard_attention_backend_selection(
|
||||
env_vars,
|
||||
selected_backend,
|
||||
expected_backend_path,
|
||||
mock_vllm_config,
|
||||
mock_on_gfx9,
|
||||
mock_on_mi3xx,
|
||||
monkeypatch,
|
||||
):
|
||||
"""Test standard attention backend selection with various configurations."""
|
||||
# Set environment variables
|
||||
for key, value in env_vars.items():
|
||||
monkeypatch.setenv(key, value)
|
||||
|
||||
# Import after setting env vars to ensure they're picked up
|
||||
# Reload envs to pick up new environment variables
|
||||
import importlib
|
||||
|
||||
import vllm.envs as envs
|
||||
|
||||
importlib.reload(envs)
|
||||
|
||||
# Convert string backend to enum if provided
|
||||
backend_enum = None
|
||||
if selected_backend:
|
||||
backend_enum = getattr(AttentionBackendEnum, selected_backend)
|
||||
|
||||
# Get the backend class path
|
||||
from vllm.platforms.rocm import RocmPlatform
|
||||
|
||||
attn_selector_config = AttentionSelectorConfig(
|
||||
head_size=128,
|
||||
dtype=torch.float16,
|
||||
kv_cache_dtype="auto",
|
||||
block_size=16,
|
||||
use_mla=False,
|
||||
has_sink=False,
|
||||
use_sparse=False,
|
||||
)
|
||||
|
||||
backend_path = RocmPlatform.get_attn_backend_cls(
|
||||
selected_backend=backend_enum, attn_selector_config=attn_selector_config
|
||||
)
|
||||
|
||||
assert backend_path == expected_backend_path
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"env_vars, selected_backend, block_size, expected_backend_path, should_raise",
|
||||
[
|
||||
# Test Case 1: TRITON_MLA with block_size != 1
|
||||
(
|
||||
{},
|
||||
"TRITON_MLA",
|
||||
16,
|
||||
AttentionBackendEnum.TRITON_MLA.get_path(),
|
||||
False,
|
||||
),
|
||||
# Test Case 2: TRITON_MLA with block_size == 1 (should raise)
|
||||
(
|
||||
{},
|
||||
"TRITON_MLA",
|
||||
1,
|
||||
None,
|
||||
True,
|
||||
),
|
||||
# Test Case 3: ROCM_AITER_MLA with block_size == 1
|
||||
(
|
||||
{},
|
||||
"ROCM_AITER_MLA",
|
||||
1,
|
||||
AttentionBackendEnum.ROCM_AITER_MLA.get_path(),
|
||||
False,
|
||||
),
|
||||
# Test Case 4: ROCM_AITER_MLA with block_size != 1 (should raise)
|
||||
(
|
||||
{},
|
||||
"ROCM_AITER_MLA",
|
||||
16,
|
||||
AttentionBackendEnum.ROCM_AITER_MLA.get_path(),
|
||||
False,
|
||||
),
|
||||
# Test Case 5: VLLM_ROCM_USE_AITER=1 with block_size == 1
|
||||
(
|
||||
{"VLLM_ROCM_USE_AITER": "1"},
|
||||
None,
|
||||
1,
|
||||
AttentionBackendEnum.ROCM_AITER_MLA.get_path(),
|
||||
False,
|
||||
),
|
||||
# Test Case 6: VLLM_ROCM_USE_AITER=1 with block_size == 16
|
||||
# (should use ROCM_AITER_MLA now, as it supports block_size 16)
|
||||
(
|
||||
{"VLLM_ROCM_USE_AITER": "1"},
|
||||
None,
|
||||
16,
|
||||
AttentionBackendEnum.ROCM_AITER_MLA.get_path(),
|
||||
False,
|
||||
),
|
||||
# Test Case 7: VLLM_ROCM_USE_AITER=1 + explicit TRITON_MLA
|
||||
(
|
||||
{"VLLM_ROCM_USE_AITER": "1"},
|
||||
"TRITON_MLA",
|
||||
16,
|
||||
AttentionBackendEnum.TRITON_MLA.get_path(),
|
||||
False,
|
||||
),
|
||||
# Test Case 8: Explicit ROCM_AITER_TRITON_MLA
|
||||
(
|
||||
{},
|
||||
"ROCM_AITER_TRITON_MLA",
|
||||
16,
|
||||
AttentionBackendEnum.ROCM_AITER_TRITON_MLA.get_path(),
|
||||
False,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_mla_backend_selection(
|
||||
env_vars,
|
||||
selected_backend,
|
||||
block_size,
|
||||
expected_backend_path,
|
||||
should_raise,
|
||||
mock_vllm_config,
|
||||
monkeypatch,
|
||||
):
|
||||
"""Test MLA backend selection with various configurations."""
|
||||
# Set environment variables
|
||||
for key, value in env_vars.items():
|
||||
monkeypatch.setenv(key, value)
|
||||
|
||||
# Import after setting env vars
|
||||
# Reload envs
|
||||
import importlib
|
||||
|
||||
import vllm.envs as envs
|
||||
|
||||
importlib.reload(envs)
|
||||
|
||||
# Mock is_aiter_mla_enabled based on env vars and block_size
|
||||
aiter_enabled = env_vars.get("VLLM_ROCM_USE_AITER") == "1"
|
||||
|
||||
mock_rocm_ops = MagicMock()
|
||||
mock_rocm_ops.is_mla_enabled.return_value = aiter_enabled
|
||||
mock_aiter_module = MagicMock()
|
||||
mock_aiter_module.rocm_aiter_ops = mock_rocm_ops
|
||||
|
||||
with patch.dict("sys.modules", {"vllm._aiter_ops": mock_aiter_module}):
|
||||
# Convert string backend to enum if provided
|
||||
backend_enum = None
|
||||
if selected_backend:
|
||||
backend_enum = getattr(AttentionBackendEnum, selected_backend)
|
||||
|
||||
from vllm.platforms.rocm import RocmPlatform
|
||||
|
||||
if should_raise:
|
||||
with pytest.raises(ValueError):
|
||||
attn_selector_config = AttentionSelectorConfig(
|
||||
head_size=128,
|
||||
dtype=torch.float16,
|
||||
kv_cache_dtype="auto",
|
||||
block_size=block_size,
|
||||
use_mla=True,
|
||||
has_sink=False,
|
||||
use_sparse=False,
|
||||
)
|
||||
attn_selector_config = AttentionSelectorConfig(
|
||||
head_size=128,
|
||||
dtype=torch.float16,
|
||||
kv_cache_dtype="auto",
|
||||
block_size=block_size,
|
||||
use_mla=True,
|
||||
has_sink=False,
|
||||
use_sparse=False,
|
||||
)
|
||||
backend_path = RocmPlatform.get_attn_backend_cls(
|
||||
selected_backend=backend_enum,
|
||||
attn_selector_config=attn_selector_config,
|
||||
)
|
||||
|
||||
else:
|
||||
attn_selector_config = AttentionSelectorConfig(
|
||||
head_size=128,
|
||||
dtype=torch.float16,
|
||||
kv_cache_dtype="auto",
|
||||
block_size=block_size,
|
||||
use_mla=True,
|
||||
has_sink=False,
|
||||
use_sparse=False,
|
||||
)
|
||||
|
||||
backend_path = RocmPlatform.get_attn_backend_cls(
|
||||
selected_backend=backend_enum, attn_selector_config=attn_selector_config
|
||||
)
|
||||
|
||||
assert backend_path == expected_backend_path
|
||||
|
||||
|
||||
def test_aiter_fa_requires_mi3xx(mock_vllm_config):
|
||||
"""Test that ROCM_AITER_FA requires mi3xx architecture."""
|
||||
from vllm.platforms.rocm import RocmPlatform
|
||||
|
||||
# Mock on_mi3xx to return False (used by supports_compute_capability)
|
||||
with (
|
||||
patch("vllm.platforms.rocm.on_mi3xx", return_value=False),
|
||||
pytest.raises(
|
||||
ValueError,
|
||||
match="compute capability not supported",
|
||||
),
|
||||
):
|
||||
attn_selector_config = AttentionSelectorConfig(
|
||||
head_size=128,
|
||||
dtype=torch.float16,
|
||||
kv_cache_dtype="auto",
|
||||
block_size=16,
|
||||
use_mla=False,
|
||||
has_sink=False,
|
||||
use_sparse=False,
|
||||
)
|
||||
|
||||
RocmPlatform.get_attn_backend_cls(
|
||||
selected_backend=AttentionBackendEnum.ROCM_AITER_FA,
|
||||
attn_selector_config=attn_selector_config,
|
||||
)
|
||||
|
||||
|
||||
def test_sparse_not_supported(mock_vllm_config):
|
||||
"""Test that sparse MLA without use_mla flag raises an error."""
|
||||
from vllm.platforms.rocm import RocmPlatform
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="No valid attention backend found",
|
||||
):
|
||||
attn_selector_config = AttentionSelectorConfig(
|
||||
head_size=128,
|
||||
dtype=torch.float16,
|
||||
kv_cache_dtype="auto",
|
||||
block_size=16,
|
||||
use_mla=False,
|
||||
has_sink=False,
|
||||
use_sparse=True,
|
||||
)
|
||||
|
||||
RocmPlatform.get_attn_backend_cls(
|
||||
selected_backend=None, attn_selector_config=attn_selector_config
|
||||
)
|
||||
769
third_party/vllm/tests/v1/attention/test_sparse_mla_backends.py
vendored
Normal file
769
third_party/vllm/tests/v1/attention/test_sparse_mla_backends.py
vendored
Normal file
@@ -0,0 +1,769 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for the sparse MLA backends and utilities."""
|
||||
|
||||
import math
|
||||
from types import MethodType, SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.v1.attention.test_mla_backends import (
|
||||
BATCH_SPECS,
|
||||
BatchSpec,
|
||||
MockSparseMLAAttentionLayer,
|
||||
create_and_prepopulate_kv_cache,
|
||||
)
|
||||
from tests.v1.attention.utils import (
|
||||
create_common_attn_metadata,
|
||||
create_standard_kv_cache_spec,
|
||||
create_vllm_config,
|
||||
)
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.config import set_current_vllm_config
|
||||
from vllm.model_executor.layers.linear import ColumnParallelLinear
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
# TODO: Integrate ROCMAiterMLASparseBackend for ROCm.
|
||||
# The ROCm sparse MLA backend (rocm_aiter_mla_sparse.py) has a compatible
|
||||
# forward_mqa interface but needs validation on ROCm hardware.
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip(
|
||||
"Sparse MLA backend tests currently only support CUDA. "
|
||||
"ROCm support requires integrating ROCMAiterMLASparseBackend.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
from vllm.utils.math_utils import cdiv
|
||||
from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
|
||||
FlashInferMLASparseBackend,
|
||||
)
|
||||
from vllm.v1.attention.backends.mla.flashmla_sparse import (
|
||||
FlashMLASparseBackend,
|
||||
triton_convert_req_index_to_global_index,
|
||||
)
|
||||
from vllm.v1.attention.backends.utils import split_prefill_chunks
|
||||
from vllm.v1.attention.ops import flashmla
|
||||
|
||||
SPARSE_BACKEND_BATCH_SPECS = {
|
||||
name: BATCH_SPECS[name]
|
||||
for name in [
|
||||
"mixed_small",
|
||||
"mixed_medium",
|
||||
"small_prefill",
|
||||
"medium_prefill",
|
||||
"single_prefill",
|
||||
]
|
||||
}
|
||||
|
||||
SPARSE_BACKEND_BATCH_SPECS["large_q_prefill"] = BatchSpec(
|
||||
seq_lens=[1024] * 2, query_lens=[256] * 2
|
||||
)
|
||||
SPARSE_BACKEND_BATCH_SPECS["large_q_pure_prefill"] = BatchSpec(
|
||||
seq_lens=[256] * 2, query_lens=[256] * 2
|
||||
)
|
||||
|
||||
|
||||
def _float_to_e8m0_truncate(f: float) -> float:
|
||||
"""Simulate SM100's float -> e8m0 -> bf16 scale conversion.
|
||||
e8m0 format only stores the exponent (power of 2).
|
||||
cudaRoundZero truncates toward zero, meaning we round down to the
|
||||
nearest power of 2.
|
||||
"""
|
||||
if f <= 0:
|
||||
return 0.0
|
||||
# e8m0 = floor(log2(f)), then 2^(e8m0)
|
||||
# This is equivalent to truncating to the nearest power of 2 below f
|
||||
exp = math.floor(math.log2(f))
|
||||
return 2.0**exp
|
||||
|
||||
|
||||
def _dequantize_fp8_ds_mla_entry(
|
||||
cache_slice: torch.Tensor,
|
||||
kv_lora_rank: int,
|
||||
rope_dim: int,
|
||||
dtype: torch.dtype,
|
||||
simulate_sm100_e8m0_scales: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Dequantize a single fp8_ds_mla cache entry back to latent + rope.
|
||||
|
||||
Args:
|
||||
simulate_sm100_e8m0_scales: If True, simulate the SM100 kernel's
|
||||
float -> e8m0 -> bf16 scale conversion path.
|
||||
"""
|
||||
|
||||
# The first kv_lora_rank bytes store FP8 latent values with one scale per
|
||||
# 128 element tile written as float32 right after the latent payload.
|
||||
scales = cache_slice.view(torch.float32)[kv_lora_rank // 4 : kv_lora_rank // 4 + 4]
|
||||
latent = torch.empty(kv_lora_rank, dtype=torch.float16, device=cache_slice.device)
|
||||
for tile_idx in range(4):
|
||||
tile_start = tile_idx * 128
|
||||
tile_end = tile_start + 128
|
||||
scale_val = float(scales[tile_idx].item())
|
||||
if simulate_sm100_e8m0_scales:
|
||||
# Simulate the lossy float -> e8m0 -> bf16 conversion
|
||||
scale_val = _float_to_e8m0_truncate(scale_val)
|
||||
ops.convert_fp8(
|
||||
latent[tile_start:tile_end],
|
||||
cache_slice[tile_start:tile_end],
|
||||
scale_val,
|
||||
kv_dtype="fp8",
|
||||
)
|
||||
latent = latent.to(dtype)
|
||||
|
||||
rope_offset = kv_lora_rank // 2 + 8
|
||||
rope_vals = cache_slice.view(dtype)[rope_offset : rope_offset + rope_dim]
|
||||
return latent, rope_vals.clone()
|
||||
|
||||
|
||||
def _quantize_dequantize_fp8_ds_mla(
|
||||
kv_c: torch.Tensor,
|
||||
k_pe: torch.Tensor,
|
||||
block_size: int,
|
||||
scale: torch.Tensor,
|
||||
simulate_sm100_e8m0_scales: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Round-trip kv_c/k_pe though the fp8_ds_mla cache layout.
|
||||
|
||||
Args:
|
||||
simulate_sm100_e8m0_scales: If True, simulate the SM100 kernel's
|
||||
float -> e8m0 -> bf16 scale conversion in dequantization.
|
||||
"""
|
||||
|
||||
if kv_c.numel() == 0:
|
||||
return kv_c.clone(), k_pe.clone()
|
||||
|
||||
kv_lora_rank = kv_c.shape[-1]
|
||||
rope_dim = k_pe.shape[-1]
|
||||
num_tokens = kv_c.shape[0]
|
||||
num_blocks = max(1, math.ceil(num_tokens / block_size))
|
||||
entry_size = kv_lora_rank + 4 * 4 + 2 * rope_dim
|
||||
|
||||
tmp_cache = torch.zeros(
|
||||
num_blocks, block_size, entry_size, dtype=torch.uint8, device=kv_c.device
|
||||
)
|
||||
slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=kv_c.device)
|
||||
|
||||
ops.concat_and_cache_mla(
|
||||
kv_c, k_pe, tmp_cache, slot_mapping, kv_cache_dtype="fp8_ds_mla", scale=scale
|
||||
)
|
||||
|
||||
dequant_kv_c = torch.empty_like(kv_c)
|
||||
dequant_k_pe = torch.empty_like(k_pe)
|
||||
|
||||
for token_idx in range(num_tokens):
|
||||
slot = slot_mapping[token_idx].item()
|
||||
block_idx = slot // block_size
|
||||
block_offset = slot % block_size
|
||||
cache_slice = tmp_cache[block_idx, block_offset]
|
||||
latent, rope_vals = _dequantize_fp8_ds_mla_entry(
|
||||
cache_slice,
|
||||
kv_lora_rank,
|
||||
rope_dim,
|
||||
kv_c.dtype,
|
||||
simulate_sm100_e8m0_scales=simulate_sm100_e8m0_scales,
|
||||
)
|
||||
dequant_kv_c[token_idx] = latent
|
||||
dequant_k_pe[token_idx] = rope_vals
|
||||
|
||||
return dequant_kv_c, dequant_k_pe
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"backend_cls",
|
||||
[FlashMLASparseBackend, FlashInferMLASparseBackend],
|
||||
ids=["FlashMLA", "FlashInfer"],
|
||||
)
|
||||
@pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys()))
|
||||
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_ds_mla"])
|
||||
@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
|
||||
@pytest.mark.parametrize("block_size", [32, 64])
|
||||
def test_sparse_backend_decode_correctness(
|
||||
default_vllm_config,
|
||||
dist_init,
|
||||
backend_cls,
|
||||
batch_name,
|
||||
kv_cache_dtype,
|
||||
tensor_parallel_size,
|
||||
block_size,
|
||||
workspace_init,
|
||||
):
|
||||
if kv_cache_dtype not in backend_cls.supported_kv_cache_dtypes:
|
||||
pytest.skip(f"{backend_cls.get_name()} does not support {kv_cache_dtype}")
|
||||
|
||||
if (
|
||||
backend_cls == FlashMLASparseBackend
|
||||
and kv_cache_dtype.startswith("fp8")
|
||||
and kv_cache_dtype != "fp8_ds_mla"
|
||||
):
|
||||
pytest.skip(
|
||||
"FlashMLA Sparse Attention backend fp8 only supports "
|
||||
"fp8_ds_mla kv-cache dtype"
|
||||
)
|
||||
|
||||
supported_block_sizes = backend_cls.get_supported_kernel_block_sizes()
|
||||
if block_size not in supported_block_sizes:
|
||||
pytest.skip(
|
||||
f"{backend_cls.get_name()} does not support block_size={block_size}"
|
||||
)
|
||||
|
||||
if backend_cls == FlashMLASparseBackend:
|
||||
ok, reason = flashmla.is_flashmla_sparse_supported()
|
||||
if not ok:
|
||||
pytest.skip(reason)
|
||||
elif backend_cls == FlashInferMLASparseBackend:
|
||||
if not current_platform.has_device_capability(100):
|
||||
pytest.skip("FlashInferMLASparseBackend requires SM 10.0 or higher")
|
||||
|
||||
batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name]
|
||||
use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla"
|
||||
|
||||
device = torch.device("cuda")
|
||||
dtype = torch.bfloat16
|
||||
|
||||
# Model hyper-parameters (kept intentionally small for the unit test)
|
||||
total_num_heads = 128
|
||||
# Compute per-rank heads for simulated TP
|
||||
num_heads = max(1, total_num_heads // tensor_parallel_size)
|
||||
|
||||
kv_lora_rank = 512
|
||||
qk_nope_head_dim = 128
|
||||
qk_rope_head_dim = 64
|
||||
v_head_dim = 128
|
||||
head_size = kv_lora_rank + qk_rope_head_dim
|
||||
topk_tokens = 128
|
||||
|
||||
max_seqlen = max(batch_spec.seq_lens)
|
||||
total_cache_tokens = sum(batch_spec.seq_lens)
|
||||
|
||||
# Note: We use TP=1 to avoid multi-GPU requirements in CI.
|
||||
# The test simulates head partitioning via mocked methods below.
|
||||
vllm_config = create_vllm_config(
|
||||
model_name="deepseek-ai/DeepSeek-V2-Lite-Chat",
|
||||
tensor_parallel_size=1,
|
||||
max_model_len=max_seqlen,
|
||||
num_gpu_blocks=max(2048, cdiv(total_cache_tokens, block_size) + 1),
|
||||
block_size=block_size,
|
||||
hf_config_override={
|
||||
"index_topk": topk_tokens,
|
||||
"attn_module_list_cfg": [{"topk_tokens": topk_tokens}],
|
||||
},
|
||||
)
|
||||
model_config = vllm_config.model_config
|
||||
model_config.hf_text_config = SimpleNamespace(
|
||||
q_lora_rank=None,
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
qk_nope_head_dim=qk_nope_head_dim,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
v_head_dim=v_head_dim,
|
||||
model_type="deepseek_v2",
|
||||
)
|
||||
model_config.dtype = dtype
|
||||
model_config.get_num_attention_heads = MethodType(
|
||||
lambda self, parallel_config: num_heads,
|
||||
model_config,
|
||||
)
|
||||
model_config.get_num_kv_heads = MethodType(
|
||||
lambda self, parallel_config: 1, model_config
|
||||
)
|
||||
model_config.get_head_size = MethodType(lambda self: head_size, model_config)
|
||||
model_config.get_sliding_window = MethodType(lambda self: None, model_config)
|
||||
|
||||
kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
|
||||
|
||||
torch.manual_seed(0)
|
||||
|
||||
scale = 1.0 / math.sqrt(head_size)
|
||||
|
||||
# Shared MLA projection weights to keep reference and backend in sync
|
||||
W_UK = torch.rand(
|
||||
kv_lora_rank, num_heads, qk_nope_head_dim, dtype=dtype, device=device
|
||||
)
|
||||
W_UV = torch.rand(kv_lora_rank, num_heads, v_head_dim, dtype=dtype, device=device)
|
||||
|
||||
# Build synthetic decode-only workload
|
||||
seq_lens = batch_spec.seq_lens
|
||||
query_lens = batch_spec.query_lens
|
||||
|
||||
# Pre-compute positions and sparse indices for all tokens.
|
||||
# We need these BEFORE computing the reference to use sparse attention masks.
|
||||
total_query_tokens = sum(query_lens)
|
||||
positions = []
|
||||
for i in range(batch_spec.batch_size):
|
||||
s_len = seq_lens[i]
|
||||
q_len = query_lens[i]
|
||||
ctx_len = s_len - q_len
|
||||
for q_idx in range(q_len):
|
||||
positions.append(ctx_len + q_idx)
|
||||
|
||||
# Create sparse indices with UNIQUE per-token offsets to catch bugs where
|
||||
# the kernel uses wrong indices for some tokens (e.g., due to incorrect
|
||||
# tensor shapes like [1, num_tokens, ...] instead of [num_tokens, 1, ...]).
|
||||
# Also include -1 masked indices to verify the kernel handles them correctly.
|
||||
sparse_indices = torch.empty(
|
||||
total_query_tokens, topk_tokens, dtype=torch.int32, device=device
|
||||
)
|
||||
for tok_idx in range(total_query_tokens):
|
||||
max_valid_idx = positions[tok_idx]
|
||||
offset = tok_idx * 7 # Prime number for varied offsets
|
||||
# Use only half the topk indices as valid, mask the rest with -1
|
||||
# This tests that the kernel correctly ignores -1 indices
|
||||
num_valid = min(topk_tokens // 2, max_valid_idx + 1)
|
||||
if num_valid > 0:
|
||||
valid_range = torch.arange(num_valid, device=device, dtype=torch.int32)
|
||||
tok_indices = (valid_range + offset) % (max_valid_idx + 1)
|
||||
# Pad with -1 for the remaining positions
|
||||
tok_indices = torch.cat(
|
||||
[
|
||||
tok_indices,
|
||||
torch.full(
|
||||
(topk_tokens - num_valid,), -1, device=device, dtype=torch.int32
|
||||
),
|
||||
]
|
||||
)
|
||||
else:
|
||||
tok_indices = torch.full(
|
||||
(topk_tokens,), -1, device=device, dtype=torch.int32
|
||||
)
|
||||
tok_indices[0] = 0 # At least one valid index
|
||||
sparse_indices[tok_idx] = tok_indices
|
||||
|
||||
all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []
|
||||
kv_c_contexts, k_pe_contexts = [], []
|
||||
reference_outputs = []
|
||||
|
||||
kv_cache_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
|
||||
global_token_idx = 0
|
||||
|
||||
for i in range(batch_spec.batch_size):
|
||||
s_len = seq_lens[i]
|
||||
q_len = query_lens[i]
|
||||
ctx_len = s_len - q_len
|
||||
|
||||
q_c = torch.rand(
|
||||
q_len,
|
||||
num_heads,
|
||||
qk_nope_head_dim + qk_rope_head_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
kv_c_full = torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device)
|
||||
k_pe_full = torch.rand(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device)
|
||||
|
||||
if use_fp8_ds_mla_quantization:
|
||||
is_sm100 = torch.cuda.get_device_capability()[0] >= 10
|
||||
kv_c_full, k_pe_squeezed = _quantize_dequantize_fp8_ds_mla(
|
||||
kv_c_full,
|
||||
k_pe_full.squeeze(1),
|
||||
block_size=block_size,
|
||||
scale=kv_cache_scale,
|
||||
simulate_sm100_e8m0_scales=is_sm100,
|
||||
)
|
||||
k_pe_full = k_pe_squeezed.unsqueeze(1)
|
||||
|
||||
q_nope, q_pe = q_c.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)
|
||||
ql_nope = torch.einsum("qnh,lnh->qnl", q_nope, W_UK)
|
||||
q_mqa = torch.cat([ql_nope, q_pe], dim=-1)
|
||||
|
||||
k_mqa = torch.cat([kv_c_full, k_pe_full.squeeze(1)], dim=-1)
|
||||
v_mqa = kv_c_full
|
||||
|
||||
# Compute sparse SDPA reference per query token using its sparse indices
|
||||
for q_idx in range(q_len):
|
||||
tok_sparse_idx = sparse_indices[global_token_idx]
|
||||
valid_mask = tok_sparse_idx >= 0
|
||||
valid_indices = tok_sparse_idx[valid_mask].long()
|
||||
|
||||
q_tok = q_mqa[q_idx : q_idx + 1] # [1, num_heads, head_dim]
|
||||
k_sparse = k_mqa[valid_indices] # [num_valid, head_dim]
|
||||
v_sparse = v_mqa[valid_indices] # [num_valid, kv_lora_rank]
|
||||
|
||||
k_sparse = k_sparse.unsqueeze(1).expand(-1, num_heads, -1)
|
||||
v_sparse = v_sparse.unsqueeze(1).expand(-1, num_heads, -1)
|
||||
|
||||
# SDPA: [1, num_heads, 1, head_dim] x [1, num_heads, num_valid, head_dim]
|
||||
q_sdpa_in = q_tok.unsqueeze(0).transpose(1, 2)
|
||||
k_sdpa_in = k_sparse.unsqueeze(0).transpose(1, 2)
|
||||
v_sdpa_in = v_sparse.unsqueeze(0).transpose(1, 2)
|
||||
|
||||
sdpa_out = torch.nn.functional.scaled_dot_product_attention(
|
||||
q_sdpa_in, k_sdpa_in, v_sdpa_in, scale=scale
|
||||
)
|
||||
sdpa_out = sdpa_out.transpose(1, 2).squeeze(
|
||||
0
|
||||
) # [1, num_heads, kv_lora_rank]
|
||||
|
||||
sdpa_out = torch.einsum("qnl,lnv->qnv", sdpa_out, W_UV)
|
||||
reference_outputs.append(sdpa_out.flatten(start_dim=-2))
|
||||
|
||||
global_token_idx += 1
|
||||
|
||||
all_q_vllm.append(q_c)
|
||||
all_kv_c_vllm.append(kv_c_full[ctx_len:])
|
||||
all_k_pe_vllm.append(k_pe_full[ctx_len:])
|
||||
kv_c_contexts.append(kv_c_full[: ctx_len + 1])
|
||||
k_pe_contexts.append(k_pe_full[: ctx_len + 1])
|
||||
|
||||
query_vllm = torch.cat(all_q_vllm, dim=0)
|
||||
kv_c_vllm = torch.cat(all_kv_c_vllm, dim=0)
|
||||
k_pe_vllm = torch.cat(all_k_pe_vllm, dim=0)
|
||||
sdpa_reference = torch.cat(reference_outputs, dim=0)
|
||||
|
||||
vllm_config.cache_config.cache_dtype = kv_cache_dtype
|
||||
vllm_config.model_config.hf_config.index_topk = topk_tokens
|
||||
|
||||
common_attn_metadata = create_common_attn_metadata(
|
||||
batch_spec,
|
||||
vllm_config.cache_config.block_size,
|
||||
device,
|
||||
arange_block_indices=True,
|
||||
)
|
||||
|
||||
kv_cache = create_and_prepopulate_kv_cache(
|
||||
kv_c_contexts=kv_c_contexts,
|
||||
k_pe_contexts=k_pe_contexts,
|
||||
block_size=vllm_config.cache_config.block_size,
|
||||
head_size=head_size,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
num_blocks=vllm_config.cache_config.num_gpu_blocks,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
randomize_blocks=False,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
scale=kv_cache_scale,
|
||||
)
|
||||
|
||||
builder_cls = backend_cls.get_builder_cls()
|
||||
builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device)
|
||||
metadata = builder.build(
|
||||
common_prefix_len=0, common_attn_metadata=common_attn_metadata
|
||||
)
|
||||
|
||||
# Use the pre-computed sparse_indices for the mock indexer
|
||||
mock_indexer = SimpleNamespace(topk_indices_buffer=sparse_indices)
|
||||
|
||||
kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1)
|
||||
kv_b_proj_weight = kv_b_proj_weight.view(
|
||||
kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim)
|
||||
)
|
||||
|
||||
mock_kv_b_proj = ColumnParallelLinear(
|
||||
input_size=kv_lora_rank,
|
||||
output_size=num_heads * (qk_nope_head_dim + v_head_dim),
|
||||
bias=False,
|
||||
).to(device=device, dtype=dtype)
|
||||
mock_kv_b_proj.weight = torch.nn.Parameter(kv_b_proj_weight.T.contiguous())
|
||||
|
||||
impl_cls = backend_cls.get_impl_cls()
|
||||
with set_current_vllm_config(vllm_config):
|
||||
impl = impl_cls(
|
||||
num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
scale=scale,
|
||||
num_kv_heads=1,
|
||||
alibi_slopes=None,
|
||||
sliding_window=None,
|
||||
kv_cache_dtype=vllm_config.cache_config.cache_dtype,
|
||||
logits_soft_cap=None,
|
||||
attn_type="decoder",
|
||||
kv_sharing_target_layer_name=None,
|
||||
q_lora_rank=None,
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
qk_nope_head_dim=qk_nope_head_dim,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,
|
||||
v_head_dim=v_head_dim,
|
||||
kv_b_proj=mock_kv_b_proj,
|
||||
indexer=mock_indexer,
|
||||
)
|
||||
|
||||
impl.process_weights_after_loading(dtype)
|
||||
|
||||
# Create mock sparse MLA layer with weight matrices
|
||||
mock_layer = MockSparseMLAAttentionLayer(
|
||||
impl=impl,
|
||||
num_heads=num_heads,
|
||||
qk_nope_head_dim=qk_nope_head_dim,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
v_head_dim=v_head_dim,
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
device=device,
|
||||
W_UK=W_UK,
|
||||
W_UV=W_UV,
|
||||
)
|
||||
|
||||
out_buffer = torch.empty(
|
||||
metadata.num_actual_tokens, num_heads * v_head_dim, dtype=dtype, device=device
|
||||
)
|
||||
|
||||
with torch.inference_mode():
|
||||
backend_output = mock_layer.forward_impl(
|
||||
query_vllm,
|
||||
kv_c_vllm,
|
||||
k_pe_vllm,
|
||||
kv_cache,
|
||||
metadata,
|
||||
out_buffer,
|
||||
)
|
||||
|
||||
assert backend_output.shape == sdpa_reference.shape
|
||||
assert backend_output.dtype == sdpa_reference.dtype
|
||||
assert torch.isfinite(backend_output).all()
|
||||
|
||||
# FP8 quantization introduces some error, but should be within reasonable bounds
|
||||
# BF16 (auto) should be very accurate, FP8 allows slightly more tolerance
|
||||
if kv_cache_dtype.startswith("fp8"):
|
||||
torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.05, atol=0.05)
|
||||
else:
|
||||
torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.01, atol=0.01)
|
||||
|
||||
|
||||
def _triton_convert_reference_impl(
|
||||
req_ids: torch.Tensor,
|
||||
block_table: torch.Tensor,
|
||||
token_indices: torch.Tensor,
|
||||
block_size: int,
|
||||
num_topk_tokens: int,
|
||||
HAS_PREFILL_WORKSPACE: bool = False,
|
||||
prefill_workspace_request_ids: torch.Tensor | None = None,
|
||||
prefill_workspace_starts: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Reference implementation for triton_convert_req_index_to_global_index."""
|
||||
num_tokens = req_ids.shape[0]
|
||||
max_blocks_per_req = block_table.shape[1]
|
||||
result = torch.empty(
|
||||
num_tokens, num_topk_tokens, dtype=torch.int32, device=req_ids.device
|
||||
)
|
||||
|
||||
for token_id in range(num_tokens):
|
||||
req_id = req_ids[token_id].item()
|
||||
|
||||
# Determine if this token uses workspace or paged cache
|
||||
use_prefill_workspace = False
|
||||
workspace_start = 0
|
||||
if HAS_PREFILL_WORKSPACE and prefill_workspace_request_ids is not None:
|
||||
assert prefill_workspace_starts is not None
|
||||
prefill_req_id = prefill_workspace_request_ids[token_id].item()
|
||||
if prefill_req_id >= 0:
|
||||
use_prefill_workspace = True
|
||||
workspace_start = prefill_workspace_starts[prefill_req_id].item()
|
||||
|
||||
for idx_id in range(num_topk_tokens):
|
||||
token_idx = token_indices[token_id, idx_id].item()
|
||||
|
||||
if token_idx == -1:
|
||||
result[token_id, idx_id] = -1
|
||||
elif use_prefill_workspace:
|
||||
# Prefill + using prefill workspace: map to workspace offset
|
||||
result[token_id, idx_id] = workspace_start + token_idx
|
||||
else:
|
||||
# Decode: map to paged cache
|
||||
block_id = token_idx // block_size
|
||||
if block_id >= max_blocks_per_req:
|
||||
result[token_id, idx_id] = -1
|
||||
else:
|
||||
block_num = block_table[req_id, block_id].item()
|
||||
offset = token_idx % block_size
|
||||
result[token_id, idx_id] = block_num * block_size + offset
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@pytest.mark.parametrize("block_size", [16, 64, 128])
|
||||
@pytest.mark.parametrize("num_topk_tokens", [128, 256, 512])
|
||||
@pytest.mark.skipif(
|
||||
torch.cuda.get_device_capability() < (9, 0),
|
||||
reason="FlashMLASparseBackend requires CUDA 9.0 or higher",
|
||||
)
|
||||
def test_triton_convert_req_index_to_global_index_decode_only(
|
||||
block_size, num_topk_tokens
|
||||
):
|
||||
device = torch.device("cuda")
|
||||
num_tokens = 8
|
||||
num_requests = 4
|
||||
max_blocks_per_req = 10
|
||||
|
||||
req_id = torch.randint(
|
||||
0, num_requests, (num_tokens,), dtype=torch.int32, device=device
|
||||
)
|
||||
block_table = torch.randint(
|
||||
0, 100, (num_requests, max_blocks_per_req), dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
token_indices = torch.randint(
|
||||
0,
|
||||
block_size * max_blocks_per_req,
|
||||
(num_tokens, num_topk_tokens),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Set some to -1 to test masking
|
||||
token_indices[0, :10] = -1
|
||||
token_indices[3, 50:60] = -1
|
||||
|
||||
# Set some to out of bounds
|
||||
token_indices[2, 100:110] = max_blocks_per_req * block_size
|
||||
token_indices[6, 150:160] = max_blocks_per_req * block_size
|
||||
|
||||
result = triton_convert_req_index_to_global_index(
|
||||
req_id,
|
||||
block_table,
|
||||
token_indices,
|
||||
BLOCK_SIZE=block_size,
|
||||
NUM_TOPK_TOKENS=num_topk_tokens,
|
||||
)
|
||||
|
||||
reference_result = _triton_convert_reference_impl(
|
||||
req_id,
|
||||
block_table,
|
||||
token_indices,
|
||||
block_size,
|
||||
num_topk_tokens,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(result, reference_result, rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("block_size", [16])
|
||||
@pytest.mark.skipif(
|
||||
torch.cuda.get_device_capability() < (9, 0),
|
||||
reason="FlashMLASparseBackend requires CUDA 9.0 or higher",
|
||||
)
|
||||
def test_triton_convert_req_index_to_global_index_with_prefill_workspace(block_size):
|
||||
device = torch.device("cuda")
|
||||
num_requests = 4
|
||||
max_blocks_per_req = 8
|
||||
num_topk_tokens = 128
|
||||
|
||||
# First 6 tokens are decode (reqs 0, 1), last 6 are prefill (reqs 2, 3)
|
||||
req_id = torch.tensor(
|
||||
[0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], dtype=torch.int32, device=device
|
||||
)
|
||||
prefill_workspace_request_ids = torch.tensor(
|
||||
[-1, -1, -1, -1, -1, -1, 0, 0, 0, 1, 1, 1], dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
# Workspace starts for the 2 prefill reqs: req 2 starts at 0, req 3 starts at 100
|
||||
prefill_workspace_starts = torch.tensor([0, 100], dtype=torch.int32, device=device)
|
||||
|
||||
block_table = torch.randint(
|
||||
0, 50, (num_requests, max_blocks_per_req), dtype=torch.int32, device=device
|
||||
)
|
||||
token_indices = torch.randint(
|
||||
0,
|
||||
block_size * max_blocks_per_req,
|
||||
(req_id.shape[0], num_topk_tokens),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Set some to -1 to test masking
|
||||
token_indices[0, :10] = -1
|
||||
token_indices[3, 50:60] = -1
|
||||
|
||||
# Set some to out of bounds
|
||||
token_indices[2, 100:110] = max_blocks_per_req * block_size
|
||||
token_indices[6, 150:160] = max_blocks_per_req * block_size
|
||||
|
||||
result = triton_convert_req_index_to_global_index(
|
||||
req_id,
|
||||
block_table,
|
||||
token_indices,
|
||||
BLOCK_SIZE=block_size,
|
||||
NUM_TOPK_TOKENS=num_topk_tokens,
|
||||
HAS_PREFILL_WORKSPACE=True,
|
||||
prefill_workspace_request_ids=prefill_workspace_request_ids,
|
||||
prefill_workspace_starts=prefill_workspace_starts,
|
||||
)
|
||||
|
||||
reference_result = _triton_convert_reference_impl(
|
||||
req_id,
|
||||
block_table,
|
||||
token_indices,
|
||||
block_size,
|
||||
num_topk_tokens,
|
||||
HAS_PREFILL_WORKSPACE=True,
|
||||
prefill_workspace_request_ids=prefill_workspace_request_ids,
|
||||
prefill_workspace_starts=prefill_workspace_starts,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(result, reference_result, rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"seq_lens,max_buf,expected",
|
||||
[
|
||||
# Basic split: totals per chunk ≤ max_buf
|
||||
(torch.tensor([2, 3, 4, 2]), 5, [(0, 2), (2, 3), (3, 4)]),
|
||||
# Exact fits should split between items when adding the next would overflow
|
||||
(torch.tensor([5, 5, 5]), 5, [(0, 1), (1, 2), (2, 3)]),
|
||||
# All requests fit in a single chunk
|
||||
(torch.tensor([1, 1, 1]), 10, [(0, 3)]),
|
||||
# Large buffer
|
||||
(torch.tensor([4, 4, 4]), 100, [(0, 3)]),
|
||||
],
|
||||
)
|
||||
def test_split_prefill_chunks(seq_lens, max_buf, expected):
|
||||
out = split_prefill_chunks(seq_lens, max_buf)
|
||||
assert out == expected
|
||||
|
||||
|
||||
def test_triton_convert_returns_valid_counts():
|
||||
"""Test that return_valid_counts correctly counts non-negative indices."""
|
||||
device = torch.device("cuda")
|
||||
num_tokens = 8
|
||||
num_requests = 2
|
||||
max_blocks_per_req = 10
|
||||
block_size = 64
|
||||
num_topk_tokens = 128
|
||||
|
||||
req_id = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1], dtype=torch.int32, device=device)
|
||||
block_table = torch.arange(
|
||||
num_requests * max_blocks_per_req, dtype=torch.int32, device=device
|
||||
).view(num_requests, max_blocks_per_req)
|
||||
|
||||
# Create token indices with varying numbers of valid entries
|
||||
# Token 0: 64 valid, 64 invalid (-1)
|
||||
# Token 1: 32 valid, 96 invalid
|
||||
# Token 2: 128 valid (all)
|
||||
# Token 3: 1 valid, 127 invalid
|
||||
# etc.
|
||||
token_indices = torch.full(
|
||||
(num_tokens, num_topk_tokens), -1, dtype=torch.int32, device=device
|
||||
)
|
||||
expected_valid = []
|
||||
for i in range(num_tokens):
|
||||
num_valid = [64, 32, 128, 1, 64, 32, 128, 1][i]
|
||||
token_indices[i, :num_valid] = torch.arange(
|
||||
num_valid, dtype=torch.int32, device=device
|
||||
) % (block_size * max_blocks_per_req)
|
||||
expected_valid.append(num_valid)
|
||||
|
||||
expected_valid_tensor = torch.tensor(
|
||||
expected_valid, dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
# Test with return_valid_counts=True
|
||||
result, valid_counts = triton_convert_req_index_to_global_index(
|
||||
req_id,
|
||||
block_table,
|
||||
token_indices,
|
||||
BLOCK_SIZE=block_size,
|
||||
NUM_TOPK_TOKENS=num_topk_tokens,
|
||||
return_valid_counts=True,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(valid_counts, expected_valid_tensor, rtol=0, atol=0)
|
||||
|
||||
# Test that return_valid_counts=False returns only the indices
|
||||
result_only = triton_convert_req_index_to_global_index(
|
||||
req_id,
|
||||
block_table,
|
||||
token_indices,
|
||||
BLOCK_SIZE=block_size,
|
||||
NUM_TOPK_TOKENS=num_topk_tokens,
|
||||
return_valid_counts=False,
|
||||
)
|
||||
assert isinstance(result_only, torch.Tensor)
|
||||
torch.testing.assert_close(result_only, result, rtol=0, atol=0)
|
||||
360
third_party/vllm/tests/v1/attention/test_trtllm_attention_integration.py
vendored
Normal file
360
third_party/vllm/tests/v1/attention/test_trtllm_attention_integration.py
vendored
Normal file
@@ -0,0 +1,360 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Integration tests for TRTLLM gen-full attention through FlashInfer."""
|
||||
|
||||
import unittest.mock
|
||||
from functools import partial
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch.nn.attention.flex_attention import create_block_mask, flex_attention
|
||||
|
||||
from tests.v1.attention.utils import (
|
||||
BatchSpec,
|
||||
create_common_attn_metadata,
|
||||
create_vllm_config,
|
||||
)
|
||||
from vllm.config import set_current_vllm_config
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.math_utils import cdiv
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
from vllm.v1.attention.backends.utils import (
|
||||
PerLayerParameters,
|
||||
get_kv_cache_layout,
|
||||
set_kv_cache_layout,
|
||||
)
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec
|
||||
|
||||
if not current_platform.is_device_capability_family(100):
|
||||
pytest.skip(
|
||||
"TRTLLM integration tests require NVIDIA Blackwell (SM100).",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
from vllm.v1.attention.backends.flashinfer import ( # noqa: E402
|
||||
FlashInferImpl,
|
||||
FlashInferMetadataBuilder,
|
||||
TRTLLMDecode,
|
||||
TRTLLMPrefill,
|
||||
)
|
||||
|
||||
|
||||
class MockAttentionLayer:
|
||||
"""Minimal mock of an attention layer for testing."""
|
||||
|
||||
def __init__(self, device: torch.device):
|
||||
self._q_scale = torch.tensor(1.0, device=device)
|
||||
self._k_scale = torch.tensor(1.0, device=device)
|
||||
self._v_scale = torch.tensor(1.0, device=device)
|
||||
self._q_scale_float = 1.0
|
||||
self._k_scale_float = 1.0
|
||||
self._v_scale_float = 1.0
|
||||
self._o_scale_float = None
|
||||
|
||||
|
||||
MODEL = "Qwen/Qwen2.5-0.5B"
|
||||
BLOCK_SIZE = 16
|
||||
NUM_GPU_BLOCKS = 8192
|
||||
|
||||
BATCH_SPECS = {
|
||||
"decode_only": BatchSpec(
|
||||
seq_lens=[128, 256, 512],
|
||||
query_lens=[1, 1, 1],
|
||||
),
|
||||
"prefill_only": BatchSpec(
|
||||
seq_lens=[64, 128, 256],
|
||||
query_lens=[16, 32, 16],
|
||||
),
|
||||
"mixed": BatchSpec(
|
||||
seq_lens=[128, 256, 512, 128],
|
||||
query_lens=[1, 1, 8, 16],
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
|
||||
head_size = vllm_config.model_config.get_head_size()
|
||||
return {
|
||||
name: PerLayerParameters(
|
||||
window_left=-1,
|
||||
logits_soft_cap=0.0,
|
||||
sm_scale=1.0 / (head_size**0.5),
|
||||
)
|
||||
for name in layer_names
|
||||
}
|
||||
|
||||
|
||||
def _create_hnd_kv_cache(
|
||||
k_contexts,
|
||||
v_contexts,
|
||||
block_size,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
dtype,
|
||||
device,
|
||||
num_blocks,
|
||||
common_attn_metadata,
|
||||
):
|
||||
"""Create and populate a KV cache with HND-compatible strides.
|
||||
|
||||
The returned tensor has logical shape
|
||||
(num_blocks, 2, block_size, num_kv_heads, head_size) but is physically
|
||||
laid out as (num_blocks, 2, num_kv_heads, block_size, head_size) so that
|
||||
``kv_cache.permute(0, 1, 3, 2, 4)`` yields a contiguous HND view.
|
||||
"""
|
||||
seq_lens = common_attn_metadata.seq_lens.cpu()
|
||||
query_lens = (
|
||||
common_attn_metadata.query_start_loc_cpu[1:]
|
||||
- common_attn_metadata.query_start_loc_cpu[:-1]
|
||||
)
|
||||
block_table = common_attn_metadata.block_table_tensor
|
||||
slot_mapping = common_attn_metadata.slot_mapping
|
||||
batch_size = len(k_contexts)
|
||||
|
||||
# Build cache in (2, num_blocks, block_size, num_kv_heads, head_size)
|
||||
# then convert to HND format (same approach as test_attention_backends.py).
|
||||
kv_cache_raw = torch.zeros(
|
||||
2,
|
||||
num_blocks,
|
||||
block_size,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
kv_cache_flat = kv_cache_raw.view(2, -1, num_kv_heads, head_size)
|
||||
|
||||
start_block_idx = 1
|
||||
for i in range(batch_size):
|
||||
k_ctx, v_ctx = k_contexts[i], v_contexts[i]
|
||||
start = start_block_idx * block_size
|
||||
end = start + k_ctx.shape[0]
|
||||
kv_cache_flat[0, start:end] = k_ctx
|
||||
kv_cache_flat[1, start:end] = v_ctx
|
||||
start_block_idx += cdiv(int(seq_lens[i]), block_size)
|
||||
|
||||
blocks_end = start_block_idx
|
||||
|
||||
# Randomly permute blocks (starting from block 1; block 0 is null).
|
||||
perm = torch.randperm(blocks_end - 1) + 1
|
||||
inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
|
||||
inv_perm[1:] = torch.argsort(perm) + 1
|
||||
kv_cache_raw[:, 1:blocks_end] = kv_cache_raw[:, perm]
|
||||
|
||||
# Build block table.
|
||||
start_block_idx = 1
|
||||
for i in range(batch_size):
|
||||
n_blocks = cdiv(int(seq_lens[i]), block_size)
|
||||
block_table[i, :n_blocks] = inv_perm[
|
||||
start_block_idx : start_block_idx + n_blocks
|
||||
]
|
||||
start_block_idx += n_blocks
|
||||
|
||||
# Build slot mapping that is consistent with the block table.
|
||||
for i in range(batch_size):
|
||||
ctx_len = int(seq_lens[i]) - int(query_lens[i])
|
||||
token_offsets = torch.arange(int(query_lens[i])) + ctx_len
|
||||
block_indices = token_offsets // block_size
|
||||
intra_block_offsets = token_offsets % block_size
|
||||
start = common_attn_metadata.query_start_loc_cpu[i]
|
||||
end = common_attn_metadata.query_start_loc_cpu[i + 1]
|
||||
slot_mapping[start:end] = block_table[
|
||||
i, block_indices
|
||||
] * block_size + intra_block_offsets.to(device)
|
||||
|
||||
# Transpose to FlashInfer logical shape then make HND-strided.
|
||||
kv_cache = kv_cache_raw.transpose(0, 1)
|
||||
kv_cache = kv_cache.transpose(2, 3).contiguous().transpose(2, 3)
|
||||
return kv_cache
|
||||
|
||||
|
||||
def _run_trtllm_integration(batch_spec):
|
||||
"""Run TRTLLM attention through the full FlashInfer pipeline
|
||||
and compare against an SDPA reference."""
|
||||
set_random_seed(42)
|
||||
device = torch.device("cuda:0")
|
||||
|
||||
vllm_config = create_vllm_config(
|
||||
model_name=MODEL,
|
||||
max_model_len=max(batch_spec.seq_lens),
|
||||
block_size=BLOCK_SIZE,
|
||||
num_gpu_blocks=NUM_GPU_BLOCKS,
|
||||
)
|
||||
vllm_config.attention_config.use_trtllm_attention = True
|
||||
|
||||
num_q_heads = vllm_config.model_config.get_num_attention_heads(
|
||||
vllm_config.parallel_config
|
||||
)
|
||||
num_kv_heads = vllm_config.model_config.get_num_kv_heads(
|
||||
vllm_config.parallel_config
|
||||
)
|
||||
head_size = vllm_config.model_config.get_head_size()
|
||||
dtype = vllm_config.model_config.dtype
|
||||
scale = 1.0 / (head_size**0.5)
|
||||
|
||||
# 1. Generate data and compute SDPA reference
|
||||
all_q, all_k, all_v = [], [], []
|
||||
all_sdpa_out = []
|
||||
k_contexts, v_contexts = [], []
|
||||
|
||||
for i in range(batch_spec.batch_size):
|
||||
s_len = batch_spec.seq_lens[i]
|
||||
q_len = batch_spec.query_lens[i]
|
||||
ctx_len = s_len - q_len
|
||||
|
||||
q = torch.randn(q_len, num_q_heads, head_size, dtype=dtype, device=device)
|
||||
k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
|
||||
v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
|
||||
|
||||
# SDPA reference (N=1, H, L, D)
|
||||
q_sdpa = q.unsqueeze(0).transpose(1, 2)
|
||||
k_sdpa = k_full.unsqueeze(0).transpose(1, 2)
|
||||
v_sdpa = v_full.unsqueeze(0).transpose(1, 2)
|
||||
|
||||
if num_q_heads != num_kv_heads:
|
||||
repeats = num_q_heads // num_kv_heads
|
||||
k_sdpa = k_sdpa.repeat_interleave(repeats, dim=1)
|
||||
v_sdpa = v_sdpa.repeat_interleave(repeats, dim=1)
|
||||
|
||||
def causal_mask_mod(b, h, q_idx, kv_idx, *, context_len):
|
||||
return (q_idx + context_len) >= kv_idx
|
||||
|
||||
mask_fn = partial(causal_mask_mod, context_len=ctx_len)
|
||||
block_mask = create_block_mask(
|
||||
mask_fn, B=None, H=None, Q_LEN=q_len, KV_LEN=s_len, device=device
|
||||
)
|
||||
sdpa_out = flex_attention(
|
||||
q_sdpa,
|
||||
k_sdpa,
|
||||
v_sdpa,
|
||||
block_mask=block_mask,
|
||||
scale=scale,
|
||||
enable_gqa=True,
|
||||
)
|
||||
all_sdpa_out.append(sdpa_out.transpose(1, 2).squeeze(0))
|
||||
|
||||
all_q.append(q)
|
||||
all_k.append(k_full[ctx_len:])
|
||||
all_v.append(v_full[ctx_len:])
|
||||
k_contexts.append(k_full[:ctx_len])
|
||||
v_contexts.append(v_full[:ctx_len])
|
||||
|
||||
query_vllm = torch.cat(all_q, dim=0)
|
||||
key_vllm = torch.cat(all_k, dim=0)
|
||||
value_vllm = torch.cat(all_v, dim=0)
|
||||
sdpa_output = torch.cat(all_sdpa_out, dim=0)
|
||||
|
||||
common_attn_metadata = create_common_attn_metadata(batch_spec, BLOCK_SIZE, device)
|
||||
|
||||
# 2. Create HND KV cache
|
||||
kv_cache = _create_hnd_kv_cache(
|
||||
k_contexts,
|
||||
v_contexts,
|
||||
BLOCK_SIZE,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
dtype,
|
||||
device,
|
||||
NUM_GPU_BLOCKS,
|
||||
common_attn_metadata,
|
||||
)
|
||||
|
||||
# 3. Run through FlashInfer with TRTLLM enabled
|
||||
set_kv_cache_layout("HND")
|
||||
get_kv_cache_layout.cache_clear()
|
||||
|
||||
try:
|
||||
kv_cache_spec = FullAttentionSpec(
|
||||
block_size=BLOCK_SIZE,
|
||||
num_kv_heads=num_kv_heads,
|
||||
head_size=head_size,
|
||||
dtype=dtype,
|
||||
)
|
||||
layer_names = ["test_layer_0"]
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
unittest.mock.patch(
|
||||
"vllm.utils.flashinfer.supports_trtllm_attention",
|
||||
return_value=True,
|
||||
),
|
||||
unittest.mock.patch(
|
||||
"vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
|
||||
_mock_get_per_layer_parameters,
|
||||
),
|
||||
):
|
||||
builder = FlashInferMetadataBuilder(
|
||||
kv_cache_spec, layer_names, vllm_config, device
|
||||
)
|
||||
attn_metadata = builder.build(
|
||||
common_prefix_len=0,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
)
|
||||
|
||||
# Verify the correct TRTLLM metadata types were produced.
|
||||
has_prefills = any(ql > 1 for ql in batch_spec.query_lens)
|
||||
has_decodes = any(ql == 1 for ql in batch_spec.query_lens)
|
||||
|
||||
if has_prefills:
|
||||
assert isinstance(attn_metadata.prefill, TRTLLMPrefill), (
|
||||
f"Expected TRTLLMPrefill, got {type(attn_metadata.prefill)}"
|
||||
)
|
||||
if has_decodes:
|
||||
assert isinstance(attn_metadata.decode, TRTLLMDecode), (
|
||||
f"Expected TRTLLMDecode, got {type(attn_metadata.decode)}"
|
||||
)
|
||||
|
||||
impl = FlashInferImpl(
|
||||
num_heads=num_q_heads,
|
||||
head_size=head_size,
|
||||
scale=scale,
|
||||
num_kv_heads=num_kv_heads,
|
||||
alibi_slopes=None,
|
||||
sliding_window=None,
|
||||
kv_cache_dtype="auto",
|
||||
)
|
||||
|
||||
mock_layer = MockAttentionLayer(device)
|
||||
output = torch.empty_like(query_vllm)
|
||||
|
||||
impl.do_kv_cache_update(
|
||||
mock_layer,
|
||||
key_vllm,
|
||||
value_vllm,
|
||||
kv_cache,
|
||||
attn_metadata.slot_mapping,
|
||||
)
|
||||
|
||||
output = impl.forward(
|
||||
mock_layer,
|
||||
query_vllm,
|
||||
key_vllm,
|
||||
value_vllm,
|
||||
kv_cache,
|
||||
attn_metadata,
|
||||
output=output,
|
||||
)
|
||||
|
||||
# 4. Compare against SDPA reference
|
||||
torch.testing.assert_close(
|
||||
output,
|
||||
sdpa_output,
|
||||
atol=1e-2,
|
||||
rtol=1e-2,
|
||||
)
|
||||
|
||||
finally:
|
||||
set_kv_cache_layout(None)
|
||||
get_kv_cache_layout.cache_clear()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_spec_name",
|
||||
list(BATCH_SPECS.keys()),
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_trtllm_gen_full_attention_integration(batch_spec_name: str):
|
||||
"""Test TRTLLM gen-full attention through the full FlashInfer
|
||||
MetadataBuilder.build() -> FlashInferImpl.forward() pipeline,
|
||||
with real TRTLLM kernels on Blackwell."""
|
||||
_run_trtllm_integration(BATCH_SPECS[batch_spec_name])
|
||||
360
third_party/vllm/tests/v1/attention/utils.py
vendored
Normal file
360
third_party/vllm/tests/v1/attention/utils.py
vendored
Normal file
@@ -0,0 +1,360 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Utility functions for attention-related v1 tests."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.config import (
|
||||
CacheConfig,
|
||||
CompilationConfig,
|
||||
DeviceConfig,
|
||||
LoadConfig,
|
||||
ModelConfig,
|
||||
ParallelConfig,
|
||||
SchedulerConfig,
|
||||
VllmConfig,
|
||||
)
|
||||
from vllm.config.model import ModelDType
|
||||
from vllm.v1.attention.backend import (
|
||||
AttentionImpl,
|
||||
AttentionMetadataBuilder,
|
||||
CommonAttentionMetadata,
|
||||
)
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchSpec:
|
||||
"""Specification for a batch configuration (workload shape only)."""
|
||||
|
||||
seq_lens: list[int]
|
||||
query_lens: list[int]
|
||||
|
||||
name: str = "unnamed"
|
||||
|
||||
@property
|
||||
def batch_size(self):
|
||||
return len(self.seq_lens)
|
||||
|
||||
def __post_init__(self):
|
||||
assert len(self.seq_lens) == len(self.query_lens)
|
||||
|
||||
def compute_num_tokens(self):
|
||||
return sum(self.query_lens)
|
||||
|
||||
|
||||
def create_common_attn_metadata(
|
||||
batch_spec: BatchSpec,
|
||||
block_size: int,
|
||||
device: torch.device,
|
||||
max_block_idx: int = 1000,
|
||||
arange_block_indices: bool = False,
|
||||
) -> CommonAttentionMetadata:
|
||||
"""Create CommonAttentionMetadata from a BatchSpec and ModelParams."""
|
||||
# Create query start locations
|
||||
query_start_loc = torch.zeros(
|
||||
batch_spec.batch_size + 1, dtype=torch.int32, device=device
|
||||
)
|
||||
query_start_loc[1:] = torch.tensor(
|
||||
batch_spec.query_lens, dtype=torch.int32, device=device
|
||||
).cumsum(0)
|
||||
query_start_loc_cpu = query_start_loc.cpu()
|
||||
num_tokens = batch_spec.compute_num_tokens()
|
||||
|
||||
# Create sequence lengths
|
||||
seq_lens = torch.tensor(batch_spec.seq_lens, dtype=torch.int32, device=device)
|
||||
seq_lens_cpu = seq_lens.cpu()
|
||||
max_seq_len = int(seq_lens_cpu.max())
|
||||
|
||||
# Create computed tokens (context length for each sequence)
|
||||
context_lens = [
|
||||
batch_spec.seq_lens[i] - batch_spec.query_lens[i]
|
||||
for i in range(batch_spec.batch_size)
|
||||
]
|
||||
num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32)
|
||||
|
||||
# Create block table and slot mapping
|
||||
max_blocks = (max(batch_spec.seq_lens) + block_size - 1) // block_size
|
||||
if arange_block_indices:
|
||||
num_blocks = batch_spec.batch_size * max_blocks
|
||||
block_table_tensor = torch.arange(
|
||||
num_blocks, dtype=torch.int32, device=device
|
||||
).view(batch_spec.batch_size, max_blocks)
|
||||
slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device).view(
|
||||
num_tokens
|
||||
)
|
||||
else:
|
||||
block_table_tensor = torch.randint(
|
||||
0,
|
||||
max_block_idx,
|
||||
(batch_spec.batch_size, max_blocks),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
slot_mapping = torch.randint(
|
||||
0, max_block_idx, (num_tokens,), dtype=torch.int64, device=device
|
||||
)
|
||||
|
||||
# Calculate max query length
|
||||
max_query_len = max(batch_spec.query_lens)
|
||||
|
||||
return CommonAttentionMetadata(
|
||||
query_start_loc=query_start_loc,
|
||||
query_start_loc_cpu=query_start_loc_cpu,
|
||||
seq_lens=seq_lens,
|
||||
_seq_lens_cpu=seq_lens_cpu,
|
||||
_num_computed_tokens_cpu=num_computed_tokens_cpu,
|
||||
num_reqs=batch_spec.batch_size,
|
||||
num_actual_tokens=num_tokens,
|
||||
max_query_len=max_query_len,
|
||||
max_seq_len=max_seq_len,
|
||||
block_table_tensor=block_table_tensor,
|
||||
slot_mapping=slot_mapping,
|
||||
causal=True,
|
||||
)
|
||||
|
||||
|
||||
def try_get_attention_backend(
|
||||
backend: AttentionBackendEnum,
|
||||
) -> tuple[type[AttentionMetadataBuilder], type[AttentionImpl]]:
|
||||
"""Try to get the attention backend class, skipping test if not found."""
|
||||
try:
|
||||
backend_class = backend.get_class()
|
||||
return backend_class.get_builder_cls(), backend_class.get_impl_cls()
|
||||
except ImportError as e:
|
||||
pytest.skip(f"{backend.name} not available: {e}")
|
||||
raise AssertionError("unreachable") from None
|
||||
|
||||
|
||||
def try_backend_includes_kv_cache_update(
|
||||
backend: AttentionBackendEnum,
|
||||
) -> bool:
|
||||
"""Try to get the attention backend class, skipping test if not found."""
|
||||
try:
|
||||
backend_class = backend.get_class()
|
||||
return backend_class.forward_includes_kv_cache_update
|
||||
except ImportError as e:
|
||||
pytest.skip(f"{backend.name} not available: {e}")
|
||||
raise AssertionError("unreachable") from None
|
||||
|
||||
|
||||
def create_standard_kv_cache_spec(vllm_config: VllmConfig) -> FullAttentionSpec:
|
||||
"""Create a FullAttentionSpec from ModelParams only."""
|
||||
return FullAttentionSpec(
|
||||
block_size=vllm_config.cache_config.block_size,
|
||||
num_kv_heads=vllm_config.model_config.get_num_kv_heads(
|
||||
vllm_config.parallel_config
|
||||
),
|
||||
head_size=vllm_config.model_config.get_head_size(),
|
||||
dtype=vllm_config.model_config.dtype,
|
||||
sliding_window=vllm_config.model_config.get_sliding_window(),
|
||||
)
|
||||
|
||||
|
||||
def create_vllm_config(
|
||||
model_name: str = "meta-llama/Meta-Llama-3-8B",
|
||||
tensor_parallel_size: int = 1,
|
||||
max_model_len: int = 1024,
|
||||
dtype: ModelDType | torch.dtype = "auto",
|
||||
num_gpu_blocks: int = 1000,
|
||||
block_size: int = 16,
|
||||
max_num_seqs: int = 256,
|
||||
max_num_batched_tokens: int = 8192,
|
||||
enable_chunked_prefill: bool = True,
|
||||
add_mock_model_methods: bool = True,
|
||||
hf_config_override: dict | None = None,
|
||||
) -> VllmConfig:
|
||||
"""Create a VllmConfig for testing with reasonable defaults."""
|
||||
|
||||
model_config = ModelConfig(
|
||||
model=model_name,
|
||||
tokenizer=model_name,
|
||||
trust_remote_code=False,
|
||||
dtype=dtype,
|
||||
seed=0,
|
||||
max_model_len=max_model_len,
|
||||
)
|
||||
|
||||
cache_config = CacheConfig(
|
||||
block_size=block_size,
|
||||
cache_dtype="auto",
|
||||
)
|
||||
# Set cache blocks for testing
|
||||
# (these may be set during initialization normally)
|
||||
cache_config.num_gpu_blocks = num_gpu_blocks
|
||||
cache_config.num_cpu_blocks = 0
|
||||
|
||||
parallel_config = ParallelConfig(
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
)
|
||||
|
||||
scheduler_config = SchedulerConfig(
|
||||
max_num_seqs=max_num_seqs,
|
||||
max_num_batched_tokens=max_num_batched_tokens,
|
||||
enable_chunked_prefill=enable_chunked_prefill,
|
||||
max_model_len=model_config.max_model_len,
|
||||
is_encoder_decoder=model_config.is_encoder_decoder,
|
||||
)
|
||||
|
||||
device_config = DeviceConfig()
|
||||
load_config = LoadConfig()
|
||||
compilation_config = CompilationConfig()
|
||||
|
||||
if add_mock_model_methods:
|
||||
# Add mock methods to satisfy backends that need them
|
||||
# This is a workaround because tests don't build full, real models,
|
||||
# but some backends expect to query the model for layer-specific
|
||||
# parameters
|
||||
import types
|
||||
|
||||
model_config.get_num_layers = types.MethodType(lambda self: 1, model_config)
|
||||
model_config.get_sliding_window_for_layer = types.MethodType(
|
||||
lambda self, i: None, model_config
|
||||
)
|
||||
model_config.get_logits_soft_cap_for_layer = types.MethodType(
|
||||
lambda self, i: 0.0, model_config
|
||||
)
|
||||
model_config.get_sm_scale_for_layer = types.MethodType(
|
||||
lambda self, i: 1.0 / model_config.get_head_size() ** 0.5, model_config
|
||||
)
|
||||
|
||||
if hf_config_override:
|
||||
model_config.hf_config.update(hf_config_override)
|
||||
|
||||
return VllmConfig(
|
||||
model_config=model_config,
|
||||
cache_config=cache_config,
|
||||
parallel_config=parallel_config,
|
||||
scheduler_config=scheduler_config,
|
||||
device_config=device_config,
|
||||
load_config=load_config,
|
||||
compilation_config=compilation_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
|
||||
attention_config: dict
|
||||
comp_config: dict
|
||||
specific_gpu_arch: tuple | None = None
|
||||
|
||||
|
||||
# Define all backend configurations of full cudagraph to be tested
|
||||
full_cg_backend_configs = {
|
||||
# FA3 on Hopper
|
||||
"FA3": BackendConfig(
|
||||
name="FA3",
|
||||
attention_config={
|
||||
"backend": "FLASH_ATTN",
|
||||
"flash_attn_version": 3,
|
||||
"flash_attn_max_num_splits_for_cuda_graph": 16,
|
||||
},
|
||||
comp_config={
|
||||
"cudagraph_mode": "FULL",
|
||||
},
|
||||
specific_gpu_arch=(9, 0),
|
||||
),
|
||||
# FlashMLA on Hopper
|
||||
"FlashMLA": BackendConfig(
|
||||
name="FlashMLA",
|
||||
attention_config={"backend": "FLASHMLA"},
|
||||
comp_config={
|
||||
"cudagraph_mode": "FULL_AND_PIECEWISE",
|
||||
},
|
||||
specific_gpu_arch=(9, 0),
|
||||
),
|
||||
# Cutlass MLA on Blackwell
|
||||
"CutlassMLA": BackendConfig(
|
||||
name="CutlassMLA",
|
||||
attention_config={"backend": "CUTLASS_MLA"},
|
||||
comp_config={
|
||||
"cudagraph_mode": "FULL_AND_PIECEWISE",
|
||||
},
|
||||
specific_gpu_arch=(10, 0),
|
||||
),
|
||||
# FlashInfer MLA on Blackwell
|
||||
"FlashInferMLA": BackendConfig(
|
||||
name="FlashInferMLA",
|
||||
attention_config={"backend": "FLASHINFER_MLA"},
|
||||
comp_config={
|
||||
"cudagraph_mode": "FULL_AND_PIECEWISE",
|
||||
},
|
||||
specific_gpu_arch=(10, 0),
|
||||
),
|
||||
# FlashAttention MLA on Hopper
|
||||
"FlashAttentionMLA": BackendConfig(
|
||||
name="FlashAttentionMLA",
|
||||
attention_config={
|
||||
"backend": "FLASH_ATTN_MLA",
|
||||
"flash_attn_max_num_splits_for_cuda_graph": 16,
|
||||
},
|
||||
comp_config={
|
||||
"cudagraph_mode": "FULL_DECODE_ONLY",
|
||||
},
|
||||
specific_gpu_arch=(9, 0),
|
||||
),
|
||||
# FA2
|
||||
"FA2": BackendConfig(
|
||||
name="FA2",
|
||||
attention_config={
|
||||
"backend": "FLASH_ATTN",
|
||||
"flash_attn_version": 2,
|
||||
"flash_attn_max_num_splits_for_cuda_graph": 16,
|
||||
},
|
||||
comp_config={
|
||||
"cudagraph_mode": "FULL_AND_PIECEWISE",
|
||||
},
|
||||
),
|
||||
# Triton Attention
|
||||
"TritonAttn": BackendConfig(
|
||||
name="TritonAttn",
|
||||
attention_config={"backend": "TRITON_ATTN"},
|
||||
comp_config={
|
||||
"cudagraph_mode": "FULL_AND_PIECEWISE",
|
||||
},
|
||||
),
|
||||
# FlashInfer
|
||||
"FlashInfer": BackendConfig(
|
||||
name="FlashInfer",
|
||||
attention_config={"backend": "FLASHINFER"},
|
||||
comp_config={
|
||||
"cudagraph_mode": "FULL_AND_PIECEWISE",
|
||||
},
|
||||
),
|
||||
"RocmAttn": BackendConfig(
|
||||
name="RocmAttn",
|
||||
attention_config={
|
||||
"backend": "ROCM_ATTN",
|
||||
"use_prefill_decode_attention": True,
|
||||
},
|
||||
comp_config={
|
||||
"cudagraph_mode": "FULL",
|
||||
},
|
||||
),
|
||||
}
|
||||
0
third_party/vllm/tests/v1/core/__init__.py
vendored
Normal file
0
third_party/vllm/tests/v1/core/__init__.py
vendored
Normal file
249
third_party/vllm/tests/v1/core/test_async_scheduler.py
vendored
Normal file
249
third_party/vllm/tests/v1/core/test_async_scheduler.py
vendored
Normal file
@@ -0,0 +1,249 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from collections import deque
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
from vllm.v1.request import RequestStatus
|
||||
from vllm.v1.utils import ConstantList
|
||||
|
||||
from .utils import create_requests, create_scheduler
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
|
||||
|
||||
def _make_model_runner_output(
|
||||
scheduler_output: SchedulerOutput,
|
||||
) -> ModelRunnerOutput:
|
||||
req_ids = list(scheduler_output.num_scheduled_tokens.keys())
|
||||
return ModelRunnerOutput(
|
||||
req_ids=req_ids,
|
||||
req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)},
|
||||
sampled_token_ids=[[i] for i in range(len(req_ids))],
|
||||
logprobs=None,
|
||||
prompt_logprobs_dict={},
|
||||
pooler_output=[],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("max_tokens", [1, 2, 3, 5])
|
||||
def test_stop_by_max_tokens(max_tokens: int):
|
||||
scheduler = create_scheduler(async_scheduling=True)
|
||||
requests = create_requests(num_requests=2, max_tokens=max_tokens)
|
||||
req0, req1 = requests
|
||||
|
||||
expected_total_num_scheduled_tokens = 0
|
||||
sched_outputs: deque[SchedulerOutput] = deque()
|
||||
scheduler.add_request(req0)
|
||||
sched_outputs.append(scheduler.schedule())
|
||||
expected_total_num_scheduled_tokens += req0.num_prompt_tokens + max_tokens - 1
|
||||
|
||||
scheduler.add_request(req1)
|
||||
sched_outputs.append(scheduler.schedule())
|
||||
expected_total_num_scheduled_tokens += req1.num_prompt_tokens + max_tokens - 1
|
||||
|
||||
total_num_scheduled_tokens = 0
|
||||
while sched_outputs:
|
||||
sched_output = sched_outputs.popleft()
|
||||
total_num_scheduled_tokens += sched_output.total_num_scheduled_tokens
|
||||
model_runner_output = _make_model_runner_output(sched_output)
|
||||
scheduler.update_from_output(sched_output, model_runner_output)
|
||||
|
||||
sched_output = scheduler.schedule()
|
||||
if sched_output.num_scheduled_tokens:
|
||||
sched_outputs.append(sched_output)
|
||||
|
||||
assert scheduler.get_num_unfinished_requests() == 0
|
||||
assert req0.num_output_tokens == max_tokens
|
||||
assert req1.num_output_tokens == max_tokens
|
||||
# Ensure we aren't scheduling more tokens than necessary.
|
||||
assert total_num_scheduled_tokens == expected_total_num_scheduled_tokens
|
||||
|
||||
|
||||
def test_abort():
|
||||
scheduler = create_scheduler(async_scheduling=True)
|
||||
requests = create_requests(num_requests=10, max_tokens=20)
|
||||
|
||||
for req in requests:
|
||||
scheduler.add_request(req)
|
||||
|
||||
sched_outputs: deque[SchedulerOutput] = deque()
|
||||
sched_outputs.append(scheduler.schedule())
|
||||
sched_outputs.append(scheduler.schedule())
|
||||
|
||||
abort_order = [0, 8, 3, 1, 6, 4, 2, 5, 7, 9]
|
||||
abort_order_copy = abort_order.copy()
|
||||
|
||||
def abort_request():
|
||||
if not abort_order:
|
||||
return
|
||||
req = requests[abort_order.pop(0)]
|
||||
scheduler.finish_requests(req.request_id, RequestStatus.FINISHED_ABORTED)
|
||||
|
||||
while sched_outputs:
|
||||
# Abort a scheduled request.
|
||||
abort_request()
|
||||
sched_output = sched_outputs.popleft()
|
||||
model_runner_output = _make_model_runner_output(sched_output)
|
||||
scheduler.update_from_output(sched_output, model_runner_output)
|
||||
|
||||
sched_output = scheduler.schedule()
|
||||
if sched_output.num_scheduled_tokens:
|
||||
sched_outputs.append(sched_output)
|
||||
|
||||
for i, req in enumerate(requests):
|
||||
assert req.status == RequestStatus.FINISHED_ABORTED
|
||||
assert req.num_output_tokens == abort_order_copy.index(i)
|
||||
|
||||
|
||||
def test_preempt():
|
||||
scheduler = create_scheduler(async_scheduling=True)
|
||||
requests = create_requests(num_requests=10, max_tokens=20)
|
||||
|
||||
for req in requests:
|
||||
scheduler.add_request(req)
|
||||
|
||||
sched_outputs: deque[SchedulerOutput] = deque()
|
||||
sched_outputs.append(scheduler.schedule())
|
||||
sched_outputs.append(scheduler.schedule())
|
||||
|
||||
abort_order = [0, 8, 3, 1, 6, 4, 2, 5, 7, 9]
|
||||
abort_order_copy = abort_order.copy()
|
||||
|
||||
def abort_request():
|
||||
if not abort_order:
|
||||
return
|
||||
req = requests[abort_order.pop(0)]
|
||||
scheduler.finish_requests(req.request_id, RequestStatus.FINISHED_ABORTED)
|
||||
|
||||
while sched_outputs:
|
||||
# Abort a scheduled request.
|
||||
abort_request()
|
||||
sched_output = sched_outputs.popleft()
|
||||
model_runner_output = _make_model_runner_output(sched_output)
|
||||
scheduler.update_from_output(sched_output, model_runner_output)
|
||||
|
||||
sched_output = scheduler.schedule()
|
||||
if sched_output.num_scheduled_tokens:
|
||||
sched_outputs.append(sched_output)
|
||||
|
||||
for i, req in enumerate(requests):
|
||||
assert req.status == RequestStatus.FINISHED_ABORTED
|
||||
assert req.num_output_tokens == abort_order_copy.index(i)
|
||||
|
||||
|
||||
def test_prefix_caching_for_prefill_dedup():
|
||||
CHUNK_SIZE = 1000
|
||||
BLOCK_SIZE = 16
|
||||
num_prompt_tokens = 100
|
||||
scheduler = create_scheduler(
|
||||
async_scheduling=True,
|
||||
max_num_batched_tokens=CHUNK_SIZE,
|
||||
enable_prefix_caching=True,
|
||||
block_size=BLOCK_SIZE,
|
||||
)
|
||||
requests = create_requests(
|
||||
num_requests=5,
|
||||
num_tokens=num_prompt_tokens,
|
||||
max_tokens=3,
|
||||
same_prompt=True,
|
||||
block_size=BLOCK_SIZE,
|
||||
)
|
||||
requests_copy = requests.copy()
|
||||
|
||||
# Two requests with the same prompt.
|
||||
req0 = requests.pop(0)
|
||||
req1 = requests.pop(0)
|
||||
scheduler.add_request(req0)
|
||||
scheduler.add_request(req1)
|
||||
|
||||
sched_outputs: deque[SchedulerOutput] = deque()
|
||||
sched_output = scheduler.schedule()
|
||||
sched_outputs.append(sched_output)
|
||||
# Make sure prefix caching de-duplicates the prompts in the same step,
|
||||
# so all the blocks except the last are shared between the two requests.
|
||||
assert len(sched_output.num_scheduled_tokens) == 2
|
||||
num_blocks = num_prompt_tokens // BLOCK_SIZE
|
||||
assert req0.num_cached_tokens == 0
|
||||
assert req1.num_cached_tokens >= num_blocks * BLOCK_SIZE
|
||||
|
||||
sched_outputs.append(scheduler.schedule())
|
||||
while sched_outputs:
|
||||
if requests:
|
||||
scheduler.add_request(requests.pop(0))
|
||||
sched_output = sched_outputs.popleft()
|
||||
model_runner_output = _make_model_runner_output(sched_output)
|
||||
scheduler.update_from_output(sched_output, model_runner_output)
|
||||
sched_output = scheduler.schedule()
|
||||
if sched_output.num_scheduled_tokens:
|
||||
sched_outputs.append(sched_output)
|
||||
|
||||
# Other requests scheduled after the two requests should also get
|
||||
# prefix cache hit.
|
||||
assert scheduler.get_num_unfinished_requests() == 0
|
||||
for req in requests_copy[1:]:
|
||||
assert req.num_cached_tokens >= num_blocks * BLOCK_SIZE
|
||||
|
||||
|
||||
def test_prefix_caching_for_multi_turn():
|
||||
CHUNK_SIZE = 1000
|
||||
BLOCK_SIZE = 16
|
||||
num_prompt_tokens = 100
|
||||
num_output_tokens = 200
|
||||
scheduler = create_scheduler(
|
||||
async_scheduling=True,
|
||||
max_num_batched_tokens=CHUNK_SIZE,
|
||||
enable_prefix_caching=True,
|
||||
block_size=BLOCK_SIZE,
|
||||
)
|
||||
requests = create_requests(
|
||||
num_requests=5,
|
||||
num_tokens=num_prompt_tokens,
|
||||
max_tokens=num_output_tokens,
|
||||
block_size=BLOCK_SIZE,
|
||||
)
|
||||
|
||||
for req in requests:
|
||||
scheduler.add_request(req)
|
||||
sched_outputs: deque[SchedulerOutput] = deque()
|
||||
sched_outputs.append(scheduler.schedule())
|
||||
sched_outputs.append(scheduler.schedule())
|
||||
|
||||
# Process the requests.
|
||||
while sched_outputs:
|
||||
sched_output = sched_outputs.popleft()
|
||||
model_runner_output = _make_model_runner_output(sched_output)
|
||||
scheduler.update_from_output(sched_output, model_runner_output)
|
||||
sched_output = scheduler.schedule()
|
||||
if sched_output.num_scheduled_tokens:
|
||||
sched_outputs.append(sched_output)
|
||||
assert scheduler.get_num_unfinished_requests() == 0
|
||||
|
||||
# Create next-turn requests whose prompts are the full output of the
|
||||
# previous turn.
|
||||
next_turn_requests = create_requests(
|
||||
num_requests=5,
|
||||
num_tokens=num_prompt_tokens + num_output_tokens,
|
||||
max_tokens=num_output_tokens,
|
||||
block_size=BLOCK_SIZE,
|
||||
)
|
||||
for i, req in enumerate(next_turn_requests):
|
||||
req.prompt_token_ids = requests[i].prompt_token_ids + list(
|
||||
requests[i].output_token_ids
|
||||
)
|
||||
req._all_token_ids = req.prompt_token_ids.copy()
|
||||
req.all_token_ids = ConstantList(req._all_token_ids)
|
||||
req.block_hashes = []
|
||||
req.update_block_hashes()
|
||||
|
||||
# Schedule the next-turn requests.
|
||||
for req in next_turn_requests:
|
||||
scheduler.add_request(req)
|
||||
sched_outputs.append(scheduler.schedule())
|
||||
|
||||
# Make sure the next-turn requests get prefix cache hit by the previous
|
||||
# requests.
|
||||
for req in next_turn_requests:
|
||||
assert req.num_cached_tokens == req.num_prompt_tokens // BLOCK_SIZE * BLOCK_SIZE
|
||||
337
third_party/vllm/tests/v1/core/test_encoder_cache_manager.py
vendored
Normal file
337
third_party/vllm/tests/v1/core/test_encoder_cache_manager.py
vendored
Normal file
@@ -0,0 +1,337 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
|
||||
from vllm.v1.core.encoder_cache_manager import (
|
||||
EncoderCacheManager,
|
||||
EncoderDecoderCacheManager,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
|
||||
|
||||
# ------------------ Mock Classes ------------------ #
|
||||
class MockRequest:
|
||||
def __init__(self, request_id, mm_hashes, token_counts):
|
||||
self.request_id = request_id
|
||||
self._token_counts = token_counts
|
||||
self.mm_features = []
|
||||
for i, mm_hash in enumerate(mm_hashes):
|
||||
feature = MultiModalFeatureSpec(
|
||||
data=None,
|
||||
modality="image",
|
||||
identifier=mm_hash,
|
||||
mm_position=PlaceholderRange(offset=0, length=self._token_counts[i]),
|
||||
)
|
||||
self.mm_features.append(feature)
|
||||
|
||||
def get_num_encoder_embeds(self, input_id: int) -> int:
|
||||
return self._token_counts[input_id]
|
||||
|
||||
|
||||
# ------------------ Unit Tests ------------------ #
|
||||
def test_basic_allocate_and_reuse():
|
||||
cache = EncoderCacheManager(cache_size=10)
|
||||
req = MockRequest("r1", ["imgA"], [4])
|
||||
|
||||
assert not cache.check_and_update_cache(req, 0)
|
||||
assert cache.can_allocate(req, 0, int(1e9), 0)
|
||||
|
||||
cache.allocate(req, 0)
|
||||
|
||||
assert cache.check_and_update_cache(req, 0)
|
||||
assert "r1" in cache.cached["imgA"]
|
||||
assert cache.num_free_slots == 6
|
||||
|
||||
# Free twice to bring refcount to 0.
|
||||
cache.free_encoder_input(req, 0)
|
||||
cache.free_encoder_input(req, 0)
|
||||
|
||||
assert not cache.cached["imgA"]
|
||||
assert "imgA" in cache.freeable
|
||||
assert cache.num_freeable_slots == 10
|
||||
assert cache.num_free_slots == 6
|
||||
|
||||
|
||||
def test_freeing_decreases_refcount_and_moves_to_freeable():
|
||||
manager = EncoderCacheManager(cache_size=10)
|
||||
req = MockRequest("req2", ["img3"], [5])
|
||||
|
||||
assert manager.can_allocate(req, 0, int(1e9), 0)
|
||||
manager.allocate(req, 0)
|
||||
|
||||
assert len(manager.cached["img3"]) == 1
|
||||
|
||||
manager.free_encoder_input(req, 0)
|
||||
|
||||
assert not manager.cached["img3"]
|
||||
assert "img3" in manager.freeable
|
||||
assert manager.num_freeable_slots == 10
|
||||
|
||||
|
||||
def test_free_request_frees_all_inputs():
|
||||
manager = EncoderCacheManager(cache_size=10)
|
||||
req = MockRequest("req3", ["a", "b"], [2, 3])
|
||||
|
||||
assert manager.can_allocate(req, 0, int(1e9), 0)
|
||||
manager.allocate(req, 0)
|
||||
|
||||
assert manager.can_allocate(req, 1, int(1e9), 0)
|
||||
manager.allocate(req, 1)
|
||||
|
||||
assert len(manager.cached["a"]) == 1
|
||||
assert len(manager.cached["b"]) == 1
|
||||
|
||||
manager.free(req)
|
||||
|
||||
assert not manager.cached["a"]
|
||||
assert not manager.cached["b"]
|
||||
assert "a" in manager.freeable
|
||||
assert "b" in manager.freeable
|
||||
assert manager.num_freeable_slots == 10
|
||||
|
||||
|
||||
def test_eviction_when_cache_is_full():
|
||||
manager = EncoderCacheManager(cache_size=10)
|
||||
|
||||
req1 = MockRequest("req1", ["x"], [6])
|
||||
req2 = MockRequest("req2", ["y"], [5])
|
||||
|
||||
assert manager.can_allocate(req1, 0, int(1e9), 0)
|
||||
manager.allocate(req1, 0)
|
||||
manager.free_encoder_input(req1, 0)
|
||||
|
||||
assert manager.can_allocate(req2, 0, int(1e9), 0)
|
||||
manager.allocate(req2, 0)
|
||||
|
||||
# 'x' should have been evicted.
|
||||
assert "x" not in manager.cached
|
||||
assert "x" in manager.get_freed_mm_hashes()
|
||||
|
||||
|
||||
def test_get_cached_input_ids():
|
||||
manager = EncoderCacheManager(cache_size=10)
|
||||
req = MockRequest("reqX", ["m", "n", "o"], [2, 4, 3])
|
||||
|
||||
assert manager.can_allocate(req, 0, int(1e9), 0)
|
||||
manager.allocate(req, 0)
|
||||
|
||||
assert manager.can_allocate(req, 2, int(1e9), 0)
|
||||
manager.allocate(req, 2)
|
||||
|
||||
cached_ids = manager.get_cached_input_ids(req)
|
||||
assert cached_ids == {0, 2}
|
||||
|
||||
|
||||
def test_has_cache_restores_from_freeable():
|
||||
manager = EncoderCacheManager(cache_size=10)
|
||||
req = MockRequest("reqY", ["imgZ"], [4])
|
||||
|
||||
assert manager.can_allocate(req, 0, int(1e9), 0)
|
||||
manager.allocate(req, 0)
|
||||
|
||||
manager.free_encoder_input(req, 0)
|
||||
|
||||
# Should restore from freeable.
|
||||
assert manager.check_and_update_cache(req, 0)
|
||||
assert len(manager.cached["imgZ"]) == 1
|
||||
assert "imgZ" not in manager.freeable
|
||||
assert manager.num_freeable_slots == 6
|
||||
|
||||
|
||||
def test_get_freed_mm_hashes_clears_freed_list():
|
||||
manager = EncoderCacheManager(cache_size=10)
|
||||
req1 = MockRequest("reqA", ["a"], [5])
|
||||
req2 = MockRequest("reqB", ["b"], [6])
|
||||
|
||||
assert manager.can_allocate(req1, 0, int(1e9), 0)
|
||||
manager.allocate(req1, 0)
|
||||
manager.free_encoder_input(req1, 0)
|
||||
|
||||
# Should trigger eviction of 'a'.
|
||||
assert manager.can_allocate(req2, 0, int(1e9), 0)
|
||||
manager.allocate(req2, 0)
|
||||
|
||||
freed = manager.get_freed_mm_hashes()
|
||||
assert "a" in freed
|
||||
assert manager.get_freed_mm_hashes() == []
|
||||
|
||||
|
||||
def test_schedule_request_multi_images_respect_space_limit():
|
||||
manager = EncoderCacheManager(cache_size=10)
|
||||
req = MockRequest("reqA", ["a", "b"], [5, 6])
|
||||
compute_budget = 100
|
||||
|
||||
num_tokens_to_schedule = 0
|
||||
assert manager.can_allocate(req, 0, compute_budget, num_tokens_to_schedule)
|
||||
num_tokens_to_schedule += req.get_num_encoder_embeds(0)
|
||||
compute_budget -= req.get_num_encoder_embeds(0)
|
||||
|
||||
assert not manager.can_allocate(req, 1, compute_budget, num_tokens_to_schedule)
|
||||
|
||||
|
||||
def test_schedule_request_multi_images_respect_compute_limit():
|
||||
manager = EncoderCacheManager(cache_size=100)
|
||||
req = MockRequest("reqA", ["a", "b"], [5, 6])
|
||||
compute_budget = 10
|
||||
num_tokens_to_schedule = 0
|
||||
assert manager.can_allocate(req, 0, compute_budget, num_tokens_to_schedule)
|
||||
num_tokens_to_schedule += req.get_num_encoder_embeds(0)
|
||||
compute_budget -= req.get_num_encoder_embeds(0)
|
||||
|
||||
assert not manager.can_allocate(req, 1, compute_budget, num_tokens_to_schedule)
|
||||
|
||||
|
||||
def test_encoder_cache_with_is_embed_mask():
|
||||
class MockRequestWithMask(MockRequest):
|
||||
def get_num_encoder_embeds(self, input_id: int) -> int:
|
||||
return self.mm_features[input_id].mm_position.get_num_embeds()
|
||||
|
||||
is_embed = torch.zeros(100, dtype=torch.bool)
|
||||
is_embed[torch.tensor([5, 15, 25, 35, 45, 55, 65, 75])] = True
|
||||
|
||||
request = MockRequestWithMask("r1", ["img1"], [100])
|
||||
request.mm_features[0] = MultiModalFeatureSpec(
|
||||
data=None,
|
||||
modality="image",
|
||||
identifier="img1",
|
||||
mm_position=PlaceholderRange(offset=0, length=100, is_embed=is_embed),
|
||||
)
|
||||
|
||||
manager = EncoderCacheManager(cache_size=100)
|
||||
manager.allocate(request, 0)
|
||||
|
||||
assert manager.num_free_slots == 92
|
||||
assert "img1" in manager.cached
|
||||
|
||||
old_size = 100
|
||||
new_size = request.mm_features[0].mm_position.get_num_embeds()
|
||||
assert new_size == 8
|
||||
savings_ratio = old_size / new_size
|
||||
assert savings_ratio == 12.5
|
||||
|
||||
|
||||
def test_encoder_cache_mask_based_retrieval():
|
||||
class MockRequestWithMask(MockRequest):
|
||||
def get_num_encoder_embeds(self, input_id: int) -> int:
|
||||
return self.mm_features[input_id].mm_position.get_num_embeds()
|
||||
|
||||
is_embed = torch.tensor(
|
||||
[False, False, True, True, False, True, True, True, False, False]
|
||||
)
|
||||
|
||||
request = MockRequestWithMask("r1", ["img1"], [10])
|
||||
request.mm_features[0] = MultiModalFeatureSpec(
|
||||
data=None,
|
||||
modality="image",
|
||||
identifier="img1",
|
||||
mm_position=PlaceholderRange(offset=0, length=10, is_embed=is_embed),
|
||||
)
|
||||
|
||||
manager = EncoderCacheManager(cache_size=50)
|
||||
manager.allocate(request, 0)
|
||||
|
||||
assert request.mm_features[0].mm_position.get_num_embeds() == 5
|
||||
|
||||
start_idx = 2
|
||||
end_idx = 8
|
||||
num_embeds_before = is_embed[:start_idx].sum().item()
|
||||
num_embeds_in_range = is_embed[start_idx:end_idx].sum().item()
|
||||
|
||||
assert num_embeds_before == 0
|
||||
assert num_embeds_in_range == 5
|
||||
|
||||
start_idx = 0
|
||||
end_idx = 5
|
||||
num_embeds_before = is_embed[:start_idx].sum().item() if start_idx > 0 else 0
|
||||
num_embeds_in_range = is_embed[start_idx:end_idx].sum().item()
|
||||
|
||||
assert num_embeds_before == 0
|
||||
assert num_embeds_in_range == 2
|
||||
|
||||
|
||||
def test_reset_clears_all_state():
|
||||
"""Test that reset() clears all cached entries and restores capacity."""
|
||||
manager = EncoderCacheManager(cache_size=20)
|
||||
|
||||
req1 = MockRequest("req1", ["img1", "img2"], [5, 3])
|
||||
req2 = MockRequest("req2", ["img3"], [4])
|
||||
|
||||
manager.allocate(req1, 0)
|
||||
manager.allocate(req1, 1)
|
||||
manager.allocate(req2, 0)
|
||||
manager.free_encoder_input(req1, 0)
|
||||
|
||||
req3 = MockRequest("req3", ["img4"], [10])
|
||||
manager.free_encoder_input(req1, 1)
|
||||
manager.free_encoder_input(req2, 0)
|
||||
manager.can_allocate(req3, 0, int(1e9), 0)
|
||||
manager.allocate(req3, 0)
|
||||
|
||||
assert len(manager.cached) > 0
|
||||
assert manager.num_free_slots < 20
|
||||
|
||||
manager.reset()
|
||||
|
||||
assert len(manager.cached) == 0
|
||||
assert len(manager.freeable) == 0
|
||||
assert len(manager.freed) == 0
|
||||
assert manager.num_free_slots == 20
|
||||
assert manager.num_freeable_slots == 20
|
||||
|
||||
|
||||
def test_reset_allows_fresh_allocations():
|
||||
manager = EncoderCacheManager(cache_size=10)
|
||||
|
||||
req1 = MockRequest("req1", ["img1"], [10])
|
||||
manager.allocate(req1, 0)
|
||||
assert manager.num_free_slots == 0
|
||||
|
||||
manager.reset()
|
||||
|
||||
req2 = MockRequest("req2", ["img2"], [8])
|
||||
assert manager.can_allocate(req2, 0, int(1e9), 0)
|
||||
manager.allocate(req2, 0)
|
||||
|
||||
assert manager.num_free_slots == 2
|
||||
assert "img2" in manager.cached
|
||||
assert "img1" not in manager.cached
|
||||
|
||||
|
||||
def test_encoder_decoder_cache_manager_reset():
|
||||
manager = EncoderDecoderCacheManager(cache_size=20)
|
||||
|
||||
req1 = MockRequest("req1", ["img1"], [5])
|
||||
req2 = MockRequest("req2", ["img2"], [3])
|
||||
|
||||
manager.allocate(req1, 0)
|
||||
manager.allocate(req2, 0)
|
||||
manager.free(req1)
|
||||
manager.get_freed_mm_hashes()
|
||||
|
||||
assert manager.num_free_slots < 20
|
||||
|
||||
manager.reset()
|
||||
|
||||
assert len(manager.allocated) == 0
|
||||
assert len(manager.to_free) == 0
|
||||
assert manager.num_free_slots == 20
|
||||
|
||||
|
||||
def test_encoder_decoder_cache_manager_reset_allows_fresh_allocations():
|
||||
manager = EncoderDecoderCacheManager(cache_size=10)
|
||||
|
||||
req1 = MockRequest("req1", ["img1"], [10])
|
||||
manager.allocate(req1, 0)
|
||||
assert manager.num_free_slots == 0
|
||||
|
||||
manager.reset()
|
||||
|
||||
req2 = MockRequest("req2", ["img2"], [8])
|
||||
assert manager.can_allocate(req2, 0, int(1e9), 0)
|
||||
manager.allocate(req2, 0)
|
||||
|
||||
assert manager.num_free_slots == 2
|
||||
assert "img2" in manager.allocated
|
||||
224
third_party/vllm/tests/v1/core/test_kv_cache_metrics.py
vendored
Normal file
224
third_party/vllm/tests/v1/core/test_kv_cache_metrics.py
vendored
Normal file
@@ -0,0 +1,224 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.v1.core.kv_cache_metrics import (
|
||||
BlockMetricsState,
|
||||
KVCacheMetricsCollector,
|
||||
)
|
||||
from vllm.v1.core.kv_cache_utils import KVCacheBlock
|
||||
|
||||
|
||||
class TestBlockMetricsState:
|
||||
def test_init(self):
|
||||
with patch("time.monotonic_ns", return_value=1000000000):
|
||||
state = BlockMetricsState()
|
||||
assert state.birth_time_ns == 1000000000
|
||||
assert state.last_access_ns == 1000000000
|
||||
assert len(state.access_history) == 0
|
||||
|
||||
def test_access_tracking(self):
|
||||
with patch("time.monotonic_ns", return_value=1000000000):
|
||||
state = BlockMetricsState()
|
||||
|
||||
with patch("time.monotonic_ns", return_value=2000000000):
|
||||
state.record_access()
|
||||
|
||||
assert state.last_access_ns == 2000000000
|
||||
assert list(state.access_history) == [2000000000]
|
||||
|
||||
def test_ring_buffer_wraps_at_4(self):
|
||||
with patch("time.monotonic_ns", return_value=1000000000):
|
||||
state = BlockMetricsState()
|
||||
|
||||
for i in range(5):
|
||||
t = 1000000000 + (i + 1) * 1000000000
|
||||
with patch("time.monotonic_ns", return_value=t):
|
||||
state.record_access()
|
||||
|
||||
assert len(state.access_history) == 4
|
||||
assert list(state.access_history) == [
|
||||
3000000000,
|
||||
4000000000,
|
||||
5000000000,
|
||||
6000000000,
|
||||
]
|
||||
|
||||
def test_lifetime(self):
|
||||
with patch("time.monotonic_ns", return_value=1000000000):
|
||||
state = BlockMetricsState()
|
||||
with patch("time.monotonic_ns", return_value=6500000000):
|
||||
assert abs(state.get_lifetime_seconds() - 5.5) < 0.001
|
||||
|
||||
def test_idle_time(self):
|
||||
with patch("time.monotonic_ns", return_value=1000000000):
|
||||
state = BlockMetricsState()
|
||||
state.last_access_ns = 2000000000
|
||||
with patch("time.monotonic_ns", return_value=5200000000):
|
||||
assert abs(state.get_idle_time_seconds() - 3.2) < 0.001
|
||||
|
||||
def test_reuse_gaps(self):
|
||||
with patch("time.monotonic_ns", return_value=1000000000):
|
||||
state = BlockMetricsState()
|
||||
|
||||
base = 1000000000
|
||||
for offset in [0, 1.5, 3.0, 5.5]:
|
||||
state.access_history.append(base + int(offset * 1e9))
|
||||
|
||||
gaps = state.get_reuse_gaps_seconds()
|
||||
assert len(gaps) == 3
|
||||
assert gaps[0] == 1.5 and gaps[1] == 1.5 and gaps[2] == 2.5
|
||||
|
||||
def test_ring_wrap_only_gives_3_gaps(self):
|
||||
# 5 accesses in size-4 buffer = 3 gaps
|
||||
with patch("time.monotonic_ns", return_value=1000000000):
|
||||
state = BlockMetricsState()
|
||||
|
||||
for i in range(5):
|
||||
state.access_history.append(1000000000 + i * 1000000000)
|
||||
|
||||
assert len(state.get_reuse_gaps_seconds()) == 3
|
||||
|
||||
|
||||
class TestKVCacheMetricsCollector:
|
||||
def test_sample_rate_validation(self):
|
||||
with pytest.raises(AssertionError):
|
||||
KVCacheMetricsCollector(sample_rate=-0.1)
|
||||
with pytest.raises(AssertionError):
|
||||
KVCacheMetricsCollector(sample_rate=1.5)
|
||||
with pytest.raises(AssertionError):
|
||||
KVCacheMetricsCollector(sample_rate=0.0)
|
||||
|
||||
def test_sampling(self):
|
||||
c = KVCacheMetricsCollector(sample_rate=1.0)
|
||||
assert sum(1 for _ in range(100) if c.should_sample_block()) == 100
|
||||
|
||||
c = KVCacheMetricsCollector(sample_rate=0.5)
|
||||
samples = sum(1 for _ in range(1000) if c.should_sample_block())
|
||||
assert 400 < samples < 600
|
||||
|
||||
def test_alloc(self):
|
||||
c = KVCacheMetricsCollector(sample_rate=1.0)
|
||||
|
||||
blocks = [KVCacheBlock(block_id=i) for i in range(5)]
|
||||
with patch("time.monotonic_ns", return_value=1000000000):
|
||||
for block in blocks:
|
||||
c.on_block_allocated(block)
|
||||
|
||||
assert len(c.block_metrics) == 5
|
||||
|
||||
def test_access(self):
|
||||
c = KVCacheMetricsCollector(sample_rate=1.0)
|
||||
block = KVCacheBlock(block_id=0)
|
||||
|
||||
with patch("time.monotonic_ns", return_value=1000000000):
|
||||
c.on_block_allocated(block)
|
||||
|
||||
for i in range(3):
|
||||
t = 1000000000 + (i + 1) * 1000000000
|
||||
with patch("time.monotonic_ns", return_value=t):
|
||||
c.on_block_accessed(block)
|
||||
|
||||
assert len(c.block_metrics[0].access_history) == 3
|
||||
|
||||
def test_evict_no_accesses(self):
|
||||
# lifetime should equal idle if never accessed
|
||||
c = KVCacheMetricsCollector(sample_rate=1.0)
|
||||
|
||||
block = KVCacheBlock(block_id=0)
|
||||
with patch("time.monotonic_ns", return_value=1000000000):
|
||||
c.on_block_allocated(block)
|
||||
|
||||
with patch("time.monotonic_ns", return_value=6000000000):
|
||||
c.on_block_evicted(block)
|
||||
|
||||
events = c.drain_events()
|
||||
assert len(events) == 1
|
||||
assert abs(events[0].lifetime_seconds - 5.0) < 0.001
|
||||
assert abs(events[0].idle_seconds - 5.0) < 0.001
|
||||
|
||||
def test_evict(self):
|
||||
c = KVCacheMetricsCollector(sample_rate=1.0)
|
||||
|
||||
block = KVCacheBlock(block_id=0)
|
||||
with patch("time.monotonic_ns", return_value=1000000000):
|
||||
c.on_block_allocated(block)
|
||||
|
||||
with patch("time.monotonic_ns", return_value=2000000000):
|
||||
c.on_block_accessed(block)
|
||||
with patch("time.monotonic_ns", return_value=3000000000):
|
||||
c.on_block_accessed(block)
|
||||
|
||||
with patch("time.monotonic_ns", return_value=4000000000):
|
||||
c.on_block_evicted(block)
|
||||
|
||||
events = c.drain_events()
|
||||
assert len(events) == 1
|
||||
sample = events[0]
|
||||
assert abs(sample.lifetime_seconds - 3.0) < 0.001
|
||||
assert abs(sample.idle_seconds - 1.0) < 0.001
|
||||
assert sample.reuse_gaps_seconds == (1.0,)
|
||||
assert 0 not in c.block_metrics
|
||||
|
||||
def test_reset(self):
|
||||
c = KVCacheMetricsCollector(sample_rate=1.0)
|
||||
|
||||
with patch("time.monotonic_ns", return_value=1000000000):
|
||||
for i in range(5):
|
||||
c.on_block_allocated(KVCacheBlock(block_id=i))
|
||||
|
||||
assert len(c.block_metrics) == 5
|
||||
c.reset()
|
||||
assert len(c.block_metrics) == 0
|
||||
|
||||
with patch("time.monotonic_ns", return_value=2000000000):
|
||||
c.on_block_allocated(KVCacheBlock(block_id=10))
|
||||
assert 10 in c.block_metrics
|
||||
|
||||
def test_huge_time_jump(self):
|
||||
c = KVCacheMetricsCollector(sample_rate=1.0)
|
||||
|
||||
block = KVCacheBlock(block_id=0)
|
||||
with patch("time.monotonic_ns", return_value=1000000000):
|
||||
c.on_block_allocated(block)
|
||||
|
||||
with patch("time.monotonic_ns", return_value=9999999999999999):
|
||||
c.on_block_evicted(block)
|
||||
|
||||
events = c.drain_events()
|
||||
assert len(events) == 1
|
||||
assert events[0].lifetime_seconds > 0
|
||||
|
||||
|
||||
def test_kv_cache_metrics_collector_smoke() -> None:
|
||||
"""Simple smoke test for KVCacheMetricsCollector on CPU."""
|
||||
collector = KVCacheMetricsCollector(sample_rate=1.0)
|
||||
block = KVCacheBlock(block_id=123)
|
||||
|
||||
# Allocate at t = 1.0s.
|
||||
with patch("time.monotonic_ns", return_value=1_000_000_000):
|
||||
collector.on_block_allocated(block)
|
||||
|
||||
# Access at t = 2.0s and t = 3.0s.
|
||||
with patch("time.monotonic_ns", return_value=2_000_000_000):
|
||||
collector.on_block_accessed(block)
|
||||
with patch("time.monotonic_ns", return_value=3_000_000_000):
|
||||
collector.on_block_accessed(block)
|
||||
|
||||
# Evict at t = 4.0s.
|
||||
with patch("time.monotonic_ns", return_value=4_000_000_000):
|
||||
collector.on_block_evicted(block)
|
||||
|
||||
events = collector.drain_events()
|
||||
assert len(events) == 1
|
||||
|
||||
event = events[0]
|
||||
# Lifetime: 1.0s → 4.0s.
|
||||
assert abs(event.lifetime_seconds - 3.0) < 1e-6
|
||||
# Idle: last access at 3.0s, evicted at 4.0s.
|
||||
assert abs(event.idle_seconds - 1.0) < 1e-6
|
||||
# One reuse gap between the two accesses.
|
||||
assert event.reuse_gaps_seconds == (1.0,)
|
||||
2096
third_party/vllm/tests/v1/core/test_kv_cache_utils.py
vendored
Normal file
2096
third_party/vllm/tests/v1/core/test_kv_cache_utils.py
vendored
Normal file
File diff suppressed because it is too large
Load Diff
105
third_party/vllm/tests/v1/core/test_kv_sharing.py
vendored
Normal file
105
third_party/vllm/tests/v1/core/test_kv_sharing.py
vendored
Normal file
@@ -0,0 +1,105 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec
|
||||
from vllm.v1.worker.utils import add_kv_sharing_layers_to_kv_cache_groups
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
|
||||
|
||||
def new_kv_cache_spec():
|
||||
return FullAttentionSpec(
|
||||
block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32
|
||||
)
|
||||
|
||||
|
||||
def test_initialize_kv_cache_for_kv_sharing_different_attn_groups():
|
||||
"""
|
||||
Test initializing KV cache sharing with different attention groups.
|
||||
Layers in the same KV cache group might be placed in different attn groups
|
||||
if they have different attention backends.
|
||||
"""
|
||||
shared_kv_cache_layers = {
|
||||
"model.layers.2": "model.layers.0",
|
||||
"model.layers.3": "model.layers.1",
|
||||
}
|
||||
|
||||
# Layers 0 and 1 both belong in KV cache group 0
|
||||
# However, if they have different attention backends, they will be
|
||||
# placed in different attention groups for KV cache group 0
|
||||
kv_cache_groups = [
|
||||
KVCacheGroupSpec(["model.layers.0", "model.layers.1"], new_kv_cache_spec()),
|
||||
]
|
||||
|
||||
add_kv_sharing_layers_to_kv_cache_groups(
|
||||
shared_kv_cache_layers=shared_kv_cache_layers,
|
||||
kv_cache_groups=kv_cache_groups,
|
||||
)
|
||||
|
||||
# Check that the layers were added to the correct KV cache group
|
||||
assert len(kv_cache_groups) == 1
|
||||
assert kv_cache_groups[0].layer_names == [
|
||||
"model.layers.0",
|
||||
"model.layers.1",
|
||||
"model.layers.2",
|
||||
"model.layers.3",
|
||||
]
|
||||
|
||||
|
||||
def test_initialize_kv_cache_for_kv_sharing_same_attn_groups():
|
||||
"""
|
||||
Test case assuming that all layers in the same KV cache group have the same
|
||||
attention backends. This is true for most models.
|
||||
"""
|
||||
shared_kv_cache_layers = {
|
||||
"model.layers.2": "model.layers.0",
|
||||
"model.layers.3": "model.layers.1",
|
||||
}
|
||||
|
||||
kv_cache_groups = [
|
||||
KVCacheGroupSpec(["model.layers.0", "model.layers.1"], new_kv_cache_spec()),
|
||||
]
|
||||
|
||||
add_kv_sharing_layers_to_kv_cache_groups(
|
||||
shared_kv_cache_layers=shared_kv_cache_layers,
|
||||
kv_cache_groups=kv_cache_groups,
|
||||
)
|
||||
|
||||
# Check that the layers were added to the correct KV cache group
|
||||
assert len(kv_cache_groups) == 1
|
||||
assert kv_cache_groups[0].layer_names == [
|
||||
"model.layers.0",
|
||||
"model.layers.1",
|
||||
"model.layers.2",
|
||||
"model.layers.3",
|
||||
]
|
||||
|
||||
|
||||
def test_initialize_kv_cache_for_kv_sharing_no_attn_groups():
|
||||
"""
|
||||
Test KV sharing set up when no attention groups are provided.
|
||||
This is the case for the TPU model runner, which doesn't have
|
||||
support for attention groups yet.
|
||||
"""
|
||||
shared_kv_cache_layers = {
|
||||
"model.layers.2": "model.layers.0",
|
||||
"model.layers.3": "model.layers.1",
|
||||
}
|
||||
|
||||
kv_cache_groups = [
|
||||
KVCacheGroupSpec(["model.layers.0"], new_kv_cache_spec()),
|
||||
KVCacheGroupSpec(["model.layers.1"], new_kv_cache_spec()),
|
||||
]
|
||||
|
||||
add_kv_sharing_layers_to_kv_cache_groups(
|
||||
shared_kv_cache_layers=shared_kv_cache_layers,
|
||||
kv_cache_groups=kv_cache_groups,
|
||||
)
|
||||
|
||||
# Check that the layers were added to the correct KV cache group
|
||||
assert len(kv_cache_groups) == 2
|
||||
assert kv_cache_groups[0].layer_names == ["model.layers.0", "model.layers.2"]
|
||||
assert kv_cache_groups[1].layer_names == ["model.layers.1", "model.layers.3"]
|
||||
36
third_party/vllm/tests/v1/core/test_output.py
vendored
Normal file
36
third_party/vllm/tests/v1/core/test_output.py
vendored
Normal file
@@ -0,0 +1,36 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import torch
|
||||
|
||||
from vllm.v1.core.sched.output import NewRequestData
|
||||
|
||||
|
||||
def _create_new_requests_data(prompt_embeds: torch.Tensor | None) -> NewRequestData:
|
||||
return NewRequestData(
|
||||
req_id="test_req",
|
||||
prompt_token_ids=None,
|
||||
mm_features=[],
|
||||
sampling_params=None,
|
||||
pooling_params=None,
|
||||
block_ids=([],),
|
||||
num_computed_tokens=0,
|
||||
lora_request=None,
|
||||
prompt_embeds=prompt_embeds,
|
||||
)
|
||||
|
||||
|
||||
def test_repr_with_none() -> None:
|
||||
"""Test repr when prompt_embeds is None."""
|
||||
new_requests_data = _create_new_requests_data(None)
|
||||
|
||||
assert "prompt_embeds_shape=None" in repr(new_requests_data)
|
||||
assert "prompt_embeds_shape=None" in new_requests_data.anon_repr()
|
||||
|
||||
|
||||
def test_repr_with_multi_element_tensor() -> None:
|
||||
"""Test repr when prompt_embeds is a multi-element tensor."""
|
||||
prompt_embeds = torch.randn(10, 768)
|
||||
new_requests_data = _create_new_requests_data(prompt_embeds)
|
||||
|
||||
assert "prompt_embeds_shape=torch.Size([10, 768])" in repr(new_requests_data)
|
||||
assert "prompt_embeds_shape=torch.Size([10, 768])" in new_requests_data.anon_repr()
|
||||
2358
third_party/vllm/tests/v1/core/test_prefix_caching.py
vendored
Normal file
2358
third_party/vllm/tests/v1/core/test_prefix_caching.py
vendored
Normal file
File diff suppressed because it is too large
Load Diff
260
third_party/vllm/tests/v1/core/test_priority_scheduler_random.py
vendored
Normal file
260
third_party/vllm/tests/v1/core/test_priority_scheduler_random.py
vendored
Normal file
@@ -0,0 +1,260 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import random
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.multimodal.inputs import (
|
||||
MultiModalFeatureSpec,
|
||||
MultiModalKwargsItem,
|
||||
PlaceholderRange,
|
||||
)
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.utils.hashing import get_hash_fn_by_name
|
||||
from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
|
||||
from vllm.v1.request import Request
|
||||
|
||||
from .test_scheduler import create_scheduler_with_priority
|
||||
from .utils import EOS_TOKEN_ID
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
|
||||
|
||||
def _create_random_request(
|
||||
max_tokens_range: tuple[int, int],
|
||||
num_tokens_range: tuple[int, int],
|
||||
arrival_time_range: tuple[float, float],
|
||||
priority_range: tuple[int, int],
|
||||
num_mm_item_range: tuple[int, int],
|
||||
vllm_config: VllmConfig,
|
||||
):
|
||||
max_tokens = random.randint(*max_tokens_range)
|
||||
num_tokens = random.randint(*num_tokens_range)
|
||||
priority = random.randint(*priority_range)
|
||||
arrival_time = random.uniform(*arrival_time_range)
|
||||
num_mm_item = random.randint(*num_mm_item_range)
|
||||
|
||||
mm_positions: list[PlaceholderRange] = []
|
||||
for mm_start in sorted(
|
||||
random.sample(range(num_tokens), min(num_mm_item, num_tokens))
|
||||
):
|
||||
if mm_start + 10 > num_tokens:
|
||||
continue
|
||||
mm_positions.append(PlaceholderRange(offset=mm_start, length=10))
|
||||
|
||||
request_id = uuid.uuid4().hex
|
||||
|
||||
sampling_params = SamplingParams(ignore_eos=False, max_tokens=max_tokens)
|
||||
sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
|
||||
|
||||
mm_features = []
|
||||
for j, position in enumerate(mm_positions):
|
||||
identifier = f"{request_id}_hash_{j}"
|
||||
mm_feature = MultiModalFeatureSpec(
|
||||
data=MultiModalKwargsItem.dummy(),
|
||||
mm_position=position,
|
||||
identifier=identifier,
|
||||
modality="image",
|
||||
)
|
||||
mm_features.append(mm_feature)
|
||||
|
||||
prompt_token_ids = random.choices(range(100), k=num_tokens)
|
||||
|
||||
caching_hash_fn = get_hash_fn_by_name(
|
||||
vllm_config.cache_config.prefix_caching_hash_algo
|
||||
)
|
||||
init_none_hash(caching_hash_fn)
|
||||
block_hasher = get_request_block_hasher(
|
||||
vllm_config.cache_config.block_size, caching_hash_fn
|
||||
)
|
||||
|
||||
request = Request(
|
||||
request_id=request_id,
|
||||
prompt_token_ids=prompt_token_ids,
|
||||
sampling_params=sampling_params,
|
||||
pooling_params=None,
|
||||
mm_features=mm_features if mm_features else None,
|
||||
arrival_time=arrival_time,
|
||||
priority=priority,
|
||||
block_hasher=block_hasher,
|
||||
)
|
||||
return request
|
||||
|
||||
|
||||
def _mock_execute_model(
|
||||
scheduler_output: SchedulerOutput, num_output_tokens_range: tuple[int, int]
|
||||
) -> ModelRunnerOutput:
|
||||
request_ids: list[str] = []
|
||||
request_ids.extend(req.req_id for req in scheduler_output.scheduled_new_reqs)
|
||||
request_ids.extend(scheduler_output.scheduled_cached_reqs.req_ids)
|
||||
random.shuffle(request_ids)
|
||||
|
||||
num_output_tokens = [
|
||||
random.randint(*num_output_tokens_range) for _ in range(len(request_ids))
|
||||
]
|
||||
sampled_token_ids = [
|
||||
[random.randint(0, 100) for _ in range(num_tokens)]
|
||||
for num_tokens in num_output_tokens
|
||||
]
|
||||
|
||||
return ModelRunnerOutput(
|
||||
req_ids=request_ids,
|
||||
req_id_to_index={req_id: i for i, req_id in enumerate(request_ids)},
|
||||
sampled_token_ids=sampled_token_ids,
|
||||
logprobs=None,
|
||||
prompt_logprobs_dict={},
|
||||
pooler_output=[],
|
||||
)
|
||||
|
||||
|
||||
def _mock_draft_token_ids(
|
||||
scheduler_output: SchedulerOutput,
|
||||
num_output_tokens_range: tuple[int, int],
|
||||
seen_request_prompt_length: dict[str, int],
|
||||
) -> DraftTokenIds:
|
||||
request_ids: list[str] = []
|
||||
sampled_token_ids: list[list[int]] = []
|
||||
for request in scheduler_output.scheduled_new_reqs:
|
||||
assert request.req_id not in seen_request_prompt_length
|
||||
seen_request_prompt_length[request.req_id] = len(request.prompt_token_ids or [])
|
||||
if request.num_computed_tokens >= seen_request_prompt_length[request.req_id]:
|
||||
num_tokens = random.randint(*num_output_tokens_range)
|
||||
request_ids.append(request.req_id)
|
||||
sampled_token_ids.append(
|
||||
[random.randint(0, 100) for _ in range(num_tokens)]
|
||||
)
|
||||
for req_id, num_computed_tokens in zip(
|
||||
scheduler_output.scheduled_cached_reqs.req_ids,
|
||||
scheduler_output.scheduled_cached_reqs.num_computed_tokens,
|
||||
):
|
||||
if num_computed_tokens >= seen_request_prompt_length[req_id]:
|
||||
num_tokens = random.randint(*num_output_tokens_range)
|
||||
request_ids.append(req_id)
|
||||
sampled_token_ids.append(
|
||||
[random.randint(0, 100) for _ in range(num_tokens)]
|
||||
)
|
||||
return DraftTokenIds(req_ids=request_ids, draft_token_ids=sampled_token_ids)
|
||||
|
||||
|
||||
def _check_valid_scheduler_output(
|
||||
scheduler_output: SchedulerOutput,
|
||||
seen_request_ids: set[str],
|
||||
seen_mm_hashes: set[str],
|
||||
):
|
||||
for req in scheduler_output.scheduled_new_reqs:
|
||||
assert req.req_id not in seen_request_ids
|
||||
seen_request_ids.add(req.req_id)
|
||||
for req_id in scheduler_output.scheduled_cached_reqs.req_ids:
|
||||
assert req_id in seen_request_ids
|
||||
|
||||
req_ids = set[str]()
|
||||
req_ids.update(req.req_id for req in scheduler_output.scheduled_new_reqs)
|
||||
req_ids.update(scheduler_output.scheduled_cached_reqs.req_ids)
|
||||
|
||||
assert set(scheduler_output.num_scheduled_tokens.keys()) == req_ids
|
||||
assert (
|
||||
sum(scheduler_output.num_scheduled_tokens.values())
|
||||
== scheduler_output.total_num_scheduled_tokens
|
||||
)
|
||||
|
||||
assert set(scheduler_output.scheduled_spec_decode_tokens.keys()) <= req_ids
|
||||
assert set(scheduler_output.scheduled_encoder_inputs.keys()) <= req_ids
|
||||
|
||||
for req in scheduler_output.scheduled_new_reqs:
|
||||
for mm_feature in req.mm_features:
|
||||
seen_mm_hashes.add(mm_feature.identifier)
|
||||
for mm_hash in scheduler_output.free_encoder_mm_hashes:
|
||||
assert mm_hash in seen_mm_hashes
|
||||
|
||||
assert scheduler_output.finished_req_ids <= seen_request_ids
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_prefix_caching", [True, False])
|
||||
@pytest.mark.parametrize("num_speculative_tokens", [None, 1, 5])
|
||||
@pytest.mark.parametrize(
|
||||
("max_input_tokens", "max_output_tokens", "max_num_seqs", "num_blocks"),
|
||||
[
|
||||
# Standard profile
|
||||
(5000, 500, 256, 10000),
|
||||
# Generation heavy + high max_num_seqs + low num_blocks -> Many preemptions
|
||||
(500, 5000, 1024, 1000),
|
||||
],
|
||||
ids=["standard", "preemption"],
|
||||
)
|
||||
def test_priority_scheduling_blast(
|
||||
enable_prefix_caching: bool,
|
||||
num_speculative_tokens: int | None,
|
||||
max_input_tokens: int,
|
||||
max_output_tokens: int,
|
||||
max_num_seqs: int,
|
||||
num_blocks: int,
|
||||
):
|
||||
random.seed(42)
|
||||
seen_request_prompt_length = dict[str, int]()
|
||||
seen_request_ids = set[str]()
|
||||
seen_mm_hashes = set[str]()
|
||||
|
||||
scheduler = create_scheduler_with_priority(
|
||||
model="Qwen/Qwen2.5-VL-3B-Instruct",
|
||||
max_num_seqs=max_num_seqs,
|
||||
enable_prefix_caching=enable_prefix_caching,
|
||||
num_blocks=num_blocks,
|
||||
num_speculative_tokens=num_speculative_tokens,
|
||||
)
|
||||
|
||||
num_initial_requests = 10
|
||||
for _ in range(num_initial_requests):
|
||||
req = _create_random_request(
|
||||
max_tokens_range=(1, max_output_tokens),
|
||||
num_tokens_range=(1, max_input_tokens),
|
||||
arrival_time_range=(0, 1),
|
||||
priority_range=(-3, 3),
|
||||
num_mm_item_range=(0, 2),
|
||||
vllm_config=scheduler.vllm_config,
|
||||
)
|
||||
scheduler.add_request(req)
|
||||
num_initial_requests = 2
|
||||
for _ in range(num_initial_requests):
|
||||
req = _create_random_request(
|
||||
max_tokens_range=(1, max_output_tokens),
|
||||
num_tokens_range=(1, max_input_tokens),
|
||||
arrival_time_range=(0, 0),
|
||||
priority_range=(4, 4),
|
||||
num_mm_item_range=(0, 2),
|
||||
vllm_config=scheduler.vllm_config,
|
||||
)
|
||||
scheduler.add_request(req)
|
||||
for _ in range(20000):
|
||||
if len(scheduler.waiting) == 0:
|
||||
num_new_requests = random.randint(0, 2)
|
||||
for _ in range(num_new_requests):
|
||||
req = _create_random_request(
|
||||
max_tokens_range=(1, max_output_tokens),
|
||||
num_tokens_range=(1, max_input_tokens),
|
||||
arrival_time_range=(0, 1),
|
||||
priority_range=(-3, 3),
|
||||
num_mm_item_range=(0, 2),
|
||||
vllm_config=scheduler.vllm_config,
|
||||
)
|
||||
scheduler.add_request(req)
|
||||
scheduler_output = scheduler.schedule()
|
||||
_check_valid_scheduler_output(
|
||||
scheduler_output, seen_request_ids, seen_mm_hashes
|
||||
)
|
||||
model_output = _mock_execute_model(
|
||||
scheduler_output,
|
||||
num_output_tokens_range=(1, 1 + (num_speculative_tokens or 0)),
|
||||
)
|
||||
scheduler.update_from_output(scheduler_output, model_output)
|
||||
if num_speculative_tokens is not None:
|
||||
scheduler.update_draft_token_ids(
|
||||
_mock_draft_token_ids(
|
||||
scheduler_output,
|
||||
(0, num_speculative_tokens),
|
||||
seen_request_prompt_length,
|
||||
)
|
||||
)
|
||||
290
third_party/vllm/tests/v1/core/test_repetition_detection.py
vendored
Normal file
290
third_party/vllm/tests/v1/core/test_repetition_detection.py
vendored
Normal file
@@ -0,0 +1,290 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
|
||||
from vllm.sampling_params import RepetitionDetectionParams, SamplingParams
|
||||
from vllm.v1.core.sched.utils import check_sequence_repetition, check_stop
|
||||
from vllm.v1.request import Request, RequestStatus
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
|
||||
# ============================================================================
|
||||
# UNIT TESTS - check_sequence_repetition function
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestCheckSequenceRepetition:
|
||||
"""Unit tests for the check_sequence_repetition function"""
|
||||
|
||||
def test_simple_repetition_detected(self):
|
||||
"""Test detection of simple repetitive patterns"""
|
||||
token_ids = [1, 2, 3, 1, 2, 3, 1, 2, 3]
|
||||
params = RepetitionDetectionParams(
|
||||
max_pattern_size=3,
|
||||
min_pattern_size=2,
|
||||
min_count=3,
|
||||
)
|
||||
assert check_sequence_repetition(token_ids, params)
|
||||
|
||||
def test_repetition_below_min_count(self):
|
||||
"""Test that pattern below min_count is not detected"""
|
||||
token_ids = [1, 2, 3, 1, 2, 3]
|
||||
params = RepetitionDetectionParams(
|
||||
max_pattern_size=3,
|
||||
min_pattern_size=2,
|
||||
min_count=3,
|
||||
)
|
||||
assert not check_sequence_repetition(token_ids, params)
|
||||
|
||||
def test_two_token_pattern(self):
|
||||
"""Test detection of 2-token patterns"""
|
||||
token_ids = [1, 2, 1, 2, 1, 2, 1, 2]
|
||||
params = RepetitionDetectionParams(
|
||||
max_pattern_size=5,
|
||||
min_pattern_size=2,
|
||||
min_count=4,
|
||||
)
|
||||
assert check_sequence_repetition(token_ids, params)
|
||||
|
||||
def test_no_repetition_varied_sequence(self):
|
||||
"""Test that non-repetitive sequences are not flagged"""
|
||||
token_ids = [1, 2, 3, 4, 5, 6, 7, 8, 9]
|
||||
params = RepetitionDetectionParams(
|
||||
max_pattern_size=5,
|
||||
min_pattern_size=2,
|
||||
min_count=2,
|
||||
)
|
||||
assert not check_sequence_repetition(token_ids, params)
|
||||
|
||||
def test_partial_repetition_not_detected(self):
|
||||
"""Test that incomplete repetitions are not detected"""
|
||||
token_ids = [1, 2, 3, 1, 2, 3, 1, 2, 4]
|
||||
params = RepetitionDetectionParams(
|
||||
max_pattern_size=3,
|
||||
min_pattern_size=2,
|
||||
min_count=3,
|
||||
)
|
||||
assert not check_sequence_repetition(token_ids, params)
|
||||
|
||||
def test_empty_token_list(self):
|
||||
"""Test with empty token list"""
|
||||
params = RepetitionDetectionParams(
|
||||
max_pattern_size=3,
|
||||
min_pattern_size=2,
|
||||
min_count=2,
|
||||
)
|
||||
assert not check_sequence_repetition([], params)
|
||||
|
||||
def test_detection_disabled_max_size_zero(self):
|
||||
"""Test that zero max_pattern_size disables detection"""
|
||||
token_ids = [1, 2, 1, 2, 1, 2]
|
||||
params = RepetitionDetectionParams()
|
||||
assert not check_sequence_repetition(token_ids, params)
|
||||
|
||||
def test_invalid_min_count(self):
|
||||
"""Test that min_count < 2 returns False"""
|
||||
token_ids = [1, 2, 1, 2]
|
||||
params = RepetitionDetectionParams()
|
||||
assert not check_sequence_repetition(token_ids, params)
|
||||
|
||||
def test_repetition_at_end_of_sequence(self):
|
||||
"""Test detection when repetition occurs at the end"""
|
||||
token_ids = [1, 2, 3, 4, 5, 6, 5, 6, 5, 6]
|
||||
params = RepetitionDetectionParams(
|
||||
max_pattern_size=3,
|
||||
min_pattern_size=2,
|
||||
min_count=3,
|
||||
)
|
||||
assert check_sequence_repetition(token_ids, params)
|
||||
|
||||
def test_large_pattern_many_repetitions(self):
|
||||
"""Test large pattern repeated many times"""
|
||||
token_ids = [1, 2, 3, 4, 5, 6, 7, 8] * 5
|
||||
params = RepetitionDetectionParams(
|
||||
max_pattern_size=10,
|
||||
min_pattern_size=2,
|
||||
min_count=3,
|
||||
)
|
||||
assert check_sequence_repetition(token_ids, params)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# INTEGRATION TESTS - check_stop with repetition detection
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestRepetitionDetectionIntegration:
|
||||
"""Integration tests for repetition detection in check_stop"""
|
||||
|
||||
def test_basic_repetition_stops_generation(self):
|
||||
"""Test that repetition is detected and stops generation"""
|
||||
params = SamplingParams(
|
||||
max_tokens=100,
|
||||
repetition_detection=RepetitionDetectionParams(
|
||||
max_pattern_size=5,
|
||||
min_pattern_size=2,
|
||||
min_count=3,
|
||||
),
|
||||
)
|
||||
request = Request(
|
||||
request_id="test",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
sampling_params=params,
|
||||
pooling_params=None,
|
||||
)
|
||||
request.append_output_token_ids([10, 20, 10, 20, 10, 20])
|
||||
assert check_stop(request, max_model_len=1024)
|
||||
assert request.status == RequestStatus.FINISHED_REPETITION
|
||||
assert request.stop_reason == "repetition_detected"
|
||||
|
||||
def test_detection_disabled_no_stop(self):
|
||||
"""Test that disabled detection doesn't stop generation"""
|
||||
params = SamplingParams(
|
||||
max_tokens=100,
|
||||
)
|
||||
request = Request(
|
||||
request_id="test",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
sampling_params=params,
|
||||
pooling_params=None,
|
||||
)
|
||||
request.append_output_token_ids([10, 20, 10, 20, 10, 20])
|
||||
assert not check_stop(request, max_model_len=1024)
|
||||
|
||||
def test_repetition_respects_min_tokens(self):
|
||||
"""Test that repetition detection respects min_tokens"""
|
||||
params = SamplingParams(
|
||||
min_tokens=10,
|
||||
max_tokens=100,
|
||||
repetition_detection=RepetitionDetectionParams(
|
||||
max_pattern_size=5,
|
||||
min_pattern_size=2,
|
||||
min_count=3,
|
||||
),
|
||||
)
|
||||
request = Request(
|
||||
request_id="test",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
sampling_params=params,
|
||||
pooling_params=None,
|
||||
)
|
||||
request.append_output_token_ids([10, 20, 10, 20, 10, 20])
|
||||
assert not check_stop(request, max_model_len=1024)
|
||||
|
||||
def test_no_repetition_continues_generation(self):
|
||||
"""Test that non-repetitive tokens don't stop generation"""
|
||||
params = SamplingParams(
|
||||
max_tokens=100,
|
||||
repetition_detection=RepetitionDetectionParams(
|
||||
max_pattern_size=5,
|
||||
min_pattern_size=2,
|
||||
min_count=3,
|
||||
),
|
||||
)
|
||||
request = Request(
|
||||
request_id="test",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
sampling_params=params,
|
||||
pooling_params=None,
|
||||
)
|
||||
request.append_output_token_ids([10, 20, 30, 40, 50, 60])
|
||||
assert not check_stop(request, max_model_len=1024)
|
||||
|
||||
def test_pattern_at_size_boundary(self):
|
||||
"""Test detection at exact pattern size boundary"""
|
||||
params = SamplingParams(
|
||||
max_tokens=100,
|
||||
repetition_detection=RepetitionDetectionParams(
|
||||
max_pattern_size=3,
|
||||
min_pattern_size=3,
|
||||
min_count=2,
|
||||
),
|
||||
)
|
||||
request = Request(
|
||||
request_id="test",
|
||||
prompt_token_ids=[1, 2],
|
||||
sampling_params=params,
|
||||
pooling_params=None,
|
||||
)
|
||||
request.append_output_token_ids([10, 20, 30, 10, 20, 30])
|
||||
assert check_stop(request, max_model_len=1024)
|
||||
assert request.status == RequestStatus.FINISHED_REPETITION
|
||||
|
||||
def test_multiple_pattern_sizes_checked(self):
|
||||
"""Test that function checks pattern sizes in range"""
|
||||
params = SamplingParams(
|
||||
max_tokens=100,
|
||||
repetition_detection=RepetitionDetectionParams(
|
||||
max_pattern_size=5,
|
||||
min_pattern_size=2,
|
||||
min_count=3,
|
||||
),
|
||||
)
|
||||
request = Request(
|
||||
request_id="test",
|
||||
prompt_token_ids=[1],
|
||||
sampling_params=params,
|
||||
pooling_params=None,
|
||||
)
|
||||
request.append_output_token_ids([7, 8, 9, 10, 7, 8, 9, 10, 7, 8, 9, 10])
|
||||
assert check_stop(request, max_model_len=1024)
|
||||
assert request.status == RequestStatus.FINISHED_REPETITION
|
||||
|
||||
def test_eos_takes_precedence_over_repetition(self):
|
||||
"""Test that EOS token stops before repetition check"""
|
||||
params = SamplingParams(
|
||||
max_tokens=100,
|
||||
stop_token_ids=[999],
|
||||
repetition_detection=RepetitionDetectionParams(
|
||||
max_pattern_size=5,
|
||||
min_pattern_size=2,
|
||||
min_count=3,
|
||||
),
|
||||
)
|
||||
request = Request(
|
||||
request_id="test",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
sampling_params=params,
|
||||
pooling_params=None,
|
||||
)
|
||||
request.append_output_token_ids([10, 20, 10, 20, 999])
|
||||
assert check_stop(request, max_model_len=1024)
|
||||
assert request.status == RequestStatus.FINISHED_STOPPED
|
||||
|
||||
def test_min_pattern_size_filters_small_patterns(self):
|
||||
"""Test that min_pattern_size filters out smaller patterns"""
|
||||
params = SamplingParams(
|
||||
max_tokens=100,
|
||||
repetition_detection=RepetitionDetectionParams(
|
||||
max_pattern_size=5,
|
||||
min_pattern_size=3,
|
||||
min_count=3,
|
||||
),
|
||||
)
|
||||
request = Request(
|
||||
request_id="test",
|
||||
prompt_token_ids=[1],
|
||||
sampling_params=params,
|
||||
pooling_params=None,
|
||||
)
|
||||
request.append_output_token_ids([10, 20, 10, 20, 10, 20])
|
||||
assert not check_stop(request, max_model_len=1024)
|
||||
|
||||
def test_high_repetition_threshold(self):
|
||||
"""Test that high min_count requires many repetitions"""
|
||||
params = SamplingParams(
|
||||
max_tokens=100,
|
||||
repetition_detection=RepetitionDetectionParams(
|
||||
max_pattern_size=5,
|
||||
min_pattern_size=2,
|
||||
min_count=5,
|
||||
),
|
||||
)
|
||||
request = Request(
|
||||
request_id="test",
|
||||
prompt_token_ids=[1],
|
||||
sampling_params=params,
|
||||
pooling_params=None,
|
||||
)
|
||||
request.append_output_token_ids([10, 20, 10, 20, 10, 20])
|
||||
assert not check_stop(request, max_model_len=1024)
|
||||
69
third_party/vllm/tests/v1/core/test_reset_prefix_cache_e2e.py
vendored
Normal file
69
third_party/vllm/tests/v1/core/test_reset_prefix_cache_e2e.py
vendored
Normal file
@@ -0,0 +1,69 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from vllm import EngineArgs, LLMEngine, SamplingParams
|
||||
|
||||
PROMPTS = [
|
||||
"A robot may not injure a human being ",
|
||||
"To be or not to be,",
|
||||
"What is the meaning of life?",
|
||||
"What does the fox say? " * 20, # Test long prompt
|
||||
]
|
||||
|
||||
|
||||
def test_reset_prefix_cache_e2e(monkeypatch):
|
||||
# "spawn" is required for test to be deterministic
|
||||
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
engine_args = EngineArgs(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
gpu_memory_utilization=0.2,
|
||||
async_scheduling=True,
|
||||
max_num_batched_tokens=32,
|
||||
max_model_len=2048,
|
||||
compilation_config={"mode": 0},
|
||||
dtype="float16",
|
||||
)
|
||||
engine = LLMEngine.from_engine_args(engine_args)
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.0,
|
||||
max_tokens=16,
|
||||
)
|
||||
|
||||
# No preempt case:
|
||||
for i, prompt in enumerate(PROMPTS):
|
||||
engine.add_request("ground_truth_" + str(i), prompt, sampling_params)
|
||||
|
||||
ground_truth_results = {}
|
||||
while engine.has_unfinished_requests():
|
||||
request_outputs = engine.step()
|
||||
for request_output in request_outputs:
|
||||
if request_output.finished:
|
||||
ground_truth_results[request_output.request_id] = request_output
|
||||
|
||||
# Preempt case:
|
||||
for i, prompt in enumerate(PROMPTS):
|
||||
engine.add_request("preempted_" + str(i), prompt, sampling_params)
|
||||
|
||||
step_id = 0
|
||||
preempted_results = {}
|
||||
while engine.has_unfinished_requests():
|
||||
if step_id == 10:
|
||||
engine.reset_prefix_cache(reset_running_requests=True)
|
||||
|
||||
request_outputs = engine.step()
|
||||
|
||||
for request_output in request_outputs:
|
||||
if request_output.finished:
|
||||
preempted_results[request_output.request_id] = request_output
|
||||
step_id += 1
|
||||
|
||||
for i in range(len(PROMPTS)):
|
||||
assert (
|
||||
ground_truth_results["ground_truth_" + str(i)].outputs[0].text
|
||||
== preempted_results["preempted_" + str(i)].outputs[0].text
|
||||
), (
|
||||
f"ground_truth_results['ground_truth_{i}'].outputs[0].text="
|
||||
f"{ground_truth_results['ground_truth_' + str(i)].outputs[0].text} "
|
||||
f"preempted_results['preempted_{i}'].outputs[0].text="
|
||||
f"{preempted_results['preempted_' + str(i)].outputs[0].text}"
|
||||
)
|
||||
4142
third_party/vllm/tests/v1/core/test_scheduler.py
vendored
Normal file
4142
third_party/vllm/tests/v1/core/test_scheduler.py
vendored
Normal file
File diff suppressed because it is too large
Load Diff
37
third_party/vllm/tests/v1/core/test_scheduler_e2e.py
vendored
Normal file
37
third_party/vllm/tests/v1/core/test_scheduler_e2e.py
vendored
Normal file
@@ -0,0 +1,37 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm import LLM
|
||||
|
||||
MODEL = "hmellor/tiny-random-LlamaForCausalLM"
|
||||
PROMPT = "Hello my name is Robert and I"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def llm() -> LLM:
|
||||
return LLM(
|
||||
MODEL,
|
||||
enforce_eager=True,
|
||||
enable_prefix_caching=True,
|
||||
long_prefill_token_threshold=2,
|
||||
max_num_batched_tokens=6,
|
||||
max_num_seqs=3,
|
||||
block_size=16,
|
||||
)
|
||||
|
||||
|
||||
def test_concurrent_partial_prefill(llm):
|
||||
outputs = llm.generate([PROMPT] * 3)
|
||||
assert len(outputs) == 3
|
||||
for output in outputs:
|
||||
assert len(output.outputs) == 1
|
||||
|
||||
|
||||
def test_prefix_cache_stats_is_recorded(llm):
|
||||
# 17 tokens will make sure first 16 tokens are cached in a block
|
||||
input_tokens = {"prompt_token_ids": [101] * 17}
|
||||
_ = llm.generate([input_tokens])
|
||||
outputs = llm.generate([input_tokens])
|
||||
assert outputs[0].num_cached_tokens == 16
|
||||
431
third_party/vllm/tests/v1/core/test_single_type_kv_cache_manager.py
vendored
Normal file
431
third_party/vllm/tests/v1/core/test_single_type_kv_cache_manager.py
vendored
Normal file
@@ -0,0 +1,431 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.v1.core.block_pool import BlockPool
|
||||
from vllm.v1.core.kv_cache_utils import (
|
||||
BlockHash,
|
||||
KVCacheBlock,
|
||||
make_block_hash_with_group_id,
|
||||
)
|
||||
from vllm.v1.core.single_type_kv_cache_manager import (
|
||||
ChunkedLocalAttentionManager,
|
||||
SlidingWindowManager,
|
||||
)
|
||||
from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec, SlidingWindowSpec
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
|
||||
|
||||
def get_sliding_window_manager(sliding_window_spec, block_pool, enable_caching=True):
|
||||
return SlidingWindowManager(
|
||||
sliding_window_spec,
|
||||
block_pool=block_pool,
|
||||
enable_caching=enable_caching,
|
||||
kv_cache_group_id=0,
|
||||
)
|
||||
|
||||
|
||||
def get_chunked_local_attention_manager(
|
||||
chunked_local_attention_spec, block_pool, enable_caching=True
|
||||
):
|
||||
return ChunkedLocalAttentionManager(
|
||||
chunked_local_attention_spec,
|
||||
block_pool=block_pool,
|
||||
enable_caching=enable_caching,
|
||||
kv_cache_group_id=0,
|
||||
)
|
||||
|
||||
|
||||
def test_chunked_local_attention_possible_cached_prefix():
|
||||
block_size = 2
|
||||
chunked_local_attention_spec = ChunkedLocalAttentionSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
attention_chunk_size=4,
|
||||
)
|
||||
|
||||
block_pool = BlockPool(
|
||||
num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size
|
||||
)
|
||||
manager = get_chunked_local_attention_manager(
|
||||
chunked_local_attention_spec, block_pool
|
||||
)
|
||||
|
||||
def run_one_case(block_is_cached, tail_token, expect_length):
|
||||
block_hash_list = [
|
||||
BlockHash(str(i).encode()) for i in range(len(block_is_cached))
|
||||
]
|
||||
|
||||
block_pool.cached_block_hash_to_block._cache.clear()
|
||||
|
||||
# Mock the block pool with the cached blocks
|
||||
for i, (block_hash, is_cached) in enumerate(
|
||||
zip(block_hash_list, block_is_cached)
|
||||
):
|
||||
if is_cached:
|
||||
block_pool.cached_block_hash_to_block.insert(
|
||||
make_block_hash_with_group_id(block_hash, 0),
|
||||
block_pool.blocks[i + 10],
|
||||
)
|
||||
|
||||
computed_blocks = manager.find_longest_cache_hit(
|
||||
block_hashes=block_hash_list,
|
||||
max_length=len(block_hash_list) * block_size + tail_token,
|
||||
kv_cache_group_ids=[0],
|
||||
block_pool=block_pool,
|
||||
kv_cache_spec=chunked_local_attention_spec,
|
||||
use_eagle=False,
|
||||
alignment_tokens=block_size,
|
||||
)[0]
|
||||
assert len(computed_blocks) == expect_length
|
||||
|
||||
assert all(
|
||||
block == block_pool.null_block
|
||||
for block in computed_blocks[: (expect_length - 1) // 2]
|
||||
)
|
||||
|
||||
run_one_case([True], 0, 1)
|
||||
run_one_case([True], 1, 1)
|
||||
run_one_case([True, False], 0, 2)
|
||||
run_one_case([True, False], 1, 2)
|
||||
run_one_case([True, True], 0, 2)
|
||||
run_one_case([True, True], 1, 2)
|
||||
run_one_case([True, True, False], 0, 2)
|
||||
run_one_case([True, True, False], 1, 2)
|
||||
run_one_case([True, True, True], 0, 3)
|
||||
run_one_case([True, True, True], 1, 3)
|
||||
run_one_case([True, True, True, False], 0, 4)
|
||||
run_one_case([True, True, True, False], 1, 4)
|
||||
run_one_case([random.choice([True, False])] * 8 + [True], 1, 9)
|
||||
run_one_case([random.choice([True, False])] * 8 + [False], 1, 8)
|
||||
run_one_case([random.choice([True, False])] * 8 + [True, True], 1, 10)
|
||||
run_one_case([random.choice([True, False])] * 8 + [True, False], 0, 10)
|
||||
run_one_case([random.choice([True, False])] * 8 + [True, False], 1, 10)
|
||||
run_one_case([random.choice([True, False])] * 8 + [False, True], 0, 10)
|
||||
run_one_case([random.choice([True, False])] * 8 + [False, True], 1, 10)
|
||||
run_one_case([random.choice([True, False])] * 8 + [False, False], 0, 10)
|
||||
run_one_case([random.choice([True, False])] * 8 + [False, False], 1, 10)
|
||||
|
||||
|
||||
def test_sliding_window_possible_cached_prefix():
|
||||
block_size = 2
|
||||
sliding_window_spec = SlidingWindowSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
sliding_window=4,
|
||||
)
|
||||
|
||||
block_pool = BlockPool(
|
||||
num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size
|
||||
)
|
||||
manager = get_sliding_window_manager(sliding_window_spec, block_pool)
|
||||
|
||||
def run_one_case(block_is_cached, expect_length):
|
||||
block_hash_list = [
|
||||
BlockHash(str(i).encode()) for i in range(len(block_is_cached))
|
||||
]
|
||||
|
||||
block_pool.cached_block_hash_to_block._cache.clear()
|
||||
|
||||
# Mock the block pool with the cached blocks
|
||||
for i, (block_hash, is_cached) in enumerate(
|
||||
zip(block_hash_list, block_is_cached)
|
||||
):
|
||||
if is_cached:
|
||||
block_pool.cached_block_hash_to_block.insert(
|
||||
make_block_hash_with_group_id(block_hash, 0),
|
||||
block_pool.blocks[i + 10],
|
||||
)
|
||||
|
||||
computed_blocks = manager.find_longest_cache_hit(
|
||||
block_hashes=block_hash_list,
|
||||
max_length=len(block_hash_list) * block_size,
|
||||
kv_cache_group_ids=[0],
|
||||
block_pool=block_pool,
|
||||
kv_cache_spec=sliding_window_spec,
|
||||
use_eagle=False,
|
||||
alignment_tokens=block_size,
|
||||
)[0]
|
||||
assert len(computed_blocks) == expect_length
|
||||
|
||||
assert all(
|
||||
block == block_pool.null_block
|
||||
for block in computed_blocks[: expect_length - 2]
|
||||
)
|
||||
for i in range(2):
|
||||
if i < expect_length:
|
||||
block_index = expect_length - i - 1
|
||||
assert computed_blocks[block_index].block_id == block_index + 10
|
||||
|
||||
run_one_case([False] * 10, 0)
|
||||
run_one_case([True], 1)
|
||||
run_one_case([True, False], 1)
|
||||
run_one_case([True, True], 2)
|
||||
run_one_case([True, True, False], 2)
|
||||
run_one_case([True, True, True], 3)
|
||||
run_one_case([True, True, True, False], 3)
|
||||
run_one_case(
|
||||
[True, True, False, True, False, False, True, True, False, True, True, True], 12
|
||||
)
|
||||
run_one_case(
|
||||
[True, True, False, True, False, False, True, True, False, False, False], 8
|
||||
)
|
||||
run_one_case(
|
||||
[True, True, False, True, False, False, True, True, False, False, False, True],
|
||||
8,
|
||||
)
|
||||
|
||||
|
||||
def test_chunked_local_attention_remove_skipped_blocks():
|
||||
attention_spec = ChunkedLocalAttentionSpec(
|
||||
block_size=2,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
attention_chunk_size=4,
|
||||
)
|
||||
|
||||
block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True, hash_block_size=2)
|
||||
|
||||
manager = get_chunked_local_attention_manager(attention_spec, block_pool)
|
||||
|
||||
null_block_id = block_pool.null_block.block_id
|
||||
|
||||
def id_to_block_table(ids) -> list[KVCacheBlock]:
|
||||
return [
|
||||
KVCacheBlock(id_) if id_ != null_block_id else block_pool.null_block
|
||||
for id_ in ids
|
||||
]
|
||||
|
||||
def assert_block_id(block_table: list[KVCacheBlock], ids: list[int]):
|
||||
for block, id_ in zip(block_table, ids):
|
||||
if id_ == null_block_id:
|
||||
assert block == block_pool.null_block
|
||||
else:
|
||||
assert block.block_id == id_
|
||||
|
||||
original_block_ids = [
|
||||
1000,
|
||||
1001,
|
||||
1002,
|
||||
1003,
|
||||
1004,
|
||||
1005,
|
||||
1006,
|
||||
1007,
|
||||
1008,
|
||||
1009,
|
||||
1010,
|
||||
]
|
||||
block_table = id_to_block_table(original_block_ids)
|
||||
manager.req_to_blocks["test"] = block_table
|
||||
|
||||
manager.remove_skipped_blocks("test", 0)
|
||||
assert_block_id(block_table, original_block_ids)
|
||||
|
||||
# For 4th token (0-indexed), token 0-3 is out of the local attention window.
|
||||
manager.remove_skipped_blocks("test", 4)
|
||||
assert_block_id(block_table, [null_block_id] * 2)
|
||||
|
||||
# For 6th token (0-indexed), token 4 - 6 are in local attention window,
|
||||
# token 0 - 3 are out, 2 blocks can be removed.
|
||||
manager.remove_skipped_blocks("test", 6)
|
||||
assert_block_id(block_table, [null_block_id] * 2 + original_block_ids[2:])
|
||||
# For 12th token (0-indexed),
|
||||
# token 0-11 are out, 6 block can be removed.
|
||||
manager.remove_skipped_blocks("test", 12)
|
||||
assert_block_id(block_table, [null_block_id] * 6)
|
||||
|
||||
|
||||
def test_sliding_window_remove_skipped_blocks():
|
||||
sliding_window_spec = SlidingWindowSpec(
|
||||
block_size=2,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
sliding_window=4,
|
||||
)
|
||||
|
||||
block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True, hash_block_size=2)
|
||||
|
||||
manager = get_sliding_window_manager(sliding_window_spec, block_pool)
|
||||
|
||||
null_block_id = block_pool.null_block.block_id
|
||||
|
||||
def id_to_block_table(ids) -> list[KVCacheBlock]:
|
||||
return [
|
||||
KVCacheBlock(id_) if id_ != null_block_id else block_pool.null_block
|
||||
for id_ in ids
|
||||
]
|
||||
|
||||
def assert_block_id(block_table: list[KVCacheBlock], ids: list[int]):
|
||||
for block, id_ in zip(block_table, ids):
|
||||
if id_ == null_block_id:
|
||||
assert block == block_pool.null_block
|
||||
else:
|
||||
assert block.block_id == id_
|
||||
|
||||
original_block_ids = [
|
||||
1000,
|
||||
1001,
|
||||
1002,
|
||||
1003,
|
||||
1004,
|
||||
1005,
|
||||
1006,
|
||||
1007,
|
||||
1008,
|
||||
1009,
|
||||
1010,
|
||||
]
|
||||
block_table = id_to_block_table(original_block_ids)
|
||||
manager.req_to_blocks["test"] = block_table
|
||||
|
||||
manager.remove_skipped_blocks("test", 0)
|
||||
assert_block_id(block_table, original_block_ids)
|
||||
|
||||
# 4 tokens are computed. Only token 0 is out of the sliding window. As
|
||||
# block 1000 also contains token 1 that is in the sliding window, block 1000
|
||||
# cannot be removed.
|
||||
manager.remove_skipped_blocks("test", 4)
|
||||
assert_block_id(block_table, original_block_ids)
|
||||
|
||||
# 5 tokens are computed. Token 0 & 1 are out of the sliding window.
|
||||
# Block 1000 can be removed.
|
||||
manager.remove_skipped_blocks("test", 5)
|
||||
assert_block_id(block_table, [null_block_id] + original_block_ids[1:])
|
||||
|
||||
# 6 tokens are computed. Token 0-2 are out of the sliding window.
|
||||
# Cannot remove new block as the block 1001 is still used by token 3.
|
||||
manager.remove_skipped_blocks("test", 6)
|
||||
assert_block_id(block_table, [null_block_id] + original_block_ids[1:])
|
||||
|
||||
# 7 tokens are computed. Token 0-3 are out of the sliding window.
|
||||
# Block 1001 can be removed and block 1000 is already removed.
|
||||
manager.remove_skipped_blocks("test", 7)
|
||||
assert_block_id(block_table, [null_block_id] * 2 + original_block_ids[2:])
|
||||
|
||||
# 11 tokens are computed. Token 0-7 are out of the sliding window.
|
||||
# Block 1002 & 1003 can be removed now. Block 1003 represents a longer
|
||||
# sequence, and is expected to be evicted earlier than 1002, so the order
|
||||
# of removed blocks should be [1003, 1002].
|
||||
manager.remove_skipped_blocks("test", 11)
|
||||
assert_block_id(block_table, [null_block_id] * 4 + original_block_ids[4:])
|
||||
|
||||
|
||||
def test_get_num_blocks_to_allocate():
|
||||
block_size = 2
|
||||
sliding_window_spec = SlidingWindowSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
sliding_window=4, # Placeholder value, not related to test result
|
||||
)
|
||||
|
||||
block_pool = BlockPool(
|
||||
num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size
|
||||
)
|
||||
manager = get_sliding_window_manager(sliding_window_spec, block_pool)
|
||||
cached_blocks_1 = [KVCacheBlock(i + 1) for i in range(10)]
|
||||
cached_blocks_2 = [block_pool.null_block for _ in range(5)] + [
|
||||
KVCacheBlock(i + 1) for i in range(5)
|
||||
]
|
||||
|
||||
assert (
|
||||
manager.get_num_blocks_to_allocate(
|
||||
"1", 20 * block_size, cached_blocks_1, 0, 20 * block_size
|
||||
)
|
||||
== 20
|
||||
)
|
||||
assert (
|
||||
manager.get_num_blocks_to_allocate(
|
||||
"2", 20 * block_size, cached_blocks_2, 0, 20 * block_size
|
||||
)
|
||||
== 15
|
||||
)
|
||||
|
||||
|
||||
def test_evictable_cached_blocks_not_double_allocated():
|
||||
block_size = 2
|
||||
sliding_window_length = 2 * block_size
|
||||
sliding_window_spec = SlidingWindowSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
sliding_window=sliding_window_length,
|
||||
)
|
||||
|
||||
block_pool = BlockPool(
|
||||
num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size
|
||||
)
|
||||
manager = get_sliding_window_manager(sliding_window_spec, block_pool)
|
||||
|
||||
request_id = "req"
|
||||
evictable_block = block_pool.blocks[1] # ref_cnt == 0, eviction candidate
|
||||
|
||||
num_blocks_to_allocate = manager.get_num_blocks_to_allocate(
|
||||
request_id=request_id,
|
||||
num_tokens=2 * block_size,
|
||||
new_computed_blocks=[evictable_block],
|
||||
total_computed_tokens=block_size,
|
||||
num_tokens_main_model=2 * block_size,
|
||||
)
|
||||
# Free capacity check should count evictable cached blocks, but allocation
|
||||
# should only allocate the truly new block.
|
||||
assert num_blocks_to_allocate == 2
|
||||
|
||||
manager.allocate_new_computed_blocks(
|
||||
request_id,
|
||||
[evictable_block],
|
||||
num_local_computed_tokens=block_size,
|
||||
num_external_computed_tokens=0,
|
||||
)
|
||||
new_blocks = manager.allocate_new_blocks(
|
||||
request_id, num_tokens=4, num_tokens_main_model=4
|
||||
)
|
||||
assert len(new_blocks) == 1
|
||||
assert len(manager.req_to_blocks[request_id]) == 2
|
||||
|
||||
|
||||
def test_chunked_local_attention_get_num_blocks_to_allocate():
|
||||
block_size = 2
|
||||
attention_spec = ChunkedLocalAttentionSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
attention_chunk_size=4, # Placeholder value, not related to test result
|
||||
)
|
||||
|
||||
block_pool = BlockPool(
|
||||
num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size
|
||||
)
|
||||
manager = get_chunked_local_attention_manager(attention_spec, block_pool)
|
||||
cached_blocks_1 = [KVCacheBlock(i + 1) for i in range(10)]
|
||||
cached_blocks_2 = [block_pool.null_block for _ in range(5)] + [
|
||||
KVCacheBlock(i + 1) for i in range(5)
|
||||
]
|
||||
|
||||
assert (
|
||||
manager.get_num_blocks_to_allocate(
|
||||
"1", 20 * block_size, cached_blocks_1, 0, 20 * block_size
|
||||
)
|
||||
== 20
|
||||
)
|
||||
assert (
|
||||
manager.get_num_blocks_to_allocate(
|
||||
"2", 20 * block_size, cached_blocks_2, 0, 20 * block_size
|
||||
)
|
||||
== 15
|
||||
)
|
||||
257
third_party/vllm/tests/v1/core/utils.py
vendored
Normal file
257
third_party/vllm/tests/v1/core/utils.py
vendored
Normal file
@@ -0,0 +1,257 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
|
||||
from tests.v1.kv_connector.unit.utils import MockKVConfig
|
||||
from vllm.config import (
|
||||
CacheConfig,
|
||||
ECTransferConfig,
|
||||
KVTransferConfig,
|
||||
ModelConfig,
|
||||
ParallelConfig,
|
||||
SchedulerConfig,
|
||||
SpeculativeConfig,
|
||||
VllmConfig,
|
||||
)
|
||||
from vllm.multimodal.inputs import (
|
||||
MultiModalFeatureSpec,
|
||||
MultiModalKwargsItem,
|
||||
PlaceholderRange,
|
||||
)
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.utils.hashing import sha256
|
||||
from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
|
||||
from vllm.v1.core.sched.async_scheduler import AsyncScheduler
|
||||
from vllm.v1.core.sched.scheduler import Scheduler
|
||||
from vllm.v1.kv_cache_interface import (
|
||||
FullAttentionSpec,
|
||||
KVCacheConfig,
|
||||
KVCacheGroupSpec,
|
||||
)
|
||||
from vllm.v1.request import Request
|
||||
from vllm.v1.structured_output import StructuredOutputManager
|
||||
|
||||
EOS_TOKEN_ID = 50256
|
||||
|
||||
|
||||
def mock_kv(matched_tokens: int, is_async: bool):
|
||||
return MockKVConfig(matched_tokens=matched_tokens, is_async=is_async)
|
||||
|
||||
|
||||
def create_scheduler(
|
||||
model: str = "facebook/opt-125m",
|
||||
max_num_seqs: int = 16,
|
||||
max_num_batched_tokens: int = 8192,
|
||||
enable_chunked_prefill: bool = True,
|
||||
enable_prefix_caching: bool = False,
|
||||
long_prefill_token_threshold: int = 0,
|
||||
disable_chunked_mm_input: bool = False,
|
||||
use_kv_connector: None | bool | MockKVConfig = None,
|
||||
num_blocks: int = 10000,
|
||||
block_size: int = 16,
|
||||
max_model_len: int | None = None,
|
||||
num_speculative_tokens: int | None = None,
|
||||
skip_tokenizer_init: bool = False,
|
||||
async_scheduling: bool = False,
|
||||
pipeline_parallel_size: int = 1,
|
||||
use_ec_connector: bool = False,
|
||||
ec_role: str | None = None,
|
||||
) -> Scheduler | AsyncScheduler:
|
||||
"""Create scheduler under test.
|
||||
|
||||
Args:
|
||||
model: model under test
|
||||
max_num_seqs: max sequences to schedule
|
||||
max_num_batch_tokens: max num tokens to batch
|
||||
enable_prefix_caching: optionally force APC config
|
||||
(True/False) or use default
|
||||
(False)
|
||||
|
||||
Returns:
|
||||
{class}`Scheduler` instance
|
||||
"""
|
||||
model_config = ModelConfig(
|
||||
model=model,
|
||||
trust_remote_code=True,
|
||||
dtype="float16",
|
||||
seed=42,
|
||||
skip_tokenizer_init=skip_tokenizer_init,
|
||||
)
|
||||
if max_model_len is None:
|
||||
max_model_len = max_num_batched_tokens
|
||||
scheduler_config = SchedulerConfig(
|
||||
max_num_seqs=max_num_seqs,
|
||||
max_num_batched_tokens=max_num_batched_tokens,
|
||||
max_model_len=max_model_len,
|
||||
long_prefill_token_threshold=long_prefill_token_threshold,
|
||||
disable_chunked_mm_input=disable_chunked_mm_input,
|
||||
enable_chunked_prefill=enable_chunked_prefill,
|
||||
async_scheduling=async_scheduling,
|
||||
is_encoder_decoder=model_config.is_encoder_decoder,
|
||||
)
|
||||
# Cache config, optionally force APC
|
||||
cache_config = CacheConfig(
|
||||
block_size=block_size,
|
||||
gpu_memory_utilization=0.9,
|
||||
cache_dtype="auto",
|
||||
enable_prefix_caching=enable_prefix_caching,
|
||||
)
|
||||
kv_transfer_config = None
|
||||
if isinstance(use_kv_connector, MockKVConfig):
|
||||
kv_transfer_config = KVTransferConfig(
|
||||
kv_connector="MockKVConnector",
|
||||
kv_role="kv_both",
|
||||
kv_connector_extra_config={
|
||||
"matched_tokens": use_kv_connector.matched_tokens,
|
||||
"is_async": use_kv_connector.is_async,
|
||||
},
|
||||
)
|
||||
elif use_kv_connector:
|
||||
kv_transfer_config = KVTransferConfig(
|
||||
kv_connector="ExampleConnector",
|
||||
kv_role="kv_both",
|
||||
kv_connector_extra_config={"shared_storage_path": "local_storage"},
|
||||
)
|
||||
|
||||
speculative_config: SpeculativeConfig | None = None
|
||||
if num_speculative_tokens is not None:
|
||||
speculative_config = SpeculativeConfig(
|
||||
model="ngram", num_speculative_tokens=num_speculative_tokens
|
||||
)
|
||||
|
||||
ec_transfer_config = (
|
||||
ECTransferConfig(
|
||||
ec_connector="ECExampleConnector",
|
||||
ec_role=ec_role,
|
||||
ec_connector_extra_config={"shared_storage_path": "/tmp/ec_test"},
|
||||
)
|
||||
if use_ec_connector
|
||||
else None
|
||||
)
|
||||
|
||||
vllm_config = VllmConfig(
|
||||
scheduler_config=scheduler_config,
|
||||
model_config=model_config,
|
||||
cache_config=cache_config,
|
||||
parallel_config=ParallelConfig(pipeline_parallel_size=pipeline_parallel_size),
|
||||
kv_transfer_config=kv_transfer_config,
|
||||
speculative_config=speculative_config,
|
||||
ec_transfer_config=ec_transfer_config,
|
||||
)
|
||||
kv_cache_config = KVCacheConfig(
|
||||
num_blocks=num_blocks, # A large number of blocks to hold all requests
|
||||
kv_cache_tensors=[],
|
||||
kv_cache_groups=[
|
||||
KVCacheGroupSpec(
|
||||
["layer"],
|
||||
FullAttentionSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
)
|
||||
],
|
||||
)
|
||||
cache_config.num_gpu_blocks = num_blocks
|
||||
scheduler_cls = AsyncScheduler if async_scheduling else Scheduler
|
||||
return scheduler_cls(
|
||||
vllm_config=vllm_config,
|
||||
kv_cache_config=kv_cache_config,
|
||||
block_size=block_size,
|
||||
log_stats=True,
|
||||
structured_output_manager=StructuredOutputManager(vllm_config),
|
||||
)
|
||||
|
||||
|
||||
_none_hash_initialized = False
|
||||
|
||||
|
||||
def create_requests(
|
||||
num_requests: int,
|
||||
num_tokens: int = 10,
|
||||
mm_hashes_list: list[list[str]] | None = None,
|
||||
mm_positions: list[list[PlaceholderRange]] | None = None,
|
||||
ignore_eos: bool = False,
|
||||
max_tokens: int = 16,
|
||||
stop_token_ids: list[int] | None = None,
|
||||
prompt_logprobs: int | None = None,
|
||||
same_prompt: bool = False,
|
||||
block_size: int = 16,
|
||||
req_ids: list[str] | None = None,
|
||||
) -> list[Request]:
|
||||
global _none_hash_initialized
|
||||
if not _none_hash_initialized:
|
||||
init_none_hash(sha256)
|
||||
_none_hash_initialized = True
|
||||
|
||||
block_hasher = get_request_block_hasher(block_size, sha256)
|
||||
sampling_params = SamplingParams(
|
||||
ignore_eos=ignore_eos,
|
||||
max_tokens=max_tokens,
|
||||
stop_token_ids=stop_token_ids,
|
||||
prompt_logprobs=prompt_logprobs,
|
||||
)
|
||||
sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
|
||||
requests = []
|
||||
|
||||
if mm_hashes_list is not None:
|
||||
# NOTE: allow manual input; some mm items can have the same identifier
|
||||
# no. of mm_hashes and mm_positions for each request should be identical
|
||||
assert mm_positions is not None, (
|
||||
"mm_positions must be provided when mm_hashes_list is provided"
|
||||
)
|
||||
assert len(mm_hashes_list) == len(mm_positions) == num_requests
|
||||
assert [len(h) for h in mm_hashes_list] == [len(p) for p in mm_positions]
|
||||
|
||||
# Since same identifier would imply they are identical encoder output
|
||||
# Verify mm items with identical identifier are having mm_position.length
|
||||
seen_hashes: dict[str, int] = {}
|
||||
|
||||
if req_ids:
|
||||
assert len(req_ids) == num_requests
|
||||
else:
|
||||
req_ids = [f"{i}" for i in range(num_requests)]
|
||||
|
||||
for i in range(num_requests):
|
||||
mm_features = []
|
||||
|
||||
for j, position in enumerate(
|
||||
mm_positions[i] if mm_positions is not None else []
|
||||
):
|
||||
if mm_hashes_list is not None:
|
||||
identifier = mm_hashes_list[i][j]
|
||||
|
||||
# Verify if position length is identical
|
||||
position_length = position.length
|
||||
if identifier in seen_hashes:
|
||||
assert seen_hashes[identifier] == position_length, (
|
||||
f"mm_hash '{identifier}' has inconsistent position lengths: "
|
||||
f"previously {seen_hashes[identifier]}, now {position_length} "
|
||||
f"at request {i}, position {j}"
|
||||
)
|
||||
else:
|
||||
seen_hashes[identifier] = position_length
|
||||
else:
|
||||
# Unique dummy hash for each mm item
|
||||
identifier = f"hash{i}_{j}"
|
||||
mm_feature = MultiModalFeatureSpec(
|
||||
data=MultiModalKwargsItem.dummy(),
|
||||
mm_position=position,
|
||||
identifier=identifier,
|
||||
modality="image",
|
||||
)
|
||||
mm_features.append(mm_feature)
|
||||
|
||||
prompt_token_ids = [0] * num_tokens if same_prompt else [i] * num_tokens
|
||||
request = Request(
|
||||
request_id=req_ids[i],
|
||||
prompt_token_ids=prompt_token_ids,
|
||||
sampling_params=sampling_params,
|
||||
pooling_params=None,
|
||||
mm_features=mm_features if mm_features else None,
|
||||
block_hasher=block_hasher,
|
||||
)
|
||||
requests.append(request)
|
||||
return requests
|
||||
0
third_party/vllm/tests/v1/cudagraph/__init__.py
vendored
Normal file
0
third_party/vllm/tests/v1/cudagraph/__init__.py
vendored
Normal file
561
third_party/vllm/tests/v1/cudagraph/test_cudagraph_dispatch.py
vendored
Normal file
561
third_party/vllm/tests/v1/cudagraph/test_cudagraph_dispatch.py
vendored
Normal file
@@ -0,0 +1,561 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from dataclasses import replace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from tests.utils import create_new_process_for_each_test
|
||||
from vllm.compilation.cuda_graph import CUDAGraphWrapper
|
||||
from vllm.compilation.monitor import set_cudagraph_capturing_enabled
|
||||
from vllm.config import (
|
||||
CompilationConfig,
|
||||
CompilationMode,
|
||||
CUDAGraphMode,
|
||||
ParallelConfig,
|
||||
SchedulerConfig,
|
||||
VllmConfig,
|
||||
)
|
||||
from vllm.config.lora import LoRAConfig
|
||||
from vllm.forward_context import BatchDescriptor, set_forward_context
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
|
||||
|
||||
|
||||
# Helper MLP for testing
|
||||
class SimpleMLP(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.fc1 = nn.Linear(10, 10)
|
||||
self.fc2 = nn.Linear(10, 10)
|
||||
|
||||
def forward(self, x):
|
||||
return self.fc2(self.fc1(x))
|
||||
|
||||
|
||||
def _create_vllm_config(
|
||||
compilation_config: CompilationConfig,
|
||||
max_num_seqs: int = 8,
|
||||
lora_config: bool = False,
|
||||
) -> MagicMock:
|
||||
mock_config = MagicMock(spec=VllmConfig)
|
||||
mock_config.compilation_config = compilation_config
|
||||
mock_config.scheduler_config = SchedulerConfig.default_factory(
|
||||
max_num_seqs=max_num_seqs,
|
||||
)
|
||||
mock_config.parallel_config = ParallelConfig()
|
||||
mock_config.speculative_config = None # No speculative decoding
|
||||
if not lora_config:
|
||||
mock_config.lora_config = None
|
||||
else:
|
||||
# Create a real LoRAConfig with specialize_active_lora enabled
|
||||
mock_config.lora_config = LoRAConfig(
|
||||
max_loras=4,
|
||||
specialize_active_lora=True,
|
||||
)
|
||||
# Mimic the behavior of VllmConfig.__post_init__()
|
||||
if compilation_config.mode == CompilationMode.VLLM_COMPILE:
|
||||
compilation_config.set_splitting_ops_for_v1(
|
||||
all2all_backend=mock_config.parallel_config.all2all_backend,
|
||||
data_parallel_size=mock_config.parallel_config.data_parallel_size,
|
||||
)
|
||||
|
||||
# mimic VllmConfig.__post_init__
|
||||
if compilation_config.cudagraph_capture_sizes:
|
||||
compilation_config.max_cudagraph_capture_size = (
|
||||
compilation_config.cudagraph_capture_sizes[-1]
|
||||
)
|
||||
|
||||
compilation_config.post_init_cudagraph_sizes()
|
||||
|
||||
return mock_config
|
||||
|
||||
|
||||
class TestCudagraphDispatcher:
|
||||
@pytest.mark.parametrize(
|
||||
"cudagraph_mode_str,compilation_mode,lora_config",
|
||||
[
|
||||
# Test case 0: Full CG for mixed batches, no separate routine
|
||||
("FULL", CompilationMode.NONE, False),
|
||||
# Test case 1: Full CG for uniform batches, piecewise for mixed
|
||||
("FULL_AND_PIECEWISE", CompilationMode.NONE, False),
|
||||
# Test case 2: Full CG for uniform batches, no CG for mixed
|
||||
("FULL_DECODE_ONLY", CompilationMode.NONE, False),
|
||||
# Test case 3: PIECEWISE for all
|
||||
("PIECEWISE", CompilationMode.VLLM_COMPILE, False),
|
||||
# Test case 4: PIECEWISE for all, specialize LoRA cases
|
||||
("PIECEWISE", CompilationMode.VLLM_COMPILE, True),
|
||||
],
|
||||
)
|
||||
def test_dispatcher(self, cudagraph_mode_str, compilation_mode, lora_config):
|
||||
# Setup dispatcher
|
||||
comp_config = CompilationConfig(
|
||||
cudagraph_mode=cudagraph_mode_str,
|
||||
mode=compilation_mode,
|
||||
cudagraph_capture_sizes=[1, 8],
|
||||
)
|
||||
|
||||
config = _create_vllm_config(
|
||||
comp_config, max_num_seqs=8, lora_config=lora_config
|
||||
)
|
||||
if (
|
||||
cudagraph_mode_str == "FULL_AND_PIECEWISE"
|
||||
and compilation_mode == CompilationMode.NONE
|
||||
):
|
||||
with pytest.raises(AssertionError):
|
||||
dispatcher = CudagraphDispatcher(config)
|
||||
return
|
||||
|
||||
dispatcher = CudagraphDispatcher(config)
|
||||
dispatcher.initialize_cudagraph_keys(
|
||||
cudagraph_mode=comp_config.cudagraph_mode, uniform_decode_query_len=1
|
||||
)
|
||||
|
||||
# Verify the key is initialized correctly
|
||||
# With LoRA specialization (max_loras=4, specialize_active_lora=True):
|
||||
# - lora_cases = [0, 1, 2, 4, 5] (no-lora + powers of 2 up to 4 + max_loras+1)
|
||||
# - capture_sizes = [1, 8]
|
||||
# - Total keys = 2 sizes × 5 lora_cases = 10
|
||||
if cudagraph_mode_str in ["FULL_AND_PIECEWISE", "PIECEWISE"]:
|
||||
assert len(dispatcher.cudagraph_keys[CUDAGraphMode.PIECEWISE]) == (
|
||||
10 if lora_config else 2
|
||||
)
|
||||
else:
|
||||
assert len(dispatcher.cudagraph_keys[CUDAGraphMode.PIECEWISE]) == 0
|
||||
if cudagraph_mode_str not in ["NONE", "PIECEWISE"]:
|
||||
assert len(dispatcher.cudagraph_keys[CUDAGraphMode.FULL]) == (
|
||||
10 if lora_config else 2
|
||||
)
|
||||
else:
|
||||
assert len(dispatcher.cudagraph_keys[CUDAGraphMode.FULL]) == 0
|
||||
|
||||
# Test dispatch logic
|
||||
# 1. non-uniform batch, size in cudagraph size list
|
||||
# FULL mode uses exact keys with num_reqs set
|
||||
desc_full_with_reqs = BatchDescriptor(num_tokens=8, num_reqs=8, uniform=False)
|
||||
# PIECEWISE mode uses relaxed keys with num_reqs=None
|
||||
desc_piecewise = BatchDescriptor(num_tokens=8, num_reqs=None, uniform=False)
|
||||
rt_mode, key = dispatcher.dispatch(
|
||||
num_tokens=8, uniform_decode=False, has_lora=False
|
||||
)
|
||||
if cudagraph_mode_str == "FULL":
|
||||
assert rt_mode == CUDAGraphMode.FULL
|
||||
assert key == desc_full_with_reqs
|
||||
elif cudagraph_mode_str in ["FULL_AND_PIECEWISE", "PIECEWISE"]:
|
||||
assert rt_mode == CUDAGraphMode.PIECEWISE
|
||||
assert key == desc_piecewise
|
||||
else:
|
||||
assert rt_mode == CUDAGraphMode.NONE
|
||||
|
||||
# 2. uniform decode batch, size in cudagraph size list
|
||||
desc_uniform_exact = BatchDescriptor(num_tokens=8, num_reqs=8, uniform=True)
|
||||
desc_non_uniform = BatchDescriptor(num_tokens=8, num_reqs=8, uniform=False)
|
||||
rt_mode, key = dispatcher.dispatch(
|
||||
num_tokens=8, uniform_decode=True, has_lora=False
|
||||
)
|
||||
if cudagraph_mode_str == "FULL":
|
||||
# Pure FULL mode uses non-uniform keys for all batches
|
||||
assert rt_mode == CUDAGraphMode.FULL
|
||||
assert key == desc_non_uniform
|
||||
elif cudagraph_mode_str in ["FULL_DECODE_ONLY", "FULL_AND_PIECEWISE"]:
|
||||
# These modes have separate uniform decode keys
|
||||
assert rt_mode == CUDAGraphMode.FULL
|
||||
assert key == desc_uniform_exact
|
||||
elif cudagraph_mode_str == "PIECEWISE":
|
||||
assert rt_mode == CUDAGraphMode.PIECEWISE
|
||||
assert key == replace(desc_uniform_exact, num_reqs=None, uniform=False)
|
||||
else:
|
||||
assert rt_mode == CUDAGraphMode.NONE
|
||||
|
||||
# 3. No key match
|
||||
rt_mode, key = dispatcher.dispatch(
|
||||
num_tokens=15, uniform_decode=False, has_lora=False
|
||||
)
|
||||
assert rt_mode == CUDAGraphMode.NONE
|
||||
assert key == BatchDescriptor(num_tokens=15)
|
||||
|
||||
# 4. invalid_modes={FULL} should have a fall back mode
|
||||
# (e.g., cascade attention)
|
||||
desc_full_exact = BatchDescriptor(num_tokens=8, uniform=False)
|
||||
rt_mode, key = dispatcher.dispatch(
|
||||
num_tokens=8,
|
||||
uniform_decode=False,
|
||||
has_lora=False,
|
||||
invalid_modes={CUDAGraphMode.FULL},
|
||||
)
|
||||
|
||||
if "PIECEWISE" in cudagraph_mode_str: # string contains check
|
||||
assert rt_mode == CUDAGraphMode.PIECEWISE
|
||||
assert key == replace(desc_full_exact, num_reqs=None, uniform=False)
|
||||
else:
|
||||
assert rt_mode == CUDAGraphMode.NONE
|
||||
|
||||
# 5. valid_modes={NONE} always returns NONE even when keys exist
|
||||
rt_mode, key = dispatcher.dispatch(
|
||||
num_tokens=8,
|
||||
uniform_decode=False,
|
||||
has_lora=False,
|
||||
valid_modes={CUDAGraphMode.NONE},
|
||||
)
|
||||
assert rt_mode == CUDAGraphMode.NONE
|
||||
assert key == BatchDescriptor(num_tokens=8)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cudagraph_mode_str,compilation_mode,expected_modes",
|
||||
[
|
||||
# FULL mode: only FULL keys, no PIECEWISE
|
||||
("FULL", CompilationMode.NONE, [CUDAGraphMode.FULL]),
|
||||
# PIECEWISE mode: only PIECEWISE keys
|
||||
("PIECEWISE", CompilationMode.VLLM_COMPILE, [CUDAGraphMode.PIECEWISE]),
|
||||
# FULL_DECODE_ONLY: only FULL keys for uniform decode
|
||||
("FULL_DECODE_ONLY", CompilationMode.NONE, [CUDAGraphMode.FULL]),
|
||||
# NONE mode: no keys
|
||||
("NONE", CompilationMode.NONE, []),
|
||||
],
|
||||
)
|
||||
def test_get_capture_descs(
|
||||
self, cudagraph_mode_str, compilation_mode, expected_modes
|
||||
):
|
||||
"""Test get_capture_descs returns correctly grouped and ordered descs."""
|
||||
comp_config = CompilationConfig(
|
||||
cudagraph_mode=cudagraph_mode_str,
|
||||
mode=compilation_mode,
|
||||
cudagraph_capture_sizes=[1, 4, 8, 16],
|
||||
)
|
||||
|
||||
config = _create_vllm_config(comp_config, max_num_seqs=16)
|
||||
dispatcher = CudagraphDispatcher(config)
|
||||
dispatcher.initialize_cudagraph_keys(
|
||||
cudagraph_mode=comp_config.cudagraph_mode, uniform_decode_query_len=1
|
||||
)
|
||||
|
||||
capture_descs = dispatcher.get_capture_descs()
|
||||
|
||||
# Verify we get the expected modes
|
||||
actual_modes = [mode for mode, _ in capture_descs]
|
||||
assert actual_modes == expected_modes
|
||||
|
||||
# Verify each group is sorted largest-first
|
||||
for mode, descs in capture_descs:
|
||||
assert len(descs) > 0, "Each group should have at least one descriptor"
|
||||
num_tokens_list = [d.num_tokens for d in descs]
|
||||
assert num_tokens_list == sorted(num_tokens_list, reverse=True), (
|
||||
f"Descriptors for {mode} should be sorted largest-first"
|
||||
)
|
||||
|
||||
# All descriptors in a group should have same uniform value
|
||||
uniform_values = [d.uniform for d in descs]
|
||||
assert len(set(uniform_values)) == 1, (
|
||||
"All descriptors in a group should have the same uniform value"
|
||||
)
|
||||
|
||||
def test_get_capture_descs_empty_when_not_initialized(self):
|
||||
"""Test that get_capture_descs returns empty list when keys not initialized."""
|
||||
comp_config = CompilationConfig(
|
||||
cudagraph_mode="FULL",
|
||||
mode=CompilationMode.NONE,
|
||||
cudagraph_capture_sizes=[1, 8],
|
||||
)
|
||||
config = _create_vllm_config(comp_config, max_num_seqs=8)
|
||||
dispatcher = CudagraphDispatcher(config)
|
||||
# Don't initialize keys
|
||||
|
||||
assert dispatcher.get_capture_descs() == []
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
|
||||
class TestCUDAGraphWrapper:
|
||||
def setup_method(self):
|
||||
self.vllm_config = _create_vllm_config(CompilationConfig())
|
||||
self.model = SimpleMLP().to("cuda")
|
||||
self.persistent_input_buffer = torch.zeros(1, 10, device="cuda")
|
||||
self.input_tensor = torch.randn(1, 10, device="cuda")
|
||||
|
||||
def test_capture_and_replay(self):
|
||||
wrapper = CUDAGraphWrapper(
|
||||
self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
|
||||
)
|
||||
batch_descriptor = BatchDescriptor(num_tokens=10)
|
||||
|
||||
# 0. global warmup
|
||||
with set_forward_context(
|
||||
attn_metadata=None,
|
||||
vllm_config=self.vllm_config,
|
||||
cudagraph_runtime_mode=CUDAGraphMode.NONE,
|
||||
batch_descriptor=None,
|
||||
):
|
||||
wrapper(self.input_tensor)
|
||||
|
||||
# 1. Capture
|
||||
with (
|
||||
set_forward_context(
|
||||
attn_metadata=None,
|
||||
vllm_config=self.vllm_config,
|
||||
cudagraph_runtime_mode=CUDAGraphMode.FULL,
|
||||
batch_descriptor=batch_descriptor,
|
||||
),
|
||||
patch("torch.cuda.graph", wraps=torch.cuda.graph) as mock_cuda_graph,
|
||||
):
|
||||
output1 = wrapper(self.input_tensor)
|
||||
# capturing phase should generate a zero output
|
||||
assert torch.allclose(output1, torch.zeros_like(output1))
|
||||
mock_cuda_graph.assert_called_once()
|
||||
|
||||
assert batch_descriptor in wrapper.concrete_cudagraph_entries
|
||||
entry = wrapper.concrete_cudagraph_entries[batch_descriptor]
|
||||
assert entry.cudagraph is not None
|
||||
|
||||
# 2. Replay
|
||||
with (
|
||||
set_forward_context(
|
||||
attn_metadata=None,
|
||||
vllm_config=self.vllm_config,
|
||||
cudagraph_runtime_mode=CUDAGraphMode.FULL,
|
||||
batch_descriptor=batch_descriptor,
|
||||
),
|
||||
patch.object(
|
||||
entry.cudagraph, "replay", wraps=entry.cudagraph.replay
|
||||
) as mock_replay,
|
||||
):
|
||||
output2 = wrapper(self.input_tensor)
|
||||
mock_replay.assert_called_once()
|
||||
|
||||
# Compare with eager output
|
||||
eager_output = self.model(self.input_tensor)
|
||||
torch.testing.assert_close(eager_output, output2)
|
||||
|
||||
def test_bypass_on_mode_mismatch(self):
|
||||
wrapper = CUDAGraphWrapper(
|
||||
self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
|
||||
)
|
||||
batch_descriptor = BatchDescriptor(num_tokens=10)
|
||||
|
||||
with (
|
||||
set_forward_context(
|
||||
attn_metadata=None,
|
||||
vllm_config=self.vllm_config,
|
||||
cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
|
||||
batch_descriptor=batch_descriptor,
|
||||
),
|
||||
patch("torch.cuda.graph", wraps=torch.cuda.graph) as mock_cuda_graph,
|
||||
patch.object(
|
||||
self.model, "forward", wraps=self.model.forward
|
||||
) as mock_forward,
|
||||
):
|
||||
wrapper(self.input_tensor)
|
||||
mock_cuda_graph.assert_not_called()
|
||||
mock_forward.assert_called_once()
|
||||
assert not wrapper.concrete_cudagraph_entries
|
||||
|
||||
def test_bypass_on_mode_none(self):
|
||||
wrapper = CUDAGraphWrapper(
|
||||
self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
|
||||
)
|
||||
batch_descriptor = BatchDescriptor(num_tokens=10)
|
||||
|
||||
with (
|
||||
set_forward_context(
|
||||
attn_metadata=None,
|
||||
vllm_config=self.vllm_config,
|
||||
cudagraph_runtime_mode=CUDAGraphMode.NONE,
|
||||
batch_descriptor=batch_descriptor,
|
||||
),
|
||||
patch("torch.cuda.graph", wraps=torch.cuda.graph) as mock_cuda_graph,
|
||||
):
|
||||
wrapper(self.input_tensor)
|
||||
mock_cuda_graph.assert_not_called()
|
||||
assert not wrapper.concrete_cudagraph_entries
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
|
||||
class TestCudagraphIntegration:
|
||||
def setup_method(self):
|
||||
# only FULL mode for non-uniform batches
|
||||
self.comp_config = CompilationConfig(
|
||||
mode=CompilationMode.VLLM_COMPILE,
|
||||
cudagraph_mode="FULL",
|
||||
cudagraph_capture_sizes=[10, 20],
|
||||
)
|
||||
self.vllm_config = _create_vllm_config(self.comp_config)
|
||||
self.dispatcher = CudagraphDispatcher(self.vllm_config)
|
||||
self.dispatcher.initialize_cudagraph_keys(
|
||||
self.comp_config.cudagraph_mode, uniform_decode_query_len=1
|
||||
)
|
||||
|
||||
def _run_and_monitor_call(
|
||||
self, wrapper, input_tensor, runtime_mode, batch_descriptor
|
||||
):
|
||||
"""Helper to run a single call and monitor the action."""
|
||||
|
||||
with (
|
||||
patch("torch.cuda.graph", wraps=torch.cuda.graph) as mock_graph_context,
|
||||
patch.object(wrapper, "runnable", wraps=wrapper.runnable) as mock_runnable,
|
||||
):
|
||||
entry = wrapper.concrete_cudagraph_entries.get(batch_descriptor, None)
|
||||
|
||||
context = set_forward_context(
|
||||
attn_metadata=None,
|
||||
vllm_config=self.vllm_config,
|
||||
cudagraph_runtime_mode=runtime_mode,
|
||||
batch_descriptor=batch_descriptor,
|
||||
)
|
||||
mock_replay = MagicMock()
|
||||
if entry and entry.cudagraph:
|
||||
with (
|
||||
context,
|
||||
patch.object(
|
||||
entry.cudagraph, "replay", new_callable=MagicMock
|
||||
) as mock_replay,
|
||||
):
|
||||
wrapper(input_tensor)
|
||||
else:
|
||||
with context:
|
||||
wrapper(input_tensor)
|
||||
|
||||
if mock_graph_context.called:
|
||||
# note that this is globally mocked, so it will be detected
|
||||
# even whether called by the inner or outer wrapper
|
||||
return "capture_global"
|
||||
if mock_replay.called:
|
||||
# only for outer wrapper
|
||||
return "replay"
|
||||
if mock_runnable.call_count > 0:
|
||||
# only for outer wrapper
|
||||
return "bypass"
|
||||
return "unknown"
|
||||
|
||||
@create_new_process_for_each_test("spawn")
|
||||
def test_capture_replay_bypass_logic(self):
|
||||
model = SimpleMLP().to("cuda")
|
||||
full_wrapper = CUDAGraphWrapper(model, self.vllm_config, CUDAGraphMode.FULL)
|
||||
max_bs = 16
|
||||
persistent_input_buffer = torch.zeros(max_bs, 10, device="cuda")
|
||||
input_1 = persistent_input_buffer[:1]
|
||||
input_2 = persistent_input_buffer[:2]
|
||||
input_3 = persistent_input_buffer[:3]
|
||||
|
||||
desc_1 = BatchDescriptor(num_tokens=1)
|
||||
desc_2 = BatchDescriptor(num_tokens=2)
|
||||
desc_3_unseen = BatchDescriptor(num_tokens=3)
|
||||
|
||||
# 0. global warmup
|
||||
with set_forward_context(
|
||||
attn_metadata=None,
|
||||
vllm_config=self.vllm_config,
|
||||
cudagraph_runtime_mode=CUDAGraphMode.NONE,
|
||||
batch_descriptor=None,
|
||||
):
|
||||
full_wrapper(input_1)
|
||||
|
||||
rt_mode, key = self.dispatcher.dispatch(num_tokens=desc_1.num_tokens)
|
||||
# 1. Capture first shape
|
||||
action = self._run_and_monitor_call(full_wrapper, input_1, rt_mode, key)
|
||||
assert action == "capture_global"
|
||||
|
||||
# 2. Replay first shape
|
||||
action = self._run_and_monitor_call(full_wrapper, input_1, rt_mode, key)
|
||||
assert action == "replay"
|
||||
|
||||
rt_mode, key = self.dispatcher.dispatch(num_tokens=desc_2.num_tokens)
|
||||
# 3. Capture second shape
|
||||
action = self._run_and_monitor_call(full_wrapper, input_2, rt_mode, key)
|
||||
assert action == "capture_global"
|
||||
|
||||
# 4. Replay second shape
|
||||
action = self._run_and_monitor_call(
|
||||
full_wrapper, input_2, CUDAGraphMode.FULL, desc_2
|
||||
)
|
||||
assert action == "replay"
|
||||
|
||||
# 5. Bypass if no key match
|
||||
rt_mode, key = self.dispatcher.dispatch(num_tokens=desc_3_unseen.num_tokens)
|
||||
assert rt_mode == CUDAGraphMode.NONE
|
||||
action = self._run_and_monitor_call(full_wrapper, input_3, rt_mode, key)
|
||||
assert action == "bypass"
|
||||
|
||||
# capture unseen shape is not allowed after disable
|
||||
set_cudagraph_capturing_enabled(False)
|
||||
with pytest.raises(RuntimeError):
|
||||
self._run_and_monitor_call(
|
||||
full_wrapper, input_3, CUDAGraphMode.FULL, desc_3_unseen
|
||||
)
|
||||
set_cudagraph_capturing_enabled(True)
|
||||
|
||||
@create_new_process_for_each_test("spawn")
|
||||
def test_nested_wrappers(self):
|
||||
"""Tests a scenario with a PIECEWISE wrapper inside a FULL one."""
|
||||
model = SimpleMLP().to("cuda")
|
||||
full_wrapper = CUDAGraphWrapper(model, self.vllm_config, CUDAGraphMode.FULL)
|
||||
input_1 = torch.randn(1, 10, device="cuda")
|
||||
|
||||
# Setup: Inner model is wrapped with PIECEWISE, outer with FULL
|
||||
inner_model = SimpleMLP().to("cuda")
|
||||
piecewise_wrapper = CUDAGraphWrapper(
|
||||
inner_model, self.vllm_config, CUDAGraphMode.PIECEWISE
|
||||
)
|
||||
inner_model.forward = MagicMock(wraps=inner_model.forward)
|
||||
outer_model = SimpleMLP().to("cuda")
|
||||
# When outer model is called, it calls the piecewise_wrapper
|
||||
outer_model.forward = MagicMock(
|
||||
wraps=outer_model.forward, side_effect=piecewise_wrapper
|
||||
)
|
||||
full_wrapper = CUDAGraphWrapper(
|
||||
outer_model, self.vllm_config, CUDAGraphMode.FULL
|
||||
)
|
||||
|
||||
desc_1 = BatchDescriptor(num_tokens=1)
|
||||
|
||||
# 0. global warmup
|
||||
with set_forward_context(
|
||||
attn_metadata=None,
|
||||
vllm_config=self.vllm_config,
|
||||
cudagraph_runtime_mode=CUDAGraphMode.NONE,
|
||||
batch_descriptor=None,
|
||||
):
|
||||
full_wrapper(input_1)
|
||||
|
||||
# --- Test runtime mode FULL---
|
||||
# Run with FULL mode context. Expect outer wrapper to capture.
|
||||
# The inner mock should be called once inside the graph capture.
|
||||
outer_model.forward.reset_mock()
|
||||
inner_model.forward.reset_mock()
|
||||
action = self._run_and_monitor_call(
|
||||
full_wrapper, input_1, CUDAGraphMode.FULL, desc_1
|
||||
)
|
||||
assert action == "capture_global"
|
||||
assert outer_model.forward.call_count == 1
|
||||
assert inner_model.forward.call_count == 1
|
||||
|
||||
# Run again. Expect outer wrapper to replay.
|
||||
# The outer model should NOT be called because the whole graph
|
||||
# is replayed.
|
||||
action = self._run_and_monitor_call(
|
||||
full_wrapper, input_1, CUDAGraphMode.FULL, desc_1
|
||||
)
|
||||
assert action == "replay"
|
||||
assert outer_model.forward.call_count == 1 # No new call
|
||||
assert inner_model.forward.call_count == 1
|
||||
|
||||
# --- Test runtime mode PIECEWISE ---
|
||||
outer_model.forward.reset_mock()
|
||||
inner_model.forward.reset_mock()
|
||||
# Run with PIECEWISE mode context.
|
||||
# Expect outer wrapper to bypass and call inner wrapper.
|
||||
# Inner wrapper should capture.
|
||||
action = self._run_and_monitor_call(
|
||||
full_wrapper, input_1, CUDAGraphMode.PIECEWISE, desc_1
|
||||
)
|
||||
assert action == "capture_global"
|
||||
assert outer_model.forward.call_count == 1
|
||||
assert inner_model.forward.call_count == 1
|
||||
|
||||
# Run again with PIECEWISE.
|
||||
# Outer bypasses, inner replays.
|
||||
action = self._run_and_monitor_call(
|
||||
full_wrapper, input_1, CUDAGraphMode.PIECEWISE, desc_1
|
||||
)
|
||||
assert action == "bypass"
|
||||
assert outer_model.forward.call_count == 2
|
||||
assert inner_model.forward.call_count == 1
|
||||
133
third_party/vllm/tests/v1/cudagraph/test_cudagraph_mode.py
vendored
Normal file
133
third_party/vllm/tests/v1/cudagraph/test_cudagraph_mode.py
vendored
Normal file
@@ -0,0 +1,133 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import weakref
|
||||
from contextlib import ExitStack
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.utils import wait_for_gpu_memory_to_clear
|
||||
from tests.v1.attention.utils import full_cg_backend_configs as backend_configs
|
||||
from vllm import LLM
|
||||
from vllm.config import CompilationConfig, CompilationMode
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
# test attention backend and cudagraph_mode combo
|
||||
# (backend_name, cudagraph_mode, supported)
|
||||
if current_platform.is_rocm():
|
||||
combo_cases_1 = [
|
||||
("RocmAttn", "FULL", True),
|
||||
("RocmAttn", "FULL_AND_PIECEWISE", True),
|
||||
("TritonAttn", "FULL", True),
|
||||
("TritonAttn", "FULL_AND_PIECEWISE", True),
|
||||
]
|
||||
else:
|
||||
combo_cases_1 = [
|
||||
("FA3", "FULL", True),
|
||||
("FA3", "FULL_AND_PIECEWISE", True),
|
||||
("FA2", "FULL", True), # Should fallback to FULL_AND_PIECEWISE
|
||||
("FA2", "FULL_AND_PIECEWISE", True),
|
||||
("FlashInfer", "FULL", True), # Should fallback to FULL_AND_PIECEWISE
|
||||
("FlashInfer", "FULL_AND_PIECEWISE", True),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend_name, cudagraph_mode, supported", combo_cases_1)
|
||||
def test_backend_and_cudagraph_mode_combo(backend_name, cudagraph_mode, supported):
|
||||
if backend_name == "FlashInfer":
|
||||
try:
|
||||
import flashinfer # noqa: F401
|
||||
except ImportError:
|
||||
pytest.skip("FlashInfer is not installed")
|
||||
backend_config = backend_configs[backend_name]
|
||||
# Dynamically skip test if GPU capability is not met
|
||||
if (
|
||||
backend_config.specific_gpu_arch
|
||||
and backend_config.specific_gpu_arch != current_platform.get_device_capability()
|
||||
):
|
||||
pytest.skip("Only Hopper GPUs support FA3 and FlashMLA")
|
||||
|
||||
attention_config = backend_config.attention_config
|
||||
|
||||
with ExitStack() as stack:
|
||||
if not supported:
|
||||
stack.enter_context(pytest.raises(Exception))
|
||||
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen2-1.5B-Instruct",
|
||||
max_num_seqs=256,
|
||||
trust_remote_code=True,
|
||||
gpu_memory_utilization=0.45,
|
||||
max_model_len=1024,
|
||||
attention_config=attention_config,
|
||||
compilation_config=CompilationConfig(
|
||||
mode=CompilationMode.VLLM_COMPILE, cudagraph_mode=cudagraph_mode
|
||||
),
|
||||
)
|
||||
llm.generate(["Hello, my name is"] * 10)
|
||||
# when above code raises, `llm` may be undefined, so we need to catch that
|
||||
try:
|
||||
llm = weakref.proxy(llm)
|
||||
del llm
|
||||
except UnboundLocalError:
|
||||
pass
|
||||
|
||||
wait_for_gpu_memory_to_clear(
|
||||
devices=[0],
|
||||
threshold_ratio=0.1,
|
||||
)
|
||||
|
||||
|
||||
# test cudagraph_mode with different compilation mode.
|
||||
# (backend_name, cudagraph_mode, compilation_mode, supported)
|
||||
attn_backend = "RocmAttn" if current_platform.is_rocm() else "FA2"
|
||||
|
||||
combo_cases_2 = [
|
||||
(attn_backend, "FULL", CompilationMode.NONE, True),
|
||||
(attn_backend, "FULL", CompilationMode.VLLM_COMPILE, True),
|
||||
(attn_backend, "PIECEWISE", CompilationMode.NONE, True),
|
||||
(attn_backend, "PIECEWISE", CompilationMode.VLLM_COMPILE, True),
|
||||
(attn_backend, "FULL_AND_PIECEWISE", CompilationMode.NONE, True),
|
||||
(attn_backend, "FULL_AND_PIECEWISE", CompilationMode.VLLM_COMPILE, True),
|
||||
(attn_backend, "FULL_DECODE_ONLY", CompilationMode.NONE, True),
|
||||
(attn_backend, "FULL_DECODE_ONLY", CompilationMode.VLLM_COMPILE, True),
|
||||
(attn_backend, "NONE", CompilationMode.NONE, True),
|
||||
(attn_backend, "NONE", CompilationMode.VLLM_COMPILE, True),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"backend_name,cudagraph_mode,compilation_mode,supported", combo_cases_2
|
||||
)
|
||||
def test_cudagraph_compilation_combo(
|
||||
backend_name, cudagraph_mode, compilation_mode, supported
|
||||
):
|
||||
backend_config = backend_configs[backend_name]
|
||||
attention_config = backend_config.attention_config
|
||||
|
||||
with ExitStack() as stack:
|
||||
if not supported:
|
||||
stack.enter_context(pytest.raises(Exception))
|
||||
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen2-1.5B-Instruct",
|
||||
max_num_seqs=256,
|
||||
trust_remote_code=True,
|
||||
gpu_memory_utilization=0.45,
|
||||
max_model_len=1024,
|
||||
attention_config=attention_config,
|
||||
compilation_config=CompilationConfig(
|
||||
mode=compilation_mode, cudagraph_mode=cudagraph_mode
|
||||
),
|
||||
)
|
||||
llm.generate(["Hello, my name is"] * 10)
|
||||
# when above code raises, `llm` may be undefined, so we need to catch that
|
||||
try:
|
||||
llm = weakref.proxy(llm)
|
||||
del llm
|
||||
except UnboundLocalError:
|
||||
pass
|
||||
finally:
|
||||
wait_for_gpu_memory_to_clear(
|
||||
devices=[0],
|
||||
threshold_ratio=0.1,
|
||||
)
|
||||
12
third_party/vllm/tests/v1/determinism/conftest.py
vendored
Normal file
12
third_party/vllm/tests/v1/determinism/conftest.py
vendored
Normal file
@@ -0,0 +1,12 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
|
||||
import vllm.model_executor.layers.batch_invariant as batch_invariant
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def enable_batch_invariant_mode(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Automatically enable batch invariant kernel overrides for all tests."""
|
||||
monkeypatch.setattr(batch_invariant, "VLLM_BATCH_INVARIANT", True)
|
||||
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
|
||||
948
third_party/vllm/tests/v1/determinism/test_batch_invariance.py
vendored
Normal file
948
third_party/vllm/tests/v1/determinism/test_batch_invariance.py
vendored
Normal file
@@ -0,0 +1,948 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import contextlib
|
||||
import os
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from utils import (
|
||||
BACKENDS,
|
||||
_extract_step_logprobs,
|
||||
_random_prompt,
|
||||
is_device_capability_below_90,
|
||||
resolve_model_name,
|
||||
skip_unsupported,
|
||||
)
|
||||
|
||||
import vllm.model_executor.layers.batch_invariant as batch_invariant
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
IS_DEVICE_CAPABILITY_BELOW_90 = is_device_capability_below_90()
|
||||
|
||||
|
||||
@skip_unsupported
|
||||
@pytest.mark.timeout(1000)
|
||||
@pytest.mark.parametrize(
|
||||
"backend",
|
||||
BACKENDS,
|
||||
)
|
||||
def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
|
||||
backend,
|
||||
):
|
||||
"""
|
||||
Ensures that the same request (the 'needle' prompt) yields identical output
|
||||
whether run alone (bs=1) or mixed into a larger batch (e.g., bs=64),
|
||||
using the high-level v1 LLM() API only (no manual batching).
|
||||
|
||||
Strategy:
|
||||
- Create two LLM engines with identical config except max_num_seqs: 1 vs N.
|
||||
- Compute a baseline output for the needle prompt with the bs=1 engine.
|
||||
- For many trials, generate a batch (size N) where the needle appears at a
|
||||
random position among random filler prompts using the bs=N engine.
|
||||
- Track how many trials match vs mismatch, and report totals at the end.
|
||||
The test fails if any mismatches occur, but we still dump pass/fail
|
||||
counts.
|
||||
|
||||
Notes:
|
||||
- Use seeded stochastic sampling with a fixed seed to test determinism.
|
||||
- Outputs are intentionally longer and sampled at higher temperature/top_p
|
||||
to produce a more random-sounding phrase, yet remain deterministic by
|
||||
seed.
|
||||
- Keep max_tokens and max_model_len bounded for speed and memory use.
|
||||
"""
|
||||
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
|
||||
random.seed(seed)
|
||||
|
||||
attention_config = {"backend": backend}
|
||||
# Allow overrides from environment (useful for CI tuning)
|
||||
# "facebook/opt-125m" is too small, doesn't reliably test determinism
|
||||
model = resolve_model_name(backend)
|
||||
num_trials = int(os.getenv("VLLM_NEEDLE_TRIALS", "5"))
|
||||
max_batch_size = int(os.getenv("VLLM_NEEDLE_BATCH_SIZE", "128"))
|
||||
min_random_prompt = int(os.getenv("VLLM_MIN_PROMPT", "1024"))
|
||||
max_random_prompt = int(os.getenv("VLLM_MAX_PROMPT", "2048"))
|
||||
assert max_batch_size >= 2, "Batch size should be >= 2 to mix needle."
|
||||
|
||||
# Keep GPU memory usage low to avoid startup allocation failures.
|
||||
gpu_mem_util = float(os.getenv("VLLM_GPU_MEMORY_UTILIZATION", "0.4"))
|
||||
max_model_len = int(os.getenv("VLLM_MAX_MODEL_LEN", "5120"))
|
||||
|
||||
# Sampling parameters: longer outputs with a more random-sounding
|
||||
# continuation,but still deterministic due to fixed seed.
|
||||
temperature = float(os.getenv("VLLM_NEEDLE_TEMPERATURE", "0.0"))
|
||||
top_p = float(os.getenv("VLLM_NEEDLE_TOP_P", "0.95"))
|
||||
max_tokens = int(os.getenv("VLLM_NEEDLE_MAX_TOKENS", "128"))
|
||||
|
||||
sampling = SamplingParams(
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
max_tokens=max_tokens,
|
||||
seed=20240919,
|
||||
)
|
||||
|
||||
needle_prompt = "There once was a "
|
||||
|
||||
llm_bs1 = None
|
||||
llm_bsN = None
|
||||
try:
|
||||
# Engine with bs=1 behavior
|
||||
llm_bs1 = LLM_with_max_seqs(
|
||||
model=model,
|
||||
max_num_seqs=max_batch_size,
|
||||
gpu_memory_utilization=gpu_mem_util,
|
||||
max_model_len=max_model_len,
|
||||
attention_config=attention_config,
|
||||
)
|
||||
|
||||
# Baseline generation for the needle prompt alone.
|
||||
baseline_out = llm_bs1.generate([needle_prompt], sampling)
|
||||
assert len(baseline_out) == 1
|
||||
assert len(baseline_out[0].outputs) >= 1
|
||||
baseline_text = baseline_out[0].outputs[0].text
|
||||
|
||||
# Engine with larger batch limit (e.g., 64)
|
||||
llm_bsN = LLM_with_max_seqs(
|
||||
model=model,
|
||||
max_num_seqs=max_batch_size,
|
||||
gpu_memory_utilization=gpu_mem_util,
|
||||
max_model_len=max_model_len,
|
||||
attention_config=attention_config,
|
||||
)
|
||||
|
||||
mismatches = 0
|
||||
|
||||
for trial in range(num_trials):
|
||||
# Create a batch of size `max_batch_size` and insert the needle at
|
||||
# a random index
|
||||
prompts: list[str] = []
|
||||
batch_size = random.randint(max_batch_size // 2, max_batch_size)
|
||||
needle_pos = random.randint(0, batch_size - 1)
|
||||
for i in range(batch_size):
|
||||
if i == needle_pos:
|
||||
prompts.append(needle_prompt)
|
||||
else:
|
||||
prompts.append(_random_prompt(min_random_prompt, max_random_prompt))
|
||||
|
||||
# Generate with the larger-batch engine
|
||||
outputs = llm_bsN.generate(prompts, sampling)
|
||||
# Find the needle output by position
|
||||
needle_output = outputs[needle_pos]
|
||||
assert needle_output.prompt == needle_prompt
|
||||
assert len(needle_output.outputs) >= 1
|
||||
text = needle_output.outputs[0].text
|
||||
|
||||
if text != baseline_text:
|
||||
print(f"{text}\n\n== Not the same as ==\n\n{baseline_text}\n\n")
|
||||
mismatches += 1
|
||||
|
||||
passes = num_trials - mismatches
|
||||
# Dump how many passed vs failed
|
||||
print(
|
||||
f"[determinism] total={num_trials}, passed={passes}, "
|
||||
f"failed={mismatches}, max_batch_size={max_batch_size}"
|
||||
)
|
||||
|
||||
if mismatches > 0:
|
||||
pytest.fail(
|
||||
f"Nondeterministic outputs detected: {mismatches} failed out "
|
||||
f"of {num_trials} trials (max_batch_size={max_batch_size})."
|
||||
)
|
||||
|
||||
finally:
|
||||
# Ensure engines are shutdown to free GPU/VRAM across test sessions
|
||||
if llm_bs1 is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
llm_bs1.shutdown()
|
||||
if llm_bsN is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
llm_bsN.shutdown()
|
||||
|
||||
|
||||
@skip_unsupported
|
||||
@pytest.mark.parametrize(
|
||||
"backend",
|
||||
BACKENDS,
|
||||
)
|
||||
def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
|
||||
backend,
|
||||
):
|
||||
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
|
||||
random.seed(seed)
|
||||
model_name = resolve_model_name(backend)
|
||||
tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
|
||||
|
||||
# For batch invariance, disable custom all-reduce to ensure deterministic
|
||||
# all-reduce operations (custom all-reduce may not be deterministic)
|
||||
from vllm.model_executor.layers.batch_invariant import (
|
||||
vllm_is_batch_invariant,
|
||||
)
|
||||
|
||||
disable_custom_ar = vllm_is_batch_invariant()
|
||||
|
||||
if disable_custom_ar:
|
||||
print(f"\n{'=' * 80}")
|
||||
print(f"BATCH INVARIANCE MODE: Disabling custom all-reduce (TP={tp_size})")
|
||||
print(f"{'=' * 80}\n")
|
||||
|
||||
llm = LLM(
|
||||
model=model_name,
|
||||
tensor_parallel_size=tp_size,
|
||||
max_num_seqs=128,
|
||||
max_model_len=8192,
|
||||
dtype="bfloat16", # not everything is supported
|
||||
gpu_memory_utilization=0.9,
|
||||
enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
|
||||
attention_config={"backend": backend},
|
||||
)
|
||||
|
||||
# Use more realistic prompts for better token generation
|
||||
prompts = [_random_prompt(10, 50) for _ in range(32)]
|
||||
|
||||
# TODO: Update prompts to have ragged lengths in order to test chunked prefill
|
||||
# The above tests are not currently long enough to exercise chunking.
|
||||
# prompts = (
|
||||
# [_random_prompt(10, 50) for _ in range(28)]
|
||||
# + [_random_prompt(256, 512) for _ in range(50)]
|
||||
# + [_random_prompt(2048, 4096) for _ in range(50)]
|
||||
# )
|
||||
|
||||
sp = SamplingParams(
|
||||
temperature=0.6,
|
||||
top_p=1.0,
|
||||
max_tokens=16,
|
||||
seed=1234,
|
||||
logprobs=5,
|
||||
)
|
||||
|
||||
# BS=1: run prompts individually and collect logprobs per step.
|
||||
print("\n" + "=" * 80)
|
||||
print("STARTING BS=1 RUNS (each prompt individually)")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
bs1_logprobs_per_prompt = []
|
||||
bs1_tokens_per_prompt = []
|
||||
for idx, p in enumerate(prompts):
|
||||
print(f"\n[BS=1] Running prompt {idx}/{len(prompts)} - Preview: {p[:80]}...")
|
||||
outs = llm.generate([p], sp, use_tqdm=False)
|
||||
assert len(outs) == 1
|
||||
step_logprobs, token_ids = _extract_step_logprobs(outs[0])
|
||||
if step_logprobs is None:
|
||||
pytest.skip(
|
||||
"Logits are not available on RequestOutput; "
|
||||
"enable logprobs return to run this test."
|
||||
)
|
||||
bs1_logprobs_per_prompt.append(step_logprobs)
|
||||
bs1_tokens_per_prompt.append(token_ids)
|
||||
print(f"[BS=1] Prompt {idx} generated tokens: {token_ids}")
|
||||
|
||||
# BS=N: run prompts in a batch and collect logprobs per step for each
|
||||
# prompt.
|
||||
print("\n" + "=" * 80)
|
||||
print(f"STARTING BS={len(prompts)} RUN (all prompts batched)")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
outs_batched = llm.generate(prompts, sp, use_tqdm=False)
|
||||
assert len(outs_batched) == len(prompts)
|
||||
bsN_logprobs_per_prompt = []
|
||||
bsN_tokens_per_prompt = []
|
||||
|
||||
print(f"\n[BS={len(prompts)}] Processing batched outputs...")
|
||||
for idx, o in enumerate(outs_batched):
|
||||
tokens = o.outputs[0].token_ids if o.outputs else "N/A"
|
||||
print(f"[BS={len(prompts)}] Prompt {idx} generated tokens: {tokens}")
|
||||
step_logprobs, token_ids = _extract_step_logprobs(o)
|
||||
if step_logprobs is None:
|
||||
pytest.skip(
|
||||
"Logits are not available on RequestOutput; "
|
||||
"enable logprobs return to run this test."
|
||||
)
|
||||
bsN_logprobs_per_prompt.append(step_logprobs)
|
||||
bsN_tokens_per_prompt.append(token_ids)
|
||||
|
||||
# Compare step-by-step logprobs for each prompt between BS=1 and BS=N runs.
|
||||
failed_prompts = []
|
||||
for i, (logprobs_bs1, logprobs_bsN, tokens_bs1, tokens_bsN) in enumerate(
|
||||
zip(
|
||||
bs1_logprobs_per_prompt,
|
||||
bsN_logprobs_per_prompt,
|
||||
bs1_tokens_per_prompt,
|
||||
bsN_tokens_per_prompt,
|
||||
)
|
||||
):
|
||||
if len(logprobs_bs1) != len(logprobs_bsN):
|
||||
reason = (
|
||||
f"Different number of steps: {len(logprobs_bs1)} (BS=1) "
|
||||
f"vs {len(logprobs_bsN)} (BS=N)"
|
||||
)
|
||||
failed_prompts.append(
|
||||
{
|
||||
"prompt_idx": i,
|
||||
"step": "all",
|
||||
"reason": reason,
|
||||
"prompt_preview": prompts[i][:100],
|
||||
"bs1_tokens": tokens_bs1,
|
||||
"bsN_tokens": tokens_bsN,
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
# Check if tokens match first
|
||||
if tokens_bs1 != tokens_bsN:
|
||||
failed_prompts.append(
|
||||
{
|
||||
"prompt_idx": i,
|
||||
"step": "sampling",
|
||||
"reason": "Different tokens sampled",
|
||||
"prompt_preview": prompts[i][:100],
|
||||
"bs1_tokens": tokens_bs1,
|
||||
"bsN_tokens": tokens_bsN,
|
||||
"bs1_all_logprobs": [
|
||||
logprobs_bs1[s].tolist() for s in range(len(logprobs_bs1))
|
||||
],
|
||||
"bsN_all_logprobs": [
|
||||
logprobs_bsN[s].tolist() for s in range(len(logprobs_bsN))
|
||||
],
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
for t, (a, b) in enumerate(zip(logprobs_bs1, logprobs_bsN)):
|
||||
if a.shape != b.shape:
|
||||
failed_prompts.append(
|
||||
{
|
||||
"prompt_idx": i,
|
||||
"step": t,
|
||||
"reason": f"Shape mismatch: {a.shape} vs {b.shape}",
|
||||
"prompt_preview": prompts[i][:100],
|
||||
"bs1_tokens": tokens_bs1,
|
||||
"bsN_tokens": tokens_bsN,
|
||||
}
|
||||
)
|
||||
break
|
||||
|
||||
if not torch.equal(a, b):
|
||||
max_diff = torch.abs(a - b).max().item()
|
||||
# Print which token failed
|
||||
print(f"\n[DIVERGENCE] Prompt {i}, Token {t}: max_diff={max_diff:.6e}")
|
||||
bs1_tok = tokens_bs1[t] if t < len(tokens_bs1) else "N/A"
|
||||
bsN_tok = tokens_bsN[t] if t < len(tokens_bsN) else "N/A"
|
||||
print(f" Token IDs: bs1={bs1_tok}, bsN={bsN_tok}")
|
||||
print(f" BS=1 logprob: {a.tolist()}")
|
||||
print(f" BS=N logprob: {b.tolist()}")
|
||||
failed_prompts.append(
|
||||
{
|
||||
"prompt_idx": i,
|
||||
"step": t,
|
||||
"reason": f"Bitwise mismatch (max_diff={max_diff:.6e})",
|
||||
"prompt_preview": prompts[i][:100],
|
||||
"bs1_tokens": tokens_bs1,
|
||||
"bsN_tokens": tokens_bsN,
|
||||
"bs1_all_logprobs": [
|
||||
logprobs_bs1[s].tolist() for s in range(len(logprobs_bs1))
|
||||
],
|
||||
"bsN_all_logprobs": [
|
||||
logprobs_bsN[s].tolist() for s in range(len(logprobs_bsN))
|
||||
],
|
||||
}
|
||||
)
|
||||
break
|
||||
|
||||
# Print summary of all failures
|
||||
if failed_prompts:
|
||||
print(f"\n{'=' * 80}")
|
||||
fail_msg = (
|
||||
f"BATCH INVARIANCE FAILURES: {len(failed_prompts)}/"
|
||||
f"{len(prompts)} prompts failed"
|
||||
)
|
||||
print(fail_msg)
|
||||
print(f"{'=' * 80}")
|
||||
for fail in failed_prompts:
|
||||
print(f"\nPrompt {fail['prompt_idx']} (step {fail['step']}):")
|
||||
print(f" Reason: {fail['reason']}")
|
||||
print(f" Preview: {fail['prompt_preview']}...")
|
||||
|
||||
# Always show the tokens
|
||||
if "bs1_tokens" in fail:
|
||||
print(f" BS=1 tokens: {fail['bs1_tokens']}")
|
||||
if "bsN_tokens" in fail:
|
||||
print(f" BS=N tokens: {fail['bsN_tokens']}")
|
||||
|
||||
if "bs1_all_logprobs" in fail:
|
||||
print(f" BS=1 logprobs for all {len(fail['bs1_all_logprobs'])} steps:")
|
||||
for step_idx, logprobs in enumerate(fail["bs1_all_logprobs"]):
|
||||
print(f" Step {step_idx}: {logprobs}")
|
||||
print(f" BS=N logprobs for all {len(fail['bsN_all_logprobs'])} steps:")
|
||||
for step_idx, logprobs in enumerate(fail["bsN_all_logprobs"]):
|
||||
print(f" Step {step_idx}: {logprobs}")
|
||||
print(f"{'=' * 80}\n")
|
||||
|
||||
# Fail the test with summary
|
||||
msg = (
|
||||
f"Batch invariance violated in {len(failed_prompts)}/"
|
||||
f"{len(prompts)} prompts. See output above for details."
|
||||
)
|
||||
pytest.fail(msg)
|
||||
|
||||
|
||||
@skip_unsupported
|
||||
@pytest.mark.parametrize(
|
||||
"backend",
|
||||
BACKENDS,
|
||||
)
|
||||
def test_simple_generation(backend):
|
||||
"""
|
||||
Simple test that runs the model with a basic prompt and prints the output.
|
||||
Useful for quick smoke testing and debugging.
|
||||
"""
|
||||
model = resolve_model_name(backend)
|
||||
|
||||
llm = LLM(
|
||||
model=model,
|
||||
max_num_seqs=1,
|
||||
tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),
|
||||
gpu_memory_utilization=0.9,
|
||||
max_model_len=2048,
|
||||
dtype="bfloat16",
|
||||
enable_prefix_caching=False,
|
||||
enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
|
||||
attention_config={"backend": backend},
|
||||
)
|
||||
|
||||
prompt = "the capital of france is"
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.0,
|
||||
max_tokens=20,
|
||||
)
|
||||
|
||||
print(f"\n{'=' * 80}")
|
||||
print("Running simple generation test")
|
||||
print(f"Prompt: '{prompt}'")
|
||||
print(f"{'=' * 80}\n")
|
||||
|
||||
try:
|
||||
outputs = llm.generate([prompt], sampling_params)
|
||||
|
||||
assert len(outputs) == 1
|
||||
output_text = outputs[0].outputs[0].text
|
||||
|
||||
print(f"Output: '{output_text}'")
|
||||
print(f"\n{'=' * 80}")
|
||||
print(f"Full completion: '{prompt}{output_text}'")
|
||||
print(f"{'=' * 80}\n")
|
||||
|
||||
finally:
|
||||
with contextlib.suppress(Exception):
|
||||
llm.shutdown()
|
||||
|
||||
|
||||
@skip_unsupported
|
||||
@pytest.mark.parametrize(
|
||||
"backend",
|
||||
BACKENDS,
|
||||
)
|
||||
def test_logprobs_without_batch_invariance_should_fail(
|
||||
backend, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
"""
|
||||
This test is the inverse of test_logprobs_bitwise_batch_invariance_bs1_vs_bsN.
|
||||
It DISABLES batch invariance mode and expects to see non-deterministic behavior
|
||||
between BS=1 and BS=N runs. This demonstrates that batch invariance is actually
|
||||
doing something useful.
|
||||
|
||||
The test will PASS if we detect differences (proving batch invariance matters).
|
||||
The test will FAIL if everything matches (suggesting batch invariance isn't needed).
|
||||
"""
|
||||
# CRITICAL: Disable batch invariance for this test
|
||||
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0")
|
||||
monkeypatch.setattr(batch_invariant, "VLLM_BATCH_INVARIANT", False)
|
||||
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
|
||||
random.seed(seed)
|
||||
model_name = resolve_model_name(backend)
|
||||
tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
|
||||
|
||||
print(f"\n{'=' * 80}")
|
||||
print("BATCH INVARIANCE DISABLED: Expecting non-deterministic behavior")
|
||||
print(f"{'=' * 80}\n")
|
||||
|
||||
llm = LLM(
|
||||
model=model_name,
|
||||
tensor_parallel_size=tp_size,
|
||||
max_num_seqs=32,
|
||||
max_model_len=8192,
|
||||
dtype="bfloat16",
|
||||
enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
|
||||
attention_config={"backend": backend},
|
||||
)
|
||||
|
||||
# build ragged prompts to change shapes significantly across BS=1 vs BS=N
|
||||
long_min = int(os.getenv("VLLM_MIN_PROMPT", "768"))
|
||||
long_max = int(os.getenv("VLLM_MAX_PROMPT", "2048"))
|
||||
prompts: list[str] = []
|
||||
options = [
|
||||
(max(long_min, 1536), max(long_max, 3072)), # very long
|
||||
(max(1024, long_min), max(2048, long_max)), # long
|
||||
(256, 512), # mid
|
||||
(10, 20), # short
|
||||
]
|
||||
|
||||
for _ in range(32):
|
||||
lo, hi = random.choice(options)
|
||||
prompts.append(_random_prompt(lo, hi))
|
||||
|
||||
sp = SamplingParams(
|
||||
temperature=0.6,
|
||||
top_p=1.0,
|
||||
max_tokens=8,
|
||||
seed=1234,
|
||||
logprobs=5,
|
||||
)
|
||||
|
||||
# BS=1: run prompts individually and collect logprobs per step.
|
||||
print("\n" + "=" * 80)
|
||||
print("STARTING BS=1 RUNS (each prompt individually)")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
bs1_logprobs_per_prompt = []
|
||||
bs1_tokens_per_prompt = []
|
||||
for idx, p in enumerate(prompts):
|
||||
print(f"\n[BS=1] Running prompt {idx}/{len(prompts)} - Preview: {p[:80]}...")
|
||||
outs = llm.generate([p], sp, use_tqdm=False)
|
||||
assert len(outs) == 1
|
||||
step_logprobs, token_ids = _extract_step_logprobs(outs[0])
|
||||
if step_logprobs is None:
|
||||
pytest.skip(
|
||||
"Logits are not available on RequestOutput; "
|
||||
"enable logprobs return to run this test."
|
||||
)
|
||||
bs1_logprobs_per_prompt.append(step_logprobs)
|
||||
bs1_tokens_per_prompt.append(token_ids)
|
||||
print(f"[BS=1] Prompt {idx} generated tokens: {token_ids}")
|
||||
|
||||
# BS=N: run prompts in a batch and collect logprobs per step for each prompt.
|
||||
print("\n" + "=" * 80)
|
||||
print(f"STARTING BS={len(prompts)} RUN (all prompts batched)")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
outs_batched = llm.generate(prompts, sp, use_tqdm=False)
|
||||
assert len(outs_batched) == len(prompts)
|
||||
bsN_logprobs_per_prompt = []
|
||||
bsN_tokens_per_prompt = []
|
||||
|
||||
print(f"\n[BS={len(prompts)}] Processing batched outputs...")
|
||||
for idx, o in enumerate(outs_batched):
|
||||
tokens = o.outputs[0].token_ids if o.outputs else "N/A"
|
||||
print(f"[BS={len(prompts)}] Prompt {idx} generated tokens: {tokens}")
|
||||
step_logprobs, token_ids = _extract_step_logprobs(o)
|
||||
if step_logprobs is None:
|
||||
pytest.skip(
|
||||
"Logits are not available on RequestOutput; "
|
||||
"enable logprobs return to run this test."
|
||||
)
|
||||
bsN_logprobs_per_prompt.append(step_logprobs)
|
||||
bsN_tokens_per_prompt.append(token_ids)
|
||||
|
||||
# Compare step-by-step logprobs for each prompt between BS=1 and BS=N runs.
|
||||
differences_found = []
|
||||
for i, (logprobs_bs1, logprobs_bsN, tokens_bs1, tokens_bsN) in enumerate(
|
||||
zip(
|
||||
bs1_logprobs_per_prompt,
|
||||
bsN_logprobs_per_prompt,
|
||||
bs1_tokens_per_prompt,
|
||||
bsN_tokens_per_prompt,
|
||||
)
|
||||
):
|
||||
if len(logprobs_bs1) != len(logprobs_bsN):
|
||||
reason = (
|
||||
f"Different number of steps: {len(logprobs_bs1)} (BS=1) "
|
||||
f"vs {len(logprobs_bsN)} (BS=N)"
|
||||
)
|
||||
differences_found.append(
|
||||
{
|
||||
"prompt_idx": i,
|
||||
"step": "all",
|
||||
"reason": reason,
|
||||
"prompt_preview": prompts[i][:100],
|
||||
"bs1_tokens": tokens_bs1,
|
||||
"bsN_tokens": tokens_bsN,
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
# Check if tokens match first
|
||||
if tokens_bs1 != tokens_bsN:
|
||||
differences_found.append(
|
||||
{
|
||||
"prompt_idx": i,
|
||||
"step": "sampling",
|
||||
"reason": "Different tokens sampled",
|
||||
"prompt_preview": prompts[i][:100],
|
||||
"bs1_tokens": tokens_bs1,
|
||||
"bsN_tokens": tokens_bsN,
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
for t, (a, b) in enumerate(zip(logprobs_bs1, logprobs_bsN)):
|
||||
if a.shape != b.shape:
|
||||
differences_found.append(
|
||||
{
|
||||
"prompt_idx": i,
|
||||
"step": t,
|
||||
"reason": f"Shape mismatch: {a.shape} vs {b.shape}",
|
||||
"prompt_preview": prompts[i][:100],
|
||||
"bs1_tokens": tokens_bs1,
|
||||
"bsN_tokens": tokens_bsN,
|
||||
}
|
||||
)
|
||||
break
|
||||
|
||||
if not torch.equal(a, b):
|
||||
max_diff = torch.abs(a - b).max().item()
|
||||
print(
|
||||
f"\n[EXPECTED DIVERGENCE FOUND] Prompt {i}, "
|
||||
f"Token {t}: max_diff={max_diff:.6e}"
|
||||
)
|
||||
bs1_tok = tokens_bs1[t] if t < len(tokens_bs1) else "N/A"
|
||||
bsN_tok = tokens_bsN[t] if t < len(tokens_bsN) else "N/A"
|
||||
print(f" Token IDs: bs1={bs1_tok}, bsN={bsN_tok}")
|
||||
print(f" BS=1 logprob: {a.tolist()}")
|
||||
print(f" BS=N logprob: {b.tolist()}")
|
||||
differences_found.append(
|
||||
{
|
||||
"prompt_idx": i,
|
||||
"step": t,
|
||||
"reason": f"Bitwise mismatch (max_diff={max_diff:.6e})",
|
||||
"prompt_preview": prompts[i][:100],
|
||||
"bs1_tokens": tokens_bs1,
|
||||
"bsN_tokens": tokens_bsN,
|
||||
}
|
||||
)
|
||||
break
|
||||
|
||||
# Print summary
|
||||
print(f"\n{'=' * 80}")
|
||||
if differences_found:
|
||||
success_msg = (
|
||||
f"✓ SUCCESS: Batch invariance is doing something! "
|
||||
f"Found {len(differences_found)}/{len(prompts)} prompts "
|
||||
f"with differences when batch invariance was DISABLED."
|
||||
)
|
||||
print(success_msg)
|
||||
print(f"{'=' * 80}")
|
||||
for diff in differences_found:
|
||||
print(f"\nPrompt {diff['prompt_idx']} (step {diff['step']}):")
|
||||
print(f" Reason: {diff['reason']}")
|
||||
print(f" Preview: {diff['prompt_preview']}...")
|
||||
if "bs1_tokens" in diff:
|
||||
print(f" BS=1 tokens: {diff['bs1_tokens']}")
|
||||
if "bsN_tokens" in diff:
|
||||
print(f" BS=N tokens: {diff['bsN_tokens']}")
|
||||
print(f"{'=' * 80}\n")
|
||||
# Test PASSES because we found differences (batch invariance matters!)
|
||||
return
|
||||
else:
|
||||
# Test FAILS because everything matched even without batch invariance
|
||||
fail_msg = (
|
||||
f"✗ UNEXPECTED: All {len(prompts)} prompts matched "
|
||||
f"between BS=1 and BS=N even with batch invariance DISABLED. "
|
||||
f"This suggests batch invariance might not be necessary, "
|
||||
f"or the test needs more sensitive prompts."
|
||||
)
|
||||
print(fail_msg)
|
||||
print(f"{'=' * 80}\n")
|
||||
pytest.fail(fail_msg)
|
||||
|
||||
|
||||
@skip_unsupported
|
||||
@pytest.mark.parametrize("backend", ["FLASH_ATTN"])
|
||||
def test_decode_logprobs_match_prefill_logprobs(
|
||||
backend,
|
||||
):
|
||||
"""
|
||||
Test that verifies decode logprobs match prefill logprobs.
|
||||
|
||||
For each decoded token at position i:
|
||||
1. Run decode to generate N tokens and collect their logprobs
|
||||
2. For each position i in [0, N):
|
||||
- Take prefix = prompt + tokens[0:i]
|
||||
- Run prefill(prefix + tokens[i]) to get logprob of tokens[i]
|
||||
- Verify prefill logprob matches decode logprob bitwise
|
||||
|
||||
This ensures that the logprobs from decode are consistent with what
|
||||
we would get if we ran prefill on each prefix.
|
||||
"""
|
||||
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
|
||||
random.seed(seed)
|
||||
model_name = resolve_model_name(backend)
|
||||
tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
|
||||
|
||||
from vllm.model_executor.layers.batch_invariant import (
|
||||
vllm_is_batch_invariant,
|
||||
)
|
||||
|
||||
disable_custom_ar = vllm_is_batch_invariant()
|
||||
|
||||
if disable_custom_ar:
|
||||
print(f"\n{'=' * 80}")
|
||||
print(f"BATCH INVARIANCE MODE: Disabling custom all-reduce (TP={tp_size})")
|
||||
print(f"{'=' * 80}\n")
|
||||
|
||||
llm = LLM(
|
||||
model=model_name,
|
||||
tensor_parallel_size=tp_size,
|
||||
max_num_seqs=32,
|
||||
max_model_len=8192,
|
||||
dtype="bfloat16",
|
||||
enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
|
||||
attention_config={"backend": backend},
|
||||
)
|
||||
|
||||
# Use a few test prompts
|
||||
num_test_prompts = int(os.getenv("VLLM_DECODE_PREFILL_NUM_PROMPTS", "4"))
|
||||
prompts = [_random_prompt(10, 50) for _ in range(num_test_prompts)]
|
||||
|
||||
# Generate longer sequences to test multiple decode steps
|
||||
max_tokens = int(os.getenv("VLLM_DECODE_PREFILL_MAX_TOKENS", "16"))
|
||||
|
||||
sp = SamplingParams(
|
||||
temperature=0.0, # Greedy for determinism
|
||||
max_tokens=max_tokens,
|
||||
logprobs=5,
|
||||
)
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("STEP 1: Running decode to generate tokens and collect logprobs")
|
||||
print("=" * 80 + "\n")
|
||||
|
||||
# Step 1: Run decode and collect logprobs
|
||||
decode_outputs = llm.generate(prompts, sp, use_tqdm=False)
|
||||
|
||||
failed_comparisons = []
|
||||
|
||||
for prompt_idx, (prompt, decode_output) in enumerate(zip(prompts, decode_outputs)):
|
||||
print(f"\n[Prompt {prompt_idx}] Testing: {prompt[:80]}...")
|
||||
|
||||
# Extract decode logprobs and tokens
|
||||
decode_logprobs, token_ids = _extract_step_logprobs(decode_output)
|
||||
if decode_logprobs is None:
|
||||
pytest.skip(
|
||||
"Logprobs are not available on RequestOutput; "
|
||||
"enable logprobs return to run this test."
|
||||
)
|
||||
|
||||
print(f"[Prompt {prompt_idx}] Generated {len(token_ids)} tokens: {token_ids}")
|
||||
print(f"[Prompt {prompt_idx}] Decode logprobs: {decode_logprobs.tolist()}")
|
||||
|
||||
# Step 2: For each token position, run prefill and compare
|
||||
print(f"\n[Prompt {prompt_idx}] Verifying each token via prefill...")
|
||||
|
||||
for token_idx in range(len(token_ids)):
|
||||
# Construct the prefix up to (but not including) this token
|
||||
current_token = token_ids[token_idx]
|
||||
|
||||
# We need to detokenize to get the text prefix
|
||||
# For this, we'll use the tokenizer from the LLM
|
||||
# However, the LLM API doesn't expose tokenizer easily, so we'll
|
||||
# construct the prefix by decoding from the original prompt
|
||||
|
||||
# Get text up to this point by using the output text
|
||||
# This is approximate but should work for verification
|
||||
if token_idx == 0:
|
||||
prefix_prompt = prompt
|
||||
else:
|
||||
# Use the partial output text up to this token
|
||||
# We'll need to construct this from the full output
|
||||
prefix_output = decode_output.outputs[0]
|
||||
# Get the text for tokens 0 to token_idx-1
|
||||
# Unfortunately, we don't have per-token text, so we'll use
|
||||
# a different approach: run prefill with prompt + tokens[0:token_idx]
|
||||
|
||||
# Actually, we need to get the actual text. Let's use a workaround:
|
||||
# Run a generation with max_tokens = token_idx to get that prefix
|
||||
prefix_sp = SamplingParams(
|
||||
temperature=0.0,
|
||||
max_tokens=token_idx,
|
||||
logprobs=1,
|
||||
)
|
||||
prefix_output = llm.generate([prompt], prefix_sp, use_tqdm=False)[0]
|
||||
prefix_prompt = prompt + prefix_output.outputs[0].text
|
||||
|
||||
# Now run prefill with max_tokens=1 to get the logprob of the next token
|
||||
prefill_sp = SamplingParams(
|
||||
temperature=0.0,
|
||||
max_tokens=1,
|
||||
logprobs=5,
|
||||
)
|
||||
|
||||
print(
|
||||
f" [Token {token_idx}] Running prefill for prefix "
|
||||
f"(len={len(prefix_prompt)})..."
|
||||
)
|
||||
prefill_output = llm.generate([prefix_prompt], prefill_sp, use_tqdm=False)[
|
||||
0
|
||||
]
|
||||
prefill_logprobs, prefill_token_ids = _extract_step_logprobs(prefill_output)
|
||||
|
||||
if prefill_logprobs is None:
|
||||
print(f" [Token {token_idx}] Warning: No prefill logprobs available")
|
||||
continue
|
||||
|
||||
# The first token from prefill should match the current token
|
||||
prefill_token = prefill_token_ids[0]
|
||||
prefill_logprob = prefill_logprobs[0].item()
|
||||
decode_logprob = decode_logprobs[token_idx].item()
|
||||
|
||||
print(
|
||||
f" [Token {token_idx}] Decode token: {current_token}, "
|
||||
f"logprob: {decode_logprob:.8f}"
|
||||
)
|
||||
print(
|
||||
f" [Token {token_idx}] Prefill token: {prefill_token}, "
|
||||
f"logprob: {prefill_logprob:.8f}"
|
||||
)
|
||||
|
||||
# Check if tokens match
|
||||
if current_token != prefill_token:
|
||||
failed_comparisons.append(
|
||||
{
|
||||
"prompt_idx": prompt_idx,
|
||||
"token_idx": token_idx,
|
||||
"reason": "Token mismatch",
|
||||
"decode_token": current_token,
|
||||
"prefill_token": prefill_token,
|
||||
"decode_logprob": decode_logprob,
|
||||
"prefill_logprob": prefill_logprob,
|
||||
"prompt_text": prompt[:100],
|
||||
"prefix_text": prefix_prompt[:100],
|
||||
}
|
||||
)
|
||||
print(f" [Token {token_idx}] ✗ TOKEN MISMATCH!")
|
||||
continue
|
||||
|
||||
# Check if logprobs match bitwise
|
||||
if decode_logprob != prefill_logprob:
|
||||
diff = abs(decode_logprob - prefill_logprob)
|
||||
failed_comparisons.append(
|
||||
{
|
||||
"prompt_idx": prompt_idx,
|
||||
"token_idx": token_idx,
|
||||
"reason": "Logprob mismatch",
|
||||
"decode_token": current_token,
|
||||
"prefill_token": prefill_token,
|
||||
"decode_logprob": decode_logprob,
|
||||
"prefill_logprob": prefill_logprob,
|
||||
"diff": diff,
|
||||
"prompt_text": prompt[:100],
|
||||
"prefix_text": prefix_prompt[:100],
|
||||
"decode_all_tokens": token_ids,
|
||||
"decode_all_logprobs": decode_logprobs.tolist(),
|
||||
}
|
||||
)
|
||||
print(f" [Token {token_idx}] ✗ LOGPROB MISMATCH! diff={diff:.8e}")
|
||||
else:
|
||||
print(f" [Token {token_idx}] ✓ Match (bitwise equal)")
|
||||
|
||||
# Print summary
|
||||
print(f"\n{'=' * 80}")
|
||||
if failed_comparisons:
|
||||
print(f"DECODE-PREFILL MISMATCH: {len(failed_comparisons)} failures detected")
|
||||
print(f"{'=' * 80}")
|
||||
|
||||
# Group failures by prompt for better readability
|
||||
failures_by_prompt: dict[int, list[dict]] = {}
|
||||
for fail in failed_comparisons:
|
||||
pid = fail["prompt_idx"]
|
||||
if pid not in failures_by_prompt:
|
||||
failures_by_prompt[pid] = []
|
||||
failures_by_prompt[pid].append(fail)
|
||||
|
||||
for prompt_idx, failures in failures_by_prompt.items():
|
||||
print(f"\n{'=' * 80}")
|
||||
print(f"PROMPT {prompt_idx}: {failures[0]['prompt_text']}...")
|
||||
print(f"{'=' * 80}")
|
||||
print(f"Total failures for this prompt: {len(failures)}")
|
||||
|
||||
# Show where mismatches occur (which token positions)
|
||||
mismatch_positions = [f["token_idx"] for f in failures]
|
||||
print(f"Mismatch at token positions: {mismatch_positions}")
|
||||
|
||||
# Show first few failures in detail
|
||||
for i, fail in enumerate(failures[:5]): # Show first 5 failures per prompt
|
||||
print(f"\n [Failure {i + 1}] Token position {fail['token_idx']}:")
|
||||
print(f" Reason: {fail['reason']}")
|
||||
print(f" Prefix text: '{fail['prefix_text']}...'")
|
||||
print(
|
||||
f" Decode: token={fail['decode_token']}, "
|
||||
f"logprob={fail['decode_logprob']:.10f}"
|
||||
)
|
||||
print(
|
||||
f" Prefill: token={fail['prefill_token']}, "
|
||||
f"logprob={fail['prefill_logprob']:.10f}"
|
||||
)
|
||||
if "diff" in fail:
|
||||
print(f" Difference: {fail['diff']:.10e}")
|
||||
# Show in hex to see bitwise difference
|
||||
import struct
|
||||
|
||||
decode_hex = struct.pack("f", fail["decode_logprob"]).hex()
|
||||
prefill_hex = struct.pack("f", fail["prefill_logprob"]).hex()
|
||||
print(f" Decode logprob (hex): 0x{decode_hex}")
|
||||
print(f" Prefill logprob (hex): 0x{prefill_hex}")
|
||||
|
||||
# If we have all tokens/logprobs, show the context
|
||||
if "decode_all_tokens" in fail and "decode_all_logprobs" in fail:
|
||||
token_idx = fail["token_idx"]
|
||||
all_tokens = fail["decode_all_tokens"]
|
||||
all_logprobs = fail["decode_all_logprobs"]
|
||||
|
||||
# Show context: 2 tokens before and after
|
||||
start = max(0, token_idx - 2)
|
||||
end = min(len(all_tokens), token_idx + 3)
|
||||
|
||||
print(f" Context (tokens {start} to {end - 1}):")
|
||||
for j in range(start, end):
|
||||
marker = " <-- MISMATCH" if j == token_idx else ""
|
||||
print(
|
||||
f" [{j}] token={all_tokens[j]}, "
|
||||
f"logprob={all_logprobs[j]:.8f}{marker}"
|
||||
)
|
||||
|
||||
if len(failures) > 5:
|
||||
print(f"\n ... and {len(failures) - 5} more failures for this prompt")
|
||||
|
||||
print(f"\n{'=' * 80}\n")
|
||||
|
||||
pytest.fail(
|
||||
f"Decode logprobs do not match prefill logprobs: "
|
||||
f"{len(failed_comparisons)} mismatches found."
|
||||
)
|
||||
else:
|
||||
print("✓ SUCCESS: All decode logprobs match prefill logprobs bitwise!")
|
||||
print(f"{'=' * 80}\n")
|
||||
|
||||
|
||||
def LLM_with_max_seqs(
|
||||
model: str,
|
||||
max_num_seqs: int,
|
||||
gpu_memory_utilization: float,
|
||||
max_model_len: int,
|
||||
attention_config: dict | None = None,
|
||||
) -> LLM:
|
||||
"""
|
||||
Helper to construct an LLM with a specific max_num_seqs (batch-size limit)
|
||||
using the high-level v1 LLM API, while constraining memory usage.
|
||||
"""
|
||||
return LLM(
|
||||
model=model,
|
||||
max_num_seqs=max_num_seqs,
|
||||
gpu_memory_utilization=gpu_memory_utilization,
|
||||
max_model_len=max_model_len,
|
||||
dtype="bfloat16",
|
||||
tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),
|
||||
enable_prefix_caching=False,
|
||||
enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
|
||||
attention_config=attention_config,
|
||||
# Enable for MOE models
|
||||
# enable_expert_parallel=True,
|
||||
)
|
||||
169
third_party/vllm/tests/v1/determinism/test_online_batch_invariance.py
vendored
Normal file
169
third_party/vllm/tests/v1/determinism/test_online_batch_invariance.py
vendored
Normal file
@@ -0,0 +1,169 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
HTTP-based batch invariance test: send requests to a running
|
||||
vLLM server and compare BS=1 vs BS=N results (tokens and per-step logprobs).
|
||||
|
||||
Environment variables:
|
||||
- VLLM_TEST_MODEL: served model name (e.g., Qwen/Qwen3-1.7B / DeepSeek-R1)
|
||||
- VLLM_TP_SIZE: tensor parallelism size (e.g., 4)
|
||||
|
||||
"""
|
||||
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
from utils import BACKENDS, _random_prompt, resolve_model_name, skip_unsupported
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
|
||||
def _request_completion(
|
||||
client: openai.OpenAI,
|
||||
model: str,
|
||||
prompt: Any,
|
||||
sp: dict[str, Any],
|
||||
max_retries: int = 3,
|
||||
retry_backoff: float = 0.5,
|
||||
) -> dict[str, Any] | None:
|
||||
payload: dict[str, Any] = {"model": model, "prompt": prompt}
|
||||
payload.update(sp)
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
completion = client.completions.create(**payload)
|
||||
# Convert to plain dict so downstream logic can keep using
|
||||
# dict-style access just like with raw HTTP JSON.
|
||||
return completion.model_dump()
|
||||
except Exception as e: # pragma: no cover
|
||||
if attempt < max_retries:
|
||||
import time as _t
|
||||
|
||||
_t.sleep(retry_backoff * (2**attempt))
|
||||
continue
|
||||
sys.stderr.write(f"Error: {e}\n")
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _extract_tokens_and_logprobs(
|
||||
choice: dict[str, Any],
|
||||
) -> tuple[list[Any], list[float] | None]:
|
||||
tokens: list[Any] = []
|
||||
token_logprobs: list[float] | None = None
|
||||
lp = choice.get("logprobs")
|
||||
if lp and isinstance(lp, dict):
|
||||
tokens = lp.get("token_ids") or lp.get("tokens") or []
|
||||
token_logprobs = lp.get("token_logprobs", None)
|
||||
return tokens, token_logprobs
|
||||
|
||||
|
||||
def _compare_bs1_vs_bsn_single_process(
|
||||
prompts: list[str],
|
||||
sp_kwargs: dict[str, Any],
|
||||
client: openai.OpenAI,
|
||||
model_name: str,
|
||||
) -> None:
|
||||
# BS=1
|
||||
bs1_tokens_per_prompt: list[list[Any]] = []
|
||||
bs1_logprobs_per_prompt: list[list[float] | None] = []
|
||||
for p in prompts:
|
||||
resp = _request_completion(client, model_name, p, sp_kwargs)
|
||||
if resp is None or not resp.get("choices"):
|
||||
raise AssertionError("BS=1 empty/failed response")
|
||||
choice = resp["choices"][0]
|
||||
toks, lps = _extract_tokens_and_logprobs(choice)
|
||||
if lps is None:
|
||||
raise AssertionError(
|
||||
"logprobs not returned; ensure server supports 'logprobs'"
|
||||
)
|
||||
bs1_tokens_per_prompt.append(list(toks))
|
||||
bs1_logprobs_per_prompt.append(list(lps))
|
||||
|
||||
# BS=N
|
||||
bsN_tokens_per_prompt: list[list[Any]] = [None] * len(prompts) # type: ignore[list-item]
|
||||
bsN_logprobs_per_prompt: list[list[float] | None] = [None] * len(prompts)
|
||||
resp = _request_completion(client, model_name, prompts, sp_kwargs)
|
||||
if resp is None or not resp.get("choices"):
|
||||
raise AssertionError("BS=N empty/failed batched response")
|
||||
choices = resp.get("choices", [])
|
||||
if len(choices) != len(prompts):
|
||||
raise AssertionError(
|
||||
f"BS=N choices length {len(choices)} != num prompts {len(prompts)}"
|
||||
)
|
||||
for idx, choice in enumerate(choices):
|
||||
toks, lps = _extract_tokens_and_logprobs(choice)
|
||||
if lps is None:
|
||||
raise AssertionError(f"BS=N missing logprobs for prompt {idx}")
|
||||
bsN_tokens_per_prompt[idx] = list(toks)
|
||||
bsN_logprobs_per_prompt[idx] = list(lps)
|
||||
|
||||
# compare
|
||||
for i, (tokens_bs1, tokens_bsN, logprobs_bs1, logprobs_bsN) in enumerate(
|
||||
zip(
|
||||
bs1_tokens_per_prompt,
|
||||
bsN_tokens_per_prompt,
|
||||
bs1_logprobs_per_prompt,
|
||||
bsN_logprobs_per_prompt,
|
||||
)
|
||||
):
|
||||
if tokens_bs1 != tokens_bsN:
|
||||
raise AssertionError(
|
||||
f"Prompt {i} (sampling): Different tokens sampled. "
|
||||
f"BS=1 tokens: {tokens_bs1} BS=N tokens: {tokens_bsN}"
|
||||
)
|
||||
if logprobs_bs1 is None or logprobs_bsN is None:
|
||||
raise AssertionError(f"Prompt {i}: Missing logprobs in one of the runs")
|
||||
if len(logprobs_bs1) != len(logprobs_bsN):
|
||||
raise AssertionError(
|
||||
f"Prompt {i}: Different number of steps: "
|
||||
f"{len(logprobs_bs1)} (BS=1) vs {len(logprobs_bsN)} (BS=N)."
|
||||
)
|
||||
for t, (a, b) in enumerate(zip(logprobs_bs1, logprobs_bsN)):
|
||||
if a != b:
|
||||
diff = abs(a - b)
|
||||
raise AssertionError(
|
||||
f"Prompt {i} Step {t}: Bitwise mismatch "
|
||||
f"(abs diff={diff:.6e}). "
|
||||
f"BS=1 tokens: {tokens_bs1} BS=N tokens: {tokens_bsN}"
|
||||
)
|
||||
|
||||
|
||||
@skip_unsupported
|
||||
@pytest.mark.parametrize("backend", BACKENDS)
|
||||
def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
|
||||
backend: str,
|
||||
) -> None:
|
||||
random.seed(int(os.getenv("VLLM_TEST_SEED", "12345")))
|
||||
model_name = resolve_model_name(backend)
|
||||
prompts_all = [_random_prompt(10, 50) for _ in range(32)]
|
||||
|
||||
sp_kwargs: dict[str, Any] = {
|
||||
"temperature": 0.6,
|
||||
"top_p": 1.0,
|
||||
"max_tokens": 8,
|
||||
"seed": 42,
|
||||
"logprobs": 5,
|
||||
}
|
||||
|
||||
tp_size = os.getenv("VLLM_TP_SIZE", "1")
|
||||
server_args: list[str] = [
|
||||
"--max-model-len=8192",
|
||||
"--max-num-seqs=32",
|
||||
f"--attention-backend={backend}",
|
||||
]
|
||||
if tp_size:
|
||||
server_args += ["-tp", tp_size]
|
||||
|
||||
with RemoteOpenAIServer(model_name, server_args) as server:
|
||||
client = server.get_client()
|
||||
_compare_bs1_vs_bsn_single_process(
|
||||
prompts=prompts_all,
|
||||
sp_kwargs=sp_kwargs,
|
||||
client=client,
|
||||
model_name=model_name,
|
||||
)
|
||||
316
third_party/vllm/tests/v1/determinism/test_rms_norm_batch_invariant.py
vendored
Normal file
316
third_party/vllm/tests/v1/determinism/test_rms_norm_batch_invariant.py
vendored
Normal file
@@ -0,0 +1,316 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Test batch-invariant RMS normalization against standard implementations.
|
||||
|
||||
This test compares the Triton-based batch-invariant RMS norm implementation
|
||||
with the standard CUDA-based implementation to ensure numerical accuracy.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from utils import skip_unsupported
|
||||
|
||||
from vllm.model_executor.layers.batch_invariant import rms_norm as triton_rms_norm
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
|
||||
|
||||
@skip_unsupported
|
||||
@pytest.mark.parametrize("batch_size", [1, 4, 16, 64])
|
||||
@pytest.mark.parametrize("hidden_size", [512, 2048, 4096, 8192])
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize("eps", [1e-6, 1e-5])
|
||||
def test_rms_norm_batch_invariant_vs_standard(
|
||||
default_vllm_config,
|
||||
batch_size: int,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
eps: float,
|
||||
):
|
||||
"""
|
||||
Compare batch-invariant Triton RMS norm against standard CUDA implementation.
|
||||
|
||||
Tests that the Triton-based batch-invariant RMS norm produces numerically
|
||||
equivalent results to the standard CUDA implementation across various
|
||||
configurations.
|
||||
"""
|
||||
device = torch.device("cuda")
|
||||
|
||||
# Create test input and weight
|
||||
torch.manual_seed(42)
|
||||
input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
|
||||
weight = torch.randn(hidden_size, dtype=dtype, device=device)
|
||||
|
||||
# Standard implementation (CUDA ops)
|
||||
rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)
|
||||
rms_norm_layer.weight.data = weight.clone()
|
||||
|
||||
standard_output = rms_norm_layer.forward_cuda(input_tensor)
|
||||
|
||||
# Batch-invariant implementation (Triton)
|
||||
triton_output = triton_rms_norm(input_tensor, weight, eps=eps)
|
||||
|
||||
# Compare outputs
|
||||
# Use looser tolerance for bfloat16 due to its lower precision
|
||||
if dtype == torch.bfloat16:
|
||||
rtol, atol = 1e-1, 1e-1 # 10% relative tolerance for bfloat16
|
||||
else:
|
||||
rtol, atol = 1e-2, 1e-2 # 1% for float16/float32
|
||||
|
||||
torch.testing.assert_close(
|
||||
triton_output,
|
||||
standard_output,
|
||||
rtol=rtol,
|
||||
atol=atol,
|
||||
msg=f"RMS norm mismatch for batch_size={batch_size}, "
|
||||
f"hidden_size={hidden_size}, "
|
||||
f"dtype={dtype}, eps={eps}",
|
||||
)
|
||||
|
||||
|
||||
@skip_unsupported
|
||||
@pytest.mark.parametrize("batch_size", [1, 16, 128])
|
||||
@pytest.mark.parametrize("seq_len", [1, 32, 512])
|
||||
@pytest.mark.parametrize("hidden_size", [2048, 4096])
|
||||
def test_rms_norm_3d_input(
|
||||
default_vllm_config, batch_size: int, seq_len: int, hidden_size: int
|
||||
):
|
||||
"""
|
||||
Test RMS norm with 3D input tensors (batch, seq_len, hidden_size).
|
||||
|
||||
Ensures that the batch-invariant RMS norm correctly handles multi-dimensional
|
||||
inputs that are common in transformer models.
|
||||
"""
|
||||
device = torch.device("cuda")
|
||||
dtype = torch.bfloat16
|
||||
eps = 1e-6
|
||||
|
||||
torch.manual_seed(42)
|
||||
input_tensor = torch.randn(
|
||||
batch_size, seq_len, hidden_size, dtype=dtype, device=device
|
||||
)
|
||||
weight = torch.randn(hidden_size, dtype=dtype, device=device)
|
||||
|
||||
# Standard implementation
|
||||
rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)
|
||||
rms_norm_layer.weight.data = weight.clone()
|
||||
standard_output = rms_norm_layer.forward_cuda(input_tensor)
|
||||
|
||||
# Batch-invariant implementation
|
||||
triton_output = triton_rms_norm(input_tensor, weight, eps=eps)
|
||||
|
||||
# Use looser tolerance for bfloat16
|
||||
rtol, atol = 1e-1, 1e-1 # 10% tolerance for bfloat16
|
||||
|
||||
torch.testing.assert_close(
|
||||
triton_output,
|
||||
standard_output,
|
||||
rtol=rtol,
|
||||
atol=atol,
|
||||
msg=f"RMS norm mismatch for 3D input with batch_size={batch_size}, "
|
||||
f"seq_len={seq_len}, hidden_size={hidden_size}",
|
||||
)
|
||||
|
||||
|
||||
@skip_unsupported
|
||||
def test_rms_norm_numerical_stability(default_vllm_config):
|
||||
"""
|
||||
Test RMS norm numerical stability with extreme values.
|
||||
|
||||
Ensures that both implementations handle edge cases like very small or large
|
||||
values without producing NaN or Inf.
|
||||
"""
|
||||
device = torch.device("cuda")
|
||||
dtype = torch.float16
|
||||
eps = 1e-6
|
||||
hidden_size = 2048
|
||||
|
||||
# Test cases with extreme values
|
||||
test_cases = [
|
||||
# Very small values
|
||||
torch.ones(4, hidden_size, dtype=dtype, device=device) * 1e-5,
|
||||
# Very large values
|
||||
torch.ones(4, hidden_size, dtype=dtype, device=device) * 1e4,
|
||||
# Mixed small and large
|
||||
torch.randn(4, hidden_size, dtype=dtype, device=device) * 100,
|
||||
# Values near zero
|
||||
torch.randn(4, hidden_size, dtype=dtype, device=device) * 1e-6,
|
||||
]
|
||||
|
||||
weight = torch.ones(hidden_size, dtype=dtype, device=device)
|
||||
|
||||
for idx, input_tensor in enumerate(test_cases):
|
||||
# Standard implementation
|
||||
rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)
|
||||
rms_norm_layer.weight.data = weight.clone()
|
||||
standard_output = rms_norm_layer.forward_cuda(input_tensor)
|
||||
|
||||
# Batch-invariant implementation
|
||||
triton_output = triton_rms_norm(input_tensor, weight, eps=eps)
|
||||
|
||||
# Check for NaN or Inf
|
||||
assert not torch.isnan(standard_output).any(), (
|
||||
f"Standard RMS norm produced NaN for test case {idx}"
|
||||
)
|
||||
assert not torch.isinf(standard_output).any(), (
|
||||
f"Standard RMS norm produced Inf for test case {idx}"
|
||||
)
|
||||
assert not torch.isnan(triton_output).any(), (
|
||||
f"Triton RMS norm produced NaN for test case {idx}"
|
||||
)
|
||||
assert not torch.isinf(triton_output).any(), (
|
||||
f"Triton RMS norm produced Inf for test case {idx}"
|
||||
)
|
||||
|
||||
# Compare outputs - very lenient for extreme values with float16
|
||||
torch.testing.assert_close(
|
||||
triton_output,
|
||||
standard_output,
|
||||
rtol=2e-1, # 20% tolerance for extreme values
|
||||
atol=2e-1,
|
||||
msg=f"RMS norm mismatch for extreme value test case {idx}",
|
||||
)
|
||||
|
||||
|
||||
@skip_unsupported
|
||||
def test_rms_norm_formula(default_vllm_config):
|
||||
"""
|
||||
Test that RMS norm follows the correct mathematical formula.
|
||||
|
||||
Verifies: output = input / sqrt(mean(input^2) + eps) * weight
|
||||
"""
|
||||
device = torch.device("cuda")
|
||||
dtype = torch.float32 # Use float32 for higher precision in formula check
|
||||
eps = 1e-6
|
||||
hidden_size = 1024
|
||||
|
||||
torch.manual_seed(42)
|
||||
input_tensor = torch.randn(8, hidden_size, dtype=dtype, device=device)
|
||||
weight = torch.randn(hidden_size, dtype=dtype, device=device)
|
||||
|
||||
# Compute expected output using the formula
|
||||
variance = (input_tensor.pow(2).mean(dim=-1, keepdim=True)).to(dtype)
|
||||
expected_output = input_tensor * torch.rsqrt(variance + eps) * weight
|
||||
|
||||
# Batch-invariant implementation
|
||||
triton_output = triton_rms_norm(input_tensor, weight, eps=eps)
|
||||
|
||||
# Compare against formula
|
||||
torch.testing.assert_close(
|
||||
triton_output,
|
||||
expected_output,
|
||||
rtol=1e-4,
|
||||
atol=1e-4,
|
||||
msg="Triton RMS norm doesn't match expected formula",
|
||||
)
|
||||
|
||||
|
||||
@skip_unsupported
|
||||
@pytest.mark.parametrize("hidden_size", [128, 1024, 4096, 16384])
|
||||
def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int):
|
||||
"""
|
||||
Test RMS norm with various hidden sizes to ensure block size handling.
|
||||
|
||||
The Triton kernel uses a fixed BLOCK_SIZE=1024, so this tests that it
|
||||
correctly handles hidden sizes both smaller and larger than the block size.
|
||||
"""
|
||||
device = torch.device("cuda")
|
||||
dtype = torch.bfloat16
|
||||
eps = 1e-6
|
||||
batch_size = 16
|
||||
|
||||
torch.manual_seed(42)
|
||||
input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
|
||||
weight = torch.randn(hidden_size, dtype=dtype, device=device)
|
||||
|
||||
# Standard implementation
|
||||
rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)
|
||||
rms_norm_layer.weight.data = weight.clone()
|
||||
standard_output = rms_norm_layer.forward_cuda(input_tensor)
|
||||
|
||||
# Batch-invariant implementation
|
||||
triton_output = triton_rms_norm(input_tensor, weight, eps=eps)
|
||||
|
||||
# Use looser tolerance for bfloat16
|
||||
rtol, atol = 1e-1, 1e-1 # 10% tolerance for bfloat16
|
||||
|
||||
torch.testing.assert_close(
|
||||
triton_output,
|
||||
standard_output,
|
||||
rtol=rtol,
|
||||
atol=atol,
|
||||
msg=f"RMS norm mismatch for hidden_size={hidden_size}",
|
||||
)
|
||||
|
||||
|
||||
@skip_unsupported
|
||||
def test_rms_norm_determinism(default_vllm_config):
|
||||
"""
|
||||
Test that batch-invariant RMS norm produces deterministic results.
|
||||
|
||||
Runs the same input through the kernel multiple times and verifies
|
||||
identical outputs.
|
||||
"""
|
||||
device = torch.device("cuda")
|
||||
dtype = torch.bfloat16
|
||||
eps = 1e-6
|
||||
hidden_size = 4096
|
||||
batch_size = 32
|
||||
|
||||
torch.manual_seed(42)
|
||||
input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
|
||||
weight = torch.randn(hidden_size, dtype=dtype, device=device)
|
||||
|
||||
# Run multiple times
|
||||
outputs = []
|
||||
for _ in range(5):
|
||||
output = triton_rms_norm(input_tensor.clone(), weight, eps=eps)
|
||||
outputs.append(output)
|
||||
|
||||
# All outputs should be identical
|
||||
reference = outputs[0]
|
||||
for idx, output in enumerate(outputs[1:], start=1):
|
||||
torch.testing.assert_close(
|
||||
output,
|
||||
reference,
|
||||
rtol=0.0,
|
||||
atol=0.0,
|
||||
msg=f"RMS norm not deterministic: run {idx} differs from reference",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run a quick smoke test
|
||||
print("Running quick smoke test of RMS norm implementations...")
|
||||
|
||||
device = torch.device("cuda")
|
||||
batch_size = 8
|
||||
hidden_size = 4096
|
||||
dtype = torch.bfloat16
|
||||
eps = 1e-6
|
||||
|
||||
torch.manual_seed(42)
|
||||
input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
|
||||
weight = torch.randn(hidden_size, dtype=dtype, device=device)
|
||||
|
||||
# Standard implementation
|
||||
rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)
|
||||
rms_norm_layer.weight.data = weight.clone()
|
||||
standard_output = rms_norm_layer.forward_cuda(input_tensor)
|
||||
|
||||
# Batch-invariant implementation
|
||||
triton_output = triton_rms_norm(input_tensor, weight, eps=eps)
|
||||
|
||||
# Compare
|
||||
max_diff = (triton_output - standard_output).abs().max().item()
|
||||
mean_diff = (triton_output - standard_output).abs().mean().item()
|
||||
|
||||
print(f"Max difference: {max_diff:.6e}")
|
||||
print(f"Mean difference: {mean_diff:.6e}")
|
||||
print(f"Standard output sample: {standard_output[0, :5].tolist()}")
|
||||
print(f"Triton output sample: {triton_output[0, :5].tolist()}")
|
||||
|
||||
if max_diff < 1e-3:
|
||||
print("✓ Smoke test passed!")
|
||||
else:
|
||||
print("✗ Smoke test failed - differences too large")
|
||||
108
third_party/vllm/tests/v1/determinism/utils.py
vendored
Normal file
108
third_party/vllm/tests/v1/determinism/utils.py
vendored
Normal file
@@ -0,0 +1,108 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import os
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
|
||||
|
||||
skip_unsupported = pytest.mark.skipif(
|
||||
not (current_platform.is_cuda() and current_platform.has_device_capability(80)),
|
||||
# Supports testing on Ampere and Ada Lovelace devices.
|
||||
# Note: For devices with SM < 90, batch invariance does not support CUDA Graphs.
|
||||
reason="Requires CUDA and >= Ampere (SM80)",
|
||||
)
|
||||
|
||||
BACKENDS: list[str] = [
|
||||
"FLASH_ATTN",
|
||||
"TRITON_ATTN",
|
||||
"TRITON_MLA",
|
||||
]
|
||||
|
||||
# FlashInfer temporarily disabled due to invariant CTA sizes.
|
||||
# See FlashInfer issue #2424
|
||||
# if has_flashinfer():
|
||||
# BACKENDS.append("FLASHINFER")
|
||||
|
||||
if flash_attn_supports_mla():
|
||||
BACKENDS.append("FLASH_ATTN_MLA")
|
||||
|
||||
DEFAULT_MODEL = "Qwen/Qwen3-1.7B"
|
||||
MLA_MODEL = "deepseek-ai/DeepSeek-V2-Lite-Chat"
|
||||
|
||||
|
||||
def resolve_model_name(backend: str) -> str:
|
||||
"""Resolve the model name for the given backend."""
|
||||
model = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL)
|
||||
if backend.endswith("MLA") and model == DEFAULT_MODEL:
|
||||
return MLA_MODEL
|
||||
return model
|
||||
|
||||
|
||||
def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str:
|
||||
# Generate more realistic prompts that will actually produce varied tokens
|
||||
# Use a mix of common English text patterns
|
||||
|
||||
prompt_templates = [
|
||||
# Question-answer style
|
||||
"Question: What is the capital of France?\nAnswer: The capital of France is",
|
||||
"Q: How does photosynthesis work?\nA: Photosynthesis is the process by which",
|
||||
"User: Can you explain quantum mechanics?\nAssistant: Quantum mechanics is",
|
||||
# Story/narrative style
|
||||
"Once upon a time in a distant galaxy, there lived",
|
||||
"The old man walked slowly down the street, remembering",
|
||||
"In the year 2157, humanity finally discovered",
|
||||
# Technical/code style
|
||||
"To implement a binary search tree in Python, first we need to",
|
||||
"The algorithm works by iterating through the array and",
|
||||
"Here's how to optimize database queries using indexing:",
|
||||
# Factual/informative style
|
||||
"The Renaissance was a period in European history that",
|
||||
"Climate change is caused by several factors including",
|
||||
"The human brain contains approximately 86 billion neurons which",
|
||||
# Conversational style
|
||||
"I've been thinking about getting a new laptop because",
|
||||
"Yesterday I went to the store and bought",
|
||||
"My favorite thing about summer is definitely",
|
||||
]
|
||||
|
||||
# Pick a random template
|
||||
base_prompt = random.choice(prompt_templates)
|
||||
|
||||
if max_words < min_words:
|
||||
max_words = min_words
|
||||
target_words = random.randint(min_words, max_words)
|
||||
|
||||
if target_words > 50:
|
||||
# For longer prompts, repeat context
|
||||
padding_text = (
|
||||
" This is an interesting topic that deserves more explanation. "
|
||||
# TODO: Update to * (target_words // 10) to better align with word ratio
|
||||
* (target_words // 50)
|
||||
)
|
||||
base_prompt = padding_text + base_prompt
|
||||
|
||||
return base_prompt
|
||||
|
||||
|
||||
def _extract_step_logprobs(request_output):
|
||||
if getattr(request_output, "outputs", None):
|
||||
inner = request_output.outputs[0]
|
||||
if hasattr(inner, "logprobs") and inner.logprobs is not None:
|
||||
t = torch.tensor(
|
||||
[
|
||||
inner.logprobs[i][tid].logprob
|
||||
for i, tid in enumerate(inner.token_ids)
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
return t, inner.token_ids
|
||||
|
||||
return None, None
|
||||
|
||||
|
||||
def is_device_capability_below_90() -> bool:
|
||||
return not current_platform.has_device_capability(90)
|
||||
0
third_party/vllm/tests/v1/distributed/__init__.py
vendored
Normal file
0
third_party/vllm/tests/v1/distributed/__init__.py
vendored
Normal file
400
third_party/vllm/tests/v1/distributed/test_async_llm_dp.py
vendored
Normal file
400
third_party/vllm/tests/v1/distributed/test_async_llm_dp.py
vendored
Normal file
@@ -0,0 +1,400 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
from contextlib import ExitStack
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm import SamplingParams
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
from vllm.inputs import PromptType
|
||||
from vllm.outputs import RequestOutput
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
from vllm.v1.engine.core_client import DPAsyncMPClient
|
||||
from vllm.v1.metrics.loggers import StatLoggerBase
|
||||
from vllm.v1.metrics.stats import IterationStats, MultiModalCacheStats, SchedulerStats
|
||||
|
||||
DP_SIZE = int(os.getenv("DP_SIZE", 2))
|
||||
|
||||
|
||||
async def generate(
|
||||
engine: AsyncLLM,
|
||||
request_id: str,
|
||||
prompt: PromptType,
|
||||
output_kind: RequestOutputKind,
|
||||
max_tokens: int,
|
||||
prompt_logprobs: int | None = None,
|
||||
data_parallel_rank: int | None = None,
|
||||
) -> tuple[int, str]:
|
||||
# Ensure generate doesn't complete too fast for cancellation test.
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
count = 0
|
||||
sampling_params = SamplingParams(
|
||||
max_tokens=max_tokens,
|
||||
ignore_eos=True,
|
||||
output_kind=output_kind,
|
||||
temperature=0,
|
||||
prompt_logprobs=prompt_logprobs,
|
||||
)
|
||||
async for out in engine.generate(
|
||||
request_id=request_id,
|
||||
prompt=prompt,
|
||||
sampling_params=sampling_params,
|
||||
data_parallel_rank=data_parallel_rank,
|
||||
):
|
||||
num_tokens = len(out.outputs[0].token_ids)
|
||||
if output_kind == RequestOutputKind.DELTA:
|
||||
count += num_tokens
|
||||
else:
|
||||
count = num_tokens
|
||||
|
||||
await asyncio.sleep(0.0)
|
||||
|
||||
return count, request_id
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"ibm-research/PowerMoE-3b",
|
||||
"hmellor/tiny-random-LlamaForCausalLM",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"output_kind",
|
||||
[
|
||||
RequestOutputKind.DELTA,
|
||||
RequestOutputKind.FINAL_ONLY,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("data_parallel_backend", ["mp", "ray"])
|
||||
@pytest.mark.parametrize("async_scheduling", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_load(
|
||||
model: str,
|
||||
output_kind: RequestOutputKind,
|
||||
data_parallel_backend: str,
|
||||
async_scheduling: bool,
|
||||
):
|
||||
if async_scheduling and data_parallel_backend == "ray":
|
||||
# TODO(NickLucche) Re-enable when async scheduling is supported
|
||||
pytest.skip("Async scheduling is not supported with ray")
|
||||
elif data_parallel_backend == "ray" and current_platform.is_rocm():
|
||||
pytest.skip(
|
||||
"Ray as the distributed executor backend is not supported with ROCm."
|
||||
)
|
||||
stats_loggers = {}
|
||||
|
||||
@dataclass
|
||||
class SimpleStatsLogger(StatLoggerBase):
|
||||
init_count: int = 0
|
||||
finished_req_count: int = 0
|
||||
|
||||
def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):
|
||||
stats_loggers[engine_index] = self
|
||||
|
||||
def record(
|
||||
self,
|
||||
scheduler_stats: SchedulerStats | None,
|
||||
iteration_stats: IterationStats | None,
|
||||
mm_cache_stats: MultiModalCacheStats | None = None,
|
||||
engine_idx: int = 0,
|
||||
):
|
||||
if iteration_stats:
|
||||
self.finished_req_count += len(iteration_stats.finished_requests)
|
||||
|
||||
def log_engine_initialized(self):
|
||||
self.init_count += 1
|
||||
|
||||
with ExitStack() as after:
|
||||
prompt = "This is a test of data parallel"
|
||||
|
||||
engine_args = AsyncEngineArgs(
|
||||
model=model,
|
||||
enforce_eager=True,
|
||||
tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
|
||||
data_parallel_size=DP_SIZE,
|
||||
data_parallel_backend=data_parallel_backend,
|
||||
async_scheduling=async_scheduling,
|
||||
)
|
||||
engine = AsyncLLM.from_engine_args(
|
||||
engine_args, stat_loggers=[SimpleStatsLogger]
|
||||
)
|
||||
after.callback(engine.shutdown)
|
||||
|
||||
NUM_REQUESTS = 100
|
||||
NUM_EXPECTED_TOKENS = 10
|
||||
|
||||
request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)]
|
||||
|
||||
# Create concurrent requests.
|
||||
tasks = []
|
||||
for request_id in request_ids:
|
||||
tasks.append(
|
||||
asyncio.create_task(
|
||||
generate(
|
||||
engine, request_id, prompt, output_kind, NUM_EXPECTED_TOKENS
|
||||
)
|
||||
)
|
||||
)
|
||||
# Short sleep to ensure that requests are distributed.
|
||||
await asyncio.sleep(0.01)
|
||||
# Confirm that we got all the EXPECTED tokens from the requests.
|
||||
done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
for task in done:
|
||||
num_generated_tokens, request_id = await task
|
||||
assert num_generated_tokens == NUM_EXPECTED_TOKENS, (
|
||||
f"{request_id} generated {num_generated_tokens} but "
|
||||
f"expected {NUM_EXPECTED_TOKENS}"
|
||||
)
|
||||
|
||||
assert not engine.output_processor.has_unfinished_requests()
|
||||
|
||||
# testing internals here which may break
|
||||
core_client: DPAsyncMPClient = engine.engine_core
|
||||
# the engines only synchronize stopping every N steps so
|
||||
# allow a small amount of time here.
|
||||
for _ in range(10):
|
||||
if not core_client.engines_running:
|
||||
break
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
assert not core_client.engines_running
|
||||
assert not core_client.reqs_in_flight
|
||||
|
||||
# Check that requests were distributed between the engines
|
||||
print(f"Stats loggers after test: {stats_loggers}")
|
||||
assert len(stats_loggers) == DP_SIZE
|
||||
assert stats_loggers[0].init_count == 1
|
||||
|
||||
for sl in stats_loggers.values():
|
||||
slogger: SimpleStatsLogger = sl
|
||||
|
||||
assert slogger.finished_req_count > NUM_REQUESTS // (DP_SIZE + 1), (
|
||||
f"requests are imbalanced: {stats_loggers}"
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# DP Pause/Resume Tests
|
||||
# =============================================================================
|
||||
# When expert_parallel=False: uses non-MoE model (DP replicas as separate engines).
|
||||
# When expert_parallel=True: uses MoE model + EP (DPEngineCoreProc, sync pause path).
|
||||
|
||||
DP_PAUSE_MODEL = "hmellor/tiny-random-LlamaForCausalLM"
|
||||
DP_PAUSE_MODEL_MOE = "ibm-research/PowerMoE-3b"
|
||||
DP_PAUSE_PROMPT = "This is a test of data parallel pause"
|
||||
|
||||
|
||||
def _get_dp_pause_engine_args(expert_parallel: bool) -> AsyncEngineArgs:
|
||||
"""Engine args for DP pause tests: MoE+EP when expert_parallel else small Llama."""
|
||||
model = DP_PAUSE_MODEL_MOE if expert_parallel else DP_PAUSE_MODEL
|
||||
return AsyncEngineArgs(
|
||||
model=model,
|
||||
enforce_eager=True,
|
||||
tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
|
||||
data_parallel_size=DP_SIZE,
|
||||
data_parallel_backend="mp",
|
||||
enable_expert_parallel=expert_parallel,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("expert_parallel", [False, True])
|
||||
async def test_dp_pause_resume_basic(expert_parallel: bool):
|
||||
"""Pausing from the client (one call) pauses all DP ranks; resume clears it."""
|
||||
with ExitStack() as after:
|
||||
engine_args = _get_dp_pause_engine_args(expert_parallel)
|
||||
engine = AsyncLLM.from_engine_args(engine_args)
|
||||
after.callback(engine.shutdown)
|
||||
|
||||
assert not await engine.is_paused()
|
||||
await engine.pause_generation(mode="abort")
|
||||
assert await engine.is_paused()
|
||||
await engine.resume_generation()
|
||||
assert not await engine.is_paused()
|
||||
|
||||
# Engine still works after resume
|
||||
sampling_params = SamplingParams(max_tokens=5)
|
||||
async for out in engine.generate(
|
||||
request_id="after-resume",
|
||||
prompt=DP_PAUSE_PROMPT,
|
||||
sampling_params=sampling_params,
|
||||
):
|
||||
pass
|
||||
assert out.finished
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("expert_parallel", [False, True])
|
||||
async def test_dp_pause_abort(expert_parallel: bool):
|
||||
"""Pause with abort from one client aborts in-flight requests on all DP ranks."""
|
||||
with ExitStack() as after:
|
||||
engine_args = _get_dp_pause_engine_args(expert_parallel)
|
||||
engine = AsyncLLM.from_engine_args(engine_args)
|
||||
after.callback(engine.shutdown)
|
||||
|
||||
# Start several requests so they are distributed across ranks
|
||||
sampling_params = SamplingParams(max_tokens=500, ignore_eos=True)
|
||||
num_requests = 4
|
||||
outputs_by_id: dict[str, list[RequestOutput]] = {}
|
||||
|
||||
async def gen(rid: str):
|
||||
out_list: list[RequestOutput] = []
|
||||
outputs_by_id[rid] = out_list
|
||||
async for out in engine.generate(
|
||||
request_id=rid,
|
||||
prompt=DP_PAUSE_PROMPT,
|
||||
sampling_params=sampling_params,
|
||||
):
|
||||
out_list.append(out)
|
||||
return out_list[-1] if out_list else None
|
||||
|
||||
tasks = [asyncio.create_task(gen(f"req-{i}")) for i in range(num_requests)]
|
||||
# Wait for some tokens on at least one request
|
||||
while not any(len(o) >= 2 for o in outputs_by_id.values()):
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
await engine.pause_generation(mode="abort")
|
||||
|
||||
finals = await asyncio.gather(*tasks)
|
||||
for i, final in enumerate(finals):
|
||||
assert final is not None, f"req-{i} had no output"
|
||||
assert final.finished
|
||||
assert final.outputs[0].finish_reason == "abort"
|
||||
|
||||
assert await engine.is_paused()
|
||||
await engine.resume_generation()
|
||||
assert not await engine.is_paused()
|
||||
|
||||
# New request completes after resume
|
||||
async for out in engine.generate(
|
||||
request_id="after-abort",
|
||||
prompt=DP_PAUSE_PROMPT,
|
||||
sampling_params=SamplingParams(max_tokens=5),
|
||||
):
|
||||
pass
|
||||
assert out.finished
|
||||
assert not engine.output_processor.has_unfinished_requests()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("expert_parallel", [False, True])
|
||||
async def test_dp_pause_keep_then_resume(expert_parallel: bool):
|
||||
"""Start generation, pause after a few tokens (keep mode), resume; verify gap."""
|
||||
|
||||
pause_duration = 2.0
|
||||
min_tokens_before_pause = 3
|
||||
|
||||
with ExitStack() as after:
|
||||
engine_args = _get_dp_pause_engine_args(expert_parallel)
|
||||
engine = AsyncLLM.from_engine_args(engine_args)
|
||||
after.callback(engine.shutdown)
|
||||
|
||||
sampling_params = SamplingParams(max_tokens=15, ignore_eos=True)
|
||||
token_times: list[tuple[int, float]] = []
|
||||
pause_token_idx = 0
|
||||
|
||||
async def generator_task():
|
||||
nonlocal pause_token_idx
|
||||
out = None
|
||||
async for output in engine.generate(
|
||||
request_id="keep-resume-req",
|
||||
prompt=DP_PAUSE_PROMPT,
|
||||
sampling_params=sampling_params,
|
||||
):
|
||||
token_count = len(output.outputs[0].token_ids)
|
||||
token_times.append((token_count, time.monotonic()))
|
||||
out = output
|
||||
return out
|
||||
|
||||
async def controller_task():
|
||||
nonlocal pause_token_idx
|
||||
while len(token_times) < min_tokens_before_pause:
|
||||
await asyncio.sleep(0.01)
|
||||
await engine.pause_generation(mode="keep")
|
||||
await asyncio.sleep(pause_duration)
|
||||
pause_token_idx = len(token_times)
|
||||
await engine.resume_generation()
|
||||
|
||||
gen_task = asyncio.create_task(generator_task())
|
||||
ctrl_task = asyncio.create_task(controller_task())
|
||||
final_output, _ = await asyncio.gather(gen_task, ctrl_task)
|
||||
|
||||
assert final_output is not None and final_output.finished
|
||||
assert await engine.is_paused() is False
|
||||
assert pause_token_idx >= min_tokens_before_pause
|
||||
if pause_token_idx > 0 and pause_token_idx < len(token_times):
|
||||
pause_gap = (
|
||||
token_times[pause_token_idx][1] - token_times[pause_token_idx - 1][1]
|
||||
)
|
||||
assert pause_gap >= pause_duration * 0.8, (
|
||||
f"Expected gap ~{pause_duration}s after pause, got {pause_gap:.3f}s"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dp_pause_keep_race_staggered_engines():
|
||||
"""Race: send pause(keep) to engine 0, then add two requests,
|
||||
then pause(keep) to engine 1. Ensures no deadlock when pause
|
||||
requests are staggered and requests arrive in between."""
|
||||
if DP_SIZE != 2:
|
||||
pytest.skip("test_dp_pause_keep_race_staggered_engines requires DP_SIZE=2")
|
||||
|
||||
with ExitStack() as after:
|
||||
engine_args = _get_dp_pause_engine_args(expert_parallel=True)
|
||||
engine = AsyncLLM.from_engine_args(engine_args)
|
||||
after.callback(engine.shutdown)
|
||||
|
||||
client = engine.engine_core
|
||||
|
||||
original_call_utility = client.call_utility_async
|
||||
mid_pause_tasks: list[asyncio.Task] = []
|
||||
|
||||
async def staggered_pause_keep(method: str, *args) -> Any:
|
||||
if method != "pause_scheduler" or not args or args[0] != "keep":
|
||||
return await original_call_utility(method, *args)
|
||||
# Send pause(keep) to engine 0 first
|
||||
await client._call_utility_async(
|
||||
method, *args, engine=client.core_engines[0]
|
||||
)
|
||||
# In the middle: send two requests (race window)
|
||||
sp = SamplingParams(max_tokens=5, ignore_eos=True)
|
||||
|
||||
async def consume_gen(req_id: str) -> None:
|
||||
async for _ in engine.generate(
|
||||
request_id=req_id,
|
||||
prompt=DP_PAUSE_PROMPT,
|
||||
sampling_params=sp,
|
||||
):
|
||||
pass
|
||||
|
||||
t1 = asyncio.create_task(consume_gen("race-1"))
|
||||
t2 = asyncio.create_task(consume_gen("race-2"))
|
||||
mid_pause_tasks.extend([t1, t2])
|
||||
await asyncio.sleep(3)
|
||||
# Then send pause(keep) to engine 1
|
||||
result = await client._call_utility_async(
|
||||
method, *args, engine=client.core_engines[1]
|
||||
)
|
||||
return result
|
||||
|
||||
client.call_utility_async = staggered_pause_keep
|
||||
|
||||
await engine.pause_generation(mode="keep")
|
||||
assert await engine.is_paused()
|
||||
await engine.resume_generation()
|
||||
assert not await engine.is_paused()
|
||||
# Let the two requests we sent mid-pause complete
|
||||
await asyncio.gather(*mid_pause_tasks)
|
||||
109
third_party/vllm/tests/v1/distributed/test_dbo.py
vendored
Normal file
109
third_party/vllm/tests/v1/distributed/test_dbo.py
vendored
Normal file
@@ -0,0 +1,109 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Test Dual Batch Overlap (DBO) with Data Parallelism + Expert Parallelism.
|
||||
|
||||
DBO is specifically designed for DP+EP scenarios to hide communication latency
|
||||
by overlapping computation of two batches. This test validates that DBO works
|
||||
correctly with the DeepSeek-V2-Lite model using GSM8K evaluation.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.utils.import_utils import has_deep_ep
|
||||
|
||||
# Detect Blackwell / B200 (compute capability 10.x)
|
||||
try:
|
||||
if torch.cuda.is_available():
|
||||
cap = torch.cuda.get_device_capability(0)
|
||||
IS_BLACKWELL = cap[0] >= 10
|
||||
else:
|
||||
IS_BLACKWELL = False
|
||||
except Exception:
|
||||
# Be conservative: if we can't detect, don't xfail by default
|
||||
IS_BLACKWELL = False
|
||||
|
||||
MODEL_NAME = "deepseek-ai/DeepSeek-V2-Lite-Chat"
|
||||
DP_SIZE = 2
|
||||
|
||||
# GSM8K eval configuration
|
||||
NUM_QUESTIONS = 256 # Fast eval for CI; but must be large enough to hit dbo thresholds
|
||||
NUM_SHOTS = 5 # Few-shot examples
|
||||
MIN_ACCURACY = 0.62 # Expected 0.64 with 2% buffer (based on vLLM test data)
|
||||
|
||||
# Increase max_num_seqs to trigger DBO for decode batches
|
||||
# With 64 seqs, decode batches should exceed the 32 token threshold
|
||||
MAX_NUM_SEQS = 64 # Increased from 16 to trigger decode DBO
|
||||
|
||||
# DeepEP backends to test
|
||||
DEEPEP_BACKENDS = [
|
||||
"deepep_low_latency",
|
||||
"deepep_high_throughput",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not has_deep_ep(), reason="These tests require deep_ep to run")
|
||||
@pytest.mark.parametrize("all2all_backend", DEEPEP_BACKENDS)
|
||||
@pytest.mark.xfail(
|
||||
IS_BLACKWELL,
|
||||
reason=(
|
||||
"Temporary: DBO accuracy unstable on Blackwell "
|
||||
"(doesn't meet expectation of MIN_ACCURACY = 0.62)"
|
||||
),
|
||||
)
|
||||
def test_dbo_dp_ep_gsm8k(all2all_backend: str, num_gpus_available):
|
||||
"""
|
||||
Test DBO with DP+EP using GSM8K evaluation.
|
||||
"""
|
||||
required_gpus = DP_SIZE
|
||||
|
||||
if num_gpus_available < required_gpus:
|
||||
pytest.skip(f"Need at least {required_gpus} GPUs (DP={DP_SIZE})")
|
||||
|
||||
# Server arguments for DBO + DP + EP
|
||||
server_args = [
|
||||
"--max-model-len",
|
||||
"4096",
|
||||
"--max-num-seqs",
|
||||
str(MAX_NUM_SEQS), # Use larger batch to trigger decode DBO
|
||||
"--trust-remote-code",
|
||||
# Note: Not using --enforce-eager to test DBO's alternate CUDA graph dispatching
|
||||
"--data-parallel-size",
|
||||
str(DP_SIZE),
|
||||
"--enable-expert-parallel",
|
||||
"--enable-dbo",
|
||||
# Fix threshold so we know we trigger DBO
|
||||
"--dbo-decode-token-threshold",
|
||||
"16",
|
||||
"--dbo-prefill-token-threshold",
|
||||
"256",
|
||||
"--all2all-backend",
|
||||
all2all_backend,
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME,
|
||||
server_args,
|
||||
max_wait_seconds=600, # Allow time for model loading with DP+EP
|
||||
) as remote_server:
|
||||
# Use host and port directly from RemoteOpenAIServer
|
||||
host = f"http://{remote_server.host}"
|
||||
port = remote_server.port
|
||||
|
||||
# Run GSM8K evaluation
|
||||
results = evaluate_gsm8k(
|
||||
num_questions=NUM_QUESTIONS,
|
||||
num_shots=NUM_SHOTS,
|
||||
host=host,
|
||||
port=port,
|
||||
)
|
||||
|
||||
# Validate accuracy is reasonable
|
||||
accuracy = results["accuracy"]
|
||||
assert accuracy >= MIN_ACCURACY, (
|
||||
f"DBO+DP+EP accuracy too low ({all2all_backend}): "
|
||||
f"{accuracy:.3f} < {MIN_ACCURACY:.3f} "
|
||||
)
|
||||
106
third_party/vllm/tests/v1/distributed/test_eagle_dp.py
vendored
Normal file
106
third_party/vllm/tests/v1/distributed/test_eagle_dp.py
vendored
Normal file
@@ -0,0 +1,106 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import asyncio
|
||||
import os
|
||||
from contextlib import AsyncExitStack
|
||||
from dataclasses import replace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm import SamplingParams
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
|
||||
DP_SIZE = int(os.getenv("DP_SIZE", 2))
|
||||
|
||||
if current_platform.is_rocm():
|
||||
ATTN_BACKENDS = ["ROCM_ATTN", "TRITON_ATTN", "FLEX_ATTENTION"]
|
||||
else:
|
||||
ATTN_BACKENDS = ["FLASH_ATTN"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("attn_backend", ATTN_BACKENDS)
|
||||
@pytest.mark.xfail(
|
||||
current_platform.is_rocm(),
|
||||
reason="Test may fail on ROCm until batch invariance is enabled."
|
||||
"See: https://github.com/vllm-project/vllm/issues/27433",
|
||||
strict=False,
|
||||
)
|
||||
async def test_run_eagle_dp(monkeypatch: pytest.MonkeyPatch, attn_backend: str):
|
||||
if not current_platform.is_rocm():
|
||||
# This test checks that running a model with and without eagle
|
||||
# leads to identical tokens.
|
||||
#
|
||||
# NOTE: This is only true in batch invariant mode
|
||||
# (because the target model verifies all draft tokens in one big
|
||||
# forward pass)
|
||||
#
|
||||
# TODO[ROCm]: Test is passing on ROCm CI but may break in future.
|
||||
# Enable batch invariance for ROCm when possible. See:
|
||||
# https://github.com/vllm-project/vllm/issues/27433
|
||||
|
||||
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
|
||||
|
||||
target_model = "meta-llama/Llama-3.1-8B-Instruct"
|
||||
draft_model = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
|
||||
|
||||
engine_args = AsyncEngineArgs(
|
||||
model=target_model,
|
||||
tokenizer_mode="auto",
|
||||
enforce_eager=False,
|
||||
tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
|
||||
data_parallel_size=DP_SIZE,
|
||||
data_parallel_backend="mp", # ray takes more time
|
||||
trust_remote_code=True,
|
||||
max_model_len=16384,
|
||||
attention_config={"backend": attn_backend},
|
||||
)
|
||||
|
||||
eagle_engine_args = replace(
|
||||
engine_args,
|
||||
speculative_config={
|
||||
"model": draft_model,
|
||||
"method": "eagle",
|
||||
"num_speculative_tokens": 3,
|
||||
},
|
||||
)
|
||||
|
||||
prompt = "This is a test of data parallel with eagle"
|
||||
# This test might be flaky, see
|
||||
# https://github.com/vllm-project/vllm/issues/31913
|
||||
num_expected_tokens = 20
|
||||
sampling_params = SamplingParams(
|
||||
max_tokens=num_expected_tokens,
|
||||
ignore_eos=True,
|
||||
output_kind=RequestOutputKind.FINAL_ONLY,
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
async def generate_with_timeout(given_engine: AsyncLLM):
|
||||
async for out in given_engine.generate(
|
||||
request_id="test-eagle-dp", prompt=prompt, sampling_params=sampling_params
|
||||
):
|
||||
token_ids = out.outputs[0].token_ids
|
||||
assert len(token_ids) == num_expected_tokens
|
||||
return token_ids
|
||||
|
||||
async def engine_create_and_generate(engine_args: AsyncEngineArgs):
|
||||
async with AsyncExitStack() as after:
|
||||
engine = AsyncLLM.from_engine_args(engine_args)
|
||||
after.callback(engine.shutdown)
|
||||
|
||||
token_ids = await asyncio.wait_for(
|
||||
generate_with_timeout(engine), timeout=30
|
||||
)
|
||||
|
||||
assert not engine.output_processor.has_unfinished_requests()
|
||||
return token_ids
|
||||
|
||||
token_ids_with_eagle = await engine_create_and_generate(eagle_engine_args)
|
||||
token_ids_no_eagle = await engine_create_and_generate(engine_args)
|
||||
|
||||
# Test for correctness
|
||||
assert token_ids_with_eagle == token_ids_no_eagle
|
||||
357
third_party/vllm/tests/v1/distributed/test_external_lb_dp.py
vendored
Normal file
357
third_party/vllm/tests/v1/distributed/test_external_lb_dp.py
vendored
Normal file
@@ -0,0 +1,357 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import asyncio
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from contextlib import AsyncExitStack
|
||||
|
||||
import openai # use the official client for correctness check
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import requests
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
MODEL_NAME = "ibm-research/PowerMoE-3b"
|
||||
|
||||
# Number of data parallel ranks for external LB testing
|
||||
DP_SIZE = int(os.getenv("DP_SIZE", "2"))
|
||||
# Default tensor parallel size to use
|
||||
TP_SIZE = int(os.getenv("TP_SIZE", "1"))
|
||||
|
||||
|
||||
class ExternalLBServerManager:
|
||||
"""Manages data parallel vLLM server instances for external
|
||||
load balancer testing."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
dp_size: int,
|
||||
api_server_count: int,
|
||||
base_server_args: list,
|
||||
tp_size: int = TP_SIZE,
|
||||
):
|
||||
self.model_name = model_name
|
||||
self.dp_size = dp_size
|
||||
self.tp_size = tp_size
|
||||
self.api_server_count = api_server_count
|
||||
self.base_server_args = base_server_args
|
||||
self.servers: list[tuple[RemoteOpenAIServer, list[str]]] = []
|
||||
self.server_threads: list[threading.Thread] = []
|
||||
|
||||
def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
|
||||
"""Start all server instances for external LB mode."""
|
||||
for rank in range(self.dp_size):
|
||||
# Create server args for this specific rank
|
||||
server_args = self.base_server_args.copy()
|
||||
|
||||
# Add external LB specific arguments
|
||||
server_args.extend(
|
||||
[
|
||||
"--data-parallel-size",
|
||||
str(self.dp_size),
|
||||
"--data-parallel-rank",
|
||||
str(rank),
|
||||
"--data-parallel-size-local",
|
||||
"1",
|
||||
"--tensor-parallel-size",
|
||||
str(self.tp_size),
|
||||
"--port",
|
||||
str(8000 + rank), # Different port for each rank
|
||||
"--api-server-count",
|
||||
str(self.api_server_count),
|
||||
]
|
||||
)
|
||||
|
||||
# Use a thread to start each server to allow parallel initialization
|
||||
def start_server(r: int, sargs: list[str]):
|
||||
try:
|
||||
# Start the server
|
||||
server = RemoteOpenAIServer(
|
||||
self.model_name,
|
||||
sargs,
|
||||
auto_port=False,
|
||||
env_dict={
|
||||
"VLLM_SERVER_DEV_MODE": "1",
|
||||
current_platform.device_control_env_var: ",".join(
|
||||
str(current_platform.device_id_to_physical_device_id(i))
|
||||
for i in range(r * TP_SIZE, (r + 1) * TP_SIZE)
|
||||
),
|
||||
},
|
||||
)
|
||||
server.__enter__()
|
||||
print(
|
||||
f"Server rank {r} started successfully with "
|
||||
f"{self.api_server_count} API servers"
|
||||
)
|
||||
self.servers.append((server, sargs))
|
||||
except Exception as e:
|
||||
print(f"Failed to start server rank {r}: {e}")
|
||||
raise
|
||||
|
||||
thread = threading.Thread(target=start_server, args=(rank, server_args))
|
||||
thread.start()
|
||||
|
||||
self.server_threads.append(thread)
|
||||
|
||||
# Wait for all servers to start
|
||||
for thread in self.server_threads:
|
||||
thread.join()
|
||||
|
||||
# Give servers additional time to fully initialize and coordinate
|
||||
time.sleep(2)
|
||||
|
||||
if len(self.servers) != self.dp_size:
|
||||
raise Exception("Servers failed to start")
|
||||
|
||||
return self.servers
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Stop all server instances."""
|
||||
while self.servers:
|
||||
try:
|
||||
self.servers.pop()[0].__exit__(exc_type, exc_val, exc_tb)
|
||||
except Exception as e:
|
||||
print(f"Error stopping server: {e}")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def default_server_args():
|
||||
return [
|
||||
# use half precision for speed and memory savings in CI environment
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-model-len",
|
||||
"2048",
|
||||
"--max-num-seqs",
|
||||
"128",
|
||||
"--enforce-eager",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=[1, 4])
|
||||
def server_manager(request, default_server_args):
|
||||
api_server_count = request.param
|
||||
server_manager = ExternalLBServerManager(
|
||||
MODEL_NAME, DP_SIZE, api_server_count, default_server_args
|
||||
)
|
||||
|
||||
with server_manager:
|
||||
yield server_manager
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def servers(server_manager):
|
||||
return server_manager.servers
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def clients(servers: list[tuple[RemoteOpenAIServer, list[str]]]):
|
||||
# Create a client for each server
|
||||
async with AsyncExitStack() as stack:
|
||||
yield [
|
||||
await stack.enter_async_context(server.get_async_client())
|
||||
for server, _ in servers
|
||||
]
|
||||
|
||||
|
||||
def _get_parallel_config(server: RemoteOpenAIServer):
|
||||
response = requests.get(server.url_for("server_info?config_format=json"))
|
||||
response.raise_for_status()
|
||||
|
||||
vllm_config = response.json()["vllm_config"]
|
||||
return vllm_config["parallel_config"]
|
||||
|
||||
|
||||
def test_external_lb_server_info(server_manager):
|
||||
servers = server_manager.servers
|
||||
api_server_count = server_manager.api_server_count
|
||||
|
||||
for i, (server, _) in enumerate(servers):
|
||||
print(f"Testing {i=}")
|
||||
|
||||
# Each request will hit one of the API servers
|
||||
# `n_reqs` is set so that there is a good chance each server
|
||||
# receives at least one request
|
||||
n_reqs = 2 * api_server_count * api_server_count
|
||||
parallel_configs = [_get_parallel_config(server) for _ in range(n_reqs)]
|
||||
api_process_counts = [c["_api_process_count"] for c in parallel_configs]
|
||||
api_process_ranks = [c["_api_process_rank"] for c in parallel_configs]
|
||||
|
||||
assert all(c == api_server_count for c in api_process_counts), (
|
||||
api_process_counts
|
||||
)
|
||||
assert all(0 <= r < api_server_count for r in api_process_ranks), (
|
||||
api_process_ranks
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_external_lb_single_completion(
|
||||
clients: list[openai.AsyncOpenAI],
|
||||
servers: list[tuple[RemoteOpenAIServer, list[str]]],
|
||||
model_name: str,
|
||||
) -> None:
|
||||
async def make_request(client: openai.AsyncOpenAI):
|
||||
completion = await client.completions.create(
|
||||
model=model_name, prompt="Hello, my name is", max_tokens=10, temperature=1.0
|
||||
)
|
||||
|
||||
assert completion.id is not None
|
||||
assert completion.choices is not None and len(completion.choices) == 1
|
||||
|
||||
choice = completion.choices[0]
|
||||
# The exact number of tokens can vary slightly with temperature=1.0,
|
||||
# so we check for a reasonable minimum length.
|
||||
assert len(choice.text) >= 1
|
||||
# Finish reason might not always be 'length' if the model finishes early
|
||||
# or due to other reasons, especially with high temperature.
|
||||
# So, we'll accept 'length' or 'stop'.
|
||||
assert choice.finish_reason in ("length", "stop")
|
||||
|
||||
# Token counts can also vary, so we check they are positive.
|
||||
assert completion.usage.completion_tokens > 0
|
||||
assert completion.usage.prompt_tokens > 0
|
||||
assert completion.usage.total_tokens > 0
|
||||
return completion
|
||||
|
||||
# Test single request to each server
|
||||
for i, client in enumerate(clients):
|
||||
result = await make_request(client)
|
||||
assert result is not None
|
||||
print(f"Server {i} handled single completion request successfully")
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Send requests to all servers in round-robin fashion
|
||||
num_requests_per_server = 25 # Total 50 requests across 2 servers
|
||||
all_tasks = []
|
||||
|
||||
for i, client in enumerate(clients):
|
||||
tasks = [make_request(client) for _ in range(num_requests_per_server)]
|
||||
all_tasks.extend(tasks)
|
||||
|
||||
results = await asyncio.gather(*all_tasks)
|
||||
assert len(results) == num_requests_per_server * len(clients)
|
||||
assert all(completion is not None for completion in results)
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Second burst of requests
|
||||
all_tasks = []
|
||||
for i, client in enumerate(clients):
|
||||
tasks = [make_request(client) for _ in range(num_requests_per_server)]
|
||||
all_tasks.extend(tasks)
|
||||
|
||||
results = await asyncio.gather(*all_tasks)
|
||||
assert len(results) == num_requests_per_server * len(clients)
|
||||
assert all(completion is not None for completion in results)
|
||||
|
||||
_, server_args = servers[0]
|
||||
api_server_count = (
|
||||
server_args.count("--api-server-count")
|
||||
and server_args[server_args.index("--api-server-count") + 1]
|
||||
or 1
|
||||
)
|
||||
print(
|
||||
f"Successfully completed external LB test with {len(clients)} servers "
|
||||
f"(API server count: {api_server_count})"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_external_lb_completion_streaming(
|
||||
clients: list[openai.AsyncOpenAI],
|
||||
servers: list[tuple[RemoteOpenAIServer, list[str]]],
|
||||
model_name: str,
|
||||
) -> None:
|
||||
prompt = "What is an LLM?"
|
||||
|
||||
async def make_streaming_request(client: openai.AsyncOpenAI):
|
||||
# Perform a non-streaming request to get the expected full output
|
||||
single_completion = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
)
|
||||
single_output = single_completion.choices[0].text
|
||||
|
||||
# Perform the streaming request
|
||||
stream = await client.completions.create(
|
||||
model=model_name, prompt=prompt, max_tokens=5, temperature=0.0, stream=True
|
||||
)
|
||||
chunks: list[str] = []
|
||||
finish_reason_count = 0
|
||||
last_chunk = None
|
||||
async for chunk in stream:
|
||||
chunks.append(chunk.choices[0].text)
|
||||
if chunk.choices[0].finish_reason is not None:
|
||||
finish_reason_count += 1
|
||||
last_chunk = chunk # Keep track of the last chunk
|
||||
|
||||
# finish reason should only return in the last block for OpenAI API
|
||||
assert finish_reason_count == 1, "Finish reason should appear exactly once."
|
||||
assert last_chunk is not None, "Stream should have yielded at least one chunk."
|
||||
assert last_chunk.choices[0].finish_reason == "length", (
|
||||
"Finish reason should be 'length'."
|
||||
)
|
||||
# Check that the combined text matches the non-streamed version.
|
||||
assert "".join(chunks) == single_output, (
|
||||
"Streamed output should match non-streamed output."
|
||||
)
|
||||
return True # Indicate success for this request
|
||||
|
||||
# Test single request to each server
|
||||
for i, client in enumerate(clients):
|
||||
result = await make_streaming_request(client)
|
||||
assert result is not None
|
||||
print(f"Server {i} handled single streaming request successfully")
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Send streaming requests to all servers in round-robin fashion
|
||||
num_requests_per_server = 25 # Total 50 requests across 2 servers
|
||||
all_tasks = []
|
||||
|
||||
for i, client in enumerate(clients):
|
||||
tasks = [make_streaming_request(client) for _ in range(num_requests_per_server)]
|
||||
all_tasks.extend(tasks)
|
||||
|
||||
results = await asyncio.gather(*all_tasks)
|
||||
assert len(results) == num_requests_per_server * len(clients)
|
||||
assert all(results), "Not all streaming requests completed successfully."
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Second burst of streaming requests
|
||||
all_tasks = []
|
||||
for i, client in enumerate(clients):
|
||||
tasks = [make_streaming_request(client) for _ in range(num_requests_per_server)]
|
||||
all_tasks.extend(tasks)
|
||||
|
||||
results = await asyncio.gather(*all_tasks)
|
||||
assert len(results) == num_requests_per_server * len(clients)
|
||||
assert all(results), "Not all streaming requests completed successfully."
|
||||
|
||||
_, server_args = servers[0]
|
||||
api_server_count = (
|
||||
server_args.count("--api-server-count")
|
||||
and server_args[server_args.index("--api-server-count") + 1]
|
||||
or 1
|
||||
)
|
||||
print(
|
||||
f"Successfully completed external LB streaming test with "
|
||||
f"{len(clients)} servers (API server count: {api_server_count})"
|
||||
)
|
||||
398
third_party/vllm/tests/v1/distributed/test_hybrid_lb_dp.py
vendored
Normal file
398
third_party/vllm/tests/v1/distributed/test_hybrid_lb_dp.py
vendored
Normal file
@@ -0,0 +1,398 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import asyncio
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from contextlib import AsyncExitStack
|
||||
|
||||
import openai # use the official client for correctness check
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import requests
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from tests.v1.utils import check_request_balancing
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
MODEL_NAME = "ibm-research/PowerMoE-3b"
|
||||
|
||||
# Number of data parallel ranks for hybrid LB testing (4 total)
|
||||
DP_SIZE = int(os.getenv("DP_SIZE", "4"))
|
||||
# Default tensor parallel size to use
|
||||
TP_SIZE = int(os.getenv("TP_SIZE", "1"))
|
||||
|
||||
# Number of nodes (2 nodes, each with 2 DP ranks)
|
||||
NUM_NODES = 2
|
||||
DP_SIZE_LOCAL = DP_SIZE // NUM_NODES # 2 ranks per node
|
||||
|
||||
|
||||
class HybridLBServerManager:
|
||||
"""Manages hybrid data parallel vLLM server instances where each node
|
||||
runs a single logical API server that balances requests only to the
|
||||
DP engines running on that same node."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
dp_size: int,
|
||||
api_server_count: int,
|
||||
base_server_args: list,
|
||||
dp_size_local: int = DP_SIZE_LOCAL,
|
||||
tp_size: int = TP_SIZE,
|
||||
):
|
||||
self.model_name = model_name
|
||||
self.dp_size = dp_size
|
||||
self.dp_size_local = dp_size_local
|
||||
self.tp_size = tp_size
|
||||
self.api_server_count = api_server_count
|
||||
self.base_server_args = base_server_args
|
||||
self.servers: list[tuple[RemoteOpenAIServer, list[str]]] = []
|
||||
self.server_threads: list[threading.Thread] = []
|
||||
self.num_nodes = dp_size // dp_size_local
|
||||
|
||||
def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
|
||||
"""Start all server instances for hybrid LB mode."""
|
||||
for node_id in range(self.num_nodes):
|
||||
# Create server args for this specific node
|
||||
server_args = self.base_server_args.copy()
|
||||
|
||||
# Calculate start rank for this node
|
||||
start_rank = node_id * self.dp_size_local
|
||||
|
||||
# Add hybrid LB specific arguments
|
||||
server_args.extend(
|
||||
[
|
||||
"--data-parallel-size",
|
||||
str(self.dp_size),
|
||||
"--data-parallel-size-local",
|
||||
str(self.dp_size_local),
|
||||
"--data-parallel-start-rank",
|
||||
str(start_rank),
|
||||
"--data-parallel-hybrid-lb", # Enable hybrid LB mode
|
||||
"--tensor-parallel-size",
|
||||
str(self.tp_size),
|
||||
"--port",
|
||||
str(8000 + node_id), # Different port for each node
|
||||
"--api-server-count",
|
||||
str(self.api_server_count),
|
||||
"--data-parallel-address",
|
||||
"127.0.0.1",
|
||||
"--data-parallel-rpc-port",
|
||||
"13345",
|
||||
]
|
||||
)
|
||||
|
||||
# Use a thread to start each server to allow parallel initialization
|
||||
def start_server(node: int, sargs: list[str]):
|
||||
try:
|
||||
# Calculate GPU devices for this node
|
||||
gpus_per_node = self.dp_size_local * self.tp_size
|
||||
gpu_start = node * gpus_per_node
|
||||
gpu_end = gpu_start + gpus_per_node
|
||||
|
||||
# Start the server
|
||||
server = RemoteOpenAIServer(
|
||||
self.model_name,
|
||||
sargs,
|
||||
auto_port=False,
|
||||
env_dict={
|
||||
"VLLM_SERVER_DEV_MODE": "1",
|
||||
current_platform.device_control_env_var: ",".join(
|
||||
str(current_platform.device_id_to_physical_device_id(i))
|
||||
for i in range(gpu_start, gpu_end)
|
||||
),
|
||||
},
|
||||
)
|
||||
server.__enter__()
|
||||
print(
|
||||
f"Hybrid LB node {node} started successfully with "
|
||||
f"{self.dp_size_local} local DP ranks and "
|
||||
f"{self.api_server_count} API servers"
|
||||
)
|
||||
self.servers.append((server, sargs))
|
||||
except Exception as e:
|
||||
print(f"Failed to start hybrid LB node {node}: {e}")
|
||||
raise
|
||||
|
||||
thread = threading.Thread(target=start_server, args=(node_id, server_args))
|
||||
thread.start()
|
||||
|
||||
self.server_threads.append(thread)
|
||||
|
||||
# Wait for all servers to start
|
||||
for thread in self.server_threads:
|
||||
thread.join()
|
||||
|
||||
# Give servers additional time to fully initialize and coordinate
|
||||
time.sleep(3)
|
||||
|
||||
if len(self.servers) != self.num_nodes:
|
||||
raise Exception("Servers failed to start")
|
||||
|
||||
return self.servers
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Stop all server instances."""
|
||||
while self.servers:
|
||||
try:
|
||||
self.servers.pop()[0].__exit__(exc_type, exc_val, exc_tb)
|
||||
except Exception as e:
|
||||
print(f"Error stopping server: {e}")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def default_server_args():
|
||||
return [
|
||||
# use half precision for speed and memory savings in CI environment
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-model-len",
|
||||
"2048",
|
||||
"--max-num-seqs",
|
||||
"128",
|
||||
"--enforce-eager",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=[1, 4])
|
||||
def server_manager(request, default_server_args):
|
||||
api_server_count = request.param
|
||||
server_manager = HybridLBServerManager(
|
||||
MODEL_NAME,
|
||||
DP_SIZE,
|
||||
api_server_count,
|
||||
default_server_args,
|
||||
DP_SIZE_LOCAL,
|
||||
TP_SIZE,
|
||||
)
|
||||
|
||||
with server_manager:
|
||||
yield server_manager
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def servers(server_manager):
|
||||
return server_manager.servers
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def clients(servers: list[tuple[RemoteOpenAIServer, list[str]]]):
|
||||
# Create a client for each node (each node has its own API endpoint)
|
||||
async with AsyncExitStack() as stack:
|
||||
yield [
|
||||
await stack.enter_async_context(server.get_async_client())
|
||||
for server, _ in servers
|
||||
]
|
||||
|
||||
|
||||
def _get_parallel_config(server: RemoteOpenAIServer):
|
||||
response = requests.get(server.url_for("server_info?config_format=json"))
|
||||
response.raise_for_status()
|
||||
|
||||
vllm_config = response.json()["vllm_config"]
|
||||
return vllm_config["parallel_config"]
|
||||
|
||||
|
||||
def test_hybrid_dp_server_info(server_manager):
|
||||
servers = server_manager.servers
|
||||
api_server_count = server_manager.api_server_count
|
||||
|
||||
for i, (server, _) in enumerate(servers):
|
||||
print(f"Testing {i=}")
|
||||
|
||||
# Each request will hit one of the API servers
|
||||
# `n_reqs` is set so that there is a good chance each server
|
||||
# receives at least one request
|
||||
n_reqs = 2 * api_server_count * api_server_count
|
||||
parallel_configs = [_get_parallel_config(server) for _ in range(n_reqs)]
|
||||
api_process_counts = [c["_api_process_count"] for c in parallel_configs]
|
||||
api_process_ranks = [c["_api_process_rank"] for c in parallel_configs]
|
||||
|
||||
assert all(c == api_server_count for c in api_process_counts), (
|
||||
api_process_counts
|
||||
)
|
||||
assert all(0 <= r < api_server_count for r in api_process_ranks), (
|
||||
api_process_ranks
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_hybrid_lb_completion(
|
||||
clients: list[openai.AsyncOpenAI],
|
||||
servers: list[tuple[RemoteOpenAIServer, list[str]]],
|
||||
model_name: str,
|
||||
) -> None:
|
||||
async def make_request(client: openai.AsyncOpenAI):
|
||||
completion = await client.completions.create(
|
||||
model=model_name, prompt="Hello, my name is", max_tokens=5, temperature=1.0
|
||||
)
|
||||
|
||||
assert completion.id is not None
|
||||
assert completion.choices is not None and len(completion.choices) == 1
|
||||
|
||||
choice = completion.choices[0]
|
||||
# The exact number of tokens can vary slightly with temperature=1.0,
|
||||
# so we check for a reasonable minimum length.
|
||||
assert len(choice.text) >= 1
|
||||
# Finish reason might not always be 'length' if the model finishes early
|
||||
# or due to other reasons, especially with high temperature.
|
||||
# So, we'll accept 'length' or 'stop'.
|
||||
assert choice.finish_reason in ("length", "stop")
|
||||
|
||||
# Token counts can also vary, so we check they are positive.
|
||||
assert completion.usage.completion_tokens > 0
|
||||
assert completion.usage.prompt_tokens > 0
|
||||
assert completion.usage.total_tokens > 0
|
||||
return completion
|
||||
|
||||
# Test single request to each node
|
||||
for i, client in enumerate(clients):
|
||||
result = await make_request(client)
|
||||
assert result is not None
|
||||
print(f"Hybrid LB node {i} handled single completion request successfully")
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Send requests to all nodes - each should balance within its local DP ranks
|
||||
num_requests = 200 # Total 200 requests across 2 nodes
|
||||
all_tasks = []
|
||||
for i in range(num_requests):
|
||||
client = clients[i % len(clients)]
|
||||
all_tasks.append(asyncio.create_task(make_request(client)))
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
results = await asyncio.gather(*all_tasks)
|
||||
assert len(results) == num_requests
|
||||
assert all(completion is not None for completion in results)
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Second burst of requests
|
||||
all_tasks = []
|
||||
for i in range(num_requests):
|
||||
client = clients[i % len(clients)]
|
||||
all_tasks.append(asyncio.create_task(make_request(client)))
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
results = await asyncio.gather(*all_tasks)
|
||||
assert len(results) == num_requests
|
||||
assert all(completion is not None for completion in results)
|
||||
|
||||
_, server_args = servers[0]
|
||||
api_server_count = (
|
||||
server_args.count("--api-server-count")
|
||||
and server_args[server_args.index("--api-server-count") + 1]
|
||||
or 1
|
||||
)
|
||||
print(
|
||||
f"Successfully completed hybrid LB test with {len(clients)} nodes "
|
||||
f"({DP_SIZE_LOCAL} DP ranks each, API server count: {api_server_count})"
|
||||
)
|
||||
|
||||
# Check request balancing within each node
|
||||
for i, (server, _) in enumerate(servers):
|
||||
print(f"Checking request balancing for node {i}")
|
||||
check_request_balancing(server, DP_SIZE_LOCAL)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_hybrid_lb_completion_streaming(
|
||||
clients: list[openai.AsyncOpenAI],
|
||||
servers: list[tuple[RemoteOpenAIServer, list[str]]],
|
||||
model_name: str,
|
||||
) -> None:
|
||||
prompt = "What is an LLM?"
|
||||
|
||||
async def make_streaming_request(client: openai.AsyncOpenAI):
|
||||
# Perform a non-streaming request to get the expected full output
|
||||
single_completion = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
)
|
||||
single_output = single_completion.choices[0].text
|
||||
|
||||
# Perform the streaming request
|
||||
stream = await client.completions.create(
|
||||
model=model_name, prompt=prompt, max_tokens=5, temperature=0.0, stream=True
|
||||
)
|
||||
chunks: list[str] = []
|
||||
finish_reason_count = 0
|
||||
last_chunk = None
|
||||
async for chunk in stream:
|
||||
chunks.append(chunk.choices[0].text)
|
||||
if chunk.choices[0].finish_reason is not None:
|
||||
finish_reason_count += 1
|
||||
last_chunk = chunk # Keep track of the last chunk
|
||||
|
||||
# finish reason should only return in the last block for OpenAI API
|
||||
assert finish_reason_count == 1, "Finish reason should appear exactly once."
|
||||
assert last_chunk is not None, "Stream should have yielded at least one chunk."
|
||||
assert last_chunk.choices[0].finish_reason == "length", (
|
||||
"Finish reason should be 'length'."
|
||||
)
|
||||
# Check that the combined text matches the non-streamed version.
|
||||
assert "".join(chunks) == single_output, (
|
||||
"Streamed output should match non-streamed output."
|
||||
)
|
||||
return True # Indicate success for this request
|
||||
|
||||
# Test single request to each node
|
||||
for i, client in enumerate(clients):
|
||||
result = await make_streaming_request(client)
|
||||
assert result is not None
|
||||
print(f"Hybrid LB node {i} handled single streaming request successfully")
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Send streaming requests to all nodes
|
||||
num_requests = 200 # Total 200 requests across 2 nodes
|
||||
all_tasks = []
|
||||
for i in range(num_requests):
|
||||
client = clients[i % len(clients)]
|
||||
all_tasks.append(asyncio.create_task(make_streaming_request(client)))
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
results = await asyncio.gather(*all_tasks)
|
||||
assert len(results) == num_requests
|
||||
assert all(results), "Not all streaming requests completed successfully."
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Second burst of streaming requests
|
||||
all_tasks = []
|
||||
for i in range(num_requests):
|
||||
client = clients[i % len(clients)]
|
||||
all_tasks.append(asyncio.create_task(make_streaming_request(client)))
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
results = await asyncio.gather(*all_tasks)
|
||||
assert len(results) == num_requests
|
||||
assert all(results), "Not all streaming requests completed successfully."
|
||||
|
||||
_, server_args = servers[0]
|
||||
api_server_count = (
|
||||
server_args.count("--api-server-count")
|
||||
and server_args[server_args.index("--api-server-count") + 1]
|
||||
or 1
|
||||
)
|
||||
print(
|
||||
f"Successfully completed hybrid LB streaming test with "
|
||||
f"{len(clients)} nodes ({DP_SIZE_LOCAL} DP ranks each, "
|
||||
f"API server count: {api_server_count})"
|
||||
)
|
||||
|
||||
# Check request balancing within each node
|
||||
for i, (server, _) in enumerate(servers):
|
||||
print(f"Checking streaming request balancing for node {i}")
|
||||
check_request_balancing(server, DP_SIZE_LOCAL)
|
||||
734
third_party/vllm/tests/v1/distributed/test_internal_lb_dp.py
vendored
Normal file
734
third_party/vllm/tests/v1/distributed/test_internal_lb_dp.py
vendored
Normal file
@@ -0,0 +1,734 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import asyncio
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from typing import cast
|
||||
|
||||
import openai # use the official client for correctness check
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import requests
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from tests.v1.utils import check_request_balancing
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
MODEL_NAME = "ibm-research/PowerMoE-3b"
|
||||
|
||||
# Number of data parallel ranks for multi-node internal LB testing
|
||||
DP_SIZE = int(os.getenv("DP_SIZE", "2"))
|
||||
# Default tensor parallel size to use
|
||||
TP_SIZE = int(os.getenv("TP_SIZE", "1"))
|
||||
|
||||
# Number of nodes to simulate
|
||||
NUM_NODES = 2
|
||||
|
||||
|
||||
class MultinodeInternalLBServerManager:
|
||||
"""Manages multi-node data parallel vLLM server instances for internal
|
||||
load balancer testing using --headless mode."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
dp_size: int,
|
||||
api_server_count: int,
|
||||
base_server_args: list,
|
||||
dp_per_node: int = 1,
|
||||
tp_size: int = TP_SIZE,
|
||||
):
|
||||
self.model_name = model_name
|
||||
self.dp_size = dp_size
|
||||
self.dp_per_node = dp_per_node
|
||||
self.tp_size = tp_size
|
||||
self.api_server_count = api_server_count
|
||||
self.base_server_args = base_server_args
|
||||
self.servers: list[tuple[RemoteOpenAIServer, list[str]] | None] = [None] * (
|
||||
dp_size // dp_per_node
|
||||
)
|
||||
self.server_threads: list[threading.Thread] = []
|
||||
|
||||
def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
|
||||
"""Start all server instances for multi-node internal LB mode."""
|
||||
for server_idx, rank in enumerate(range(0, self.dp_size, self.dp_per_node)):
|
||||
# Create server args for this specific rank
|
||||
server_args = self.base_server_args.copy()
|
||||
|
||||
if rank == 0:
|
||||
# Head node - runs API server and first DP rank
|
||||
server_args.extend(
|
||||
[
|
||||
"--data-parallel-size",
|
||||
str(self.dp_size),
|
||||
"--data-parallel-size-local",
|
||||
str(self.dp_per_node),
|
||||
"--tensor-parallel-size",
|
||||
str(self.tp_size),
|
||||
"--port",
|
||||
"8000", # Single endpoint for all requests
|
||||
"--api-server-count",
|
||||
str(self.api_server_count),
|
||||
"--data-parallel-address",
|
||||
"127.0.0.1",
|
||||
"--data-parallel-rpc-port",
|
||||
"13345",
|
||||
]
|
||||
)
|
||||
else:
|
||||
# Secondary nodes - run in headless mode
|
||||
server_args.extend(
|
||||
[
|
||||
"--headless",
|
||||
"--data-parallel-size",
|
||||
str(self.dp_size),
|
||||
"--data-parallel-size-local",
|
||||
str(self.dp_per_node),
|
||||
"--data-parallel-start-rank",
|
||||
str(rank),
|
||||
"--tensor-parallel-size",
|
||||
str(self.tp_size),
|
||||
"--data-parallel-address",
|
||||
"127.0.0.1",
|
||||
"--data-parallel-rpc-port",
|
||||
"13345",
|
||||
]
|
||||
)
|
||||
|
||||
# Use a thread to start each server to allow parallel initialization
|
||||
def start_server(sidx: int, r: int, sargs: list[str]):
|
||||
gpus_per_node = self.tp_size * self.dp_per_node
|
||||
try:
|
||||
# Start the server
|
||||
server = RemoteOpenAIServer(
|
||||
self.model_name,
|
||||
sargs,
|
||||
auto_port=False,
|
||||
env_dict={
|
||||
"VLLM_SERVER_DEV_MODE": "1",
|
||||
current_platform.device_control_env_var: ",".join(
|
||||
str(current_platform.device_id_to_physical_device_id(i))
|
||||
for i in range(r, r + gpus_per_node)
|
||||
),
|
||||
},
|
||||
)
|
||||
server.__enter__()
|
||||
if r == 0:
|
||||
print(
|
||||
f"Head node (rank {r}) started successfully with "
|
||||
f"{self.api_server_count} API servers"
|
||||
)
|
||||
else:
|
||||
print(f"Headless node (rank {r}) started successfully")
|
||||
self.servers[sidx] = (server, sargs)
|
||||
except Exception as e:
|
||||
print(f"Failed to start server rank {r}: {e}")
|
||||
traceback.print_exc()
|
||||
raise
|
||||
|
||||
thread = threading.Thread(
|
||||
target=start_server, args=(server_idx, rank, server_args)
|
||||
)
|
||||
thread.start()
|
||||
|
||||
self.server_threads.append(thread)
|
||||
|
||||
# Wait for all servers to start
|
||||
for thread in self.server_threads:
|
||||
thread.join()
|
||||
|
||||
# Give servers additional time to fully initialize and coordinate
|
||||
time.sleep(3)
|
||||
|
||||
if not all(self.servers):
|
||||
raise Exception("Servers failed to start")
|
||||
|
||||
return cast(list[tuple[RemoteOpenAIServer, list[str]]], self.servers)
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Stop all server instances."""
|
||||
while self.servers:
|
||||
if server := self.servers.pop():
|
||||
try:
|
||||
server[0].__exit__(exc_type, exc_val, exc_tb)
|
||||
except Exception as e:
|
||||
print(f"Error stopping server: {e}")
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
class APIOnlyServerManager:
|
||||
"""Manages API-only server (Node 0) and headless engines server (Node 1)
|
||||
for testing separated API server and engine configuration."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
dp_size: int,
|
||||
api_server_count: int,
|
||||
base_server_args: list,
|
||||
tp_size: int = TP_SIZE,
|
||||
):
|
||||
self.model_name = model_name
|
||||
self.dp_size = dp_size
|
||||
self.tp_size = tp_size
|
||||
self.api_server_count = api_server_count
|
||||
self.base_server_args = base_server_args
|
||||
self.servers: list[tuple[RemoteOpenAIServer, list[str]] | None] = [None] * 2
|
||||
self.server_threads: list[threading.Thread] = []
|
||||
|
||||
def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
|
||||
"""Start API-only server and headless engines server."""
|
||||
|
||||
# Start API-only server (Node 0) - no engines, only API server
|
||||
api_server_args = self.base_server_args.copy()
|
||||
api_server_args.extend(
|
||||
[
|
||||
"--data-parallel-size",
|
||||
str(self.dp_size),
|
||||
"--data-parallel-size-local",
|
||||
"0", # No engines on this node
|
||||
"--tensor-parallel-size",
|
||||
str(self.tp_size),
|
||||
"--port",
|
||||
"8000",
|
||||
"--api-server-count",
|
||||
str(self.api_server_count),
|
||||
"--data-parallel-address",
|
||||
"127.0.0.1",
|
||||
"--data-parallel-rpc-port",
|
||||
"13345",
|
||||
]
|
||||
)
|
||||
|
||||
# Start headless engines server (Node 1) - all engines, no API server
|
||||
engines_server_args = self.base_server_args.copy()
|
||||
engines_server_args.extend(
|
||||
[
|
||||
"--headless",
|
||||
"--data-parallel-size",
|
||||
str(self.dp_size),
|
||||
"--data-parallel-size-local",
|
||||
str(self.dp_size), # All engines on this node
|
||||
"--tensor-parallel-size",
|
||||
str(self.tp_size),
|
||||
"--data-parallel-address",
|
||||
"127.0.0.1",
|
||||
"--data-parallel-rpc-port",
|
||||
"13345",
|
||||
]
|
||||
)
|
||||
|
||||
# Use threads to start both servers in parallel
|
||||
def start_api_server():
|
||||
try:
|
||||
server = RemoteOpenAIServer(
|
||||
self.model_name,
|
||||
api_server_args,
|
||||
auto_port=False,
|
||||
env_dict={
|
||||
"VLLM_SERVER_DEV_MODE": "1",
|
||||
# No GPUs needed for API-only server
|
||||
},
|
||||
)
|
||||
server.__enter__()
|
||||
print(
|
||||
f"API-only server started successfully with "
|
||||
f"{self.api_server_count} API servers"
|
||||
)
|
||||
self.servers[0] = (server, api_server_args)
|
||||
except Exception as e:
|
||||
print(f"Failed to start API-only server: {e}")
|
||||
raise
|
||||
|
||||
def start_engines_server():
|
||||
try:
|
||||
server = RemoteOpenAIServer(
|
||||
self.model_name,
|
||||
engines_server_args,
|
||||
auto_port=False,
|
||||
env_dict={
|
||||
current_platform.device_control_env_var: ",".join(
|
||||
str(current_platform.device_id_to_physical_device_id(i))
|
||||
for i in range(self.dp_size * self.tp_size)
|
||||
)
|
||||
},
|
||||
)
|
||||
server.__enter__()
|
||||
print(
|
||||
f"Headless engines server started successfully with "
|
||||
f"{self.dp_size} engines"
|
||||
)
|
||||
self.servers[1] = (server, engines_server_args)
|
||||
except Exception as e:
|
||||
print(f"Failed to start headless engines server: {e}")
|
||||
raise
|
||||
|
||||
# Start API server first
|
||||
api_thread = threading.Thread(target=start_api_server)
|
||||
api_thread.start()
|
||||
self.server_threads.append(api_thread)
|
||||
|
||||
# Start engines server second
|
||||
engines_thread = threading.Thread(target=start_engines_server)
|
||||
engines_thread.start()
|
||||
self.server_threads.append(engines_thread)
|
||||
|
||||
# Wait for both servers to start
|
||||
for thread in self.server_threads:
|
||||
thread.join()
|
||||
|
||||
# Give servers additional time to fully initialize and coordinate
|
||||
time.sleep(3)
|
||||
|
||||
if not all(self.servers):
|
||||
raise Exception("Both servers failed to start")
|
||||
|
||||
return cast(list[tuple[RemoteOpenAIServer, list[str]]], self.servers)
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Stop both server instances."""
|
||||
while self.servers:
|
||||
if server := self.servers.pop():
|
||||
try:
|
||||
server[0].__exit__(exc_type, exc_val, exc_tb)
|
||||
except Exception as e:
|
||||
print(f"Error stopping server: {e}")
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def default_server_args():
|
||||
return [
|
||||
# use half precision for speed and memory savings in CI environment
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-model-len",
|
||||
"2048",
|
||||
"--max-num-seqs",
|
||||
"128",
|
||||
"--enforce-eager",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=[1, 4])
|
||||
def server_manager(request, default_server_args):
|
||||
api_server_count = request.param
|
||||
server_manager = MultinodeInternalLBServerManager(
|
||||
MODEL_NAME,
|
||||
DP_SIZE,
|
||||
api_server_count,
|
||||
default_server_args,
|
||||
DP_SIZE // NUM_NODES,
|
||||
TP_SIZE,
|
||||
)
|
||||
|
||||
with server_manager:
|
||||
yield server_manager
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def servers(server_manager):
|
||||
return server_manager.servers
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=[1, 4])
|
||||
def api_only_servers(request, default_server_args):
|
||||
"""Fixture for API-only server + headless engines configuration."""
|
||||
api_server_count = request.param
|
||||
with APIOnlyServerManager(
|
||||
MODEL_NAME, DP_SIZE, api_server_count, default_server_args, TP_SIZE
|
||||
) as server_list:
|
||||
yield server_list
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(servers: list[tuple[RemoteOpenAIServer, list[str]]]):
|
||||
# For internal LB, we only connect to the head node (rank 0)
|
||||
# which provides the single API endpoint
|
||||
head_server = servers[0][0]
|
||||
async with head_server.get_async_client() as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_only_client(api_only_servers: list[tuple[RemoteOpenAIServer, list[str]]]):
|
||||
"""Client fixture for API-only server configuration."""
|
||||
# Connect to the API-only server (first server in the list)
|
||||
api_server = api_only_servers[0][0]
|
||||
async with api_server.get_async_client() as client:
|
||||
yield client
|
||||
|
||||
|
||||
def _get_parallel_config(server: RemoteOpenAIServer):
|
||||
response = requests.get(server.url_for("server_info?config_format=json"))
|
||||
response.raise_for_status()
|
||||
|
||||
vllm_config = response.json()["vllm_config"]
|
||||
return vllm_config["parallel_config"]
|
||||
|
||||
|
||||
def test_multinode_dp_server_info(server_manager):
|
||||
head_server = server_manager.servers[0][0]
|
||||
api_server_count = server_manager.api_server_count
|
||||
|
||||
# Each request will hit one of the API servers
|
||||
# `n_reqs` is set so that there is a good chance each server
|
||||
# receives at least one request
|
||||
n_reqs = 2 * api_server_count * api_server_count
|
||||
parallel_configs = [_get_parallel_config(head_server) for _ in range(n_reqs)]
|
||||
api_process_counts = [c["_api_process_count"] for c in parallel_configs]
|
||||
api_process_ranks = [c["_api_process_rank"] for c in parallel_configs]
|
||||
|
||||
assert all(c == api_server_count for c in api_process_counts), api_process_counts
|
||||
assert all(0 <= r < api_server_count for r in api_process_ranks), api_process_ranks
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_multinode_dp_completion(
|
||||
client: openai.AsyncOpenAI,
|
||||
servers: list[tuple[RemoteOpenAIServer, list[str]]],
|
||||
model_name: str,
|
||||
) -> None:
|
||||
async def make_request():
|
||||
completion = await client.completions.create(
|
||||
model=model_name, prompt="Hello, my name is", max_tokens=5, temperature=1.0
|
||||
)
|
||||
|
||||
assert completion.id is not None
|
||||
assert completion.choices is not None and len(completion.choices) == 1
|
||||
|
||||
choice = completion.choices[0]
|
||||
# The exact number of tokens can vary slightly with temperature=1.0,
|
||||
# so we check for a reasonable minimum length.
|
||||
assert len(choice.text) >= 1
|
||||
# Finish reason might not always be 'length' if the model finishes early
|
||||
# or due to other reasons, especially with high temperature.
|
||||
# So, we'll accept 'length' or 'stop'.
|
||||
assert choice.finish_reason in ("length", "stop")
|
||||
|
||||
# Token counts can also vary, so we check they are positive.
|
||||
assert completion.usage.completion_tokens > 0
|
||||
assert completion.usage.prompt_tokens > 0
|
||||
assert completion.usage.total_tokens > 0
|
||||
return completion
|
||||
|
||||
# Test single request
|
||||
result = await make_request()
|
||||
assert result is not None
|
||||
print("Multi-node internal LB handled single completion request successfully")
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Send multiple requests - internal LB should distribute across DP ranks
|
||||
num_requests = 200
|
||||
all_tasks = []
|
||||
for _ in range(num_requests):
|
||||
all_tasks.append(asyncio.create_task(make_request()))
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
results = await asyncio.gather(*all_tasks)
|
||||
assert len(results) == num_requests
|
||||
assert all(completion is not None for completion in results)
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Second burst of requests
|
||||
all_tasks = []
|
||||
for _ in range(num_requests):
|
||||
all_tasks.append(asyncio.create_task(make_request()))
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
results = await asyncio.gather(*all_tasks)
|
||||
assert len(results) == num_requests
|
||||
assert all(completion is not None for completion in results)
|
||||
|
||||
_, server_args = servers[0]
|
||||
api_server_count = (
|
||||
server_args.count("--api-server-count")
|
||||
and server_args[server_args.index("--api-server-count") + 1]
|
||||
or 1
|
||||
)
|
||||
print(
|
||||
f"Successfully completed multi-node internal LB test with "
|
||||
f"{len(servers)} DP ranks (API server count: {api_server_count})"
|
||||
)
|
||||
|
||||
# Check request balancing via Prometheus metrics
|
||||
head_server = servers[0][0]
|
||||
check_request_balancing(head_server, DP_SIZE)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_multinode_dp_completion_streaming(
|
||||
client: openai.AsyncOpenAI,
|
||||
servers: list[tuple[RemoteOpenAIServer, list[str]]],
|
||||
model_name: str,
|
||||
) -> None:
|
||||
prompt = "What is an LLM?"
|
||||
|
||||
async def make_streaming_request():
|
||||
# Perform a non-streaming request to get the expected full output
|
||||
single_completion = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
)
|
||||
single_output = single_completion.choices[0].text
|
||||
|
||||
# Perform the streaming request
|
||||
stream = await client.completions.create(
|
||||
model=model_name, prompt=prompt, max_tokens=5, temperature=0.0, stream=True
|
||||
)
|
||||
chunks: list[str] = []
|
||||
finish_reason_count = 0
|
||||
last_chunk = None
|
||||
async for chunk in stream:
|
||||
chunks.append(chunk.choices[0].text)
|
||||
if chunk.choices[0].finish_reason is not None:
|
||||
finish_reason_count += 1
|
||||
last_chunk = chunk # Keep track of the last chunk
|
||||
|
||||
# finish reason should only return in the last block for OpenAI API
|
||||
assert finish_reason_count == 1, "Finish reason should appear exactly once."
|
||||
assert last_chunk is not None, "Stream should have yielded at least one chunk."
|
||||
assert last_chunk.choices[0].finish_reason == "length", (
|
||||
"Finish reason should be 'length'."
|
||||
)
|
||||
# Check that the combined text matches the non-streamed version.
|
||||
assert "".join(chunks) == single_output, (
|
||||
"Streamed output should match non-streamed output."
|
||||
)
|
||||
return True # Indicate success for this request
|
||||
|
||||
# Test single streaming request
|
||||
result = await make_streaming_request()
|
||||
assert result is not None
|
||||
print("Multi-node internal LB handled single streaming request successfully")
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Send multiple streaming requests - internal LB should distribute across
|
||||
# DP ranks
|
||||
num_requests = 200
|
||||
all_tasks = []
|
||||
for _ in range(num_requests):
|
||||
all_tasks.append(asyncio.create_task(make_streaming_request()))
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
results = await asyncio.gather(*all_tasks)
|
||||
assert len(results) == num_requests
|
||||
assert all(results), "Not all streaming requests completed successfully."
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Second burst of streaming requests
|
||||
all_tasks = []
|
||||
for _ in range(num_requests):
|
||||
all_tasks.append(asyncio.create_task(make_streaming_request()))
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
results = await asyncio.gather(*all_tasks)
|
||||
assert len(results) == num_requests
|
||||
assert all(results), "Not all streaming requests completed successfully."
|
||||
|
||||
_, server_args = servers[0]
|
||||
api_server_count = (
|
||||
server_args.count("--api-server-count")
|
||||
and server_args[server_args.index("--api-server-count") + 1]
|
||||
or 1
|
||||
)
|
||||
print(
|
||||
f"Successfully completed multi-node internal LB streaming test with "
|
||||
f"{len(servers)} DP ranks (API server count: {api_server_count})"
|
||||
)
|
||||
|
||||
# Check request balancing via Prometheus metrics
|
||||
head_server = servers[0][0]
|
||||
check_request_balancing(head_server, DP_SIZE)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_api_only_multinode_dp_completion(
|
||||
api_only_client: openai.AsyncOpenAI,
|
||||
api_only_servers: list[tuple[RemoteOpenAIServer, list[str]]],
|
||||
model_name: str,
|
||||
) -> None:
|
||||
"""Test API-only server with all engines on separate headless server."""
|
||||
|
||||
async def make_request():
|
||||
completion = await api_only_client.completions.create(
|
||||
model=model_name, prompt="Hello, my name is", max_tokens=5, temperature=1.0
|
||||
)
|
||||
|
||||
assert completion.id is not None
|
||||
assert completion.choices is not None and len(completion.choices) == 1
|
||||
|
||||
choice = completion.choices[0]
|
||||
# The exact number of tokens can vary slightly with temperature=1.0,
|
||||
# so we check for a reasonable minimum length.
|
||||
assert len(choice.text) >= 1
|
||||
# Finish reason might not always be 'length' if the model finishes
|
||||
# early or due to other reasons, especially with high temperature.
|
||||
# So, we'll accept 'length' or 'stop'.
|
||||
assert choice.finish_reason in ("length", "stop")
|
||||
|
||||
# Token counts can also vary, so we check they are positive.
|
||||
assert completion.usage.completion_tokens > 0
|
||||
assert completion.usage.prompt_tokens > 0
|
||||
assert completion.usage.total_tokens > 0
|
||||
return completion
|
||||
|
||||
# Test single request
|
||||
result = await make_request()
|
||||
assert result is not None
|
||||
print("API-only server handled single completion request successfully")
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Send multiple requests - should be distributed across engines on
|
||||
# headless server
|
||||
num_requests = 200
|
||||
all_tasks = []
|
||||
for _ in range(num_requests):
|
||||
all_tasks.append(asyncio.create_task(make_request()))
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
results = await asyncio.gather(*all_tasks)
|
||||
assert len(results) == num_requests
|
||||
assert all(completion is not None for completion in results)
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Second burst of requests
|
||||
all_tasks = []
|
||||
for _ in range(num_requests):
|
||||
all_tasks.append(asyncio.create_task(make_request()))
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
results = await asyncio.gather(*all_tasks)
|
||||
assert len(results) == num_requests
|
||||
assert all(completion is not None for completion in results)
|
||||
|
||||
api_server, api_server_args = api_only_servers[0]
|
||||
api_server_count = (
|
||||
api_server_args.count("--api-server-count")
|
||||
and api_server_args[api_server_args.index("--api-server-count") + 1]
|
||||
or 1
|
||||
)
|
||||
print(
|
||||
f"Successfully completed API-only multi-node test with {DP_SIZE} "
|
||||
f"engines on headless server (API server count: {api_server_count})"
|
||||
)
|
||||
|
||||
# Check request balancing via Prometheus metrics
|
||||
check_request_balancing(api_server, DP_SIZE)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_api_only_multinode_dp_completion_streaming(
|
||||
api_only_client: openai.AsyncOpenAI,
|
||||
api_only_servers: list[tuple[RemoteOpenAIServer, list[str]]],
|
||||
model_name: str,
|
||||
) -> None:
|
||||
"""Test API-only server streaming with all engines on separate
|
||||
headless server."""
|
||||
prompt = "What is an LLM?"
|
||||
|
||||
async def make_streaming_request():
|
||||
# Perform a non-streaming request to get the expected full output
|
||||
single_completion = await api_only_client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
)
|
||||
single_output = single_completion.choices[0].text
|
||||
|
||||
# Perform the streaming request
|
||||
stream = await api_only_client.completions.create(
|
||||
model=model_name, prompt=prompt, max_tokens=5, temperature=0.0, stream=True
|
||||
)
|
||||
chunks: list[str] = []
|
||||
finish_reason_count = 0
|
||||
last_chunk = None
|
||||
async for chunk in stream:
|
||||
chunks.append(chunk.choices[0].text)
|
||||
if chunk.choices[0].finish_reason is not None:
|
||||
finish_reason_count += 1
|
||||
last_chunk = chunk # Keep track of the last chunk
|
||||
|
||||
# finish reason should only return in the last block for OpenAI API
|
||||
assert finish_reason_count == 1, "Finish reason should appear exactly once."
|
||||
assert last_chunk is not None, "Stream should have yielded at least one chunk."
|
||||
assert last_chunk.choices[0].finish_reason == "length", (
|
||||
"Finish reason should be 'length'."
|
||||
)
|
||||
# Check that the combined text matches the non-streamed version.
|
||||
assert "".join(chunks) == single_output, (
|
||||
"Streamed output should match non-streamed output."
|
||||
)
|
||||
return True # Indicate success for this request
|
||||
|
||||
# Test single streaming request
|
||||
result = await make_streaming_request()
|
||||
assert result is not None
|
||||
print("API-only server handled single streaming request successfully")
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Send multiple streaming requests - should be distributed across engines
|
||||
num_requests = 200
|
||||
all_tasks = []
|
||||
for _ in range(num_requests):
|
||||
all_tasks.append(asyncio.create_task(make_streaming_request()))
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
results = await asyncio.gather(*all_tasks)
|
||||
assert len(results) == num_requests
|
||||
assert all(results), "Not all streaming requests completed successfully."
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Second burst of streaming requests
|
||||
all_tasks = []
|
||||
for _ in range(num_requests):
|
||||
all_tasks.append(asyncio.create_task(make_streaming_request()))
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
results = await asyncio.gather(*all_tasks)
|
||||
assert len(results) == num_requests
|
||||
assert all(results), "Not all streaming requests completed successfully."
|
||||
|
||||
_, api_server_args = api_only_servers[0]
|
||||
api_server_count = (
|
||||
api_server_args.count("--api-server-count")
|
||||
and api_server_args[api_server_args.index("--api-server-count") + 1]
|
||||
or 1
|
||||
)
|
||||
print(
|
||||
f"Successfully completed API-only streaming test with {DP_SIZE} "
|
||||
f"engines on headless server (API server count: {api_server_count})"
|
||||
)
|
||||
|
||||
# Check request balancing via Prometheus metrics
|
||||
api_server = api_only_servers[0][0]
|
||||
check_request_balancing(api_server, DP_SIZE)
|
||||
0
third_party/vllm/tests/v1/e2e/__init__.py
vendored
Normal file
0
third_party/vllm/tests/v1/e2e/__init__.py
vendored
Normal file
0
third_party/vllm/tests/v1/e2e/general/__init__.py
vendored
Normal file
0
third_party/vllm/tests/v1/e2e/general/__init__.py
vendored
Normal file
437
third_party/vllm/tests/v1/e2e/general/test_async_scheduling.py
vendored
Normal file
437
third_party/vllm/tests/v1/e2e/general/test_async_scheduling.py
vendored
Normal file
@@ -0,0 +1,437 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import os
|
||||
from itertools import repeat
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch._dynamo.config as dynamo_config
|
||||
|
||||
from tests.utils import (
|
||||
large_gpu_mark,
|
||||
single_gpu_only,
|
||||
)
|
||||
from vllm import SamplingParams
|
||||
from vllm.logprobs import Logprob
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.sampling_params import StructuredOutputsParams
|
||||
from vllm.v1.metrics.reader import Metric
|
||||
|
||||
from ....conftest import VllmRunner
|
||||
from ....models.utils import check_outputs_equal
|
||||
|
||||
MODEL = "Qwen/Qwen3-0.6B"
|
||||
MTP_MODEL = "meta-llama/Llama-3.2-1B-Instruct"
|
||||
|
||||
# Need to enforce eager for MRV2 while we sort out cudagraph issues.
|
||||
ENFORCE_EAGER = os.getenv("ENFORCE_EAGER", "0") == "1"
|
||||
|
||||
first_prompt = (
|
||||
"The following numbers of the sequence "
|
||||
+ ", ".join(str(i) for i in range(10))
|
||||
+ " are:"
|
||||
)
|
||||
example_prompts = [first_prompt, "In one word, the capital of France is "] + [
|
||||
f"Tell me about the number {i}: " for i in range(32)
|
||||
]
|
||||
|
||||
default_params = dict(
|
||||
temperature=0.0, # greedy
|
||||
max_tokens=30,
|
||||
min_tokens=28,
|
||||
)
|
||||
|
||||
|
||||
@single_gpu_only
|
||||
def test_without_spec_decoding(
|
||||
sample_json_schema,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Test consistency of combos of async scheduling, preemption,
|
||||
uni/multiproc executor, prefill chunking."""
|
||||
struct_outputs = StructuredOutputsParams(json=sample_json_schema)
|
||||
test_sampling_params: list[dict[str, Any]] = [
|
||||
dict(),
|
||||
# dict(min_tokens=20),
|
||||
dict(frequency_penalty=-1.0),
|
||||
dict(bad_words=["the", " the"]),
|
||||
dict(logprobs=2),
|
||||
dict(logprobs=2, frequency_penalty=-1.0),
|
||||
dict(structured_outputs=struct_outputs),
|
||||
dict(
|
||||
structured_outputs=struct_outputs,
|
||||
logprobs=2,
|
||||
),
|
||||
dict(
|
||||
structured_outputs=struct_outputs,
|
||||
frequency_penalty=-1.0,
|
||||
),
|
||||
dict(
|
||||
structured_outputs=struct_outputs,
|
||||
logprobs=2,
|
||||
frequency_penalty=-1.0,
|
||||
),
|
||||
]
|
||||
|
||||
# test_preemption, executor, async_scheduling,
|
||||
# spec_config, test_prefill_chunking
|
||||
test_configs = [
|
||||
(False, "mp", False, None, False),
|
||||
(True, "mp", False, None, True),
|
||||
(False, "mp", True, None, False),
|
||||
(False, "uni", True, None, False),
|
||||
(True, "mp", True, None, False),
|
||||
(True, "uni", True, None, False),
|
||||
(False, "mp", True, None, True),
|
||||
(True, "mp", True, None, True),
|
||||
(True, "uni", True, None, True),
|
||||
]
|
||||
|
||||
if current_platform.is_rocm():
|
||||
# On ROCm, Only test with structured_outputs (deterministic)
|
||||
# and skip chunk_prefill (more variable).
|
||||
test_configs = [
|
||||
cfg
|
||||
for cfg in test_configs
|
||||
if not cfg[4] # skip chunk_prefill=True
|
||||
]
|
||||
test_sampling_params = [
|
||||
p for p in test_sampling_params if p.get("structured_outputs") is not None
|
||||
]
|
||||
|
||||
run_tests(monkeypatch, MODEL, test_configs, test_sampling_params)
|
||||
|
||||
|
||||
@single_gpu_only
|
||||
@large_gpu_mark(min_gb=16)
|
||||
def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test consistency and acceptance rates with some different combos of
|
||||
preemption, executor, async scheduling, prefill chunking,
|
||||
spec decoding model length.
|
||||
"""
|
||||
|
||||
spec_config = {
|
||||
"method": "eagle3",
|
||||
"num_speculative_tokens": 2,
|
||||
"model": "nm-testing/Llama3_2_1B_speculator.eagle3",
|
||||
}
|
||||
# Set small draft model len to force doesn't-fit-in-drafter case.
|
||||
spec_config_short = spec_config | {"max_model_len": 50}
|
||||
|
||||
struct_outputs = StructuredOutputsParams(json=sample_json_schema)
|
||||
|
||||
test_sampling_params = [
|
||||
dict(),
|
||||
dict(frequency_penalty=-1.0),
|
||||
dict(bad_words=["the", " the"]),
|
||||
dict(logprobs=2),
|
||||
dict(logprobs=2, frequency_penalty=-1.0),
|
||||
dict(structured_outputs=struct_outputs),
|
||||
dict(
|
||||
structured_outputs=struct_outputs,
|
||||
logprobs=2,
|
||||
frequency_penalty=-1.0,
|
||||
),
|
||||
]
|
||||
|
||||
# test_preemption, executor, async_scheduling,
|
||||
# spec_config, test_prefill_chunking
|
||||
test_configs = [
|
||||
(False, "mp", False, None, False),
|
||||
(False, "mp", False, spec_config, False),
|
||||
(True, "mp", False, spec_config, True),
|
||||
(True, "uni", False, spec_config_short, True),
|
||||
(False, "mp", True, spec_config, False),
|
||||
(True, "mp", True, spec_config, False),
|
||||
(False, "mp", True, spec_config_short, True),
|
||||
(True, "uni", True, spec_config, False),
|
||||
(True, "uni", True, spec_config_short, False),
|
||||
(True, "mp", True, spec_config, True),
|
||||
(True, "uni", True, spec_config_short, True),
|
||||
]
|
||||
|
||||
run_tests(monkeypatch, MTP_MODEL, test_configs, test_sampling_params)
|
||||
|
||||
|
||||
@pytest.mark.flaky(reruns=2, only_on=current_platform.is_rocm())
|
||||
def test_with_ngram_gpu_spec_decoding(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test ngram_gpu speculative decoding with different configurations.
|
||||
|
||||
This test specifically validates ngram_gpu behavior with various:
|
||||
- Number of speculative tokens (2-6)
|
||||
- Prompt lookup window sizes (min/max)
|
||||
- Async scheduling enabled (as in production)
|
||||
- Different executors and chunking settings
|
||||
"""
|
||||
|
||||
# Variant with larger speculation window
|
||||
ngram_gpu_config = {
|
||||
"method": "ngram_gpu",
|
||||
"num_speculative_tokens": 3,
|
||||
"prompt_lookup_max": 3,
|
||||
"prompt_lookup_min": 2,
|
||||
}
|
||||
|
||||
# Test configurations covering various scenarios
|
||||
# test_preemption, executor, async_scheduling,
|
||||
# spec_config, test_prefill_chunking
|
||||
test_configs = [
|
||||
(False, "mp", False, None, False),
|
||||
(False, "mp", False, ngram_gpu_config, False),
|
||||
(True, "mp", False, ngram_gpu_config, True),
|
||||
(False, "mp", True, ngram_gpu_config, False),
|
||||
(True, "mp", True, ngram_gpu_config, False),
|
||||
(True, "uni", True, ngram_gpu_config, False),
|
||||
(True, "mp", True, ngram_gpu_config, True),
|
||||
]
|
||||
|
||||
# Use MODEL (Qwen) for ngram_gpu tests as it's lighter weight
|
||||
# and ngram_gpu doesn't require a specific draft model
|
||||
run_tests(monkeypatch, MODEL, test_configs, [{}])
|
||||
|
||||
|
||||
@dynamo_config.patch(cache_size_limit=16)
|
||||
def run_tests(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
model: str,
|
||||
test_configs: list[tuple],
|
||||
test_sampling_params: list[dict[str, Any]],
|
||||
):
|
||||
"""Test consistency of combos of async scheduling, preemption,
|
||||
uni/multiproc executor with spec decoding."""
|
||||
|
||||
# Flex attention supports float32.
|
||||
attention_config = {"backend": "FLEX_ATTENTION"}
|
||||
|
||||
with monkeypatch.context() as m:
|
||||
# lock matmul precision to full FP32 (IEEE)
|
||||
m.setenv("VLLM_FLOAT32_MATMUL_PRECISION", "highest")
|
||||
outputs: list[tuple[str, list, list]] = []
|
||||
for n, (
|
||||
test_preemption,
|
||||
executor,
|
||||
async_scheduling,
|
||||
spec_config,
|
||||
test_prefill_chunking,
|
||||
) in enumerate(test_configs, 1):
|
||||
test_str = f"{n}/{len(test_configs)}"
|
||||
test_results = run_test(
|
||||
model,
|
||||
test_str,
|
||||
test_sampling_params,
|
||||
test_preemption,
|
||||
executor,
|
||||
async_scheduling,
|
||||
spec_config,
|
||||
test_prefill_chunking=test_prefill_chunking,
|
||||
attention_config=attention_config,
|
||||
)
|
||||
outputs.append(test_results)
|
||||
|
||||
baseline_config, baseline_tests, _ = outputs[0]
|
||||
_, _, baseline_acceptances = next(
|
||||
(o for o in outputs if o[2] is not None), (None, None, None)
|
||||
)
|
||||
|
||||
print(f"BASELINE: config=[{baseline_config}], accept_rates={baseline_acceptances}")
|
||||
|
||||
failure = None
|
||||
for test_config, test_outputs, test_acceptance_rates in outputs[1:]:
|
||||
for (base_outs, base_logprobs), base_acceptance_rate, (
|
||||
test_outs,
|
||||
test_logprobs,
|
||||
), test_acceptance_rate, params in zip(
|
||||
baseline_tests,
|
||||
baseline_acceptances or repeat(None),
|
||||
test_outputs,
|
||||
test_acceptance_rates or repeat(None),
|
||||
test_sampling_params,
|
||||
):
|
||||
reason = None
|
||||
try:
|
||||
check_outputs_equal(
|
||||
outputs_0_lst=base_outs,
|
||||
outputs_1_lst=test_outs,
|
||||
name_0=f"baseline=[{baseline_config}], params={params}",
|
||||
name_1=f"config=[{test_config}], params={params}",
|
||||
)
|
||||
except AssertionError as e:
|
||||
reason = "outputs ", e
|
||||
|
||||
if reason is None:
|
||||
try:
|
||||
assert _all_logprobs_match(base_logprobs, test_logprobs)
|
||||
except AssertionError as e:
|
||||
reason = "logprobs", e
|
||||
|
||||
if reason is None:
|
||||
try:
|
||||
if (
|
||||
base_acceptance_rate is not None
|
||||
and test_acceptance_rate is not None
|
||||
):
|
||||
if "spec_mml=None" in test_config:
|
||||
# Preemption causes more variance in acceptance rates
|
||||
if (
|
||||
current_platform.is_rocm()
|
||||
and "preemption=True" in test_config
|
||||
):
|
||||
tolerance = 0.10
|
||||
else:
|
||||
tolerance = 0.05
|
||||
assert (
|
||||
test_acceptance_rate > base_acceptance_rate
|
||||
or test_acceptance_rate
|
||||
== pytest.approx(base_acceptance_rate, rel=tolerance)
|
||||
)
|
||||
else:
|
||||
# Currently the reported acceptance rate is expected to be
|
||||
# lower when we sometimes skip drafting altogether.
|
||||
assert test_acceptance_rate > 0.1
|
||||
except AssertionError as e:
|
||||
reason = "accept ", e
|
||||
|
||||
if reason is None:
|
||||
print(
|
||||
f"\033[32mPASSED\033[0m: "
|
||||
f"config=[{test_config}], params={params}"
|
||||
f" accept_rate={test_acceptance_rate}"
|
||||
)
|
||||
else:
|
||||
reason_str, _ = reason
|
||||
print(
|
||||
f"\033[31mFAILED\033[0m({reason_str}): "
|
||||
f"config=[{test_config}], params={params}"
|
||||
f" accept_rate={test_acceptance_rate}"
|
||||
)
|
||||
if failure is None:
|
||||
_, failure = reason
|
||||
|
||||
if failure is not None:
|
||||
raise failure
|
||||
|
||||
|
||||
def run_test(
|
||||
model: str,
|
||||
test_str: str,
|
||||
sampling_param_tests: list[dict[str, Any]],
|
||||
test_preemption: bool,
|
||||
executor: str,
|
||||
async_scheduling: bool,
|
||||
spec_config: dict[str, Any] | None,
|
||||
test_prefill_chunking: bool,
|
||||
attention_config: dict[str, Any] | None = None,
|
||||
):
|
||||
spec_decoding = spec_config is not None
|
||||
cache_arg: dict[str, Any] = (
|
||||
# Force preemptions
|
||||
dict(num_gpu_blocks_override=32)
|
||||
if test_preemption
|
||||
else dict(gpu_memory_utilization=0.9)
|
||||
)
|
||||
spec_mml = (spec_config or {}).get("max_model_len")
|
||||
spec_method = (spec_config or {}).get("method", "none")
|
||||
test_config = (
|
||||
f"executor={executor}, preemption={test_preemption}, "
|
||||
f"async_sched={async_scheduling}, "
|
||||
f"chunk_prefill={test_prefill_chunking}, "
|
||||
f"spec_decoding={spec_decoding}, spec_method={spec_method}, spec_mml={spec_mml}"
|
||||
)
|
||||
print("-" * 80)
|
||||
print(f"---- TESTING {test_str}: {test_config}")
|
||||
print("-" * 80)
|
||||
|
||||
with VllmRunner(
|
||||
model,
|
||||
max_model_len=4096,
|
||||
enable_chunked_prefill=test_prefill_chunking,
|
||||
# Force prefill chunking
|
||||
max_num_batched_tokens=48 if test_prefill_chunking else None,
|
||||
enforce_eager=ENFORCE_EAGER,
|
||||
async_scheduling=async_scheduling,
|
||||
distributed_executor_backend=executor,
|
||||
dtype="float32",
|
||||
speculative_config=spec_config,
|
||||
disable_log_stats=False,
|
||||
attention_config=attention_config,
|
||||
enable_prefix_caching=False if current_platform.is_rocm() else None,
|
||||
**cache_arg,
|
||||
) as vllm_model:
|
||||
results = []
|
||||
acceptance_rates: list[float] | None = [] if spec_decoding else None
|
||||
for override_params in sampling_param_tests:
|
||||
metrics_before = vllm_model.llm.get_metrics()
|
||||
print(f"----------- RUNNING PARAMS: {override_params}")
|
||||
results.append(
|
||||
vllm_model.generate(
|
||||
example_prompts,
|
||||
sampling_params=SamplingParams(**default_params, **override_params),
|
||||
return_logprobs=True,
|
||||
)
|
||||
)
|
||||
metrics_after = vllm_model.llm.get_metrics()
|
||||
if acceptance_rates is not None:
|
||||
acceptance_rate = _get_acceptance_rate(metrics_before, metrics_after)
|
||||
acceptance_rates.append(acceptance_rate)
|
||||
print(f"ACCEPTANCE RATE {acceptance_rate}")
|
||||
|
||||
if test_preemption:
|
||||
preemptions = _get_count(
|
||||
metrics_before, metrics_after, "vllm:num_preemptions"
|
||||
)
|
||||
assert preemptions > 0, "preemption test had no preemptions"
|
||||
|
||||
if len(results) > 1:
|
||||
# First check that the different parameter configs
|
||||
# actually result in different output.
|
||||
for (other_test_outs, other_test_logprobs), params in zip(
|
||||
results[1:], sampling_param_tests[1:]
|
||||
):
|
||||
with pytest.raises(AssertionError):
|
||||
check_outputs_equal(
|
||||
outputs_0_lst=results[0][0],
|
||||
outputs_1_lst=other_test_outs,
|
||||
name_0=f"baseline params={params}",
|
||||
name_1=f"other params={params}",
|
||||
)
|
||||
assert _all_logprobs_match(results[0][1], other_test_logprobs)
|
||||
|
||||
return test_config, results, acceptance_rates
|
||||
|
||||
|
||||
def _all_logprobs_match(req_a, req_b) -> bool:
|
||||
return (
|
||||
req_a == req_b
|
||||
or len(req_a) == len(req_b)
|
||||
and all(
|
||||
len(seq_a) == len(seq_b)
|
||||
and all(_logprobs_match(a, b) for a, b in zip(seq_a, seq_b))
|
||||
for seq_a, seq_b in zip(req_a, req_b)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _logprobs_match(lps_a: dict[int, Logprob], lps_b: dict[int, Logprob]) -> bool:
|
||||
rel_tol, abs_tol = 1e-3, 1e-6
|
||||
return (
|
||||
len(lps_a) == len(lps_b)
|
||||
and lps_a.keys() == lps_b.keys()
|
||||
and all(
|
||||
a.decoded_token == b.decoded_token
|
||||
and a.rank == b.rank
|
||||
and a.logprob == pytest.approx(b.logprob, rel=rel_tol, abs=abs_tol)
|
||||
for a, b in ((lps_a[x], lps_b[x]) for x in lps_a)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _get_acceptance_rate(before: list[Metric], after: list[Metric]) -> float:
|
||||
draft = _get_count(before, after, "vllm:spec_decode_num_draft_tokens")
|
||||
accept = _get_count(before, after, "vllm:spec_decode_num_accepted_tokens")
|
||||
return accept / draft if draft > 0 else 0.0
|
||||
|
||||
|
||||
def _get_count(before: list[Metric], after: list[Metric], name: str) -> int:
|
||||
before_val = next(m.value for m in before if m.name == name)
|
||||
after_val = next(m.value for m in after if m.name == name)
|
||||
return after_val - before_val
|
||||
36
third_party/vllm/tests/v1/e2e/general/test_cascade_attention.py
vendored
Normal file
36
third_party/vllm/tests/v1/e2e/general/test_cascade_attention.py
vendored
Normal file
@@ -0,0 +1,36 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
from ....utils import create_new_process_for_each_test
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
@pytest.mark.parametrize("attn_backend", ["FLASH_ATTN", "FLASHINFER"])
|
||||
def test_cascade_attention(example_system_message, attn_backend):
|
||||
prompt = "\n<User>: Implement fibonacci sequence in Python.\n<Claude>:"
|
||||
|
||||
if attn_backend == "FLASHINFER":
|
||||
pytest.skip(
|
||||
"This test is failing with FlashInfer backend and "
|
||||
"needs investigation. See issue #25679."
|
||||
)
|
||||
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen2-1.5B-Instruct", attention_config={"backend": attn_backend}
|
||||
)
|
||||
sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
|
||||
|
||||
# No cascade attention.
|
||||
single_prompt = [example_system_message + prompt]
|
||||
responses = llm.generate(single_prompt, sampling_params)
|
||||
ref_output = responses[0].outputs[0].text
|
||||
|
||||
# (Probably) Use cascade attention.
|
||||
prompts = [example_system_message + prompt] * 64
|
||||
responses = llm.generate(prompts, sampling_params)
|
||||
for response in responses:
|
||||
assert response.outputs[0].text == ref_output
|
||||
63
third_party/vllm/tests/v1/e2e/general/test_context_length.py
vendored
Normal file
63
third_party/vllm/tests/v1/e2e/general/test_context_length.py
vendored
Normal file
@@ -0,0 +1,63 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for vLLM `vllm/v1/engine/processor.Processor._validate_model_input()`
|
||||
handling of maximum context length for decoder models.
|
||||
|
||||
This test ensures:
|
||||
- A prompt that is one token shorter than the model's maximum context length
|
||||
can be processed successfully when requesting one additional token.
|
||||
- A prompt that reaches the model's maximum context length throws a
|
||||
`ValueError` when requesting at least one additional token.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.conftest import VllmRunner
|
||||
from tests.utils import create_new_process_for_each_test
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
@pytest.mark.parametrize("model, max_model_len", [("JackFram/llama-160m", 2048)])
|
||||
@pytest.mark.parametrize(
|
||||
"prompt_len, max_tokens",
|
||||
[
|
||||
(2047, 1), # prompt_len = max_model_len - 1 -> allowed
|
||||
(2048, 1), # prompt_len = max_model_len -> not allowed
|
||||
],
|
||||
)
|
||||
def test_decoder_max_context_length_validation(
|
||||
model: str,
|
||||
max_model_len: int,
|
||||
vllm_runner: type[VllmRunner],
|
||||
prompt_len: int,
|
||||
max_tokens: int,
|
||||
) -> None:
|
||||
"""Check vLLM decoder model input validation for edge cases where
|
||||
the prompt length is (almost) equal to the max model length."""
|
||||
|
||||
prompt_ids = [[43] * prompt_len]
|
||||
|
||||
with vllm_runner(
|
||||
model_name=model,
|
||||
tokenizer_name=model,
|
||||
max_model_len=max_model_len,
|
||||
max_num_seqs=1,
|
||||
tensor_parallel_size=1,
|
||||
) as vllm_model:
|
||||
if prompt_len + max_tokens <= max_model_len:
|
||||
# Should succeed as constraints are met
|
||||
vllm_model.generate_greedy(prompt_ids, max_tokens)
|
||||
else:
|
||||
# Should raise the ValueError defined in
|
||||
# vllm/v1/engine/processor.Processor_validate_model_input()
|
||||
expected_msg = (
|
||||
f"The decoder prompt (length {prompt_len}) plus the number of "
|
||||
f"requested output tokens (at least 1) is longer than "
|
||||
f"the maximum model length of {max_model_len}. "
|
||||
"Make sure that `max_model_len` is no smaller than the number of "
|
||||
"text tokens (prompt + requested output tokens)."
|
||||
)
|
||||
with pytest.raises(ValueError) as excinfo:
|
||||
vllm_model.generate_greedy(prompt_ids, max_tokens)
|
||||
assert expected_msg in str(excinfo.value)
|
||||
98
third_party/vllm/tests/v1/e2e/general/test_correctness_sliding_window.py
vendored
Normal file
98
third_party/vllm/tests/v1/e2e/general/test_correctness_sliding_window.py
vendored
Normal file
@@ -0,0 +1,98 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ....utils import check_answers, prep_prompts
|
||||
|
||||
|
||||
@dataclass
|
||||
class TestConfig:
|
||||
sliding_window: int
|
||||
ln_range: tuple[int, int]
|
||||
|
||||
|
||||
model_config = {
|
||||
"bigcode/starcoder2-3b": TestConfig(4096, (800, 1100)),
|
||||
"google/gemma-3-1b-it": TestConfig(4096, (400, 800)),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"bigcode/starcoder2-3b", # sliding window only
|
||||
"google/gemma-3-1b-it", # sliding window + full attention
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("batch_size", [5])
|
||||
@pytest.mark.parametrize("seed", [1])
|
||||
@pytest.mark.parametrize("disable_hybrid_kv_cache_manager", [True, False])
|
||||
def test_sliding_window_retrieval(
|
||||
model, batch_size, seed, disable_hybrid_kv_cache_manager
|
||||
):
|
||||
"""
|
||||
The test does a bunch of assignments "x1 = 10\nx2 = 33\n..." and then
|
||||
asks for value of one of them (which is outside the sliding window).
|
||||
If we tell it upfront which we are going to be looking for, then
|
||||
it answers correctly (mostly).
|
||||
"""
|
||||
# NOTE: For ROCm, we have to enforce eager mode to use custom kernel
|
||||
# implementation of GELU with tanh approximation, as PyTorch's native
|
||||
# implementation is currently unstable with torch.compile and produces garbage.
|
||||
enforce_eager = current_platform.is_rocm()
|
||||
|
||||
test_config = model_config[model]
|
||||
|
||||
llm = LLM(
|
||||
model=model,
|
||||
disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager,
|
||||
enforce_eager=enforce_eager,
|
||||
)
|
||||
sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
|
||||
|
||||
prompts, answer, indices = prep_prompts(batch_size, ln_range=test_config.ln_range)
|
||||
|
||||
check_length(prompts, llm, test_config.sliding_window)
|
||||
|
||||
# Fresh generation
|
||||
responses = llm.generate(prompts, sampling_params)
|
||||
check_answers(
|
||||
indices,
|
||||
answer,
|
||||
[response.outputs[0].text for response in responses],
|
||||
accept_rate=1.0,
|
||||
)
|
||||
|
||||
# Re-generate with the same prompts to test prefix caching
|
||||
responses = llm.generate(prompts, sampling_params)
|
||||
check_answers(
|
||||
indices,
|
||||
answer,
|
||||
[response.outputs[0].text for response in responses],
|
||||
accept_rate=1.0,
|
||||
)
|
||||
|
||||
|
||||
def check_length(prompts: list[str], llm: LLM, sliding_window: int):
|
||||
"""
|
||||
Check if the prompt length is valid, i.e., longer than the sliding window
|
||||
size and shorter than the model's max length.
|
||||
|
||||
Args:
|
||||
prompts: list of prompts
|
||||
llm: LLM object
|
||||
sliding_window: Sliding window size
|
||||
"""
|
||||
tokenizer = llm.get_tokenizer()
|
||||
max_model_len = llm.llm_engine.model_config.max_model_len
|
||||
assert any(len(tokenizer.encode(prompt)) > sliding_window for prompt in prompts), (
|
||||
"Prompt is too short for test"
|
||||
)
|
||||
assert all(len(tokenizer.encode(prompt)) <= max_model_len for prompt in prompts), (
|
||||
"Prompt is too long for test"
|
||||
)
|
||||
102
third_party/vllm/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py
vendored
Normal file
102
third_party/vllm/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py
vendored
Normal file
@@ -0,0 +1,102 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import random
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.config import CompilationConfig, CompilationMode
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ....utils import check_answers, fork_new_process_for_each_test, prep_prompts
|
||||
|
||||
# global seed
|
||||
SEED = 42
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_prompts():
|
||||
"""
|
||||
Adapted from tests/v1/e2e/spec_decode/test_spec_decode.py
|
||||
"""
|
||||
prompt_types = ["repeat", "sentence"]
|
||||
# Setting higher num prompts increases the chance of numerics mismatch
|
||||
# due to matrix multiplication numerics depending on batch dimension
|
||||
num_prompts = 10
|
||||
prompts = []
|
||||
|
||||
random.seed(0)
|
||||
random_prompt_type_choices = random.choices(prompt_types, k=num_prompts)
|
||||
|
||||
for kind in random_prompt_type_choices:
|
||||
word_choices = ["test", "temp", "hello", "where"]
|
||||
word = random.choice(word_choices)
|
||||
if kind == "repeat":
|
||||
prompt = f"""please repeat the word '{word}' 10 times."""
|
||||
elif kind == "sentence":
|
||||
prompt = f"""please give a ten-word sentence that
|
||||
uses the word {word} at least once."""
|
||||
else:
|
||||
raise ValueError(f"Unknown prompt type: {kind}")
|
||||
prompts.append(prompt)
|
||||
|
||||
return prompts
|
||||
|
||||
|
||||
use_fork_for_test = (
|
||||
fork_new_process_for_each_test if not current_platform.is_rocm() else lambda x: x
|
||||
)
|
||||
|
||||
|
||||
@use_fork_for_test
|
||||
@pytest.mark.parametrize("kv_sharing_fast_prefill", [False, True])
|
||||
@pytest.mark.parametrize("enforce_eager", [True, False])
|
||||
def test_kv_sharing_fast_prefill(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
kv_sharing_fast_prefill: bool,
|
||||
enforce_eager: bool,
|
||||
):
|
||||
if not enforce_eager and current_platform.is_rocm():
|
||||
# Relevant context: https://github.com/vllm-project/vllm/pull/29244
|
||||
pytest.skip(
|
||||
"ROCm: torch.compile produces incorrect output for gemma-3n's GELU "
|
||||
"with tanh approximation. Use enforce_eager=True instead."
|
||||
)
|
||||
|
||||
sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
|
||||
compilation_config = CompilationConfig(
|
||||
# This allows vLLM compilation backend to handle allocating and
|
||||
# managing buffers for cudagraph
|
||||
cudagraph_copy_inputs=True,
|
||||
mode=CompilationMode.VLLM_COMPILE
|
||||
if not enforce_eager
|
||||
else CompilationMode.NONE,
|
||||
)
|
||||
batch_size = 10
|
||||
|
||||
with monkeypatch.context() as m:
|
||||
# Make scheduling deterministic for reproducibility
|
||||
if current_platform.is_rocm():
|
||||
# Use spawn to prevent cuda re-initialization error
|
||||
m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
else:
|
||||
m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
|
||||
|
||||
prompts, answer, indices = prep_prompts(batch_size)
|
||||
|
||||
llm = LLM(
|
||||
model="google/gemma-3n-E2B-it",
|
||||
enforce_eager=enforce_eager,
|
||||
compilation_config=compilation_config,
|
||||
seed=SEED,
|
||||
kv_sharing_fast_prefill=kv_sharing_fast_prefill,
|
||||
attention_backend="TRITON_ATTN",
|
||||
)
|
||||
responses = llm.generate(prompts, sampling_params)
|
||||
check_answers(
|
||||
indices,
|
||||
answer,
|
||||
[response.outputs[0].text for response in responses],
|
||||
accept_rate=1.0,
|
||||
)
|
||||
809
third_party/vllm/tests/v1/e2e/general/test_mamba_prefix_cache.py
vendored
Normal file
809
third_party/vllm/tests/v1/e2e/general/test_mamba_prefix_cache.py
vendored
Normal file
@@ -0,0 +1,809 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import traceback
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import datasets
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.utils import create_new_process_for_each_test
|
||||
from vllm import LLM, SamplingParams, TokensPrompt
|
||||
from vllm.config import CacheConfig
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc
|
||||
from vllm.sequence import IntermediateTensors
|
||||
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
|
||||
from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.engine.core_client import InprocClient
|
||||
from vllm.v1.kv_cache_interface import KVCacheConfig
|
||||
from vllm.v1.outputs import SamplerOutput
|
||||
from vllm.v1.request import Request
|
||||
from vllm.v1.sample.metadata import SamplingMetadata
|
||||
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
|
||||
from vllm.v1.worker import mamba_utils
|
||||
from vllm.v1.worker.gpu_input_batch import CachedRequestState
|
||||
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
|
||||
from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch
|
||||
from vllm.v1.worker.mamba_utils import get_mamba_groups
|
||||
|
||||
|
||||
@dataclass
|
||||
class StepAction:
|
||||
num_computed_tokens_start: int
|
||||
num_scheduled_tokens: int
|
||||
kv_cache_block_ids: list[int] # [] to follow last step
|
||||
preprocess_copy_idx: tuple[int, int] # -1, -1 for no copy
|
||||
postprocess_copy_idx: tuple[int, int] # -1, -1 for no copy
|
||||
|
||||
|
||||
num_speculative_tokens = 3
|
||||
|
||||
num_accepted_tokens = 1
|
||||
prompt_token_ids: list[int] = []
|
||||
MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8"
|
||||
BLOCK_SIZE = 560
|
||||
NUM_HIDDEN_LAYERS = 1
|
||||
cur_step_action_idx = 0
|
||||
cur_step_action: StepAction | None = None
|
||||
step_actions: list[StepAction] = []
|
||||
|
||||
|
||||
def get_fake_sample_fn() -> SamplerOutput:
|
||||
def fake_sample_fn(
|
||||
self: GPUModelRunner,
|
||||
logits: torch.Tensor | None,
|
||||
spec_decode_metadata: SpecDecodeMetadata | None,
|
||||
) -> SamplerOutput:
|
||||
assert logits is not None
|
||||
num_computed_tokens_cpu_tensor = self.input_batch.num_computed_tokens_cpu_tensor
|
||||
num_computed_tokens = num_computed_tokens_cpu_tensor[0].item()
|
||||
if num_computed_tokens < self.input_batch.num_prompt_tokens[0].item():
|
||||
first_token_id_index = self.input_batch.num_prompt_tokens[0].item()
|
||||
else:
|
||||
first_token_id_index = num_computed_tokens + 1
|
||||
if spec_decode_metadata is None:
|
||||
return SamplerOutput(
|
||||
sampled_token_ids=torch.tensor(
|
||||
[[prompt_token_ids[first_token_id_index]]],
|
||||
device="cuda",
|
||||
dtype=torch.int32,
|
||||
),
|
||||
logprobs_tensors=None,
|
||||
)
|
||||
accepted_tokens = prompt_token_ids[
|
||||
first_token_id_index : first_token_id_index
|
||||
+ min(num_accepted_tokens, logits.shape[0])
|
||||
]
|
||||
sampled_token_ids = accepted_tokens
|
||||
return SamplerOutput(
|
||||
sampled_token_ids=torch.tensor(
|
||||
[sampled_token_ids], device="cuda", dtype=torch.int32
|
||||
),
|
||||
logprobs_tensors=None,
|
||||
)
|
||||
|
||||
return fake_sample_fn
|
||||
|
||||
|
||||
def get_fake_propose_draft_token_ids_fn():
|
||||
def fake_propose_draft_token_ids_fn(
|
||||
self: GPUModelRunner,
|
||||
scheduler_output: SchedulerOutput,
|
||||
sampled_token_ids: torch.Tensor | list[list[int]],
|
||||
sampling_metadata: SamplingMetadata,
|
||||
hidden_states: torch.Tensor,
|
||||
sample_hidden_states: torch.Tensor,
|
||||
aux_hidden_states: list[torch.Tensor] | None,
|
||||
spec_decode_metadata: SpecDecodeMetadata | None,
|
||||
common_attn_metadata: CommonAttentionMetadata,
|
||||
slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None,
|
||||
) -> list[list[int]]:
|
||||
num_computed_tokens_cpu_tensor = self.input_batch.num_computed_tokens_cpu_tensor
|
||||
num_computed_tokens = num_computed_tokens_cpu_tensor[0].item()
|
||||
if (
|
||||
self.input_batch.num_tokens_no_spec[0].item()
|
||||
<= self.input_batch.num_prompt_tokens[0].item()
|
||||
):
|
||||
first_token_id_index = self.input_batch.num_prompt_tokens[0].item()
|
||||
else:
|
||||
first_token_id_index = (
|
||||
num_computed_tokens + 1
|
||||
) # bonus token isn't considered as computed
|
||||
first_token_id_index += self.input_batch.num_accepted_tokens_cpu[0].item()
|
||||
proposed_draft_token_ids = [
|
||||
prompt_token_ids[
|
||||
first_token_id_index : first_token_id_index + num_speculative_tokens
|
||||
]
|
||||
]
|
||||
|
||||
next_token_ids = torch.tensor(
|
||||
prompt_token_ids[
|
||||
first_token_id_index - 1 : first_token_id_index
|
||||
- 1
|
||||
+ num_accepted_tokens
|
||||
],
|
||||
device="cuda",
|
||||
dtype=torch.int32,
|
||||
)
|
||||
|
||||
valid_sampled_tokens_count = torch.tensor(
|
||||
[num_accepted_tokens], device="cuda", dtype=torch.int32
|
||||
)
|
||||
|
||||
self._copy_valid_sampled_token_count(next_token_ids, valid_sampled_tokens_count)
|
||||
|
||||
return torch.tensor(proposed_draft_token_ids, device="cuda", dtype=torch.int32)
|
||||
|
||||
return fake_propose_draft_token_ids_fn
|
||||
|
||||
|
||||
def get_fake_step_action_fn(original_step_action_fn: Callable):
|
||||
def fake_get_output(self: InprocClient):
|
||||
global cur_step_action_idx
|
||||
global cur_step_action
|
||||
if cur_step_action_idx < len(step_actions):
|
||||
cur_step_action = step_actions[cur_step_action_idx]
|
||||
cur_step_action_idx += 1
|
||||
else:
|
||||
cur_step_action = None
|
||||
print(f"cur_step_action: {cur_step_action_idx=} {cur_step_action=}")
|
||||
return original_step_action_fn(self)
|
||||
|
||||
return fake_get_output
|
||||
|
||||
|
||||
def get_fake_allocate_slots_fn(original_allocate_slots_fn: Callable):
|
||||
def fake_allocate_slots_fn(
|
||||
self: KVCacheManager,
|
||||
request: Request,
|
||||
num_new_tokens: int,
|
||||
num_new_computed_tokens: int = 0,
|
||||
new_computed_blocks: KVCacheBlocks | None = None,
|
||||
num_lookahead_tokens: int = 0,
|
||||
num_external_computed_tokens: int = 0,
|
||||
delay_cache_blocks: bool = False,
|
||||
num_encoder_tokens: int = 0,
|
||||
):
|
||||
ret = original_allocate_slots_fn(
|
||||
self,
|
||||
request,
|
||||
num_new_tokens,
|
||||
num_new_computed_tokens,
|
||||
new_computed_blocks,
|
||||
num_lookahead_tokens,
|
||||
num_external_computed_tokens,
|
||||
delay_cache_blocks,
|
||||
num_encoder_tokens,
|
||||
)
|
||||
if cur_step_action is not None:
|
||||
cur_block_ids = self.coordinator.single_type_managers[0].req_to_blocks[
|
||||
request.request_id
|
||||
]
|
||||
not_null_block_flags = [not block.is_null for block in cur_block_ids]
|
||||
block_ids = [1 if block else 0 for block in not_null_block_flags]
|
||||
assert block_ids == cur_step_action.kv_cache_block_ids
|
||||
return ret
|
||||
|
||||
return fake_allocate_slots_fn
|
||||
|
||||
|
||||
mamba_kv_cache_dict = {}
|
||||
|
||||
|
||||
def get_fake_execute_model_fn(original_execute_model_fn: Callable):
|
||||
last_num_computed_tokens = 0
|
||||
num_prompt_tokens = None
|
||||
|
||||
def fake_execute_model_fn(
|
||||
self: GPUModelRunner,
|
||||
scheduler_output: SchedulerOutput,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
):
|
||||
if cur_step_action is not None:
|
||||
num_scheduled_tokens = next(
|
||||
iter(scheduler_output.num_scheduled_tokens.values())
|
||||
)
|
||||
assert num_scheduled_tokens == cur_step_action.num_scheduled_tokens
|
||||
mamba_group_ids, mamba_spec = get_mamba_groups(self.kv_cache_config)
|
||||
mamba_group_id = mamba_group_ids[0]
|
||||
mamba_layer_name = self.kv_cache_config.kv_cache_groups[
|
||||
mamba_group_id
|
||||
].layer_names[0]
|
||||
nonlocal last_num_computed_tokens
|
||||
nonlocal num_prompt_tokens
|
||||
|
||||
if (
|
||||
len(scheduler_output.scheduled_new_reqs) > 0
|
||||
and scheduler_output.scheduled_new_reqs[0].prompt_token_ids is not None
|
||||
):
|
||||
# record number of prompt tokens
|
||||
num_prompt_tokens = len(
|
||||
scheduler_output.scheduled_new_reqs[0].prompt_token_ids
|
||||
)
|
||||
|
||||
if len(scheduler_output.scheduled_cached_reqs.req_ids) > 0:
|
||||
num_computed_tokens = (
|
||||
scheduler_output.scheduled_cached_reqs.num_computed_tokens[0]
|
||||
)
|
||||
if (
|
||||
self.num_spec_tokens
|
||||
and num_prompt_tokens is not None
|
||||
and num_computed_tokens > num_prompt_tokens
|
||||
):
|
||||
# NOTE (tdoublep) with async scheduling, the scheduler does not have an
|
||||
# accurate measure of the number of computed tokens; we need to subtract
|
||||
# the number of reject tokens from the previous timestep.
|
||||
num_computed_tokens -= num_speculative_tokens + 1 - num_accepted_tokens
|
||||
if (
|
||||
num_computed_tokens // BLOCK_SIZE
|
||||
> last_num_computed_tokens // BLOCK_SIZE
|
||||
):
|
||||
# generated a new aligned block in this step
|
||||
block_idx = num_computed_tokens // mamba_spec.block_size - 1
|
||||
block_id = (
|
||||
self.input_batch.block_table.block_tables[mamba_group_id]
|
||||
.block_table.cpu[0, block_idx]
|
||||
.item()
|
||||
)
|
||||
if block_id != 0:
|
||||
kv_cache = self.compilation_config.static_forward_context[
|
||||
mamba_layer_name
|
||||
].kv_cache
|
||||
mamba_kv_cache_dict[
|
||||
num_computed_tokens - num_computed_tokens % BLOCK_SIZE
|
||||
] = (
|
||||
kv_cache[0][0][block_id].clone(),
|
||||
kv_cache[0][1][block_id].clone(),
|
||||
)
|
||||
|
||||
last_num_computed_tokens = num_computed_tokens
|
||||
else:
|
||||
last_num_computed_tokens = 0
|
||||
|
||||
ret = original_execute_model_fn(self, scheduler_output, intermediate_tensors)
|
||||
|
||||
if cur_step_action is not None:
|
||||
assert (
|
||||
cur_step_action.num_computed_tokens_start
|
||||
== self.input_batch.num_computed_tokens_cpu[0].item()
|
||||
)
|
||||
|
||||
return ret
|
||||
|
||||
return fake_execute_model_fn
|
||||
|
||||
|
||||
def get_fake_process_mamba_fn(
|
||||
original_preprocess_mamba_fn: Callable,
|
||||
original_post_process_mamba_fn: Callable,
|
||||
original_copy_fn: Callable,
|
||||
):
|
||||
copy_info: tuple[list[int], list[int], list[int]] | None = None
|
||||
|
||||
def check_copy_info(
|
||||
action: tuple[int, int],
|
||||
kv_cache_config: KVCacheConfig,
|
||||
forward_context: dict[str, Any],
|
||||
input_batch: GPUInputBatch,
|
||||
):
|
||||
assert copy_info is not None
|
||||
if action == (-1, -1):
|
||||
assert len(copy_info[0]) == len(copy_info[1]) == len(copy_info[2]) == 0
|
||||
else:
|
||||
assert len(copy_info[0]) == len(copy_info[1]) == len(copy_info[2]) == 2
|
||||
mamba_group_ids, mamba_spec = get_mamba_groups(kv_cache_config)
|
||||
mamba_group_id = mamba_group_ids[0]
|
||||
mamba_layer_name = kv_cache_config.kv_cache_groups[
|
||||
mamba_group_id
|
||||
].layer_names[0]
|
||||
mamba_kv_cache = forward_context[mamba_layer_name].kv_cache[0][-1]
|
||||
mamba_block_table = input_batch.block_table.block_tables[
|
||||
mamba_group_id
|
||||
].block_table.cpu[0]
|
||||
expected_temporal_src = mamba_kv_cache[
|
||||
mamba_block_table[action[0]]
|
||||
].data_ptr()
|
||||
expected_temporal_dest = mamba_kv_cache[
|
||||
mamba_block_table[action[1]]
|
||||
].data_ptr()
|
||||
# -1 is qwen3-next's temporal. We skip checking conv as it is more complex.
|
||||
assert copy_info[0][-1] == expected_temporal_src
|
||||
assert copy_info[1][-1] == expected_temporal_dest
|
||||
|
||||
def fake_preprocess_mamba_fn(
|
||||
scheduler_output: SchedulerOutput,
|
||||
kv_cache_config: KVCacheConfig,
|
||||
cache_config: CacheConfig,
|
||||
mamba_state_idx: dict[str, int],
|
||||
input_batch: GPUInputBatch,
|
||||
requests: dict[str, CachedRequestState],
|
||||
forward_context: dict[str, Any],
|
||||
mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
|
||||
copy_bufs: mamba_utils.MambaCopyBuffers,
|
||||
):
|
||||
nonlocal copy_info
|
||||
copy_info = None
|
||||
ret = original_preprocess_mamba_fn(
|
||||
scheduler_output,
|
||||
kv_cache_config,
|
||||
cache_config,
|
||||
mamba_state_idx,
|
||||
input_batch,
|
||||
requests,
|
||||
forward_context,
|
||||
mamba_state_copy_funcs,
|
||||
copy_bufs,
|
||||
)
|
||||
if cur_step_action is not None:
|
||||
check_copy_info(
|
||||
cur_step_action.preprocess_copy_idx,
|
||||
kv_cache_config,
|
||||
forward_context,
|
||||
input_batch,
|
||||
)
|
||||
return ret
|
||||
|
||||
def fake_post_process_mamba_fn(
|
||||
scheduler_output: SchedulerOutput,
|
||||
kv_cache_config: KVCacheConfig,
|
||||
input_batch: GPUInputBatch,
|
||||
requests: dict[str, CachedRequestState],
|
||||
mamba_state_idx: dict[str, int],
|
||||
forward_context: dict[str, Any],
|
||||
mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
|
||||
copy_bufs: mamba_utils.MambaCopyBuffers,
|
||||
):
|
||||
nonlocal copy_info
|
||||
copy_info = None
|
||||
ret = original_post_process_mamba_fn(
|
||||
scheduler_output,
|
||||
kv_cache_config,
|
||||
input_batch,
|
||||
requests,
|
||||
mamba_state_idx,
|
||||
forward_context,
|
||||
mamba_state_copy_funcs,
|
||||
copy_bufs,
|
||||
)
|
||||
if cur_step_action is not None:
|
||||
check_copy_info(
|
||||
cur_step_action.postprocess_copy_idx,
|
||||
kv_cache_config,
|
||||
forward_context,
|
||||
input_batch,
|
||||
)
|
||||
return ret
|
||||
|
||||
def fake_copy_fn(copy_bufs: mamba_utils.MambaCopyBuffers):
|
||||
nonlocal copy_info
|
||||
assert copy_info is None
|
||||
n = copy_bufs.offset
|
||||
src_state_list = copy_bufs.src_ptrs.cpu[:n].tolist()
|
||||
dest_state_list = copy_bufs.dst_ptrs.cpu[:n].tolist()
|
||||
num_elements_list = copy_bufs.sizes.cpu[:n].tolist()
|
||||
copy_info = (src_state_list, dest_state_list, num_elements_list)
|
||||
return original_copy_fn(copy_bufs)
|
||||
|
||||
return fake_preprocess_mamba_fn, fake_post_process_mamba_fn, fake_copy_fn
|
||||
|
||||
|
||||
def run_ref_mamba_state_in_subprocess() -> None:
|
||||
ctx = mp.get_context("spawn")
|
||||
proc = ctx.Process(target=_run_ref_mamba_state_worker)
|
||||
proc.start()
|
||||
proc.join(timeout=600)
|
||||
if proc.exitcode != 0:
|
||||
raise RuntimeError(f"Ref mamba state process exited with code {proc.exitcode}.")
|
||||
|
||||
|
||||
def _run_ref_mamba_state_worker():
|
||||
try:
|
||||
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
|
||||
num_generated_tokens = 8000
|
||||
num_prompt_tokens = 500
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.0, max_tokens=num_generated_tokens
|
||||
)
|
||||
prompt_dataset = datasets.load_dataset("heheda/a_long_article")
|
||||
full_prompt = prompt_dataset["train"][0]["text"]
|
||||
fake_execute_model_fn = get_fake_execute_model_fn(GPUModelRunner.execute_model)
|
||||
GPUModelRunner.execute_model = fake_execute_model_fn
|
||||
fake_sample_fn = get_fake_sample_fn()
|
||||
GPUModelRunner._sample = fake_sample_fn
|
||||
engine = LLM(
|
||||
model=MODEL,
|
||||
block_size=BLOCK_SIZE,
|
||||
hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS},
|
||||
seed=42,
|
||||
)
|
||||
global prompt_token_ids
|
||||
prompt_token_ids = engine.get_tokenizer().encode(full_prompt)
|
||||
print(f"Token IDs length: {len(prompt_token_ids)}")
|
||||
|
||||
_outputs = engine.generate(
|
||||
[TokensPrompt(prompt_token_ids=prompt_token_ids[:num_prompt_tokens])],
|
||||
sampling_params,
|
||||
)
|
||||
# ref_mamba_kv_cache_dict = torch.load("mamba_kv_cache_dict.pth")
|
||||
# check_mamba_state_equal(ref_mamba_kv_cache_dict, mamba_kv_cache_dict)
|
||||
# torch.save(mamba_kv_cache_dict, "mamba_kv_cache_dict.pth")
|
||||
cpu_state_ref = {
|
||||
key: tuple(tensor.detach().cpu() for tensor in tensors)
|
||||
for key, tensors in mamba_kv_cache_dict.items()
|
||||
}
|
||||
torch.save(cpu_state_ref, "mamba_kv_cache_dict_ref.pth")
|
||||
mamba_kv_cache_dict.clear()
|
||||
del engine
|
||||
torch.accelerator.empty_cache()
|
||||
cleanup_dist_env_and_memory()
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
raise
|
||||
|
||||
|
||||
def check_mamba_state_equal(
|
||||
mamba_state_ref: dict, mamba_state_new: dict, keys_to_check: list[int]
|
||||
):
|
||||
atol = 1e-2
|
||||
rtol = 1e-2
|
||||
for key in keys_to_check:
|
||||
assert key in mamba_state_new
|
||||
assert key in mamba_state_ref
|
||||
# mamba state new is a subset of mamba state ref
|
||||
for i, (ref, new) in enumerate(zip(mamba_state_ref[key], mamba_state_new[key])):
|
||||
if ref.device != new.device:
|
||||
new = new.to(ref.device)
|
||||
new = new[: ref.shape[0]]
|
||||
if not torch.allclose(ref, new, atol=atol, rtol=rtol):
|
||||
diff_mask = ~torch.isclose(ref, new, atol=atol, rtol=rtol)
|
||||
diff_idx = torch.nonzero(diff_mask)
|
||||
if diff_idx.shape[0] * 100 < ref.numel():
|
||||
print(
|
||||
f"[WARNING] found {diff_idx.shape[0] * 100 / ref.numel()}% of the elements are different" # noqa: E501
|
||||
)
|
||||
continue
|
||||
raise ValueError(
|
||||
f"Mamba state is not equal for key: {key} at index {i}"
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
@dataclass
|
||||
class TestConfig:
|
||||
num_prompt_tokens: int
|
||||
num_generated_tokens: int
|
||||
num_accepted_tokens: int
|
||||
step_actions: list[StepAction]
|
||||
|
||||
|
||||
def apply_patch(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
|
||||
|
||||
fake_sample_fn = get_fake_sample_fn()
|
||||
monkeypatch.setattr(GPUModelRunner, "_sample", fake_sample_fn)
|
||||
|
||||
fake_propose_draft_token_ids_fn = get_fake_propose_draft_token_ids_fn()
|
||||
monkeypatch.setattr(
|
||||
GPUModelRunner, "propose_draft_token_ids", fake_propose_draft_token_ids_fn
|
||||
)
|
||||
|
||||
fake_execute_model_fn = get_fake_execute_model_fn(GPUModelRunner.execute_model)
|
||||
monkeypatch.setattr(GPUModelRunner, "execute_model", fake_execute_model_fn)
|
||||
|
||||
fake_step_action_fn = get_fake_step_action_fn(InprocClient.get_output)
|
||||
monkeypatch.setattr(InprocClient, "get_output", fake_step_action_fn)
|
||||
|
||||
fake_allocate_slots_fn = get_fake_allocate_slots_fn(KVCacheManager.allocate_slots)
|
||||
monkeypatch.setattr(KVCacheManager, "allocate_slots", fake_allocate_slots_fn)
|
||||
|
||||
fake_preprocess_mamba_fn, fake_post_process_mamba_fn, fake_copy_fn = (
|
||||
get_fake_process_mamba_fn(
|
||||
mamba_utils.preprocess_mamba,
|
||||
mamba_utils.postprocess_mamba,
|
||||
mamba_utils.do_mamba_copy_block,
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(mamba_utils, "preprocess_mamba", fake_preprocess_mamba_fn)
|
||||
monkeypatch.setattr(mamba_utils, "postprocess_mamba", fake_post_process_mamba_fn)
|
||||
monkeypatch.setattr(mamba_utils, "do_mamba_copy_block", fake_copy_fn)
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
|
||||
run_ref_mamba_state_in_subprocess()
|
||||
apply_patch(monkeypatch)
|
||||
prompt_dataset = datasets.load_dataset("heheda/a_long_article")
|
||||
full_prompt = prompt_dataset["train"][0]["text"]
|
||||
tests = {
|
||||
"accept_1": TestConfig(
|
||||
num_prompt_tokens=554,
|
||||
num_generated_tokens=20,
|
||||
num_accepted_tokens=1,
|
||||
step_actions=[
|
||||
StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(554, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(555, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(556, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(557, 4, [], (0, 1), (-1, -1)),
|
||||
StepAction(558, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(559, 4, [], (-1, -1), (1, 0)),
|
||||
StepAction(560, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(561, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
# test case 2.1: no hit, accept 2 tokens
|
||||
"accept_2_1": TestConfig(
|
||||
num_prompt_tokens=554,
|
||||
num_generated_tokens=20,
|
||||
num_accepted_tokens=2,
|
||||
step_actions=[
|
||||
StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(554, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(556, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(558, 4, [], (1, 1), (2, 0)),
|
||||
StepAction(560, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(562, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
# test case 2.2: no hit, accept 2 tokens
|
||||
"accept_2_2": TestConfig(
|
||||
num_prompt_tokens=555,
|
||||
num_generated_tokens=20,
|
||||
num_accepted_tokens=2,
|
||||
step_actions=[
|
||||
StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(555, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(557, 4, [1, 1, 1, 1, 1], (1, 1), (-1, -1)),
|
||||
StepAction(559, 4, [], (-1, -1), (1, 0)),
|
||||
StepAction(561, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"accept_3_1": TestConfig(
|
||||
num_prompt_tokens=553,
|
||||
num_generated_tokens=20,
|
||||
num_accepted_tokens=3,
|
||||
step_actions=[
|
||||
StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(553, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(556, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(559, 4, [], (2, 1), (1, 0)),
|
||||
StepAction(562, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"accept_3_2": TestConfig(
|
||||
num_prompt_tokens=554,
|
||||
num_generated_tokens=20,
|
||||
num_accepted_tokens=3,
|
||||
step_actions=[
|
||||
StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(554, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(557, 4, [1, 1, 1, 1, 1], (2, 1), (3, 0)),
|
||||
StepAction(560, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"accept_3_3": TestConfig(
|
||||
num_prompt_tokens=555,
|
||||
num_generated_tokens=20,
|
||||
num_accepted_tokens=3,
|
||||
step_actions=[
|
||||
StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(555, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(558, 4, [1, 1, 1, 1, 1], (2, 1), (2, 0)),
|
||||
StepAction(561, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(564, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"accept_4_1": TestConfig(
|
||||
num_prompt_tokens=553,
|
||||
num_generated_tokens=20,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(553, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(557, 4, [1, 1, 1, 1, 1], (3, 1), (3, 0)),
|
||||
StepAction(561, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"accept_4_2": TestConfig(
|
||||
num_prompt_tokens=554,
|
||||
num_generated_tokens=25,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(554, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(558, 4, [1, 1, 1, 1, 1], (3, 1), (2, 0)),
|
||||
StepAction(562, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(566, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"accept_4_3": TestConfig(
|
||||
num_prompt_tokens=555,
|
||||
num_generated_tokens=25,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(555, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(559, 4, [1, 1, 1, 1, 1], (3, 1), (1, 0)),
|
||||
StepAction(563, 4, [], (-1, -1), (-1, -1)),
|
||||
StepAction(567, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"accept_4_4": TestConfig(
|
||||
num_prompt_tokens=556,
|
||||
num_generated_tokens=25,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 556, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(556, 4, [], (-1, -1), (3, 0)),
|
||||
StepAction(560, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
|
||||
StepAction(564, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"prompt_block_size": TestConfig(
|
||||
num_prompt_tokens=560,
|
||||
num_generated_tokens=10,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(560, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"prompt_2_block_size": TestConfig(
|
||||
num_prompt_tokens=560 * 2,
|
||||
num_generated_tokens=10,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(560, 560, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
|
||||
StepAction(560 * 2, 4, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"prompt_2_block_size_10": TestConfig(
|
||||
num_prompt_tokens=560 * 2 + 10,
|
||||
num_generated_tokens=10,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(560, 570, [1, 0, 1, 1, 1, 1], (0, 2), (-1, -1)),
|
||||
StepAction(560 * 2 + 10, 4, [0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"prompt_3_block_size": TestConfig(
|
||||
num_prompt_tokens=560 * 3,
|
||||
num_generated_tokens=10,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(560 * 2, 560, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)),
|
||||
StepAction(560 * 3, 4, [0, 0, 1, 1, 1, 1, 1], (2, 3), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"prompt_3_block_size_10": TestConfig(
|
||||
num_prompt_tokens=560 * 3 + 10,
|
||||
num_generated_tokens=10,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(560 * 2, 570, [0, 1, 0, 1, 1, 1, 1], (1, 3), (-1, -1)),
|
||||
StepAction(560 * 3 + 10, 4, [0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
],
|
||||
),
|
||||
"prompt_10_block_size": TestConfig(
|
||||
num_prompt_tokens=560 * 10,
|
||||
num_generated_tokens=10,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 560 * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(
|
||||
560 * 5,
|
||||
560 * 4,
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1],
|
||||
(4, 8),
|
||||
(-1, -1),
|
||||
),
|
||||
StepAction(
|
||||
560 * 9,
|
||||
560,
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
|
||||
(8, 9),
|
||||
(-1, -1),
|
||||
),
|
||||
StepAction(
|
||||
560 * 10,
|
||||
4,
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
|
||||
(9, 10),
|
||||
(-1, -1),
|
||||
),
|
||||
],
|
||||
),
|
||||
"prompt_10_block_size_10": TestConfig(
|
||||
num_prompt_tokens=560 * 10 + 10,
|
||||
num_generated_tokens=10,
|
||||
num_accepted_tokens=4,
|
||||
step_actions=[
|
||||
StepAction(0, 560 * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
|
||||
StepAction(
|
||||
560 * 5,
|
||||
560 * 4,
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1],
|
||||
(4, 8),
|
||||
(-1, -1),
|
||||
),
|
||||
StepAction(
|
||||
560 * 9,
|
||||
560 + 10,
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1],
|
||||
(8, 10),
|
||||
(-1, -1),
|
||||
),
|
||||
],
|
||||
),
|
||||
}
|
||||
|
||||
engine = LLM(
|
||||
model=MODEL,
|
||||
enable_prefix_caching=True,
|
||||
block_size=BLOCK_SIZE,
|
||||
mamba_cache_mode="align",
|
||||
speculative_config={
|
||||
"method": "qwen3_next_mtp",
|
||||
"num_speculative_tokens": num_speculative_tokens,
|
||||
},
|
||||
max_num_batched_tokens=3072,
|
||||
hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS},
|
||||
seed=42,
|
||||
)
|
||||
global prompt_token_ids
|
||||
prompt_token_ids = engine.get_tokenizer().encode(full_prompt)
|
||||
print(f"Token IDs length: {len(prompt_token_ids)}")
|
||||
for test_case_name, test_config in tests.items():
|
||||
print(f"Running test case: {test_case_name}")
|
||||
num_generated_tokens = test_config.num_generated_tokens
|
||||
num_prompt_tokens = test_config.num_prompt_tokens
|
||||
global num_accepted_tokens
|
||||
num_accepted_tokens = test_config.num_accepted_tokens
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.0, max_tokens=num_generated_tokens
|
||||
)
|
||||
global cur_step_action_idx
|
||||
cur_step_action_idx = 0
|
||||
for step_action_prev, step_action_next in zip(
|
||||
test_config.step_actions[:-1], test_config.step_actions[1:]
|
||||
):
|
||||
if (
|
||||
step_action_next.kv_cache_block_ids is not None
|
||||
and len(step_action_next.kv_cache_block_ids) == 0
|
||||
):
|
||||
prev_block_ids = step_action_prev.kv_cache_block_ids
|
||||
if prev_block_ids is not None:
|
||||
step_action_next.kv_cache_block_ids = prev_block_ids.copy()
|
||||
global step_actions
|
||||
step_actions = test_config.step_actions
|
||||
_ = engine.generate(
|
||||
[TokensPrompt(prompt_token_ids=prompt_token_ids[:num_prompt_tokens])],
|
||||
sampling_params,
|
||||
)
|
||||
assert engine.llm_engine.engine_core.engine_core.scheduler.reset_prefix_cache()
|
||||
print(f"End test case: {test_case_name}")
|
||||
keys_to_check = [
|
||||
(action.postprocess_copy_idx[1] + 1) * BLOCK_SIZE
|
||||
for action in test_config.step_actions
|
||||
if action.postprocess_copy_idx and action.postprocess_copy_idx[0] != -1
|
||||
]
|
||||
mamba_state_ref = torch.load("mamba_kv_cache_dict_ref.pth")
|
||||
check_mamba_state_equal(mamba_state_ref, mamba_kv_cache_dict, keys_to_check)
|
||||
mamba_kv_cache_dict.clear()
|
||||
del engine
|
||||
torch.accelerator.empty_cache()
|
||||
cleanup_dist_env_and_memory()
|
||||
502
third_party/vllm/tests/v1/e2e/general/test_min_tokens.py
vendored
Normal file
502
third_party/vllm/tests/v1/e2e/general/test_min_tokens.py
vendored
Normal file
@@ -0,0 +1,502 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Comprehensive end-to-end tests for `min_tokens` in the V1 engine.
|
||||
|
||||
Addresses #21950: verify and add CI coverage.
|
||||
|
||||
Covers:
|
||||
1) Basic functionality
|
||||
2) Stop strings with `min_tokens` (bug #21987; fix in PR #22014)
|
||||
3) EOS behavior with `min_tokens` (potential logits-processor bug)
|
||||
4) Edge cases (min_tokens == max_tokens, min_tokens == 0)
|
||||
5) Multiple stop conditions
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.outputs import RequestOutput
|
||||
|
||||
# Test configuration
|
||||
TEST_MODEL = "facebook/opt-125m" # Small model for fast CI execution
|
||||
GREEDY = 0.0 # Deterministic generation for consistent testing
|
||||
|
||||
|
||||
class MinTokensTestCase:
|
||||
"""Data class for min_tokens test scenarios"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
min_tokens: int,
|
||||
max_tokens: int,
|
||||
stop: str | list[str] | None = None,
|
||||
expected_min_len: int | None = None,
|
||||
expected_exact_len: int | None = None,
|
||||
):
|
||||
self.name = name
|
||||
self.min_tokens = min_tokens
|
||||
self.max_tokens = max_tokens
|
||||
self.stop = stop
|
||||
self.expected_min_len = expected_min_len or min_tokens
|
||||
self.expected_exact_len = expected_exact_len
|
||||
|
||||
def __str__(self):
|
||||
return (
|
||||
f"{self.name}: min={self.min_tokens}, "
|
||||
f"max={self.max_tokens}, stop={self.stop}"
|
||||
)
|
||||
|
||||
|
||||
# Test scenarios covering all critical cases
|
||||
MIN_TOKENS_TEST_CASES = [
|
||||
# === BASIC FUNCTIONALITY (should work) ===
|
||||
MinTokensTestCase(
|
||||
name="basic_min_tokens_no_stop",
|
||||
min_tokens=8,
|
||||
max_tokens=20,
|
||||
stop=None,
|
||||
expected_min_len=8,
|
||||
),
|
||||
MinTokensTestCase(
|
||||
name="min_tokens_zero",
|
||||
min_tokens=0,
|
||||
max_tokens=10,
|
||||
stop=None,
|
||||
expected_min_len=0,
|
||||
),
|
||||
MinTokensTestCase(
|
||||
name="min_equals_max_no_stop",
|
||||
min_tokens=15,
|
||||
max_tokens=15,
|
||||
stop=None,
|
||||
expected_exact_len=15,
|
||||
),
|
||||
# === STOP STRINGS WITH MIN_TOKENS ===
|
||||
# These tests expose the detokenizer bug where stop strings
|
||||
# bypass min_tokens
|
||||
# Using mathematically guaranteed approach with wide stop nets
|
||||
pytest.param(
|
||||
MinTokensTestCase(
|
||||
name="min_tokens_with_comprehensive_stops",
|
||||
min_tokens=5,
|
||||
max_tokens=20,
|
||||
stop=[
|
||||
"a",
|
||||
"e",
|
||||
"i",
|
||||
"o",
|
||||
"u",
|
||||
"t",
|
||||
"n",
|
||||
"s",
|
||||
"r",
|
||||
"l",
|
||||
" ",
|
||||
],
|
||||
expected_min_len=5,
|
||||
),
|
||||
marks=pytest.mark.xfail(
|
||||
reason=(
|
||||
"Known bug #21987: stop strings bypass min_tokens (fixed by PR #22014)"
|
||||
),
|
||||
strict=False,
|
||||
),
|
||||
id="min_tokens_with_comprehensive_stops",
|
||||
),
|
||||
pytest.param(
|
||||
MinTokensTestCase(
|
||||
name="min_tokens_with_simple_char_stop",
|
||||
min_tokens=3,
|
||||
max_tokens=15,
|
||||
stop=["e", "a", " "],
|
||||
expected_min_len=3,
|
||||
),
|
||||
marks=pytest.mark.xfail(
|
||||
reason=(
|
||||
"Known bug #21987: stop strings bypass min_tokens (fixed by PR #22014)"
|
||||
),
|
||||
strict=False,
|
||||
),
|
||||
id="min_tokens_with_simple_char_stop",
|
||||
),
|
||||
# === EOS TOKEN WITH MIN_TOKENS (potential LogitsProcessor bug) ===
|
||||
# These test the MinTokensLogitsProcessor handling of EOS tokens
|
||||
pytest.param(
|
||||
MinTokensTestCase(
|
||||
name="min_equals_max_eos_only",
|
||||
min_tokens=20,
|
||||
max_tokens=20,
|
||||
stop=None, # Relies on default EOS token behavior
|
||||
expected_exact_len=20,
|
||||
),
|
||||
marks=pytest.mark.xfail(
|
||||
reason=("Potential logits-processor bug: EOS tokens may bypass min_tokens"),
|
||||
strict=False,
|
||||
),
|
||||
id="min_equals_max_eos_only",
|
||||
),
|
||||
# === EDGE CASES ===
|
||||
MinTokensTestCase(
|
||||
name="large_min_tokens",
|
||||
min_tokens=50,
|
||||
max_tokens=60,
|
||||
stop=None,
|
||||
expected_min_len=50,
|
||||
),
|
||||
MinTokensTestCase(
|
||||
name="min_tokens_with_empty_stop_list",
|
||||
min_tokens=5,
|
||||
max_tokens=15,
|
||||
stop=[], # Empty stop list
|
||||
expected_min_len=5,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def llm_v1():
|
||||
"""Create V1 LLM instance for testing"""
|
||||
llm = LLM(
|
||||
model=TEST_MODEL,
|
||||
tensor_parallel_size=1,
|
||||
max_model_len=1024, # Small context for fast testing
|
||||
enforce_eager=True, # Avoid graph compilation overhead
|
||||
)
|
||||
return llm
|
||||
|
||||
|
||||
def get_token_count(output: RequestOutput) -> int:
|
||||
"""Extract token count from LLM output"""
|
||||
if not output.outputs:
|
||||
return 0
|
||||
return len(output.outputs[0].token_ids)
|
||||
|
||||
|
||||
def assert_min_tokens_satisfied(
|
||||
output: RequestOutput, test_case: MinTokensTestCase
|
||||
) -> None:
|
||||
"""Assert that min_tokens requirement is satisfied"""
|
||||
token_count = get_token_count(output)
|
||||
stop_reason = output.outputs[0].stop_reason if output.outputs else "no output"
|
||||
|
||||
if test_case.expected_exact_len is not None:
|
||||
# Exact length requirement
|
||||
assert token_count == test_case.expected_exact_len, (
|
||||
f"Expected exactly {test_case.expected_exact_len} tokens, "
|
||||
f"got {token_count} tokens. "
|
||||
f"Stop reason: {stop_reason}"
|
||||
)
|
||||
else:
|
||||
# Minimum length requirement
|
||||
assert token_count >= (test_case.expected_min_len or 0), (
|
||||
f"Expected at least {test_case.expected_min_len} tokens, "
|
||||
f"got {token_count} tokens. "
|
||||
f"Stop reason: {stop_reason}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"test_case",
|
||||
MIN_TOKENS_TEST_CASES,
|
||||
ids=lambda tc: tc.name,
|
||||
)
|
||||
def test_min_tokens_comprehensive(llm_v1: LLM, test_case: MinTokensTestCase):
|
||||
"""
|
||||
Comprehensive test for min_tokens functionality in V1 engine.
|
||||
|
||||
This test covers all critical scenarios for min_tokens:
|
||||
- Basic functionality (should work)
|
||||
- Stop strings with min_tokens (known bug)
|
||||
- EOS tokens with min_tokens (potential bug)
|
||||
- Edge cases
|
||||
|
||||
Args:
|
||||
llm_v1: V1 LLM instance
|
||||
test_case: Test scenario parameters
|
||||
"""
|
||||
# Known failing cases are handled via param-level xfail marks above.
|
||||
|
||||
# Create sampling parameters
|
||||
sampling_params = SamplingParams(
|
||||
min_tokens=test_case.min_tokens,
|
||||
max_tokens=test_case.max_tokens,
|
||||
stop=test_case.stop,
|
||||
temperature=GREEDY,
|
||||
include_stop_str_in_output=True, # Include stop strings for debugging
|
||||
)
|
||||
|
||||
# Use simple prompt. Comprehensive stop lists should catch any generation
|
||||
prompt = "Hello"
|
||||
|
||||
# Generate output
|
||||
outputs = llm_v1.generate([prompt], sampling_params)
|
||||
|
||||
assert len(outputs) == 1, "Expected exactly one output"
|
||||
output = outputs[0]
|
||||
|
||||
# Debug information
|
||||
token_count = get_token_count(output)
|
||||
generated_text = output.outputs[0].text if output.outputs else ""
|
||||
stop_reason = output.outputs[0].stop_reason if output.outputs else "unknown"
|
||||
|
||||
print(f"\nTest: {test_case.name}")
|
||||
print(f"Generated {token_count} tokens")
|
||||
print(f"Stop reason: {stop_reason}")
|
||||
print(f"Generated text: {repr(generated_text)}")
|
||||
print(f"Expected min: {test_case.expected_min_len}")
|
||||
if test_case.expected_exact_len:
|
||||
print(f"Expected exact: {test_case.expected_exact_len}")
|
||||
|
||||
# Validate min_tokens requirement
|
||||
assert_min_tokens_satisfied(output, test_case)
|
||||
|
||||
|
||||
def test_min_tokens_basic_functionality(llm_v1: LLM):
|
||||
"""
|
||||
Test basic min_tokens functionality without stop conditions.
|
||||
|
||||
This is a baseline test that should always pass and validates
|
||||
that min_tokens works correctly in the simple case.
|
||||
"""
|
||||
sampling_params = SamplingParams(min_tokens=10, max_tokens=20, temperature=GREEDY)
|
||||
|
||||
prompt = "Once upon a time"
|
||||
outputs = llm_v1.generate([prompt], sampling_params)
|
||||
|
||||
assert len(outputs) == 1
|
||||
token_count = get_token_count(outputs[0])
|
||||
|
||||
assert token_count >= 10, f"Expected at least 10 tokens, got {token_count}"
|
||||
assert token_count <= 20, f"Expected at most 20 tokens, got {token_count}"
|
||||
|
||||
|
||||
@pytest.mark.xfail(
|
||||
reason=("Known bug #21987: stop strings bypass min_tokens (fixed by PR #22014)"),
|
||||
strict=False,
|
||||
)
|
||||
def test_min_tokens_stop_strings_bug(llm_v1: LLM):
|
||||
"""
|
||||
Test the specific bug where stop strings bypass min_tokens.
|
||||
|
||||
This test specifically reproduces the bug Calvin is fixing in PR #22014.
|
||||
It should fail until that fix is merged.
|
||||
|
||||
Strategy: Use guaranteed stop characters that will appear
|
||||
in any generated text.
|
||||
"""
|
||||
# If the bug is fixed upstream, this test will XPASS
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
min_tokens=15,
|
||||
max_tokens=50,
|
||||
# Common letter; likely appears early
|
||||
stop=["e"],
|
||||
temperature=GREEDY,
|
||||
include_stop_str_in_output=True,
|
||||
)
|
||||
|
||||
# Simple prompt that will generate text containing "e"
|
||||
prompt = "The quick brown fox"
|
||||
outputs = llm_v1.generate([prompt], sampling_params)
|
||||
|
||||
assert len(outputs) == 1
|
||||
token_count = get_token_count(outputs[0])
|
||||
generated_text = outputs[0].outputs[0].text if outputs[0].outputs else ""
|
||||
|
||||
# Debug info to understand what happened
|
||||
print(f"Generated text: {repr(generated_text)}")
|
||||
print(f"Token count: {token_count}")
|
||||
print(f"Contains 'e': {'e' in generated_text}")
|
||||
|
||||
# This assertion should fail due to the bug - if stop string is found early,
|
||||
# the model should still continue generating until min_tokens is reached
|
||||
stop_reason = (
|
||||
outputs[0].outputs[0].stop_reason if outputs[0].outputs else "no output"
|
||||
)
|
||||
assert token_count >= 15, (
|
||||
"Bug confirmed: "
|
||||
f"{token_count} tokens < min_tokens=15. "
|
||||
f"Reason: {stop_reason}. "
|
||||
f"Text: {repr(generated_text)}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.xfail(
|
||||
reason=("Known bug #21987: stop strings bypass min_tokens (fixed by PR #22014)"),
|
||||
strict=False,
|
||||
)
|
||||
def test_min_tokens_stop_strings_guaranteed_early_trigger(llm_v1: LLM):
|
||||
"""
|
||||
Guaranteed test for stop strings bypassing min_tokens bug.
|
||||
|
||||
Strategy: Use very low temperature and multiple common stop strings
|
||||
to virtually guarantee early detection, combined with long min_tokens
|
||||
to ensure the bug is exposed regardless of model behavior.
|
||||
"""
|
||||
# If the bug is fixed upstream, this test will XPASS
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
min_tokens=50, # Set high min_tokens to ensure bug detection
|
||||
max_tokens=200,
|
||||
# Use multiple very common patterns - at least one will appear
|
||||
stop=["e", "a", "i", "o", "u", " ", "t", "n", "s", "r"],
|
||||
temperature=GREEDY,
|
||||
include_stop_str_in_output=True,
|
||||
)
|
||||
|
||||
# Simple prompt that will generate some text
|
||||
prompt = "The cat"
|
||||
outputs = llm_v1.generate([prompt], sampling_params)
|
||||
|
||||
assert len(outputs) == 1
|
||||
token_count = get_token_count(outputs[0])
|
||||
generated_text = outputs[0].outputs[0].text if outputs[0].outputs else ""
|
||||
stop_reason = outputs[0].outputs[0].stop_reason if outputs[0].outputs else "unknown"
|
||||
|
||||
print(f"Generated text: {repr(generated_text)}")
|
||||
print(f"Token count: {token_count}")
|
||||
print(f"Stop reason: {stop_reason}")
|
||||
|
||||
# With the bug, this will fail because ANY of the common characters
|
||||
# will trigger early termination before min_tokens=50 is reached
|
||||
# It's virtually impossible to generate 50 tokens without hitting
|
||||
# at least one of: e, a, i, o, u, space, t, n, s, r
|
||||
finish_reason = (
|
||||
outputs[0].outputs[0].finish_reason if outputs[0].outputs else "unknown"
|
||||
)
|
||||
|
||||
print(f"Finish reason: {finish_reason}")
|
||||
|
||||
if finish_reason == "stop":
|
||||
assert token_count >= 50, (
|
||||
"Bug confirmed: "
|
||||
f"{token_count} tokens < min_tokens=50. "
|
||||
f"Reason: {finish_reason}. "
|
||||
f"Text: {repr(generated_text)}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.xfail(
|
||||
reason=("Potential logits-processor bug: EOS tokens may bypass min_tokens"),
|
||||
strict=False,
|
||||
)
|
||||
def test_min_tokens_eos_behavior(llm_v1: LLM):
|
||||
"""
|
||||
Verify EOS handling with and without min_tokens.
|
||||
|
||||
- Without min_tokens: expect early EOS -> finish_reason == "stop",
|
||||
stop_reason is None, and generated tokens < max_tokens (25).
|
||||
- With min_tokens: EOS should be blocked until min_tokens is reached
|
||||
(finish_reason == "length"); verify that eos_token_id does not appear
|
||||
in generated token_ids.
|
||||
"""
|
||||
# tokenizer + eos id
|
||||
tokenizer = llm_v1.get_tokenizer()
|
||||
eos_token_id = tokenizer.eos_token_id
|
||||
|
||||
prompt = "Give a file extension."
|
||||
max_toks = 32
|
||||
|
||||
# Case 1: WITHOUT min_tokens
|
||||
sp_no_min = SamplingParams(
|
||||
max_tokens=max_toks,
|
||||
temperature=GREEDY,
|
||||
)
|
||||
out_no_min = llm_v1.generate([prompt], sp_no_min)
|
||||
assert len(out_no_min) == 1
|
||||
choice_no_min = out_no_min[0].outputs[0]
|
||||
|
||||
ids_no_min = choice_no_min.token_ids or []
|
||||
finish_no_min = choice_no_min.finish_reason
|
||||
stop_no_min = choice_no_min.stop_reason
|
||||
|
||||
print(
|
||||
"[no-min] tokens=",
|
||||
len(ids_no_min),
|
||||
" finish=",
|
||||
finish_no_min,
|
||||
" stop_reason=",
|
||||
stop_no_min,
|
||||
)
|
||||
|
||||
assert finish_no_min == "stop", (
|
||||
f"Expected finish_reason 'stop' without min_tokens, got {finish_no_min}"
|
||||
)
|
||||
assert stop_no_min is None, (
|
||||
"For EOS-based stop (no user stop strings), stop_reason should be None."
|
||||
)
|
||||
assert len(ids_no_min) < max_toks, (
|
||||
f"Expected early EOS with < {max_toks} tokens, got {len(ids_no_min)}"
|
||||
)
|
||||
|
||||
# Case 2: WITH min_tokens
|
||||
sp_with_min = SamplingParams(
|
||||
min_tokens=max_toks,
|
||||
max_tokens=max_toks,
|
||||
temperature=GREEDY,
|
||||
)
|
||||
out_with_min = llm_v1.generate([prompt], sp_with_min)
|
||||
assert len(out_with_min) == 1
|
||||
choice_with_min = out_with_min[0].outputs[0]
|
||||
|
||||
ids_with_min = choice_with_min.token_ids or []
|
||||
finish_with_min = choice_with_min.finish_reason
|
||||
stop_with_min = choice_with_min.stop_reason
|
||||
|
||||
print(
|
||||
"[with-min] tokens=",
|
||||
len(ids_with_min),
|
||||
" finish=",
|
||||
finish_with_min,
|
||||
" stop_reason=",
|
||||
stop_with_min,
|
||||
)
|
||||
|
||||
# Exact length reached; EOS should have been blocked
|
||||
assert len(ids_with_min) == max_toks, (
|
||||
f"Expected exactly {max_toks} tokens with min_tokens; got {len(ids_with_min)}"
|
||||
)
|
||||
assert finish_with_min == "length", (
|
||||
f"Expected finish_reason 'length'; got {finish_with_min}"
|
||||
)
|
||||
assert eos_token_id not in ids_with_min, (
|
||||
"EOS token id should not appear when min_tokens prevents early EOS."
|
||||
)
|
||||
|
||||
|
||||
def test_min_tokens_validation():
|
||||
"""
|
||||
Test that SamplingParams correctly validates min_tokens parameters.
|
||||
|
||||
This tests the parameter validation logic in SamplingParams.
|
||||
"""
|
||||
# Valid cases
|
||||
SamplingParams(min_tokens=0, max_tokens=10)
|
||||
SamplingParams(min_tokens=5, max_tokens=10)
|
||||
SamplingParams(min_tokens=10, max_tokens=10)
|
||||
|
||||
# Invalid cases
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="min_tokens must be greater than or equal to 0",
|
||||
):
|
||||
SamplingParams(min_tokens=-1, max_tokens=10)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="min_tokens must be less than or equal to max_tokens",
|
||||
):
|
||||
SamplingParams(min_tokens=15, max_tokens=10)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
"""
|
||||
Run tests locally for development.
|
||||
|
||||
Usage:
|
||||
cd vllm/
|
||||
python -m pytest tests/v1/e2e/general/test_min_tokens.py -v
|
||||
"""
|
||||
pytest.main([__file__, "-v"])
|
||||
168
third_party/vllm/tests/v1/e2e/general/test_pooling_chunked_prefill.py
vendored
Normal file
168
third_party/vllm/tests/v1/e2e/general/test_pooling_chunked_prefill.py
vendored
Normal file
@@ -0,0 +1,168 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
prompt = """
|
||||
Generals gathered in their masses
|
||||
Just like witches at black masses
|
||||
Evil minds that plot destruction
|
||||
Sorcerer of death's construction
|
||||
In the fields, the bodies burning
|
||||
As the war machine keeps turning
|
||||
Death and hatred to mankind
|
||||
Poisoning their brainwashed minds
|
||||
Oh, Lord, yeah
|
||||
|
||||
Politicians hide themselves away
|
||||
They only started the war
|
||||
Why should they go out to fight?
|
||||
They leave that all to the poor, yeah
|
||||
Time will tell on their power minds
|
||||
Making war just for fun
|
||||
Treating people just like pawns in chess
|
||||
Wait till their judgment day comes, yeah
|
||||
|
||||
Now, in darkness, world stops turning
|
||||
Ashes where their bodies burning
|
||||
No more war pigs have the power
|
||||
Hand of God has struck the hour
|
||||
Day of Judgment, God is calling
|
||||
On their knees, the war pigs crawling
|
||||
Begging mercies for their sins
|
||||
Satan, laughing, spreads his wings
|
||||
Oh, Lord, yeah
|
||||
"""
|
||||
|
||||
|
||||
class WrapperPooler(nn.Module):
|
||||
def __init__(self, pooler):
|
||||
super().__init__()
|
||||
self.pooler = pooler
|
||||
self.chunks = []
|
||||
|
||||
def get_pooling_updates(self, task):
|
||||
return self.pooler.get_pooling_updates(task)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states,
|
||||
pooling_metadata,
|
||||
):
|
||||
self.chunks.append(hidden_states.shape[0])
|
||||
return self.pooler(hidden_states, pooling_metadata)
|
||||
|
||||
|
||||
def inject_pooler(self):
|
||||
model = self.get_model()
|
||||
wrapper = WrapperPooler(model.pooler)
|
||||
model.pooler = wrapper
|
||||
|
||||
|
||||
def retrieve_chunks(self):
|
||||
model = self.get_model()
|
||||
chunks = model.pooler.chunks
|
||||
model.pooler.chunks = []
|
||||
return chunks
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")
|
||||
def test_pooling_chunked_prefill(vllm_runner, monkeypatch):
|
||||
"""Test chunked prefill for pooling models with LastPool."""
|
||||
|
||||
with monkeypatch.context() as m:
|
||||
m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
||||
model_id = "Qwen/Qwen3-Embedding-0.6B"
|
||||
|
||||
chunk_size = 10
|
||||
|
||||
# Set chunking parameters to force chunked prefill
|
||||
# Note: Chunked prefill is automatically handled by vLLM
|
||||
# internally based on the model size and prompt
|
||||
with vllm_runner(
|
||||
model_id,
|
||||
runner="pooling",
|
||||
long_prefill_token_threshold=chunk_size,
|
||||
tensor_parallel_size=1,
|
||||
enforce_eager=True,
|
||||
enable_chunked_prefill=True,
|
||||
) as llm:
|
||||
llm.get_llm().llm_engine.collective_rpc(inject_pooler)
|
||||
|
||||
tokenizer = llm.get_llm().get_tokenizer()
|
||||
tokens = tokenizer(prompt)["input_ids"]
|
||||
prompt_len = len(tokens)
|
||||
full_chunks, last_chunk = divmod(prompt_len, chunk_size)
|
||||
expected_chunks = [chunk_size] * full_chunks
|
||||
if last_chunk:
|
||||
expected_chunks.append(last_chunk)
|
||||
llm.embed([prompt])
|
||||
chunks = llm.get_llm().llm_engine.collective_rpc(retrieve_chunks)[0]
|
||||
|
||||
# Check that PoolerWrapper was called and chunks were received
|
||||
assert len(chunks) > 1
|
||||
assert chunks == expected_chunks
|
||||
|
||||
# Disable chunked prefill
|
||||
with vllm_runner(
|
||||
model_id,
|
||||
runner="pooling",
|
||||
tensor_parallel_size=1,
|
||||
enforce_eager=True,
|
||||
) as llm:
|
||||
llm.get_llm().llm_engine.collective_rpc(inject_pooler)
|
||||
llm.embed([prompt])
|
||||
chunks = llm.get_llm().llm_engine.collective_rpc(retrieve_chunks)[0]
|
||||
|
||||
# Check that PoolerWrapper was called and no chunks were received
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0] == prompt_len
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")
|
||||
def test_pooling_prefix_cache(vllm_runner, monkeypatch):
|
||||
"""Test chunked prefill for pooling models with LastPool."""
|
||||
|
||||
verses = prompt.split("\n\n")
|
||||
|
||||
with monkeypatch.context() as m:
|
||||
m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
||||
model_id = "Qwen/Qwen3-Embedding-0.6B"
|
||||
|
||||
with vllm_runner(
|
||||
model_id,
|
||||
runner="pooling",
|
||||
enable_prefix_caching=True,
|
||||
tensor_parallel_size=1,
|
||||
enforce_eager=True,
|
||||
) as llm:
|
||||
llm.get_llm().llm_engine.collective_rpc(inject_pooler)
|
||||
tokenizer = llm.get_llm().get_tokenizer()
|
||||
|
||||
prompt1 = "\n\n".join([verses[0], verses[1]])
|
||||
prompt2 = "\n\n".join([verses[0], verses[2]])
|
||||
tokens1 = tokenizer(prompt1)["input_ids"]
|
||||
tokens2 = tokenizer(prompt2)["input_ids"]
|
||||
prompt1_len = len(tokens1)
|
||||
prompt2_len = len(tokens2)
|
||||
|
||||
llm.embed([prompt1])
|
||||
chunks = llm.get_llm().llm_engine.collective_rpc(retrieve_chunks)[0]
|
||||
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0] == prompt1_len
|
||||
|
||||
llm.embed([prompt2])
|
||||
chunks = llm.get_llm().llm_engine.collective_rpc(retrieve_chunks)[0]
|
||||
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0] <= prompt1_len
|
||||
assert chunks[0] < prompt2_len
|
||||
|
||||
vllm_config = llm.get_llm().llm_engine.vllm_config
|
||||
cache_config = vllm_config.cache_config
|
||||
print(f"{cache_config=}")
|
||||
# Prefixes are cached in blocks
|
||||
assert (prompt2_len - chunks[0]) % cache_config.block_size == 0
|
||||
657
third_party/vllm/tests/v1/e2e/general/test_streaming_input.py
vendored
Normal file
657
third_party/vllm/tests/v1/e2e/general/test_streaming_input.py
vendored
Normal file
@@ -0,0 +1,657 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
End-to-end tests for the streaming input feature in AsyncLLM.
|
||||
|
||||
These tests verify that:
|
||||
1. Streaming inputs work correctly with bunched inputs (queued)
|
||||
2. Streaming inputs work correctly with spaced out inputs
|
||||
3. Outputs are equivalent whether inputs are bunched or spaced
|
||||
4. Cancelling the output stream correctly aborts the session
|
||||
5. Closing the input stream correctly signals completion
|
||||
6. Queued inputs are cancelled when the session is aborted
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncGenerator
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from vllm import SamplingParams
|
||||
from vllm.engine.protocol import StreamingInput
|
||||
from vllm.outputs import RequestOutput
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
from vllm.utils.torch_utils import set_default_torch_num_threads
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip(reason="V1 currently only supported on CUDA.", allow_module_level=True)
|
||||
|
||||
# Use a small model that doesn't require authentication for fast tests
|
||||
MODEL = "facebook/opt-125m"
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope="module", loop_scope="module")
|
||||
async def engine():
|
||||
"""Create an AsyncLLM engine for the test.
|
||||
|
||||
Note: Using function scope because pytest_asyncio creates a new event loop
|
||||
for each test, and the output_handler task gets cancelled between tests
|
||||
with module scope.
|
||||
"""
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
|
||||
engine_args = AsyncEngineArgs(
|
||||
model=MODEL, enforce_eager=True, gpu_memory_utilization=0.7
|
||||
)
|
||||
with set_default_torch_num_threads(1):
|
||||
engine = AsyncLLM.from_engine_args(engine_args)
|
||||
try:
|
||||
yield engine
|
||||
finally:
|
||||
engine.shutdown()
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
|
||||
def get_sampling_params(max_tokens: int = 20) -> SamplingParams:
|
||||
"""Create sampling params for streaming input tests."""
|
||||
return SamplingParams(
|
||||
max_tokens=max_tokens,
|
||||
ignore_eos=True,
|
||||
output_kind=RequestOutputKind.DELTA,
|
||||
temperature=0.0, # Deterministic for reproducibility
|
||||
)
|
||||
|
||||
|
||||
async def collect_outputs(
|
||||
output_gen: AsyncGenerator[RequestOutput, None],
|
||||
) -> tuple[list[RequestOutput], str]:
|
||||
"""Collect all outputs from a generate call, return outputs and full text."""
|
||||
outputs: list[RequestOutput] = []
|
||||
full_text = ""
|
||||
async for output in output_gen:
|
||||
outputs.append(output)
|
||||
if output.outputs and output.outputs[0].text:
|
||||
full_text += output.outputs[0].text
|
||||
return outputs, full_text
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_bunched(engine: AsyncLLM):
|
||||
"""Test streaming input where all inputs are sent at once (bunched/queued).
|
||||
|
||||
This tests the case where multiple inputs arrive before any completes.
|
||||
The inputs should be queued and processed in sequence.
|
||||
"""
|
||||
request_id = "test_bunched"
|
||||
sampling_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
# Create an input generator that yields all inputs quickly
|
||||
async def bunched_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
# Send multiple inputs rapidly - they should be queued
|
||||
yield StreamingInput(prompt="Hello, my name is")
|
||||
yield StreamingInput(prompt=" Alice and I like")
|
||||
yield StreamingInput(prompt=" to code in Python")
|
||||
|
||||
outputs, full_text = await collect_outputs(
|
||||
engine.generate(
|
||||
bunched_input_generator(),
|
||||
sampling_params,
|
||||
request_id,
|
||||
)
|
||||
)
|
||||
|
||||
# Verify we got outputs
|
||||
assert len(outputs) > 0, "Should have received outputs"
|
||||
|
||||
# Verify the final output is marked as finished
|
||||
assert outputs[-1].finished, "Last output should be marked as finished"
|
||||
|
||||
# Verify intermediate outputs are not marked as finished
|
||||
for output in outputs[:-1]:
|
||||
assert not output.finished, "Intermediate outputs should not be finished"
|
||||
|
||||
# Verify we generated some text
|
||||
assert len(full_text) > 0, "Should have generated text"
|
||||
print(f"Bunched test generated: {full_text}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_spaced(engine: AsyncLLM):
|
||||
"""Test streaming input where inputs are spaced out.
|
||||
|
||||
This tests the case where each input completes processing before the
|
||||
next one is sent. Each chunk should be prefilled, generate tokens,
|
||||
then the next chunk should be processed.
|
||||
"""
|
||||
request_id = "test_spaced"
|
||||
sampling_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
# Track when each input is sent
|
||||
input_times: list[float] = []
|
||||
outputs_per_chunk: list[int] = [0, 0, 0]
|
||||
current_chunk = 0
|
||||
|
||||
async def spaced_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
nonlocal current_chunk
|
||||
import time
|
||||
|
||||
# First input
|
||||
input_times.append(time.time())
|
||||
yield StreamingInput(prompt="Hello, my name is")
|
||||
current_chunk = 0
|
||||
|
||||
# Wait for some outputs to be generated
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Second input
|
||||
input_times.append(time.time())
|
||||
current_chunk = 1
|
||||
yield StreamingInput(prompt=" Alice and I like")
|
||||
|
||||
# Wait for some outputs
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Third input
|
||||
input_times.append(time.time())
|
||||
current_chunk = 2
|
||||
yield StreamingInput(prompt=" to code in Python")
|
||||
|
||||
outputs: list[RequestOutput] = []
|
||||
full_text = ""
|
||||
|
||||
async for output in engine.generate(
|
||||
spaced_input_generator(),
|
||||
sampling_params,
|
||||
request_id,
|
||||
):
|
||||
outputs.append(output)
|
||||
if output.outputs and output.outputs[0].text:
|
||||
full_text += output.outputs[0].text
|
||||
outputs_per_chunk[current_chunk] += 1
|
||||
|
||||
# Verify we got outputs
|
||||
assert len(outputs) > 0, "Should have received outputs"
|
||||
|
||||
# Verify the final output is marked as finished
|
||||
assert outputs[-1].finished, "Last output should be marked as finished"
|
||||
|
||||
# Verify we received outputs from multiple chunks
|
||||
# (with spaced inputs, we should see outputs distributed across chunks)
|
||||
chunks_with_outputs = sum(1 for c in outputs_per_chunk if c > 0)
|
||||
assert chunks_with_outputs >= 1, "Should have outputs from at least one chunk"
|
||||
|
||||
print(f"Spaced test generated: {full_text}")
|
||||
print(f"Outputs per chunk: {outputs_per_chunk}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_output_equivalence(engine: AsyncLLM):
|
||||
"""Test that bunched and spaced inputs produce equivalent outputs.
|
||||
|
||||
When the same prompts are provided either bunched or spaced,
|
||||
the final concatenated output should be the same (with deterministic
|
||||
sampling).
|
||||
"""
|
||||
prompts = ["Hello, my name is", " Bob and I work", " at Anthropic"]
|
||||
sampling_params = get_sampling_params(max_tokens=15)
|
||||
|
||||
# Test bunched inputs
|
||||
async def bunched_gen() -> AsyncGenerator[StreamingInput, None]:
|
||||
for prompt in prompts:
|
||||
yield StreamingInput(prompt=prompt)
|
||||
|
||||
_, bunched_text = await collect_outputs(
|
||||
engine.generate(bunched_gen(), sampling_params, "equiv_bunched")
|
||||
)
|
||||
|
||||
# Test spaced inputs (same prompts, but with delays)
|
||||
async def spaced_gen() -> AsyncGenerator[StreamingInput, None]:
|
||||
for prompt in prompts:
|
||||
yield StreamingInput(prompt=prompt)
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
_, spaced_text = await collect_outputs(
|
||||
engine.generate(spaced_gen(), sampling_params, "equiv_spaced")
|
||||
)
|
||||
|
||||
# Both should produce the same output since we use temperature=0
|
||||
assert bunched_text == spaced_text, (
|
||||
f"Bunched and spaced should produce same output.\n"
|
||||
f"Bunched: {bunched_text!r}\n"
|
||||
f"Spaced: {spaced_text!r}"
|
||||
)
|
||||
|
||||
print(f"Equivalence test passed. Generated: {bunched_text}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_cancel_output_stream(engine: AsyncLLM):
|
||||
"""Test that cancelling the output stream aborts the entire session.
|
||||
|
||||
When the consumer cancels iteration over the output generator,
|
||||
the session should be aborted including any queued inputs.
|
||||
"""
|
||||
request_id = "test_cancel_output"
|
||||
sampling_params = get_sampling_params(max_tokens=1000)
|
||||
|
||||
input_completed = asyncio.Event()
|
||||
input_task_cancelled = False
|
||||
|
||||
async def slow_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
nonlocal input_task_cancelled
|
||||
try:
|
||||
yield StreamingInput(prompt="Tell me a very long story about")
|
||||
yield StreamingInput(prompt=" a dragon and a knight")
|
||||
|
||||
# This should be cancelled before we get here
|
||||
await asyncio.sleep(10)
|
||||
yield StreamingInput(prompt=" who become friends")
|
||||
input_completed.set()
|
||||
except asyncio.CancelledError:
|
||||
input_task_cancelled = True
|
||||
raise
|
||||
|
||||
outputs_received = 0
|
||||
output_gen = engine.generate(slow_input_generator(), sampling_params, request_id)
|
||||
|
||||
# Collect a few outputs then cancel
|
||||
try:
|
||||
async for output in output_gen:
|
||||
outputs_received += 1
|
||||
if outputs_received >= 5:
|
||||
# Cancel by breaking out of the loop (generator will be GC'd)
|
||||
break
|
||||
finally:
|
||||
# Explicitly close the generator to ensure cleanup
|
||||
await output_gen.aclose()
|
||||
|
||||
# Give time for cleanup
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Verify we got some outputs before cancelling
|
||||
assert outputs_received >= 5, "Should have received outputs before cancel"
|
||||
|
||||
# Verify the input task was cancelled
|
||||
assert input_task_cancelled, "Input task should have been cancelled"
|
||||
|
||||
# Verify the session is properly cleaned up
|
||||
assert not engine.output_processor.has_unfinished_requests(), (
|
||||
"Should have no unfinished requests after cancel"
|
||||
)
|
||||
|
||||
print(f"Cancel test passed. Received {outputs_received} outputs before cancel")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_close_signals_completion(engine: AsyncLLM):
|
||||
"""Test that closing the input stream signals completion.
|
||||
|
||||
When the input generator finishes (naturally or via return),
|
||||
the session should complete with finished=True on the last output.
|
||||
"""
|
||||
request_id = "test_close_completion"
|
||||
sampling_params = get_sampling_params(max_tokens=15)
|
||||
|
||||
input_generator_finished = False
|
||||
|
||||
async def limited_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
nonlocal input_generator_finished
|
||||
yield StreamingInput(prompt="What is 2 + 2? The answer is")
|
||||
# Generator finishes naturally here
|
||||
input_generator_finished = True
|
||||
|
||||
outputs, _ = await collect_outputs(
|
||||
engine.generate(limited_input_generator(), sampling_params, request_id)
|
||||
)
|
||||
|
||||
# Verify the input generator completed
|
||||
assert input_generator_finished, "Input generator should have finished"
|
||||
|
||||
# Verify we got a finished output
|
||||
assert len(outputs) > 0, "Should have received outputs"
|
||||
assert outputs[-1].finished, "Last output should be marked as finished"
|
||||
|
||||
# Verify the session is cleaned up
|
||||
assert not engine.output_processor.has_unfinished_requests(), (
|
||||
"Should have no unfinished requests"
|
||||
)
|
||||
|
||||
print("Close completion test passed")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_abort_queued_inputs(engine: AsyncLLM):
|
||||
"""Test that aborting the session cancels queued inputs.
|
||||
|
||||
When multiple inputs are queued and the session is aborted,
|
||||
all pending inputs should be cancelled.
|
||||
"""
|
||||
request_id = "test_abort_queued"
|
||||
# Use large max_tokens to ensure we have time to queue inputs
|
||||
sampling_params = get_sampling_params(max_tokens=2000)
|
||||
|
||||
inputs_sent = 0
|
||||
input_cancelled = False
|
||||
|
||||
async def many_inputs_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
nonlocal inputs_sent, input_cancelled
|
||||
try:
|
||||
# Send several inputs to fill the queue
|
||||
for i in range(10):
|
||||
yield StreamingInput(prompt=f" Part {i}: Tell me about the number {i}.")
|
||||
inputs_sent += 1
|
||||
# Small delay to interleave with output processing
|
||||
await asyncio.sleep(0.05)
|
||||
except asyncio.CancelledError:
|
||||
input_cancelled = True
|
||||
raise
|
||||
|
||||
outputs_received = 0
|
||||
output_gen = engine.generate(many_inputs_generator(), sampling_params, request_id)
|
||||
|
||||
try:
|
||||
async for output in output_gen:
|
||||
outputs_received += 1
|
||||
# Cancel after receiving some outputs
|
||||
if outputs_received >= 10:
|
||||
break
|
||||
finally:
|
||||
await output_gen.aclose()
|
||||
|
||||
# Give time for cleanup
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Verify we received some outputs
|
||||
assert outputs_received >= 10, "Should have received outputs before abort"
|
||||
|
||||
# Verify the input generator was cancelled OR finished naturally
|
||||
# (it might finish naturally if all inputs were sent before cancel)
|
||||
assert input_cancelled or inputs_sent == 10, (
|
||||
f"Input generator should have been cancelled or completed. "
|
||||
f"cancelled={input_cancelled}, inputs_sent={inputs_sent}"
|
||||
)
|
||||
|
||||
# Verify the session is cleaned up
|
||||
assert not engine.output_processor.has_unfinished_requests(), (
|
||||
"Should have no unfinished requests after abort"
|
||||
)
|
||||
|
||||
print(
|
||||
f"Abort queued test passed. Sent {inputs_sent} inputs, "
|
||||
f"received {outputs_received} outputs"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_error_propagation(engine: AsyncLLM):
|
||||
"""Test that errors in the input generator are propagated to the caller."""
|
||||
request_id = "test_error_propagation"
|
||||
sampling_params = get_sampling_params(max_tokens=20)
|
||||
|
||||
class InputError(Exception):
|
||||
pass
|
||||
|
||||
async def error_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="Start with this")
|
||||
await asyncio.sleep(0.1)
|
||||
raise InputError("Simulated input error")
|
||||
|
||||
# Note: The current implementation catches exceptions and puts them
|
||||
# in the queue, so we should get the error when iterating outputs
|
||||
with pytest.raises(InputError, match="Simulated input error"):
|
||||
async for _ in engine.generate(
|
||||
error_input_generator(), sampling_params, request_id
|
||||
):
|
||||
pass
|
||||
|
||||
# Give time for cleanup
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
# Verify the session is cleaned up
|
||||
assert not engine.output_processor.has_unfinished_requests(), (
|
||||
"Should have no unfinished requests after error"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_multiple_concurrent_sessions(engine: AsyncLLM):
|
||||
"""Test multiple concurrent streaming input sessions.
|
||||
|
||||
Multiple streaming sessions should be able to run concurrently
|
||||
without interfering with each other.
|
||||
"""
|
||||
num_sessions = 3
|
||||
results: list[tuple[str, str]] = []
|
||||
|
||||
async def run_session(session_id: int) -> tuple[str, str]:
|
||||
request_id = f"test_concurrent_{session_id}"
|
||||
sampling_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
prompts = [f"Session {session_id}: Hello", f" world from session {session_id}"]
|
||||
|
||||
async def input_gen() -> AsyncGenerator[StreamingInput, None]:
|
||||
for prompt in prompts:
|
||||
yield StreamingInput(prompt=prompt)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
_, text = await collect_outputs(
|
||||
engine.generate(input_gen(), sampling_params, request_id)
|
||||
)
|
||||
return request_id, text
|
||||
|
||||
# Run sessions concurrently
|
||||
tasks = [asyncio.create_task(run_session(i)) for i in range(num_sessions)]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
# Verify all sessions completed
|
||||
assert len(results) == num_sessions
|
||||
|
||||
for request_id, text in results:
|
||||
assert len(text) > 0, f"Session {request_id} should have generated text"
|
||||
print(f"{request_id}: {text}")
|
||||
|
||||
# Verify cleanup
|
||||
assert not engine.output_processor.has_unfinished_requests()
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_per_chunk_sampling_params(engine: AsyncLLM):
|
||||
"""Test that per-chunk sampling params are respected.
|
||||
|
||||
Each StreamingInput can have its own sampling_params.
|
||||
"""
|
||||
request_id = "test_per_chunk_params"
|
||||
base_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
async def variable_params_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
# First chunk with base params
|
||||
yield StreamingInput(prompt="Count to five:", sampling_params=base_params)
|
||||
|
||||
# Second chunk with different max_tokens
|
||||
chunk_params = get_sampling_params(max_tokens=5)
|
||||
yield StreamingInput(
|
||||
prompt=" Now count backwards:", sampling_params=chunk_params
|
||||
)
|
||||
|
||||
outputs, full_text = await collect_outputs(
|
||||
engine.generate(variable_params_generator(), base_params, request_id)
|
||||
)
|
||||
|
||||
assert len(outputs) > 0, "Should have received outputs"
|
||||
assert outputs[-1].finished, "Last output should be finished"
|
||||
assert len(full_text) > 0, "Should have generated text"
|
||||
|
||||
print(f"Per-chunk params test generated: {full_text}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_empty_generator(engine: AsyncLLM):
|
||||
"""Test behavior when the input generator yields nothing.
|
||||
|
||||
An empty generator should still produce a finished output.
|
||||
"""
|
||||
request_id = "test_empty_generator"
|
||||
sampling_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
async def empty_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
# Don't yield anything
|
||||
return
|
||||
yield # Make it a generator
|
||||
|
||||
outputs: list[RequestOutput] = []
|
||||
async for output in engine.generate(empty_generator(), sampling_params, request_id):
|
||||
outputs.append(output)
|
||||
|
||||
# Should still get a finished marker
|
||||
assert len(outputs) >= 1, "Should receive at least one output"
|
||||
assert outputs[-1].finished, "Should have a finished output"
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_single_chunk(engine: AsyncLLM):
|
||||
"""Test streaming input with a single chunk.
|
||||
|
||||
This is effectively the same as a regular non-streaming request,
|
||||
but using the streaming input API.
|
||||
"""
|
||||
request_id = "test_single_chunk"
|
||||
sampling_params = get_sampling_params(max_tokens=15)
|
||||
|
||||
async def single_chunk_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="What color is the sky? The sky is")
|
||||
|
||||
outputs, full_text = await collect_outputs(
|
||||
engine.generate(single_chunk_generator(), sampling_params, request_id)
|
||||
)
|
||||
|
||||
assert len(outputs) > 0
|
||||
assert outputs[-1].finished
|
||||
assert "blue" in full_text.lower() or len(full_text) > 0
|
||||
|
||||
print(f"Single chunk test generated: {full_text}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_reuse_request_id(engine: AsyncLLM):
|
||||
"""Test that request IDs can be reused after a session completes."""
|
||||
request_id = "test_reuse_id"
|
||||
sampling_params = get_sampling_params(max_tokens=5)
|
||||
|
||||
# First session
|
||||
async def gen1() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="First session")
|
||||
|
||||
_, text1 = await collect_outputs(
|
||||
engine.generate(gen1(), sampling_params, request_id)
|
||||
)
|
||||
|
||||
# Second session with same ID
|
||||
async def gen2() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="Second session")
|
||||
|
||||
_, text2 = await collect_outputs(
|
||||
engine.generate(gen2(), sampling_params, request_id)
|
||||
)
|
||||
|
||||
assert len(text1) > 0
|
||||
assert len(text2) > 0
|
||||
assert not engine.output_processor.has_unfinished_requests()
|
||||
|
||||
print(f"Reuse ID test: session 1: {text1}, session 2: {text2}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_validation_errors(engine: AsyncLLM):
|
||||
"""Test that invalid configurations raise appropriate errors."""
|
||||
|
||||
async def dummy_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="test")
|
||||
|
||||
# Test n > 1 is rejected
|
||||
with pytest.raises(ValueError, match="Input streaming not currently supported"):
|
||||
params_n2 = SamplingParams(max_tokens=10, n=2)
|
||||
async for _ in engine.generate(dummy_generator(), params_n2, "test_n2"):
|
||||
pass
|
||||
|
||||
# Test FINAL_ONLY is rejected
|
||||
with pytest.raises(ValueError, match="Input streaming not currently supported"):
|
||||
params_final = SamplingParams(
|
||||
max_tokens=10, output_kind=RequestOutputKind.FINAL_ONLY
|
||||
)
|
||||
async for _ in engine.generate(dummy_generator(), params_final, "test_final"):
|
||||
pass
|
||||
|
||||
# Test stop strings are rejected
|
||||
with pytest.raises(ValueError, match="Input streaming not currently supported"):
|
||||
params_stop = SamplingParams(max_tokens=10, stop=["stop"])
|
||||
async for _ in engine.generate(dummy_generator(), params_stop, "test_stop"):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_delayed_generator_exit(engine: AsyncLLM):
|
||||
"""Test that output generator exits when input generator closes after outputs.
|
||||
|
||||
This tests the case where:
|
||||
1. Multiple inputs are sent and fully processed
|
||||
2. The engine has finished
|
||||
3. The input generator doesn't exit until after the engine finishes
|
||||
4. The output generator should exit properly once the input generator exits
|
||||
"""
|
||||
request_id = "test_delayed_exit"
|
||||
sampling_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
engine_finished_event = asyncio.Event()
|
||||
input_generator_exited = False
|
||||
finish_count = 0
|
||||
|
||||
async def delayed_exit_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
nonlocal input_generator_exited
|
||||
# Send all inputs immediately
|
||||
yield StreamingInput(prompt="Hello, my name is")
|
||||
yield StreamingInput(prompt=" Alice")
|
||||
|
||||
# Wait until the engine has finished generating before exiting
|
||||
await engine_finished_event.wait()
|
||||
|
||||
# Add a small delay to ensure we're testing the "delayed exit" case
|
||||
await asyncio.sleep(0.1)
|
||||
input_generator_exited = True
|
||||
|
||||
outputs: list[RequestOutput] = []
|
||||
full_text = ""
|
||||
|
||||
async for output in engine.generate(
|
||||
delayed_exit_input_generator(), sampling_params, request_id
|
||||
):
|
||||
outputs.append(output)
|
||||
if output.outputs and output.outputs[0].text:
|
||||
full_text += output.outputs[0].text
|
||||
|
||||
# Signal when the engine finishes both input chunks (each gets a finish_reason)
|
||||
# Note: output.finished will be False while input stream is open
|
||||
if output.outputs and output.outputs[0].finish_reason is not None:
|
||||
finish_count += 1
|
||||
if finish_count == 2:
|
||||
engine_finished_event.set()
|
||||
|
||||
# Verify the input generator exited properly
|
||||
assert input_generator_exited, (
|
||||
"Input generator should have exited after engine finished"
|
||||
)
|
||||
|
||||
# Verify we got outputs
|
||||
assert len(outputs) > 0, "Should have received outputs"
|
||||
|
||||
# Verify we generated some text
|
||||
assert len(full_text) > 0, "Should have generated text"
|
||||
|
||||
# Verify the session is cleaned up
|
||||
assert not engine.output_processor.has_unfinished_requests(), (
|
||||
"Should have no unfinished requests"
|
||||
)
|
||||
|
||||
print(f"Delayed exit test passed. Generated: {full_text}")
|
||||
0
third_party/vllm/tests/v1/e2e/spec_decode/__init__.py
vendored
Normal file
0
third_party/vllm/tests/v1/e2e/spec_decode/__init__.py
vendored
Normal file
131
third_party/vllm/tests/v1/e2e/spec_decode/test_async_spec_decode.py
vendored
Normal file
131
third_party/vllm/tests/v1/e2e/spec_decode/test_async_spec_decode.py
vendored
Normal file
@@ -0,0 +1,131 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Test that verifies no implicit GPU-CPU synchronization occurs during
|
||||
speculative decoding generation under expected conditions.
|
||||
"""
|
||||
|
||||
import multiprocessing
|
||||
import sys
|
||||
import traceback
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sync_tracker():
|
||||
"""
|
||||
Fixture that patches CommonAttentionMetadata.seq_lens_cpu to detect
|
||||
lazy init syncs. Prints stack traces immediately when syncs occur.
|
||||
"""
|
||||
from vllm.v1.attention.backend import CommonAttentionMetadata
|
||||
|
||||
# Shared counter for cross-process communication (inherited by fork)
|
||||
sync_count = multiprocessing.Value("i", 0)
|
||||
|
||||
# Save original property
|
||||
original_prop = CommonAttentionMetadata.seq_lens_cpu
|
||||
original_fget = original_prop.fget
|
||||
|
||||
# Create tracking wrapper
|
||||
def tracking_seq_lens_cpu(self):
|
||||
if self._seq_lens_cpu is None:
|
||||
# Increment counter
|
||||
with sync_count.get_lock():
|
||||
sync_count.value += 1
|
||||
count = sync_count.value
|
||||
# Print stack trace immediately (shows in subprocess output)
|
||||
print(f"\n{'=' * 60}", file=sys.stderr)
|
||||
print(f"SYNC #{count}: seq_lens_cpu lazy init triggered!", file=sys.stderr)
|
||||
print(f"{'=' * 60}", file=sys.stderr)
|
||||
traceback.print_stack(file=sys.stderr)
|
||||
print(f"{'=' * 60}\n", file=sys.stderr)
|
||||
sys.stderr.flush()
|
||||
return original_fget(self)
|
||||
|
||||
# Apply patch
|
||||
CommonAttentionMetadata.seq_lens_cpu = property(tracking_seq_lens_cpu)
|
||||
|
||||
class SyncTracker:
|
||||
@property
|
||||
def count(self) -> int:
|
||||
return sync_count.value
|
||||
|
||||
def assert_no_sync(self, msg: str = ""):
|
||||
count = sync_count.value
|
||||
assert count == 0, (
|
||||
f"Unexpected GPU-CPU sync: seq_lens_cpu lazy init triggered "
|
||||
f"{count} times. See stack traces above. {msg}"
|
||||
)
|
||||
|
||||
yield SyncTracker()
|
||||
|
||||
# Restore original property
|
||||
CommonAttentionMetadata.seq_lens_cpu = original_prop
|
||||
torch._dynamo.reset()
|
||||
|
||||
|
||||
# Test configurations: (model, spec_model, method, num_spec_tokens, backend_env)
|
||||
SPEC_DECODE_CONFIGS = [
|
||||
pytest.param(
|
||||
"meta-llama/Llama-3.2-1B-Instruct",
|
||||
"nm-testing/Llama3_2_1B_speculator.eagle3",
|
||||
"eagle3",
|
||||
2,
|
||||
id="eagle3-llama",
|
||||
),
|
||||
pytest.param(
|
||||
"eagle618/deepseek-v3-random",
|
||||
"eagle618/eagle-deepseek-v3-random",
|
||||
"eagle",
|
||||
2,
|
||||
id="eagle-mla-deepseek",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model,spec_model,method,num_spec_tokens",
|
||||
SPEC_DECODE_CONFIGS,
|
||||
)
|
||||
def test_no_sync_with_spec_decode(
|
||||
sync_tracker,
|
||||
model: str,
|
||||
spec_model: str,
|
||||
method: str,
|
||||
num_spec_tokens: int,
|
||||
):
|
||||
"""
|
||||
Test that no implicit GPU-CPU sync occurs during speculative decoding
|
||||
generation.
|
||||
"""
|
||||
# Import vLLM AFTER sync_tracker fixture has applied the patch
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
|
||||
llm = LLM(
|
||||
model=model,
|
||||
max_model_len=256,
|
||||
speculative_config={
|
||||
"method": method,
|
||||
"num_speculative_tokens": num_spec_tokens,
|
||||
"model": spec_model,
|
||||
},
|
||||
enforce_eager=True,
|
||||
async_scheduling=True,
|
||||
)
|
||||
|
||||
outputs = llm.generate(
|
||||
["Hello, my name is"],
|
||||
SamplingParams(temperature=0, max_tokens=10),
|
||||
)
|
||||
|
||||
assert len(outputs) == 1
|
||||
assert len(outputs[0].outputs[0].text) > 0
|
||||
|
||||
del llm
|
||||
torch.accelerator.empty_cache()
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
sync_tracker.assert_no_sync()
|
||||
139
third_party/vllm/tests/v1/e2e/spec_decode/test_lora_with_spec_decode.py
vendored
Normal file
139
third_party/vllm/tests/v1/e2e/spec_decode/test_lora_with_spec_decode.py
vendored
Normal file
@@ -0,0 +1,139 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
This script contains:
|
||||
1. test lora with speculative decoding for batch inference
|
||||
"""
|
||||
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.lora.request import LoRARequest
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
LORA_TEST_PROMPT_MAP: dict[str, str] = {}
|
||||
|
||||
LORA_TEST_PROMPT_MAP["premjatin/qwen-linear-algebra-coder"] = """
|
||||
### INSTRUCTION:
|
||||
You are an AI assistant that generates Python code to solve linear
|
||||
algebra problems.
|
||||
|
||||
### PROBLEM:
|
||||
Find the eigenvalues and eigenvectors of the following 3x3 matrix:
|
||||
[[3, 2, 0],
|
||||
[2, 3, 0],
|
||||
[0, 0, 2]]
|
||||
|
||||
### OUTPUT FORMAT (STRICT):
|
||||
Numbers should be represented as integers only.
|
||||
|
||||
### PYTHON SOLUTION:
|
||||
"""
|
||||
|
||||
SEED = 42
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")
|
||||
@pytest.mark.parametrize(
|
||||
"model_setup",
|
||||
[
|
||||
(
|
||||
"eagle3",
|
||||
"Qwen/Qwen3-1.7B",
|
||||
"AngelSlim/Qwen3-1.7B_eagle3",
|
||||
"premjatin/qwen-linear-algebra-coder",
|
||||
1,
|
||||
)
|
||||
],
|
||||
)
|
||||
def test_batch_inference_correctness(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
model_setup: tuple[str, str, str, str, int],
|
||||
):
|
||||
"""
|
||||
Compare the outputs of a LLM with only Lora and a LLM with both SD and Lora.
|
||||
Should be the same and no failure when doing batch inference.
|
||||
model_setup: (method, model_name, spec_model_name, lora_path, tp_size)
|
||||
"""
|
||||
with monkeypatch.context() as m:
|
||||
# Disable randomness
|
||||
m.setenv("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
|
||||
torch.manual_seed(SEED)
|
||||
np.random.seed(SEED)
|
||||
random.seed(SEED)
|
||||
torch.cuda.manual_seed_all(SEED)
|
||||
torch.backends.cudnn.benchmark = False
|
||||
torch.backends.cudnn.deterministic = True
|
||||
|
||||
method, model_name, spec_model_name, lora_path, tp_size = model_setup
|
||||
|
||||
# without speculative decoding
|
||||
ref_llm = LLM(
|
||||
model=model_name,
|
||||
trust_remote_code=True,
|
||||
tensor_parallel_size=tp_size,
|
||||
max_model_len=2048,
|
||||
max_num_seqs=4,
|
||||
enable_lora=True,
|
||||
max_loras=1,
|
||||
max_cpu_loras=1,
|
||||
max_lora_rank=16,
|
||||
)
|
||||
|
||||
prompts = [LORA_TEST_PROMPT_MAP[lora_path]] * 100
|
||||
lora_request = LoRARequest("adapter", 1, lora_path)
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.0, top_p=1.0, top_k=-1, seed=SEED, max_tokens=128
|
||||
)
|
||||
|
||||
ref_outputs = ref_llm.generate(
|
||||
prompts, sampling_params, lora_request=lora_request
|
||||
)
|
||||
del ref_llm
|
||||
torch.accelerator.empty_cache()
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
lora_spec_llm = LLM(
|
||||
model=model_name,
|
||||
trust_remote_code=True,
|
||||
tensor_parallel_size=tp_size,
|
||||
speculative_config={
|
||||
"method": method,
|
||||
"model": spec_model_name,
|
||||
"num_speculative_tokens": 3,
|
||||
"max_model_len": 2048,
|
||||
},
|
||||
max_model_len=2048,
|
||||
max_num_seqs=4,
|
||||
enable_lora=True,
|
||||
max_loras=1,
|
||||
max_cpu_loras=1,
|
||||
max_lora_rank=16,
|
||||
)
|
||||
|
||||
lora_spec_outputs = lora_spec_llm.generate(
|
||||
prompts, sampling_params, lora_request=lora_request
|
||||
)
|
||||
|
||||
matches = 0
|
||||
misses = 0
|
||||
for ref_output, spec_output in zip(ref_outputs, lora_spec_outputs):
|
||||
if ref_output.outputs[0].text == spec_output.outputs[0].text:
|
||||
matches += 1
|
||||
else:
|
||||
misses += 1
|
||||
print(f"ref_output: {ref_output.outputs[0].text}")
|
||||
print(f"spec_output: {spec_output.outputs[0].text}")
|
||||
|
||||
# Heuristic: expect at least 90% of the prompts to match exactly
|
||||
# Upon failure, inspect the outputs to check for inaccuracy.
|
||||
print(f"match ratio: {matches}/{len(ref_outputs)}")
|
||||
assert matches > int(0.90 * len(ref_outputs))
|
||||
del lora_spec_llm
|
||||
torch.accelerator.empty_cache()
|
||||
cleanup_dist_env_and_memory()
|
||||
1033
third_party/vllm/tests/v1/e2e/spec_decode/test_spec_decode.py
vendored
Normal file
1033
third_party/vllm/tests/v1/e2e/spec_decode/test_spec_decode.py
vendored
Normal file
File diff suppressed because it is too large
Load Diff
171
third_party/vllm/tests/v1/ec_connector/integration/README.md
vendored
Normal file
171
third_party/vllm/tests/v1/ec_connector/integration/README.md
vendored
Normal file
@@ -0,0 +1,171 @@
|
||||
# EPD Correctness Test
|
||||
|
||||
This test verifies that EPD (Encoder-Prefill-Decode) disaggregation produces identical outputs to a baseline single instance.
|
||||
|
||||
## What It Tests
|
||||
|
||||
- **Baseline**: Single vLLM instance serving a multimodal model
|
||||
- **EPD (1E+1PD)**: 1 Encoder + 1 Prefill-Decode instance
|
||||
- **Baseline (1P+1D)**: 1 Prefill + 1 Decode instance
|
||||
- **EPD (1E+1P+1D)**: 1 Encoder + 1 Prefill + 1 Decode instance
|
||||
|
||||
The test ensures that disaggregated encoding produces **identical** outputs to the baseline.
|
||||
|
||||
Note that currently PD disaggregation set up may give slightly different results from a single instance. Therefore, we need the result from 1P+1D as the baseline for 1E+1P+1D
|
||||
|
||||
Please refer to [Disaggregated Encoder Feature](../../../docs/features/disagg_encoder.md) for the detailed explanation for the EPD features.
|
||||
|
||||
## Files
|
||||
|
||||
- `run_epd_correctness_test.sh` - Main test script (starts all instances and runs tests)
|
||||
- `test_epd_correctness.py` - Python test script (compares outputs)
|
||||
|
||||
## Usage
|
||||
|
||||
### Multimodal Prompts (Default)
|
||||
|
||||
```bash
|
||||
cd vllm
|
||||
./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
|
||||
```
|
||||
|
||||
This runs the test with actual multimodal (image) prompts.
|
||||
|
||||
### Text-Only Prompts
|
||||
|
||||
```bash
|
||||
cd vllm
|
||||
USE_MM_PROMPTS=0 ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
|
||||
```
|
||||
|
||||
This runs a quick test with text-only prompts to verify the setup works.
|
||||
|
||||
### Custom Configuration
|
||||
|
||||
```bash
|
||||
# Use specific GPUs
|
||||
GPU_E=0 GPU_PD=1 GPU_P=1 GPU_D=2 bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
|
||||
|
||||
# Use specific ports
|
||||
ENDPOINT_PORT=10001 bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
|
||||
|
||||
# Use specific model
|
||||
MODEL="Qwen/Qwen2.5-VL-3B-Instruct" bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
|
||||
|
||||
# Use specific storage path
|
||||
EC_SHARED_STORAGE_PATH="/tmp/my_ec_cache" bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
|
||||
```
|
||||
|
||||
## How It Works
|
||||
|
||||
### Step 1: Baseline
|
||||
|
||||
1. Start single vLLM instance on GPU
|
||||
2. Run test prompts (multimodal or text-only)
|
||||
3. Save outputs to `.vllm_epd_baseline.txt`
|
||||
4. Shutdown instance
|
||||
|
||||
### Step 2: EPD (1E + 1PD)
|
||||
|
||||
1. Clear encoder cache storage
|
||||
2. Start instances and proxy
|
||||
3. Run same test prompts
|
||||
4. Assert outputs match baseline exactly
|
||||
5. Shutdown instances
|
||||
|
||||
### Step 3: EPD (1E + 1P + 1D)
|
||||
|
||||
1. Clear encoder cache storage
|
||||
2. Start instances and proxy
|
||||
3. Run same test prompts
|
||||
4. Assert outputs match baseline exactly
|
||||
5. Shutdown instances
|
||||
|
||||
## Test Scenarios
|
||||
|
||||
### Multimodal Prompts (--use_mm_prompts)
|
||||
|
||||
Tests encoder cache transfer:
|
||||
|
||||
- Single image query
|
||||
- Multiple images in one request
|
||||
- Mixed image and text
|
||||
- Image with detailed questions
|
||||
|
||||
### Text-Only Prompts (default)
|
||||
|
||||
Quick sanity check:
|
||||
|
||||
- Simple text queries
|
||||
- Text-only explanations
|
||||
- Verifies proxy routing works
|
||||
|
||||
## Expected Behavior
|
||||
|
||||
### ✅ Test Passes When
|
||||
|
||||
- All disagg outputs match baseline outputs exactly
|
||||
- No errors during instance startup
|
||||
- Encoder cache is properly saved and loaded
|
||||
- Proxy correctly routes requests
|
||||
|
||||
### ❌ Test Fails When
|
||||
|
||||
- Outputs differ between baseline and disagg
|
||||
- Server startup fails
|
||||
- Encoder cache not found (should fall back to local execution)
|
||||
- Proxy routing errors
|
||||
|
||||
## Notes
|
||||
|
||||
- The test uses deterministic generation (`temperature=0.0`, `seed=42`)
|
||||
- Encoder cache should enable exact output reproduction
|
||||
- Test cleans up all instances and cache files after completion
|
||||
- Safe to run multiple times (idempotent)
|
||||
- We setup the PD disagg part with NixlConnector. Please read details about EPD in `examples/online_serving/disaggregated_encoder/README.md`
|
||||
|
||||
## Requirements
|
||||
|
||||
- Multiple GPUs (3 for 1E+1P+1D, 2 for 1E+1PD, 1 for baseline)
|
||||
- 1E+1P+1D is runnable with 2 GPU by assign E and P on the same GPU now.
|
||||
- Multimodal model (e.g., Qwen2.5-VL-3B-Instruct)
|
||||
- Internet access (for accessing vllm test images)
|
||||
|
||||
## Debugging
|
||||
|
||||
### Check Logs
|
||||
|
||||
Logs and baseline output are saved in `/tmp/` by default.
|
||||
Can be customized by changing the environment variables.
|
||||
|
||||
### Check Encoder Cache
|
||||
|
||||
```bash
|
||||
# Verify cache files are created
|
||||
ls -la $EC_SHARED_STORAGE_PATH/
|
||||
|
||||
# Should see directories with mm_hash names
|
||||
# Each containing encoder_cache.safetensors
|
||||
```
|
||||
|
||||
### Manual Testing
|
||||
|
||||
Run individual components:
|
||||
|
||||
```bash
|
||||
# Baseline only
|
||||
python test_epd_correctness.py \
|
||||
--service_url http://localhost:8000 \
|
||||
--model_name Qwen/Qwen2.5-VL-3B-Instruct \
|
||||
--mode baseline \
|
||||
--baseline_file test_output.txt \
|
||||
--use_mm_prompts
|
||||
|
||||
# Disagg only (requires baseline output file!)
|
||||
python test_epd_correctness.py \
|
||||
--service_url http://localhost:8000 \
|
||||
--model_name Qwen/Qwen2.5-VL-3B-Instruct \
|
||||
--mode disagg \
|
||||
--baseline_file test_output.txt \
|
||||
--use_mm_prompts
|
||||
```
|
||||
BIN
third_party/vllm/tests/v1/ec_connector/integration/hato.jpg
vendored
Normal file
BIN
third_party/vllm/tests/v1/ec_connector/integration/hato.jpg
vendored
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 821 KiB |
476
third_party/vllm/tests/v1/ec_connector/integration/run_epd_correctness_test.sh
vendored
Normal file
476
third_party/vllm/tests/v1/ec_connector/integration/run_epd_correctness_test.sh
vendored
Normal file
@@ -0,0 +1,476 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# EPD (Encoder-Prefill-Decode) Correctness Test
|
||||
#
|
||||
# This script tests that EPD disaggregation produces the same outputs as baseline.
|
||||
# It runs:
|
||||
# 1. Baseline: Single vLLM instance
|
||||
# 2. EPD: 1E + 1PD setup
|
||||
# 3. Baseline for (E + P + D): 1P + 1D vLLM instances disagg
|
||||
# 4. EPD: 1E + 1P + 1D setup
|
||||
|
||||
# For GPU usage
|
||||
|
||||
# set -xe
|
||||
|
||||
# Find the git repository root directory
|
||||
GIT_ROOT=$(git rev-parse --show-toplevel)
|
||||
|
||||
# Model to test
|
||||
MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"
|
||||
|
||||
# Set 1 to use multimodal prompts; else to use text-only
|
||||
USE_MM_PROMPTS="${USE_MM_PROMPTS:-1}"
|
||||
MM_FLAG=""
|
||||
if [ "$USE_MM_PROMPTS" = "1" ]; then
|
||||
MM_FLAG="--use_mm_prompts"
|
||||
fi
|
||||
|
||||
# GPU configuration
|
||||
GPU_E="${GPU_E:-0}"
|
||||
GPU_P="${GPU_P:-1}"
|
||||
GPU_D="${GPU_D:-2}"
|
||||
GPU_SINGLE="${GPU_SINGLE:-$GPU_P}"
|
||||
GPU_PD="${GPU_PD:-$GPU_P}"
|
||||
|
||||
# Port
|
||||
ENCODE_PORT="${ENCODE_PORT:-19534}"
|
||||
PREFILL_PORT="${PREFILL_PORT:-19535}"
|
||||
DECODE_PORT="${DECODE_PORT:-19536}"
|
||||
PREFILL_DECODE_PORT="${PREFILL_DECODE_PORT:-19537}"
|
||||
ENDPOINT_PORT="${ENDPOINT_PORT:-10001}"
|
||||
|
||||
# Storage path for encoder cache
|
||||
EC_SHARED_STORAGE_PATH="${EC_SHARED_STORAGE_PATH:-/tmp/ec_cache_test}"
|
||||
TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-600}"
|
||||
|
||||
# Output file for baseline comparison and logs
|
||||
LOG_PATH="${LOG_PATH:-/tmp}"
|
||||
BASELINE_FILE="${BASELINE_FILE:-/tmp/vllm_baseline.txt}"
|
||||
BASELINE_PD_FILE="${BASELINE_PD_FILE:-/tmp/vllm_epd_baseline.txt}"
|
||||
|
||||
mkdir -p "$LOG_PATH"
|
||||
|
||||
# Trap the SIGINT signal (triggered by Ctrl+C)
|
||||
trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT
|
||||
|
||||
# Wait for server to be ready
|
||||
wait_for_server() {
|
||||
local port=$1
|
||||
timeout "$TIMEOUT_SECONDS" bash -c "
|
||||
until curl -s localhost:${port}/v1/chat/completions > /dev/null; do
|
||||
sleep 1
|
||||
done" && return 0 || return 1
|
||||
}
|
||||
|
||||
# Cleanup function
|
||||
cleanup_instances() {
|
||||
echo "Cleaning up any running vLLM instances..."
|
||||
pkill -f "vllm serve" || true
|
||||
pkill -f "disagg_epd_proxy.py" || true
|
||||
sleep 2
|
||||
}
|
||||
|
||||
# Function to run baseline (single instance)
|
||||
run_baseline() {
|
||||
echo "================================"
|
||||
echo "Running BASELINE (single instance)"
|
||||
echo "================================"
|
||||
|
||||
cleanup_instances
|
||||
rm -rf "$EC_SHARED_STORAGE_PATH"
|
||||
|
||||
local PORT=$ENDPOINT_PORT
|
||||
|
||||
# Start baseline instance
|
||||
echo "Starting baseline instance on GPU $GPU_SINGLE, port $PORT"
|
||||
CUDA_VISIBLE_DEVICES="$GPU_SINGLE" vllm serve "$MODEL" \
|
||||
--port "$PORT" \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.7 \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
> "$LOG_PATH"/baseline.log 2>&1 &
|
||||
|
||||
local BASELINE_PID=$!
|
||||
|
||||
# Wait for baseline to start
|
||||
echo "Waiting for baseline instance to start..."
|
||||
wait_for_server "$PORT"
|
||||
|
||||
curl http://127.0.0.1:"$PORT"/v1/models
|
||||
echo ""
|
||||
|
||||
# Run test in baseline mode
|
||||
echo "Running baseline..."
|
||||
|
||||
python "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \
|
||||
--service_url "http://localhost:$PORT" \
|
||||
--model_name "$MODEL" \
|
||||
--mode baseline \
|
||||
--baseline_file "$BASELINE_FILE" \
|
||||
$MM_FLAG
|
||||
|
||||
# Cleanup baseline
|
||||
echo "Stopping baseline instance..."
|
||||
kill $BASELINE_PID 2>/dev/null || true
|
||||
sleep 2
|
||||
cleanup_instances
|
||||
}
|
||||
|
||||
# Function to run EPD with 1E + 1PD
|
||||
run_epd_1e_1pd() {
|
||||
echo "================================"
|
||||
echo "Running EPD (1E + 1PD)"
|
||||
echo "================================"
|
||||
|
||||
cleanup_instances
|
||||
rm -rf "$EC_SHARED_STORAGE_PATH"
|
||||
mkdir -p "$EC_SHARED_STORAGE_PATH"
|
||||
|
||||
local ENCODE_PORT=$ENCODE_PORT
|
||||
local PREFILL_DECODE_PORT=$PREFILL_DECODE_PORT
|
||||
local PROXY_PORT=$ENDPOINT_PORT
|
||||
|
||||
declare -a PIDS=()
|
||||
|
||||
# Start encoder instance
|
||||
echo "Starting encoder instance on GPU $GPU_E, port $ENCODE_PORT"
|
||||
CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \
|
||||
--port "$ENCODE_PORT" \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.01 \
|
||||
--enable-request-id-headers \
|
||||
--no-enable-prefix-caching \
|
||||
--max-num-batched-tokens 114688 \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--ec-transfer-config '{
|
||||
"ec_connector": "ECExampleConnector",
|
||||
"ec_role": "ec_producer",
|
||||
"ec_connector_extra_config": {
|
||||
"shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
|
||||
}
|
||||
}' \
|
||||
> "$LOG_PATH"/1e1pd_encoder.log 2>&1 &
|
||||
PIDS+=($!)
|
||||
|
||||
# Start prefill+decode instance
|
||||
echo "Starting PD instance on GPU $GPU_PD, port $PREFILL_DECODE_PORT"
|
||||
CUDA_VISIBLE_DEVICES="$GPU_PD" vllm serve "$MODEL" \
|
||||
--port "$PREFILL_DECODE_PORT" \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.7 \
|
||||
--enable-request-id-headers \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--ec-transfer-config '{
|
||||
"ec_connector": "ECExampleConnector",
|
||||
"ec_role": "ec_consumer",
|
||||
"ec_connector_extra_config": {
|
||||
"shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
|
||||
}
|
||||
}' \
|
||||
> "$LOG_PATH"/1e1pd_pd.log 2>&1 &
|
||||
PIDS+=($!)
|
||||
|
||||
# Wait for instances to start
|
||||
echo "Waiting for encoder instance..."
|
||||
wait_for_server "$ENCODE_PORT"
|
||||
echo "Waiting for PD instance..."
|
||||
wait_for_server "$PREFILL_DECODE_PORT"
|
||||
|
||||
# Start proxy
|
||||
echo "Starting EPD proxy on port $PROXY_PORT"
|
||||
python "${GIT_ROOT}/examples/online_serving/disaggregated_encoder/disagg_epd_proxy.py" \
|
||||
--host "0.0.0.0" \
|
||||
--port "$PROXY_PORT" \
|
||||
--encode-servers-urls "http://localhost:$ENCODE_PORT" \
|
||||
--prefill-servers-urls "disable" \
|
||||
--decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" \
|
||||
> "$LOG_PATH"/1e1pd_proxy.log 2>&1 &
|
||||
PIDS+=($!)
|
||||
|
||||
# Wait for proxy
|
||||
echo "Waiting for proxy..."
|
||||
wait_for_server "$PROXY_PORT"
|
||||
|
||||
curl http://127.0.0.1:"$PROXY_PORT"/v1/models
|
||||
curl http://127.0.0.1:"$PROXY_PORT"/health
|
||||
echo ""
|
||||
|
||||
echo "All EPD (1E+1PD) services are up!"
|
||||
|
||||
# Run test in disagg mode
|
||||
echo "Running EPD (1E+1PD) correctness test..."
|
||||
|
||||
python "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \
|
||||
--service_url "http://localhost:$PROXY_PORT" \
|
||||
--model_name "$MODEL" \
|
||||
--mode disagg \
|
||||
--baseline_file "$BASELINE_FILE" \
|
||||
$MM_FLAG
|
||||
|
||||
# Cleanup
|
||||
echo "✓✓ 1E+1PD Correctness Test finished"
|
||||
echo "Stopping EPD (1E+1PD) instances..."
|
||||
for pid in "${PIDS[@]}"; do
|
||||
kill "$pid" 2>/dev/null || true
|
||||
done
|
||||
sleep 2
|
||||
cleanup_instances
|
||||
}
|
||||
|
||||
# Function to run baseline for 1E + 1P + 1D (PD disagg)
|
||||
run_baseline_1p_1d() {
|
||||
echo "================================"
|
||||
echo "Running PD BASELINE (1P + 1D)"
|
||||
echo "================================"
|
||||
|
||||
cleanup_instances
|
||||
rm -rf "$EC_SHARED_STORAGE_PATH"
|
||||
mkdir -p "$EC_SHARED_STORAGE_PATH"
|
||||
|
||||
local PREFILL_PORT=$PREFILL_PORT
|
||||
local DECODE_PORT=$DECODE_PORT
|
||||
local PROXY_PORT=$ENDPOINT_PORT
|
||||
|
||||
declare -a PIDS=()
|
||||
|
||||
# Start prefill instance
|
||||
echo "Starting prefill instance on GPU $GPU_P, port $PREFILL_PORT"
|
||||
CUDA_VISIBLE_DEVICES="$GPU_P" \
|
||||
VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \
|
||||
vllm serve "$MODEL" \
|
||||
--port "$PREFILL_PORT" \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.7 \
|
||||
--enable-request-id-headers \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "NixlConnector",
|
||||
"kv_role": "kv_producer"
|
||||
}' \
|
||||
> "$LOG_PATH"/1p1d_prefill.log 2>&1 &
|
||||
PIDS+=($!)
|
||||
|
||||
# Start decode instance
|
||||
echo "Starting decode instance on GPU $GPU_D, port $DECODE_PORT"
|
||||
CUDA_VISIBLE_DEVICES="$GPU_D" \
|
||||
VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \
|
||||
vllm serve "$MODEL" \
|
||||
--port "$DECODE_PORT" \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.7 \
|
||||
--enable-request-id-headers \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "NixlConnector",
|
||||
"kv_role": "kv_consumer"
|
||||
}' \
|
||||
> "$LOG_PATH"/1p1d_decode.log 2>&1 &
|
||||
PIDS+=($!)
|
||||
|
||||
# Wait for instances to start
|
||||
echo "Waiting for prefill instance..."
|
||||
wait_for_server "$PREFILL_PORT"
|
||||
echo "Waiting for decode instance..."
|
||||
wait_for_server "$DECODE_PORT"
|
||||
|
||||
# Start proxy
|
||||
echo "Starting EPD proxy on port $PROXY_PORT"
|
||||
python "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py" \
|
||||
--host "0.0.0.0" \
|
||||
--port "$PROXY_PORT" \
|
||||
--prefiller-ports "$PREFILL_PORT" \
|
||||
--decoder-ports "$DECODE_PORT" \
|
||||
> "$LOG_PATH"/1p1d_proxy.log 2>&1 &
|
||||
PIDS+=($!)
|
||||
|
||||
# Wait for proxy
|
||||
echo "Waiting for proxy..."
|
||||
wait_for_server "$PROXY_PORT"
|
||||
|
||||
curl http://127.0.0.1:"$PROXY_PORT"/healthcheck
|
||||
echo ""
|
||||
|
||||
echo "All PD (1P+1D) services are up!"
|
||||
|
||||
# Run test in baseline mode
|
||||
echo "Running PD disagg baseline..."
|
||||
|
||||
python "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \
|
||||
--service_url "http://localhost:$PROXY_PORT" \
|
||||
--model_name "$MODEL" \
|
||||
--mode baseline_pd \
|
||||
--baseline_file "$BASELINE_PD_FILE" \
|
||||
$MM_FLAG
|
||||
|
||||
# Cleanup
|
||||
echo "Stopping PD (1P+1D) instances..."
|
||||
for pid in "${PIDS[@]}"; do
|
||||
kill "$pid" 2>/dev/null || true
|
||||
done
|
||||
sleep 2
|
||||
cleanup_instances
|
||||
}
|
||||
|
||||
# Function to run EPD with 1E + 1P + 1D
|
||||
run_epd_1e_1p_1d() {
|
||||
echo "================================"
|
||||
echo "Running EPD (1E + 1P + 1D)"
|
||||
echo "================================"
|
||||
|
||||
cleanup_instances
|
||||
rm -rf "$EC_SHARED_STORAGE_PATH"
|
||||
mkdir -p "$EC_SHARED_STORAGE_PATH"
|
||||
|
||||
local ENCODE_PORT=$ENCODE_PORT
|
||||
local PREFILL_PORT=$PREFILL_PORT
|
||||
local DECODE_PORT=$DECODE_PORT
|
||||
local PROXY_PORT=$ENDPOINT_PORT
|
||||
|
||||
declare -a PIDS=()
|
||||
|
||||
# Start encoder instance
|
||||
echo "Starting encoder instance on GPU $GPU_E, port $ENCODE_PORT"
|
||||
CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \
|
||||
--port "$ENCODE_PORT" \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.01 \
|
||||
--enable-request-id-headers \
|
||||
--no-enable-prefix-caching \
|
||||
--max-num-batched-tokens 114688 \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--ec-transfer-config '{
|
||||
"ec_connector": "ECExampleConnector",
|
||||
"ec_role": "ec_producer",
|
||||
"ec_connector_extra_config": {
|
||||
"shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
|
||||
}
|
||||
}' \
|
||||
> "$LOG_PATH"/1e1p1d_encoder.log 2>&1 &
|
||||
PIDS+=($!)
|
||||
|
||||
# Start prefill instance
|
||||
echo "Starting prefill instance on GPU $GPU_P, port $PREFILL_PORT"
|
||||
CUDA_VISIBLE_DEVICES="$GPU_P" \
|
||||
VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \
|
||||
vllm serve "$MODEL" \
|
||||
--port "$PREFILL_PORT" \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.7 \
|
||||
--enable-request-id-headers \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--ec-transfer-config '{
|
||||
"ec_connector": "ECExampleConnector",
|
||||
"ec_role": "ec_consumer",
|
||||
"ec_connector_extra_config": {
|
||||
"shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
|
||||
}
|
||||
}' \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "NixlConnector",
|
||||
"kv_role": "kv_producer"
|
||||
}' \
|
||||
> "$LOG_PATH"/1e1p1d_prefill.log 2>&1 &
|
||||
PIDS+=($!)
|
||||
|
||||
# Start decode instance
|
||||
echo "Starting decode instance on GPU $GPU_D, port $DECODE_PORT"
|
||||
CUDA_VISIBLE_DEVICES="$GPU_D" \
|
||||
VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \
|
||||
vllm serve "$MODEL" \
|
||||
--port "$DECODE_PORT" \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.7 \
|
||||
--enable-request-id-headers \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "NixlConnector",
|
||||
"kv_role": "kv_consumer"
|
||||
}' \
|
||||
> "$LOG_PATH"/1e1p1d_decode.log 2>&1 &
|
||||
PIDS+=($!)
|
||||
|
||||
# Wait for instances to start
|
||||
echo "Waiting for encoder instance..."
|
||||
wait_for_server "$ENCODE_PORT"
|
||||
echo "Waiting for prefill instance..."
|
||||
wait_for_server "$PREFILL_PORT"
|
||||
echo "Waiting for decode instance..."
|
||||
wait_for_server "$DECODE_PORT"
|
||||
|
||||
# Start proxy
|
||||
echo "Starting EPD proxy on port $PROXY_PORT"
|
||||
python "${GIT_ROOT}/examples/online_serving/disaggregated_encoder/disagg_epd_proxy.py" \
|
||||
--host "0.0.0.0" \
|
||||
--port "$PROXY_PORT" \
|
||||
--encode-servers-urls "http://localhost:$ENCODE_PORT" \
|
||||
--prefill-servers-urls "http://localhost:$PREFILL_PORT" \
|
||||
--decode-servers-urls "http://localhost:$DECODE_PORT" \
|
||||
> "$LOG_PATH"/1e1p1d_proxy.log 2>&1 &
|
||||
PIDS+=($!)
|
||||
|
||||
# Wait for proxy
|
||||
echo "Waiting for proxy..."
|
||||
wait_for_server "$PROXY_PORT"
|
||||
|
||||
curl http://127.0.0.1:"$PROXY_PORT"/v1/models
|
||||
curl http://127.0.0.1:"$PROXY_PORT"/health
|
||||
echo ""
|
||||
|
||||
echo "All EPD (1E+1P+1D) services are up!"
|
||||
|
||||
# Run test in disagg mode
|
||||
echo "Running EPD (1E+1P+1D) correctness test..."
|
||||
|
||||
python "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \
|
||||
--service_url "http://localhost:$PROXY_PORT" \
|
||||
--model_name "$MODEL" \
|
||||
--mode disagg \
|
||||
--baseline_file "$BASELINE_PD_FILE" \
|
||||
$MM_FLAG
|
||||
|
||||
# Cleanup
|
||||
echo "✓✓ 1E+1P+1D Correctness Test finished"
|
||||
echo "Stopping EPD (1E+1P+1D) instances..."
|
||||
for pid in "${PIDS[@]}"; do
|
||||
kill "$pid" 2>/dev/null || true
|
||||
done
|
||||
sleep 2
|
||||
cleanup_instances
|
||||
}
|
||||
|
||||
# Main execution
|
||||
echo "================================"
|
||||
echo "EPD Correctness Test Suite"
|
||||
echo "Model: $MODEL"
|
||||
echo "================================"
|
||||
|
||||
# Step 1: Run baseline
|
||||
run_baseline
|
||||
|
||||
# Step 2: Test 1E + 1PD
|
||||
run_epd_1e_1pd
|
||||
|
||||
# Step 3: Test baseline 1P + 1D
|
||||
run_baseline_1p_1d
|
||||
|
||||
# Step 4: Test 1E + 1P + 1D
|
||||
run_epd_1e_1p_1d
|
||||
|
||||
# Cleanup output file
|
||||
rm -f "$BASELINE_FILE"
|
||||
rm -f "$BASELINE_PD_FILE"
|
||||
|
||||
echo "================================"
|
||||
echo "✓✓ All EPD correctness tests finished!"
|
||||
echo "================================"
|
||||
300
third_party/vllm/tests/v1/ec_connector/integration/test_epd_correctness.py
vendored
Normal file
300
third_party/vllm/tests/v1/ec_connector/integration/test_epd_correctness.py
vendored
Normal file
@@ -0,0 +1,300 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
EPD Correctness Test
|
||||
|
||||
Tests that EPD (Encoder-Prefill-Decode) disaggregation produces the same
|
||||
outputs as a baseline single instance.
|
||||
|
||||
Usage:
|
||||
# Baseline mode (saves outputs):
|
||||
python test_epd_correctness.py \
|
||||
--service_url http://localhost:8000 \
|
||||
--model_name Qwen/Qwen2.5-VL-3B-Instruct \
|
||||
--mode baseline \
|
||||
--baseline_file .vllm_epd_baseline.txt
|
||||
|
||||
# Disagg mode (compares outputs):
|
||||
python test_epd_correctness.py \
|
||||
--service_url http://localhost:8000 \
|
||||
--model_name Qwen/Qwen2.5-VL-3B-Instruct \
|
||||
--mode disagg \
|
||||
--baseline_file .vllm_epd_baseline.txt
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
|
||||
import openai
|
||||
import requests
|
||||
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.multimodal.utils import encode_image_url
|
||||
|
||||
MAX_OUTPUT_LEN = 256
|
||||
|
||||
# Sample prompts with multimodal content
|
||||
image_1 = ImageAsset("stop_sign").pil_image.resize((1280, 720))
|
||||
image_2 = ImageAsset("cherry_blossom").pil_image.resize((1280, 720))
|
||||
|
||||
image_local_path = f"{os.path.dirname(os.path.abspath(__file__))}/hato.jpg"
|
||||
|
||||
SAMPLE_PROMPTS_MM: list[dict] = [
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": encode_image_url(image_1)},
|
||||
},
|
||||
{"type": "text", "text": "What's in this image?"},
|
||||
],
|
||||
}
|
||||
],
|
||||
"description": "Single image query",
|
||||
},
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": encode_image_url(image_2)},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"file://{image_local_path}"},
|
||||
},
|
||||
{"type": "text", "text": "Describe these 2 images in detail."},
|
||||
],
|
||||
}
|
||||
],
|
||||
"description": "2 images with detailed query",
|
||||
},
|
||||
]
|
||||
|
||||
# Text-only prompts for mixed testing
|
||||
SAMPLE_PROMPTS_TEXT: list[dict] = [
|
||||
{
|
||||
"messages": [{"role": "user", "content": "What is the capital of France?"}],
|
||||
"description": "Simple text-only query",
|
||||
},
|
||||
{
|
||||
"messages": [
|
||||
{"role": "user", "content": "Explain quantum computing in simple terms."}
|
||||
],
|
||||
"description": "Text-only explanation request",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def check_vllm_server(url: str, timeout=5, retries=10) -> bool:
|
||||
"""Check if the vLLM server is ready.
|
||||
|
||||
Args:
|
||||
url: The URL to check (usually /health or /healthcheck endpoint)
|
||||
timeout: Timeout in seconds for each request
|
||||
retries: Number of retries if the server is not ready
|
||||
|
||||
Returns:
|
||||
True if the server is ready, False otherwise
|
||||
"""
|
||||
for attempt in range(retries):
|
||||
try:
|
||||
response = requests.get(url, timeout=timeout)
|
||||
if response.status_code == 200:
|
||||
print(f"Server is ready at {url}")
|
||||
return True
|
||||
else:
|
||||
print(
|
||||
f"Attempt {attempt + 1}/{retries}: Server returned "
|
||||
f"status code {response.status_code}"
|
||||
)
|
||||
except requests.exceptions.RequestException as e:
|
||||
print(f"Attempt {attempt + 1}/{retries}: Error connecting: {e}")
|
||||
time.sleep(2) # Wait before retrying
|
||||
return False
|
||||
|
||||
|
||||
def run_chat_completion(
|
||||
base_url: str,
|
||||
model_name: str,
|
||||
messages: list,
|
||||
max_tokens: int = MAX_OUTPUT_LEN,
|
||||
) -> str:
|
||||
"""Run a chat completion request.
|
||||
|
||||
Args:
|
||||
base_url: Base URL of the vLLM server
|
||||
model_name: Name of the model
|
||||
messages: Messages for chat completion
|
||||
max_tokens: Maximum tokens to generate
|
||||
|
||||
Returns:
|
||||
Generated text content
|
||||
"""
|
||||
client = openai.OpenAI(api_key="EMPTY", base_url=base_url)
|
||||
|
||||
completion = client.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=messages,
|
||||
max_tokens=max_tokens,
|
||||
temperature=0.0,
|
||||
seed=42,
|
||||
)
|
||||
|
||||
return completion.choices[0].message.content
|
||||
|
||||
|
||||
def main():
|
||||
"""Main test function."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="EPD correctness test - compare disagg vs baseline"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--service_url",
|
||||
type=str,
|
||||
required=True,
|
||||
help="The vLLM service URL (e.g., http://localhost:8000)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--model_name",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Model name",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
type=str,
|
||||
default="baseline",
|
||||
choices=["baseline", "baseline_pd", "disagg"],
|
||||
help="Mode: baseline/baseline_pd (saves outputs) or disagg (compares outputs)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--baseline_file",
|
||||
type=str,
|
||||
default=".vllm_epd_baseline.txt",
|
||||
help="File to save/load baseline outputs",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--use_mm_prompts",
|
||||
action="store_true",
|
||||
help="Use multimodal prompts (default: use text-only for quick testing)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
print(f"Service URL: {args.service_url}")
|
||||
print(f"Model: {args.model_name}")
|
||||
print(f"Mode: {args.mode}")
|
||||
print(f"Output file: {args.baseline_file}")
|
||||
print(f"Use MM prompts: {args.use_mm_prompts}")
|
||||
|
||||
# Determine health check endpoint
|
||||
if args.mode == "baseline":
|
||||
health_check_url = f"{args.service_url}/health"
|
||||
elif args.mode == "baseline_pd":
|
||||
# Nixl toy proxy use /healthcheck
|
||||
health_check_url = f"{args.service_url}/healthcheck"
|
||||
else:
|
||||
# Disagg EPD proxy uses /health
|
||||
health_check_url = f"{args.service_url}/health"
|
||||
if not os.path.exists(args.baseline_file):
|
||||
raise ValueError(
|
||||
f"In disagg mode, the output file {args.baseline_file} from "
|
||||
"baseline does not exist. Run baseline mode first."
|
||||
)
|
||||
|
||||
# Check if server is ready
|
||||
if not check_vllm_server(health_check_url):
|
||||
raise RuntimeError(f"vLLM server at {args.service_url} is not ready!")
|
||||
|
||||
# Select prompts to use
|
||||
if args.use_mm_prompts:
|
||||
test_prompts = SAMPLE_PROMPTS_MM
|
||||
print("Using multimodal prompts")
|
||||
else:
|
||||
test_prompts = SAMPLE_PROMPTS_TEXT
|
||||
print("Using text-only prompts for quick testing")
|
||||
|
||||
# Run completions
|
||||
service_url = f"{args.service_url}/v1"
|
||||
output_strs = {}
|
||||
|
||||
for i, prompt_data in enumerate(test_prompts):
|
||||
print(
|
||||
f"\nRunning prompt {i + 1}/{len(test_prompts)}: "
|
||||
f"{prompt_data['description']}"
|
||||
)
|
||||
|
||||
output_str = run_chat_completion(
|
||||
base_url=service_url,
|
||||
model_name=args.model_name,
|
||||
messages=prompt_data["messages"],
|
||||
max_tokens=MAX_OUTPUT_LEN,
|
||||
)
|
||||
|
||||
# Use description as key for comparison
|
||||
key = prompt_data["description"]
|
||||
output_strs[key] = output_str
|
||||
print(f"Output: {output_str}")
|
||||
|
||||
if args.mode in ("baseline", "baseline_pd"):
|
||||
# Baseline mode: Save outputs
|
||||
print(f"\nSaving baseline outputs to {args.baseline_file}")
|
||||
try:
|
||||
with open(args.baseline_file, "w") as json_file:
|
||||
json.dump(output_strs, json_file, indent=4)
|
||||
print("✅ Baseline outputs saved successfully")
|
||||
except OSError as e:
|
||||
print(f"Error writing to file: {e}")
|
||||
raise
|
||||
else:
|
||||
# Disagg mode: Load and compare outputs
|
||||
print(f"\nLoading baseline outputs from {args.baseline_file}")
|
||||
baseline_outputs = None
|
||||
try:
|
||||
with open(args.baseline_file) as json_file:
|
||||
baseline_outputs = json.load(json_file)
|
||||
except OSError as e:
|
||||
print(f"Error reading from file: {e}")
|
||||
raise
|
||||
|
||||
# Verify outputs match
|
||||
print("\nComparing disagg outputs with baseline...")
|
||||
assert isinstance(baseline_outputs, dict), "Baseline outputs should be a dict"
|
||||
assert len(baseline_outputs) == len(output_strs), (
|
||||
f"Length mismatch: baseline has {len(baseline_outputs)}, "
|
||||
f"disagg has {len(output_strs)}"
|
||||
)
|
||||
|
||||
all_match = True
|
||||
for key, baseline_output in baseline_outputs.items():
|
||||
assert key in output_strs, f"{key} not in disagg outputs"
|
||||
|
||||
disagg_output = output_strs[key]
|
||||
if baseline_output == disagg_output:
|
||||
print(f"✅ {key}: MATCH")
|
||||
else:
|
||||
print(f"❌ {key}: MISMATCH")
|
||||
print(f" Baseline: {baseline_output}")
|
||||
print(f" Disagg: {disagg_output}")
|
||||
all_match = False
|
||||
|
||||
assert all_match, "❌❌Disagg outputs do not match baseline!❌❌"
|
||||
if all_match:
|
||||
print("\n✅ All outputs match! Test PASSED")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
623
third_party/vllm/tests/v1/ec_connector/unit/test_ec_example_connector.py
vendored
Normal file
623
third_party/vllm/tests/v1/ec_connector/unit/test_ec_example_connector.py
vendored
Normal file
@@ -0,0 +1,623 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Unit tests for ECExampleConnector.
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
import safetensors
|
||||
import torch
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole
|
||||
from vllm.distributed.ec_transfer.ec_connector.example_connector import (
|
||||
ECExampleConnector,
|
||||
ECExampleConnectorMetadata,
|
||||
MMMeta,
|
||||
)
|
||||
from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
|
||||
|
||||
# ------------------ Mock Classes ------------------ #
|
||||
class MockRequest:
|
||||
def __init__(self, request_id, mm_hashes: list[str], token_counts: list[int]):
|
||||
assert len(mm_hashes) == len(token_counts)
|
||||
self.request_id = request_id
|
||||
self._token_counts = token_counts
|
||||
self.mm_features = []
|
||||
for i, mm_hash in enumerate(mm_hashes):
|
||||
feature = MultiModalFeatureSpec(
|
||||
data=None,
|
||||
modality="image",
|
||||
identifier=mm_hash,
|
||||
mm_position=PlaceholderRange(offset=0, length=self._token_counts[i]),
|
||||
)
|
||||
self.mm_features.append(feature)
|
||||
|
||||
def get_num_encoder_embeds(self, input_id: int) -> int:
|
||||
assert input_id < len(self._token_counts)
|
||||
return self._token_counts[input_id]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_storage(tmp_path):
|
||||
"""Fixture providing temporary storage path."""
|
||||
return str(tmp_path)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_vllm_config_producer(temp_storage):
|
||||
"""Fixture providing mock VllmConfig for producer role."""
|
||||
config = Mock(spec=VllmConfig)
|
||||
config.ec_transfer_config = Mock()
|
||||
config.ec_transfer_config.get_from_extra_config = Mock(return_value=temp_storage)
|
||||
config.ec_transfer_config.is_ec_producer = True
|
||||
return config
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_vllm_config_consumer(temp_storage):
|
||||
"""Fixture providing mock VllmConfig for consumer role."""
|
||||
config = Mock(spec=VllmConfig)
|
||||
config.ec_transfer_config = Mock()
|
||||
config.ec_transfer_config.get_from_extra_config = Mock(return_value=temp_storage)
|
||||
config.ec_transfer_config.is_ec_producer = False
|
||||
return config
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_request_with_3_mm():
|
||||
"""Fixture providing mock Request with 3 multimodal items."""
|
||||
request_id = "test_req_123"
|
||||
mm_hashes = ["img_hash_1", "img_hash_2", "img_hash_3"]
|
||||
token_counts = [100, 150, 200]
|
||||
|
||||
request = MockRequest(request_id, mm_hashes, token_counts)
|
||||
return request
|
||||
|
||||
|
||||
# ------------------ Unit Tests ------------------ #
|
||||
class TestECExampleConnectorBasics:
|
||||
"""Test basic EC connector functionality."""
|
||||
|
||||
def test_initialization_producer(self, mock_vllm_config_producer, temp_storage):
|
||||
"""Test connector initializes correctly as producer."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.SCHEDULER,
|
||||
)
|
||||
|
||||
assert connector.role == ECConnectorRole.SCHEDULER
|
||||
assert connector.is_producer
|
||||
assert connector._storage_path == temp_storage
|
||||
assert connector._mm_datas_need_loads == {}
|
||||
|
||||
def test_initialization_consumer(self, mock_vllm_config_consumer, temp_storage):
|
||||
"""Test connector initializes correctly as consumer."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_consumer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
assert connector.role == ECConnectorRole.WORKER
|
||||
assert not connector.is_producer
|
||||
assert connector._storage_path == temp_storage
|
||||
|
||||
def test_role_assignment(self, mock_vllm_config_producer):
|
||||
"""Test role is correctly assigned."""
|
||||
scheduler_connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.SCHEDULER,
|
||||
)
|
||||
worker_connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
assert scheduler_connector.role == ECConnectorRole.SCHEDULER
|
||||
assert worker_connector.role == ECConnectorRole.WORKER
|
||||
|
||||
|
||||
class TestCacheExistence:
|
||||
"""Test cache existence checking using has_cache_item() API."""
|
||||
|
||||
def test_has_cache_item_all_exist_3_items(
|
||||
self,
|
||||
mock_vllm_config_producer,
|
||||
mock_vllm_config_consumer,
|
||||
mock_request_with_3_mm,
|
||||
):
|
||||
"""Test has_cache_item returns True when all 3 caches exist."""
|
||||
# Test for producer first
|
||||
producer = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.SCHEDULER,
|
||||
)
|
||||
|
||||
# Create cache files using save_caches (proper way)
|
||||
encoder_cache: dict[str, torch.Tensor] = {}
|
||||
|
||||
for mm_feature in mock_request_with_3_mm.mm_features:
|
||||
mm_hash = mm_feature.identifier
|
||||
encoder_cache[mm_hash] = torch.randn(10, 768)
|
||||
producer.save_caches(encoder_cache, mm_hash)
|
||||
|
||||
# Test using has_cache_item API
|
||||
producer_result = [
|
||||
producer.has_cache_item(mm_feature.identifier)
|
||||
for mm_feature in mock_request_with_3_mm.mm_features
|
||||
]
|
||||
|
||||
# Assert
|
||||
assert len(producer_result) == 3
|
||||
assert all(producer_result), f"Expected all True, got {producer_result}"
|
||||
|
||||
# Also test consumer can check if cache exists
|
||||
consumer = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_consumer,
|
||||
role=ECConnectorRole.SCHEDULER,
|
||||
)
|
||||
|
||||
# Test using has_cache_item API
|
||||
consumer_result = [
|
||||
consumer.has_cache_item(mm_feature.identifier)
|
||||
for mm_feature in mock_request_with_3_mm.mm_features
|
||||
]
|
||||
|
||||
# Assert
|
||||
assert len(consumer_result) == 3
|
||||
assert all(consumer_result), f"Expected all True, got {consumer_result}"
|
||||
|
||||
def test_has_cache_item_none_exist(
|
||||
self, mock_vllm_config_producer, mock_request_with_3_mm
|
||||
):
|
||||
"""Test has_caches returns False when no caches exist."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.SCHEDULER,
|
||||
)
|
||||
|
||||
# Test without creating any files
|
||||
result = [
|
||||
connector.has_cache_item(mm_feature.identifier)
|
||||
for mm_feature in mock_request_with_3_mm.mm_features
|
||||
]
|
||||
|
||||
# Assert
|
||||
assert len(result) == 3
|
||||
assert not any(result), f"Expected all False, got {result}"
|
||||
|
||||
def test_has_cache_item_partial_exist(
|
||||
self, mock_vllm_config_producer, mock_request_with_3_mm
|
||||
):
|
||||
"""Test has_caches with some caches existing (1 of 3)."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.SCHEDULER,
|
||||
)
|
||||
|
||||
# Create only the second cache file
|
||||
mm_hash_second = mock_request_with_3_mm.mm_features[1].identifier
|
||||
encoder_cache = {mm_hash_second: torch.randn(10, 768)}
|
||||
connector.save_caches(encoder_cache, mm_hash_second)
|
||||
|
||||
# Test
|
||||
result = [
|
||||
connector.has_cache_item(mm_feature.identifier)
|
||||
for mm_feature in mock_request_with_3_mm.mm_features
|
||||
]
|
||||
|
||||
# Assert
|
||||
assert len(result) == 3
|
||||
assert not result[0] # First doesn't exist
|
||||
assert result[1] # Second exists
|
||||
assert not result[2] # Third doesn't exist
|
||||
|
||||
|
||||
class TestStateManagement:
|
||||
"""Test connector state management."""
|
||||
|
||||
def test_update_state_after_alloc_3_items(
|
||||
self, mock_vllm_config_producer, mock_request_with_3_mm
|
||||
):
|
||||
"""Test state update after allocation for 3 MM items."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.SCHEDULER,
|
||||
)
|
||||
|
||||
# Initial state should be empty
|
||||
assert len(connector._mm_datas_need_loads) == 0
|
||||
|
||||
# Update state for all 3 items (mock cache existence)
|
||||
with patch.object(connector, "has_cache_item", return_value=True):
|
||||
for i in range(3):
|
||||
connector.update_state_after_alloc(mock_request_with_3_mm, index=i)
|
||||
|
||||
# Check state updated for all 3
|
||||
assert len(connector._mm_datas_need_loads) == 3
|
||||
assert "img_hash_1" in connector._mm_datas_need_loads
|
||||
assert "img_hash_2" in connector._mm_datas_need_loads
|
||||
assert "img_hash_3" in connector._mm_datas_need_loads
|
||||
assert connector._mm_datas_need_loads["img_hash_1"] == 100
|
||||
assert connector._mm_datas_need_loads["img_hash_2"] == 150
|
||||
assert connector._mm_datas_need_loads["img_hash_3"] == 200
|
||||
|
||||
def test_build_connector_meta_3_items(
|
||||
self, mock_vllm_config_producer, mock_request_with_3_mm
|
||||
):
|
||||
"""Test metadata building for 3 MM items."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.SCHEDULER,
|
||||
)
|
||||
|
||||
# Setup state for all 3 items (mock cache existence)
|
||||
with patch.object(connector, "has_cache_item", return_value=True):
|
||||
for i in range(3):
|
||||
connector.update_state_after_alloc(mock_request_with_3_mm, index=i)
|
||||
|
||||
# Build metadata
|
||||
scheduler_output = Mock(spec=SchedulerOutput)
|
||||
metadata = connector.build_connector_meta(scheduler_output)
|
||||
|
||||
# Assert
|
||||
assert isinstance(metadata, ECExampleConnectorMetadata)
|
||||
assert len(metadata.mm_datas) == 3
|
||||
assert metadata.mm_datas[0].mm_hash == "img_hash_1"
|
||||
assert metadata.mm_datas[0].num_token == 100
|
||||
assert metadata.mm_datas[1].mm_hash == "img_hash_2"
|
||||
assert metadata.mm_datas[1].num_token == 150
|
||||
assert metadata.mm_datas[2].mm_hash == "img_hash_3"
|
||||
assert metadata.mm_datas[2].num_token == 200
|
||||
|
||||
# State should be cleared after building
|
||||
assert len(connector._mm_datas_need_loads) == 0
|
||||
|
||||
def test_build_connector_meta_empty(self, mock_vllm_config_producer):
|
||||
"""Test metadata building with empty state."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.SCHEDULER,
|
||||
)
|
||||
|
||||
scheduler_output = Mock(spec=SchedulerOutput)
|
||||
metadata = connector.build_connector_meta(scheduler_output)
|
||||
|
||||
assert isinstance(metadata, ECExampleConnectorMetadata)
|
||||
assert len(metadata.mm_datas) == 0
|
||||
|
||||
def test_state_cleared_after_metadata_build(
|
||||
self, mock_vllm_config_producer, mock_request_with_3_mm
|
||||
):
|
||||
"""Test that state is properly cleared after building metadata."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.SCHEDULER,
|
||||
)
|
||||
|
||||
# Add state (mock cache existence)
|
||||
with patch.object(connector, "has_cache_item", return_value=True):
|
||||
for i in range(3):
|
||||
connector.update_state_after_alloc(mock_request_with_3_mm, index=i)
|
||||
assert len(connector._mm_datas_need_loads) == 3
|
||||
|
||||
# Build metadata (should clear state)
|
||||
scheduler_output = Mock(spec=SchedulerOutput)
|
||||
connector.build_connector_meta(scheduler_output)
|
||||
|
||||
# State should be empty
|
||||
assert len(connector._mm_datas_need_loads) == 0
|
||||
|
||||
# Build again should return empty metadata
|
||||
metadata2 = connector.build_connector_meta(scheduler_output)
|
||||
assert len(metadata2.mm_datas) == 0
|
||||
|
||||
|
||||
class TestCacheSaving:
|
||||
"""Test encoder cache saving (producer only)."""
|
||||
|
||||
def test_save_caches_producer_3_items(
|
||||
self, mock_vllm_config_producer, mock_request_with_3_mm, temp_storage
|
||||
):
|
||||
"""Test cache saving as producer for 3 different MM items."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
# Create and save 3 different caches
|
||||
mm_hashes = [f.identifier for f in mock_request_with_3_mm.mm_features]
|
||||
encoder_cache: dict[str, torch.Tensor] = {}
|
||||
|
||||
for mm_hash in mm_hashes:
|
||||
encoder_cache[mm_hash] = torch.randn(10, 768)
|
||||
connector.save_caches(encoder_cache, mm_hash)
|
||||
|
||||
# Verify all files exist using has_cache_item
|
||||
result = [
|
||||
connector.has_cache_item(mm_feature.identifier)
|
||||
for mm_feature in mock_request_with_3_mm.mm_features
|
||||
]
|
||||
assert all(result), f"Not all caches were saved: {result}"
|
||||
|
||||
# Verify each file's content
|
||||
for mm_hash in mm_hashes:
|
||||
filename = connector._generate_filename_debug(mm_hash)
|
||||
loaded = safetensors.torch.load_file(filename)
|
||||
assert "ec_cache" in loaded
|
||||
assert torch.allclose(loaded["ec_cache"], encoder_cache[mm_hash].cpu())
|
||||
|
||||
def test_save_caches_consumer_skips(self, mock_vllm_config_consumer):
|
||||
"""Test cache saving is skipped for consumer."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_consumer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
mm_hash = "test_hash_consumer"
|
||||
encoder_cache = {mm_hash: torch.randn(10, 768)}
|
||||
|
||||
# Save should not raise but also not create file
|
||||
connector.save_caches(encoder_cache, mm_hash)
|
||||
|
||||
# Verify file doesn't exist using has_cache_item
|
||||
result = connector.has_cache_item(mm_hash)
|
||||
assert not result, "Consumer should not save caches"
|
||||
|
||||
|
||||
class TestCacheLoading:
|
||||
"""Test encoder cache loading (consumer)."""
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
def test_start_load_caches_consumer_3_items(
|
||||
self,
|
||||
mock_vllm_config_producer,
|
||||
mock_vllm_config_consumer,
|
||||
mock_request_with_3_mm,
|
||||
temp_storage,
|
||||
):
|
||||
"""Test consumer loads 3 caches from storage."""
|
||||
# First, create producer to save caches
|
||||
producer = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
# Producer saves 3 caches
|
||||
mm_hashes = [f.identifier for f in mock_request_with_3_mm.mm_features]
|
||||
saved_caches = {}
|
||||
for mm_hash in mm_hashes:
|
||||
saved_caches[mm_hash] = torch.randn(10, 768)
|
||||
producer.save_caches(saved_caches, mm_hash)
|
||||
|
||||
# Now consumer loads
|
||||
consumer = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_consumer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
# Setup metadata for all 3
|
||||
metadata = ECExampleConnectorMetadata()
|
||||
for mm_hash in mm_hashes:
|
||||
metadata.add_mm_data(MMMeta.make_meta(mm_hash, 100))
|
||||
consumer.bind_connector_metadata(metadata)
|
||||
|
||||
# Load
|
||||
encoder_cache: dict[str, torch.Tensor] = {}
|
||||
consumer.start_load_caches(encoder_cache=encoder_cache)
|
||||
|
||||
# Verify all 3 loaded
|
||||
assert len(encoder_cache) == 3
|
||||
for mm_hash in mm_hashes:
|
||||
assert mm_hash in encoder_cache, f"{mm_hash} missing in encoder_cache"
|
||||
assert encoder_cache[mm_hash].is_cuda, (
|
||||
f"{mm_hash} cache is in {encoder_cache[mm_hash].device}"
|
||||
)
|
||||
assert torch.allclose(
|
||||
encoder_cache[mm_hash].cpu(), saved_caches[mm_hash]
|
||||
), f"{mm_hash} cache saved and loaded tesnor are not the same"
|
||||
|
||||
def test_start_load_caches_skip_existing(
|
||||
self, mock_vllm_config_producer, mock_vllm_config_consumer, temp_storage
|
||||
):
|
||||
"""Test cache loading skips already cached items."""
|
||||
# Setup: producer saves cache
|
||||
producer = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
mm_hash = "existing_hash"
|
||||
saved_cache = torch.randn(10, 768)
|
||||
producer.save_caches({mm_hash: saved_cache}, mm_hash)
|
||||
|
||||
# Consumer setup
|
||||
consumer = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_consumer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
metadata = ECExampleConnectorMetadata()
|
||||
metadata.add_mm_data(MMMeta.make_meta(mm_hash, 100))
|
||||
consumer.bind_connector_metadata(metadata)
|
||||
|
||||
# Pre-populate encoder_cache with different value
|
||||
existing_cache = torch.randn(5, 512)
|
||||
encoder_cache = {mm_hash: existing_cache}
|
||||
|
||||
# Load (should skip since already exists)
|
||||
with patch("safetensors.torch.load_file") as mock_load:
|
||||
consumer.start_load_caches(encoder_cache=encoder_cache)
|
||||
# Should not call load_file since cache exists
|
||||
mock_load.assert_not_called()
|
||||
|
||||
# Verify original cache unchanged
|
||||
assert torch.equal(encoder_cache[mm_hash], existing_cache)
|
||||
|
||||
def test_start_load_caches_empty_metadata(self, mock_vllm_config_consumer):
|
||||
"""Test loading with empty metadata does nothing."""
|
||||
consumer = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_consumer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
# Setup empty metadata
|
||||
metadata = ECExampleConnectorMetadata()
|
||||
consumer.bind_connector_metadata(metadata)
|
||||
|
||||
# Load (should not raise)
|
||||
encoder_cache: dict[str, torch.Tensor] = {}
|
||||
consumer.start_load_caches(encoder_cache=encoder_cache)
|
||||
|
||||
# Cache should remain empty
|
||||
assert len(encoder_cache) == 0
|
||||
|
||||
|
||||
class TestFilenameGeneration:
|
||||
"""Test filename and path generation."""
|
||||
|
||||
def test_generate_foldername(self, mock_vllm_config_producer, temp_storage):
|
||||
"""Test folder name generation."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
mm_hash = "test_folder_hash"
|
||||
folder = connector._generate_foldername_debug(mm_hash)
|
||||
|
||||
assert folder == os.path.join(temp_storage, mm_hash)
|
||||
assert os.path.isdir(folder) # Should be created
|
||||
|
||||
def test_generate_filename(self, mock_vllm_config_producer, temp_storage):
|
||||
"""Test filename generation."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
mm_hash = "test_file_hash"
|
||||
filename = connector._generate_filename_debug(mm_hash)
|
||||
|
||||
expected = os.path.join(temp_storage, mm_hash, "encoder_cache.safetensors")
|
||||
assert filename == expected
|
||||
assert os.path.isdir(os.path.dirname(filename)) # Folder created
|
||||
|
||||
def test_generate_filename_consistency(self, mock_vllm_config_producer):
|
||||
"""Test filename generation is consistent."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
mm_hash = "consistency_hash"
|
||||
filename1 = connector._generate_filename_debug(mm_hash)
|
||||
filename2 = connector._generate_filename_debug(mm_hash)
|
||||
|
||||
assert filename1 == filename2
|
||||
|
||||
|
||||
class TestMetadataBindingLifecycle:
|
||||
"""Test metadata binding and clearing lifecycle."""
|
||||
|
||||
def test_bind_connector_metadata(self, mock_vllm_config_consumer):
|
||||
"""Test binding connector metadata."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_consumer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
metadata = ECExampleConnectorMetadata()
|
||||
metadata.add_mm_data(MMMeta.make_meta("hash_1", 100))
|
||||
|
||||
connector.bind_connector_metadata(metadata)
|
||||
|
||||
assert connector._connector_metadata is metadata
|
||||
|
||||
def test_clear_connector_metadata(self, mock_vllm_config_consumer):
|
||||
"""Test clearing connector metadata."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_consumer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
metadata = ECExampleConnectorMetadata()
|
||||
connector.bind_connector_metadata(metadata)
|
||||
|
||||
connector.clear_connector_metadata()
|
||||
|
||||
assert connector._connector_metadata is None
|
||||
|
||||
def test_get_connector_metadata(self, mock_vllm_config_consumer):
|
||||
"""Test getting connector metadata."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_consumer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
metadata = ECExampleConnectorMetadata()
|
||||
connector.bind_connector_metadata(metadata)
|
||||
|
||||
retrieved = connector._get_connector_metadata()
|
||||
|
||||
assert retrieved is metadata
|
||||
|
||||
def test_get_connector_metadata_not_set(self, mock_vllm_config_consumer):
|
||||
"""Test getting metadata when not set raises."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_consumer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
with pytest.raises(AssertionError):
|
||||
connector._get_connector_metadata()
|
||||
|
||||
|
||||
class TestEdgeCases:
|
||||
"""Test edge cases and error handling."""
|
||||
|
||||
def test_save_empty_cache(self, mock_vllm_config_producer):
|
||||
"""Test saving empty tensor."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
mm_hash = "empty_hash"
|
||||
encoder_cache = {mm_hash: torch.empty(0)}
|
||||
|
||||
# Should not raise
|
||||
connector.save_caches(encoder_cache, mm_hash)
|
||||
|
||||
def test_load_nonexistent_cache(self, mock_vllm_config_consumer):
|
||||
"""Test loading cache that doesn't exist raises error."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_consumer,
|
||||
role=ECConnectorRole.WORKER,
|
||||
)
|
||||
|
||||
metadata = ECExampleConnectorMetadata()
|
||||
metadata.add_mm_data(MMMeta.make_meta("nonexistent_hash", 100))
|
||||
connector.bind_connector_metadata(metadata)
|
||||
|
||||
encoder_cache: dict[str, torch.Tensor] = {}
|
||||
|
||||
# Should raise FileNotFoundError
|
||||
with pytest.raises(FileNotFoundError):
|
||||
connector.start_load_caches(encoder_cache=encoder_cache)
|
||||
|
||||
def test_has_cache_item_empty_request(self, mock_vllm_config_producer):
|
||||
"""Test has_cache_item with a nonexistent identifier."""
|
||||
connector = ECExampleConnector(
|
||||
vllm_config=mock_vllm_config_producer,
|
||||
role=ECConnectorRole.SCHEDULER,
|
||||
)
|
||||
|
||||
result = connector.has_cache_item("nonexistent_hash")
|
||||
|
||||
assert result is False
|
||||
0
third_party/vllm/tests/v1/engine/__init__.py
vendored
Normal file
0
third_party/vllm/tests/v1/engine/__init__.py
vendored
Normal file
90
third_party/vllm/tests/v1/engine/conftest.py
vendored
Normal file
90
third_party/vllm/tests/v1/engine/conftest.py
vendored
Normal file
@@ -0,0 +1,90 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from tests.v1.engine.utils import (
|
||||
FULL_STRINGS,
|
||||
NUM_PROMPT_LOGPROBS_UNDER_TEST,
|
||||
NUM_SAMPLE_LOGPROBS_UNDER_TEST,
|
||||
PROMPT_LEN,
|
||||
TOKENIZER_NAME,
|
||||
DummyOutputProcessorTestVectors,
|
||||
generate_dummy_prompt_logprobs_tensors,
|
||||
generate_dummy_sample_logprobs,
|
||||
)
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
|
||||
from ...distributed.conftest import publisher_config, random_port # noqa: F401
|
||||
|
||||
EngineCoreSampleLogprobsType = list[tuple[torch.Tensor, torch.Tensor]]
|
||||
EngineCorePromptLogprobsType = tuple[torch.Tensor, torch.Tensor]
|
||||
|
||||
|
||||
def _build_test_vectors_no_logprobs() -> DummyOutputProcessorTestVectors:
|
||||
"""Generate output processor dummy test vectors, without logprobs
|
||||
|
||||
Returns:
|
||||
DummyOutputProcessorTestVectors instance with no logprobs
|
||||
"""
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
|
||||
vllm_config = EngineArgs(model=TOKENIZER_NAME).create_engine_config()
|
||||
# Tokenize prompts under test & create dummy generated tokens
|
||||
prompt_tokens = [tokenizer(text).input_ids[:PROMPT_LEN] for text in FULL_STRINGS]
|
||||
generation_tokens = [
|
||||
tokenizer(text).input_ids[PROMPT_LEN:] for text in FULL_STRINGS
|
||||
]
|
||||
# Generate prompt strings
|
||||
prompt_strings = [
|
||||
tokenizer.decode(prompt_tokens, skip_special_tokens=True)
|
||||
for prompt_tokens in prompt_tokens
|
||||
]
|
||||
prompt_strings_len = [len(prompt_string) for prompt_string in prompt_strings]
|
||||
return DummyOutputProcessorTestVectors(
|
||||
tokenizer=tokenizer,
|
||||
vllm_config=vllm_config,
|
||||
full_tokens=[tokenizer(text).input_ids for text in FULL_STRINGS],
|
||||
prompt_tokens=prompt_tokens,
|
||||
generation_tokens=generation_tokens,
|
||||
prompt_strings=prompt_strings,
|
||||
prompt_strings_len=prompt_strings_len,
|
||||
generation_strings=[
|
||||
text[prompt_len:]
|
||||
for text, prompt_len in zip(FULL_STRINGS, prompt_strings_len)
|
||||
],
|
||||
prompt_logprobs=[],
|
||||
generation_logprobs=[],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dummy_test_vectors() -> DummyOutputProcessorTestVectors:
|
||||
"""Generate output processor dummy test vectors, with logprobs
|
||||
|
||||
Returns:
|
||||
DummyOutputProcessorTestVectors instance with logprobs
|
||||
"""
|
||||
# Build dummy test vectors without logprobs
|
||||
dtv = _build_test_vectors_no_logprobs()
|
||||
# Inject logprobs into dummy test vectors
|
||||
# data structure
|
||||
dtv.generation_logprobs = [
|
||||
generate_dummy_sample_logprobs(
|
||||
sampled_tokens_list=tokens_list,
|
||||
num_logprobs=NUM_SAMPLE_LOGPROBS_UNDER_TEST,
|
||||
tokenizer=dtv.tokenizer,
|
||||
)
|
||||
for tokens_list in dtv.generation_tokens
|
||||
]
|
||||
dtv.prompt_logprobs = [
|
||||
generate_dummy_prompt_logprobs_tensors(
|
||||
prompt_tokens_list=tokens_list,
|
||||
num_logprobs=NUM_PROMPT_LOGPROBS_UNDER_TEST,
|
||||
tokenizer=dtv.tokenizer,
|
||||
)
|
||||
for tokens_list in dtv.prompt_tokens
|
||||
]
|
||||
return dtv
|
||||
311
third_party/vllm/tests/v1/engine/test_abort_final_step.py
vendored
Normal file
311
third_party/vllm/tests/v1/engine/test_abort_final_step.py
vendored
Normal file
@@ -0,0 +1,311 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Test for the fix in PR #29987: Eagerly abort cancelled final-step requests.
|
||||
|
||||
This test verifies that when a request is aborted during its final execution
|
||||
step (when it would naturally complete), it is properly marked as aborted
|
||||
rather than being treated as normally completed.
|
||||
|
||||
The test uses a dummy KV connector to verify that the connector receives
|
||||
the correct finish status (FINISHED_ABORTED, not FINISHED_LENGTH_CAPPED).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm import SamplingParams
|
||||
from vllm.config import KVTransferConfig, VllmConfig
|
||||
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
|
||||
KVConnectorBase_V1,
|
||||
KVConnectorMetadata,
|
||||
KVConnectorRole,
|
||||
)
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
from vllm.utils.torch_utils import set_default_torch_num_threads
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
from vllm.v1.kv_cache_interface import KVCacheConfig
|
||||
from vllm.v1.request import Request
|
||||
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip(reason="V1 currently only supported on CUDA.", allow_module_level=True)
|
||||
|
||||
TEXT_PROMPT = "Hello"
|
||||
|
||||
|
||||
class DummyKVConnectorMetadata(KVConnectorMetadata):
|
||||
"""Dummy metadata for the test connector."""
|
||||
|
||||
def __init__(self):
|
||||
self.requests: list = []
|
||||
|
||||
|
||||
class DummyKVConnector(KVConnectorBase_V1):
|
||||
"""
|
||||
Dummy KV connector that captures request finish statuses to a file.
|
||||
This is used to verify the fix - without the fix, a request aborted
|
||||
during its final step would be captured as FINISHED_LENGTH_CAPPED
|
||||
instead of FINISHED_ABORTED.
|
||||
|
||||
The connector runs in a separate process, so we write statuses to a file
|
||||
that can be read by the test process.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
role: KVConnectorRole,
|
||||
kv_cache_config: KVCacheConfig | None = None,
|
||||
):
|
||||
super().__init__(vllm_config, role, kv_cache_config)
|
||||
# Get the status file path from extra config
|
||||
extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config or {}
|
||||
self.status_file = extra_config.get("status_file")
|
||||
# Log that we were initialized
|
||||
if self.status_file:
|
||||
try:
|
||||
with open(self.status_file, "a") as f:
|
||||
f.write(f"INIT:{role.name}\n")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def get_num_new_matched_tokens(
|
||||
self,
|
||||
request: Request,
|
||||
num_computed_tokens: int,
|
||||
) -> tuple[int | None, bool]:
|
||||
return (0, False)
|
||||
|
||||
def update_state_after_alloc(
|
||||
self,
|
||||
request: Request,
|
||||
blocks: Any,
|
||||
num_external_tokens: int,
|
||||
):
|
||||
pass
|
||||
|
||||
def build_connector_meta(
|
||||
self, scheduler_output: SchedulerOutput
|
||||
) -> KVConnectorMetadata:
|
||||
return DummyKVConnectorMetadata()
|
||||
|
||||
def request_finished(
|
||||
self,
|
||||
request: Request,
|
||||
block_ids: list[int],
|
||||
) -> tuple[bool, dict[str, Any] | None]:
|
||||
"""Capture the request status when finished by writing to a file."""
|
||||
if self.status_file:
|
||||
try:
|
||||
with open(self.status_file, "a") as f:
|
||||
# Write the status name (e.g., "FINISHED_ABORTED")
|
||||
f.write(f"{request.status.name}\n")
|
||||
except Exception as e:
|
||||
# Log but don't fail - this is just test instrumentation
|
||||
print(f"[DummyKVConnector] Failed to write status: {e}")
|
||||
return False, None
|
||||
|
||||
def start_load_kv(self, forward_context: Any, **kwargs: Any) -> None:
|
||||
pass
|
||||
|
||||
def wait_for_layer_load(self, layer_name: str) -> None:
|
||||
pass
|
||||
|
||||
def save_kv_layer(
|
||||
self,
|
||||
layer_name: str,
|
||||
kv_layer: Any,
|
||||
attn_metadata: Any,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
def wait_for_save(self):
|
||||
pass
|
||||
|
||||
|
||||
# Register the dummy connector
|
||||
KVConnectorFactory.register_connector(
|
||||
"DummyKVConnector", __name__, DummyKVConnector.__name__
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("async_scheduling", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_abort_during_final_step(async_scheduling: bool):
|
||||
"""
|
||||
Test that a request aborted during its final execution step is treated as
|
||||
aborted rather than completed.
|
||||
|
||||
This test:
|
||||
1. Monkeypatches execute_model to wait for a file to be deleted
|
||||
2. Configures a dummy KV connector to capture finish statuses
|
||||
3. Starts a request with max_tokens=1 (will complete on first decode step)
|
||||
4. Aborts the request, then deletes the file to unblock execute_model
|
||||
5. Verifies the KV connector received FINISHED_ABORTED not FINISHED_LENGTH_CAPPED
|
||||
|
||||
See https://github.com/vllm-project/vllm/pull/29987.
|
||||
|
||||
Without the fix, the KV connector would see FINISHED_LENGTH_CAPPED because
|
||||
update_from_output() would mark the request as completed before processing
|
||||
the abort. This causes KV cache blocks to not be freed properly in
|
||||
disaggregated prefill scenarios.
|
||||
|
||||
With the fix, _process_aborts_queue() runs before update_from_output(), so the
|
||||
abort takes precedence and the KV connector sees FINISHED_ABORTED.
|
||||
"""
|
||||
|
||||
# Create three temporary files:
|
||||
# 1. ready_file: deleted by execute_model to signal it has started
|
||||
# 2. block_file: execute_model waits for this to be deleted
|
||||
# 3. status_file: KV connector writes finish statuses here
|
||||
with tempfile.NamedTemporaryFile(delete=False) as f:
|
||||
ready_file = Path(f.name)
|
||||
with tempfile.NamedTemporaryFile(delete=False) as f2:
|
||||
block_file = Path(f2.name)
|
||||
with tempfile.NamedTemporaryFile(delete=False, mode="w") as f3:
|
||||
status_file = Path(f3.name)
|
||||
|
||||
try:
|
||||
# Get the original execute_model method
|
||||
from vllm.v1.worker.gpu_worker import Worker
|
||||
|
||||
original_execute_model = Worker.execute_model
|
||||
|
||||
def execute_model_with_wait(self, scheduler_output):
|
||||
# Signal that execute_model has been called by deleting ready_file
|
||||
if ready_file.exists():
|
||||
ready_file.unlink()
|
||||
|
||||
# Wait for the block file to be deleted (triggered from test after abort)
|
||||
# This runs in the worker process (after fork), so we poll the filesystem
|
||||
while block_file.exists():
|
||||
time.sleep(0.01)
|
||||
return original_execute_model(self, scheduler_output)
|
||||
|
||||
# Patch execute_model to inject the wait
|
||||
# This happens before the worker process is forked, so the patch applies there
|
||||
with patch.object(Worker, "execute_model", execute_model_with_wait):
|
||||
request_id = "test-abort-final-step"
|
||||
|
||||
# Configure engine with dummy KV connector
|
||||
# Pass the status file path so the connector can write to it
|
||||
kv_transfer_config = KVTransferConfig(
|
||||
kv_connector="DummyKVConnector",
|
||||
kv_role="kv_both",
|
||||
kv_connector_extra_config={"status_file": str(status_file)},
|
||||
)
|
||||
engine_args = AsyncEngineArgs(
|
||||
model="meta-llama/Llama-3.2-1B-Instruct",
|
||||
enforce_eager=True,
|
||||
async_scheduling=async_scheduling,
|
||||
kv_transfer_config=kv_transfer_config,
|
||||
)
|
||||
|
||||
with set_default_torch_num_threads(1):
|
||||
engine = AsyncLLM.from_engine_args(engine_args)
|
||||
|
||||
try:
|
||||
# Create a request that will complete after just 1 token
|
||||
sampling_params = SamplingParams(
|
||||
max_tokens=1,
|
||||
ignore_eos=True,
|
||||
output_kind=RequestOutputKind.DELTA,
|
||||
)
|
||||
|
||||
# Start generation in a task
|
||||
outputs = []
|
||||
|
||||
async def generate():
|
||||
async for output in engine.generate(
|
||||
request_id=request_id,
|
||||
prompt=TEXT_PROMPT,
|
||||
sampling_params=sampling_params,
|
||||
):
|
||||
outputs.append(output)
|
||||
|
||||
gen_task = asyncio.create_task(generate())
|
||||
|
||||
# Wait for execute_model to signal it has started (with timeout)
|
||||
timeout = 5.0 # 5 second timeout
|
||||
start_time = time.time()
|
||||
while ready_file.exists():
|
||||
if time.time() - start_time > timeout:
|
||||
raise TimeoutError(
|
||||
"Timeout waiting for execute_model to start. "
|
||||
"The monkeypatch may not be working correctly, "
|
||||
"for example if spawn was used instead of fork."
|
||||
)
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
# Abort the request while execute_model is blocked
|
||||
await engine.abort(request_id)
|
||||
|
||||
# Now unblock execute_model by deleting the file
|
||||
# The abort should be processed before the model output
|
||||
block_file.unlink()
|
||||
|
||||
# Wait for generation to complete
|
||||
await gen_task
|
||||
|
||||
# Give the scheduler a moment to finish cleanup
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Verify we got output
|
||||
assert len(outputs) > 0, "Should have received at least one output"
|
||||
|
||||
# The final output should have finish_reason="abort"
|
||||
final_output = outputs[-1]
|
||||
assert final_output.finished, (
|
||||
"Final output should be marked as finished"
|
||||
)
|
||||
assert final_output.outputs[0].finish_reason == "abort", (
|
||||
f"Expected finish_reason='abort' but got "
|
||||
f"'{final_output.outputs[0].finish_reason}'. "
|
||||
)
|
||||
|
||||
with open(status_file) as f4:
|
||||
status_lines = f4.read().strip().split("\n")
|
||||
# Filter for actual finish statuses (not INIT or empty lines)
|
||||
captured_statuses = [
|
||||
line
|
||||
for line in status_lines
|
||||
if line and line.startswith("FINISHED_")
|
||||
]
|
||||
|
||||
assert len(captured_statuses) >= 1, (
|
||||
f"Expected at least 1 captured finish status, got "
|
||||
f"{len(captured_statuses)}. File content: {status_lines}"
|
||||
)
|
||||
|
||||
assert "FINISHED_ABORTED" in captured_statuses, (
|
||||
f"KV connector should see FINISHED_ABORTED but got "
|
||||
f"{captured_statuses}. "
|
||||
)
|
||||
|
||||
# Verify cleanup
|
||||
assert not engine.output_processor.has_unfinished_requests()
|
||||
|
||||
finally:
|
||||
# Shutdown the engine
|
||||
engine.shutdown()
|
||||
|
||||
finally:
|
||||
# Clean up temporary files if they still exist
|
||||
if ready_file.exists():
|
||||
ready_file.unlink()
|
||||
if block_file.exists():
|
||||
block_file.unlink()
|
||||
if status_file.exists():
|
||||
status_file.unlink()
|
||||
1016
third_party/vllm/tests/v1/engine/test_async_llm.py
vendored
Normal file
1016
third_party/vllm/tests/v1/engine/test_async_llm.py
vendored
Normal file
File diff suppressed because it is too large
Load Diff
92
third_party/vllm/tests/v1/engine/test_engine_args.py
vendored
Normal file
92
third_party/vllm/tests/v1/engine/test_engine_args.py
vendored
Normal file
@@ -0,0 +1,92 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from argparse import ArgumentError
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
from vllm.usage.usage_lib import UsageContext
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.utils.hashing import _xxhash
|
||||
|
||||
|
||||
def test_prefix_caching_from_cli():
|
||||
parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
|
||||
args = parser.parse_args([])
|
||||
vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
|
||||
assert vllm_config.cache_config.enable_prefix_caching, (
|
||||
"V1 turns on prefix caching by default."
|
||||
)
|
||||
|
||||
# Turn it off possible with flag.
|
||||
args = parser.parse_args(["--no-enable-prefix-caching"])
|
||||
vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
|
||||
assert not vllm_config.cache_config.enable_prefix_caching
|
||||
|
||||
# Turn it on with flag.
|
||||
args = parser.parse_args(["--enable-prefix-caching"])
|
||||
vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
|
||||
assert vllm_config.cache_config.enable_prefix_caching
|
||||
|
||||
# default hash algorithm is "builtin"
|
||||
assert vllm_config.cache_config.prefix_caching_hash_algo == "sha256"
|
||||
|
||||
# set hash algorithm to sha256_cbor
|
||||
args = parser.parse_args(["--prefix-caching-hash-algo", "sha256_cbor"])
|
||||
vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
|
||||
assert vllm_config.cache_config.prefix_caching_hash_algo == "sha256_cbor"
|
||||
|
||||
# set hash algorithm to sha256
|
||||
args = parser.parse_args(["--prefix-caching-hash-algo", "sha256"])
|
||||
vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
|
||||
assert vllm_config.cache_config.prefix_caching_hash_algo == "sha256"
|
||||
|
||||
# an invalid hash algorithm raises an error
|
||||
parser.exit_on_error = False
|
||||
with pytest.raises(ArgumentError):
|
||||
args = parser.parse_args(["--prefix-caching-hash-algo", "invalid"])
|
||||
|
||||
|
||||
@pytest.mark.skipif(_xxhash is None, reason="xxhash not installed")
|
||||
def test_prefix_caching_xxhash_from_cli():
|
||||
parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
|
||||
|
||||
# set hash algorithm to xxhash (pickle)
|
||||
args = parser.parse_args(["--prefix-caching-hash-algo", "xxhash"])
|
||||
vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
|
||||
assert vllm_config.cache_config.prefix_caching_hash_algo == "xxhash"
|
||||
|
||||
# set hash algorithm to xxhash_cbor
|
||||
args = parser.parse_args(["--prefix-caching-hash-algo", "xxhash_cbor"])
|
||||
vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config()
|
||||
assert vllm_config.cache_config.prefix_caching_hash_algo == "xxhash_cbor"
|
||||
|
||||
|
||||
def test_defaults_with_usage_context():
|
||||
engine_args = EngineArgs(model="facebook/opt-125m")
|
||||
vllm_config: VllmConfig = engine_args.create_engine_config(UsageContext.LLM_CLASS)
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.mem_constants import GiB_bytes
|
||||
|
||||
device_memory = current_platform.get_device_total_memory()
|
||||
device_name = current_platform.get_device_name().lower()
|
||||
if device_memory >= 70 * GiB_bytes and "a100" not in device_name:
|
||||
# For GPUs like H100, H200, and MI300x with >= 70GB memory
|
||||
default_llm_tokens = 16384
|
||||
default_server_tokens = 8192
|
||||
default_max_num_seqs = 1024
|
||||
else:
|
||||
default_llm_tokens = 8192
|
||||
default_server_tokens = 2048
|
||||
default_max_num_seqs = 256
|
||||
|
||||
assert vllm_config.scheduler_config.max_num_seqs == default_max_num_seqs
|
||||
assert vllm_config.scheduler_config.max_num_batched_tokens == default_llm_tokens # noqa: E501
|
||||
|
||||
engine_args = EngineArgs(model="facebook/opt-125m")
|
||||
vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER)
|
||||
assert vllm_config.scheduler_config.max_num_seqs == default_max_num_seqs
|
||||
assert vllm_config.scheduler_config.max_num_batched_tokens == default_server_tokens # noqa: E501
|
||||
603
third_party/vllm/tests/v1/engine/test_engine_core.py
vendored
Normal file
603
third_party/vllm/tests/v1/engine/test_engine_core.py
vendored
Normal file
@@ -0,0 +1,603 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import copy
|
||||
import time
|
||||
import uuid
|
||||
from concurrent.futures import Future, ThreadPoolExecutor
|
||||
|
||||
import pytest
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm import SamplingParams
|
||||
from vllm.config import (
|
||||
CacheConfig,
|
||||
ECTransferConfig,
|
||||
KVTransferConfig,
|
||||
ModelConfig,
|
||||
SchedulerConfig,
|
||||
VllmConfig,
|
||||
)
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_default_torch_num_threads
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
from vllm.v1.engine.core import EngineCore
|
||||
from vllm.v1.executor.abstract import Executor
|
||||
from vllm.v1.executor.uniproc_executor import UniProcExecutor
|
||||
from vllm.v1.kv_cache_interface import KVCacheConfig
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
|
||||
from ...utils import create_new_process_for_each_test, multi_gpu_test
|
||||
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip(reason="V1 currently only supported on CUDA.", allow_module_level=True)
|
||||
|
||||
MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
|
||||
TOKENIZER = AutoTokenizer.from_pretrained(MODEL_NAME)
|
||||
# test_engine_core_concurrent_batches assumes exactly 12 tokens per prompt.
|
||||
# Adjust prompt if changing model to maintain 12-token length.
|
||||
PROMPT = "I am Gyoubu Masataka Oniwa"
|
||||
PROMPT_TOKENS = TOKENIZER(PROMPT).input_ids
|
||||
|
||||
_REQUEST_COUNTER = 0
|
||||
|
||||
|
||||
def make_request() -> EngineCoreRequest:
|
||||
global _REQUEST_COUNTER
|
||||
_REQUEST_COUNTER += 1
|
||||
request_id = f"request-{_REQUEST_COUNTER}"
|
||||
return EngineCoreRequest(
|
||||
request_id=request_id,
|
||||
external_req_id=f"{request_id}-{uuid.uuid4()}",
|
||||
prompt_token_ids=PROMPT_TOKENS,
|
||||
mm_features=None,
|
||||
sampling_params=SamplingParams(),
|
||||
pooling_params=None,
|
||||
arrival_time=time.time(),
|
||||
lora_request=None,
|
||||
cache_salt=None,
|
||||
data_parallel_rank=None,
|
||||
)
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_engine_core():
|
||||
"""Setup the EngineCore."""
|
||||
engine_args = EngineArgs(model=MODEL_NAME)
|
||||
vllm_config = engine_args.create_engine_config()
|
||||
executor_class = Executor.get_class(vllm_config)
|
||||
|
||||
with set_default_torch_num_threads(1):
|
||||
engine_core = EngineCore(
|
||||
vllm_config=vllm_config, executor_class=executor_class, log_stats=True
|
||||
)
|
||||
"""Test basic request lifecycle."""
|
||||
|
||||
# First request.
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(make_request()))
|
||||
assert len(engine_core.scheduler.waiting) == 1
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
|
||||
_ = engine_core.step_fn()
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 1
|
||||
|
||||
# Second request.
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(make_request()))
|
||||
assert len(engine_core.scheduler.waiting) == 1
|
||||
assert len(engine_core.scheduler.running) == 1
|
||||
|
||||
_ = engine_core.step_fn()
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 2
|
||||
|
||||
# Add two requests in a row.
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(make_request()))
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(make_request()))
|
||||
assert len(engine_core.scheduler.waiting) == 2
|
||||
assert len(engine_core.scheduler.running) == 2
|
||||
|
||||
_ = engine_core.step_fn()
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 4
|
||||
|
||||
# Loop through until they are all done.
|
||||
while (outs := engine_core.step_fn()[0].get(0)) and outs.outputs:
|
||||
pass
|
||||
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
"""Test abort cycle."""
|
||||
|
||||
# Basic abort.
|
||||
req = make_request()
|
||||
request_id = req.request_id
|
||||
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(req))
|
||||
assert len(engine_core.scheduler.waiting) == 1
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
assert engine_core.scheduler.has_unfinished_requests()
|
||||
assert not engine_core.scheduler.has_finished_requests()
|
||||
|
||||
_ = engine_core.step_fn()
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 1
|
||||
assert engine_core.scheduler.has_unfinished_requests()
|
||||
assert not engine_core.scheduler.has_finished_requests()
|
||||
|
||||
engine_core.abort_requests([request_id])
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
assert not engine_core.scheduler.has_unfinished_requests()
|
||||
assert engine_core.scheduler.has_finished_requests()
|
||||
|
||||
_ = engine_core.step_fn()
|
||||
assert not engine_core.scheduler.has_unfinished_requests()
|
||||
assert not engine_core.scheduler.has_finished_requests()
|
||||
|
||||
# Add, step, abort 1 of the 3.
|
||||
req0 = make_request()
|
||||
req1 = make_request()
|
||||
req2 = make_request()
|
||||
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(req0))
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(req1))
|
||||
assert len(engine_core.scheduler.waiting) == 2
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
|
||||
_ = engine_core.step_fn()
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 2
|
||||
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(req2))
|
||||
assert len(engine_core.scheduler.waiting) == 1
|
||||
assert len(engine_core.scheduler.running) == 2
|
||||
|
||||
_ = engine_core.step_fn()
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 3
|
||||
|
||||
# Abort just one.
|
||||
engine_core.abort_requests([req1.request_id])
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 2
|
||||
|
||||
_ = engine_core.step_fn()
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 2
|
||||
|
||||
# Abort the other requests at the same time.
|
||||
engine_core.abort_requests([req2.request_id, req0.request_id])
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
|
||||
# Sending duplicate requests with same request_id
|
||||
req0 = make_request()
|
||||
req1 = make_request()
|
||||
req0.request_id = req1.request_id = "test"
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(req0))
|
||||
|
||||
while engine_core.scheduler.has_requests():
|
||||
engine_core.step_fn()
|
||||
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(req1))
|
||||
while engine_core.scheduler.has_requests():
|
||||
engine_core.step_fn()
|
||||
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_engine_core_advanced_sampling():
|
||||
"""
|
||||
A basic end-to-end test to verify that the engine functions correctly
|
||||
when additional sampling parameters, such as top_p, min_tokens, and
|
||||
presence_penalty, are set.
|
||||
"""
|
||||
"""Setup the EngineCore."""
|
||||
engine_args = EngineArgs(model=MODEL_NAME)
|
||||
vllm_config = engine_args.create_engine_config()
|
||||
executor_class = Executor.get_class(vllm_config)
|
||||
|
||||
with set_default_torch_num_threads(1):
|
||||
engine_core = EngineCore(
|
||||
vllm_config=vllm_config, executor_class=executor_class, log_stats=True
|
||||
)
|
||||
"""Test basic request lifecycle."""
|
||||
# First request.
|
||||
request: EngineCoreRequest = make_request()
|
||||
request.sampling_params = SamplingParams(
|
||||
min_tokens=4,
|
||||
presence_penalty=1.0,
|
||||
frequency_penalty=1.0,
|
||||
repetition_penalty=0.1,
|
||||
stop_token_ids=[1001, 1002],
|
||||
)
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(request))
|
||||
|
||||
def _check_engine_state():
|
||||
assert len(engine_core.scheduler.waiting) == 1
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
# Loop through until they are all done.
|
||||
while engine_core.scheduler.has_requests():
|
||||
engine_core.step_fn()
|
||||
assert len(engine_core.scheduler.waiting) == 0
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
|
||||
_check_engine_state()
|
||||
|
||||
# Second request.
|
||||
request2 = make_request()
|
||||
request2.sampling_params = SamplingParams(
|
||||
top_p=0.99,
|
||||
top_k=50,
|
||||
)
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(request2))
|
||||
_check_engine_state()
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_engine_core_concurrent_batches():
|
||||
"""
|
||||
Test that the engine can handle multiple concurrent batches.
|
||||
"""
|
||||
|
||||
def make_request_with_max_tokens(req_id: str, max_tokens: int) -> EngineCoreRequest:
|
||||
request = make_request()
|
||||
request.request_id = req_id
|
||||
request.sampling_params.max_tokens = max_tokens
|
||||
return request
|
||||
|
||||
class DummyExecutor(UniProcExecutor):
|
||||
def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None:
|
||||
super().initialize_from_config(kv_cache_configs)
|
||||
|
||||
# Create a thread pool with a single worker
|
||||
self.thread_pool = ThreadPoolExecutor(max_workers=1)
|
||||
|
||||
def execute_model(
|
||||
self,
|
||||
scheduler_output,
|
||||
non_block=False,
|
||||
) -> Future[ModelRunnerOutput | None]:
|
||||
"""Make execute_model non-blocking."""
|
||||
|
||||
# DummyExecutor used only for testing async case.
|
||||
assert non_block
|
||||
|
||||
def _execute():
|
||||
output = self.collective_rpc("execute_model", args=(scheduler_output,))
|
||||
# Make a copy because output[0] may be reused
|
||||
# by the next batch.
|
||||
return copy.deepcopy(output[0])
|
||||
|
||||
# Use the thread pool instead of creating a new thread
|
||||
return self.thread_pool.submit(_execute)
|
||||
|
||||
def sample_tokens(
|
||||
self, grammar_output, non_block=False
|
||||
) -> Future[ModelRunnerOutput]:
|
||||
"""Make sample_tokens non-blocking."""
|
||||
|
||||
# DummyExecutor used only for testing async case.
|
||||
assert non_block
|
||||
|
||||
def _execute():
|
||||
output = self.collective_rpc("sample_tokens", args=(grammar_output,))
|
||||
# Make a copy because output[0] may be reused
|
||||
# by the next batch.
|
||||
return copy.deepcopy(output[0])
|
||||
|
||||
# Use the thread pool instead of creating a new thread
|
||||
return self.thread_pool.submit(_execute)
|
||||
|
||||
@property
|
||||
def max_concurrent_batches(self) -> int:
|
||||
return 2
|
||||
|
||||
def shutdown(self):
|
||||
if hasattr(self, "thread_pool"):
|
||||
self.thread_pool.shutdown(wait=False)
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model=MODEL_NAME,
|
||||
# To test concurrent batches.
|
||||
max_num_seqs=2,
|
||||
# Avoid all requests being scheduled once.
|
||||
enable_prefix_caching=False,
|
||||
max_num_batched_tokens=10,
|
||||
# Reduce startup time.
|
||||
enforce_eager=True,
|
||||
# Test concurrent batch behaviour independently of async scheduling.
|
||||
async_scheduling=False,
|
||||
)
|
||||
vllm_config = engine_args.create_engine_config()
|
||||
with set_default_torch_num_threads(1):
|
||||
engine_core = EngineCore(
|
||||
vllm_config=vllm_config, log_stats=False, executor_class=DummyExecutor
|
||||
)
|
||||
assert engine_core.batch_queue is not None
|
||||
|
||||
# Add two requests in a row. Each request have 12 prompt tokens.
|
||||
req0 = make_request_with_max_tokens("0", 5)
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(req0))
|
||||
req1 = make_request_with_max_tokens("1", 5)
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(req1))
|
||||
|
||||
# Schedule Batch 1: (10, req0)
|
||||
assert engine_core.step_with_batch_queue()[0] is None
|
||||
assert len(engine_core.batch_queue) == 1
|
||||
scheduler_output = engine_core.batch_queue[-1][1]
|
||||
assert scheduler_output.num_scheduled_tokens["0"] == 10
|
||||
# num_computed_tokens should have been updated immediately.
|
||||
assert engine_core.scheduler.requests[req0.request_id].num_computed_tokens == 10
|
||||
|
||||
# Schedule Batch 2: (2, req0), (8, req1)
|
||||
assert engine_core.step_with_batch_queue()[0] == {}
|
||||
assert len(engine_core.batch_queue) == 1
|
||||
scheduler_output = engine_core.batch_queue[-1][1]
|
||||
assert scheduler_output.num_scheduled_tokens["0"] == 2
|
||||
assert scheduler_output.num_scheduled_tokens["1"] == 8
|
||||
# num_computed_tokens should have been updated immediately.
|
||||
assert engine_core.scheduler.requests["0"].num_computed_tokens == 12
|
||||
assert engine_core.scheduler.requests["1"].num_computed_tokens == 8
|
||||
|
||||
assert engine_core.scheduler.get_num_unfinished_requests() == 2
|
||||
|
||||
# Finish Batch 1 and schedule Batch 3: (4, req1).
|
||||
# Note that req0 cannot be scheduled
|
||||
# because it is in the decoding stage now.
|
||||
engine_core.step_with_batch_queue()
|
||||
assert len(engine_core.batch_queue) == 1
|
||||
scheduler_output = engine_core.batch_queue[-1][1]
|
||||
assert scheduler_output.num_scheduled_tokens["1"] == 4
|
||||
|
||||
# Finish Batch 2. Get first token of req0.
|
||||
# Schedule Batch 4: (1, req0).
|
||||
output = engine_core.step_with_batch_queue()[0].get(0)
|
||||
assert output is not None
|
||||
assert len(output.outputs) == 1
|
||||
assert engine_core.scheduler.requests[req0.request_id].num_tokens == 13
|
||||
scheduler_output = engine_core.batch_queue[-1][1]
|
||||
assert scheduler_output.num_scheduled_tokens["0"] == 1
|
||||
|
||||
# Finish Batch 3. Get first token of req1. Schedule Batch 5: (1, req1).
|
||||
output = engine_core.step_with_batch_queue()[0].get(0)
|
||||
assert output is not None
|
||||
assert len(output.outputs) == 1
|
||||
assert engine_core.scheduler.requests[req1.request_id].num_tokens == 13
|
||||
scheduler_output = engine_core.batch_queue[-1][1]
|
||||
assert scheduler_output.num_scheduled_tokens["1"] == 1
|
||||
|
||||
# Loop until req0 is finished.
|
||||
req_id = 0
|
||||
expected_num_tokens = [
|
||||
engine_core.scheduler.requests["0"].num_tokens + 1,
|
||||
engine_core.scheduler.requests["1"].num_tokens + 1,
|
||||
]
|
||||
while engine_core.scheduler.get_num_unfinished_requests() == 2:
|
||||
output = engine_core.step_with_batch_queue()[0]
|
||||
# Every step consumes an output.
|
||||
assert output is not None
|
||||
assert len(output[0].outputs) == 1
|
||||
if req_id in engine_core.scheduler.requests:
|
||||
assert (
|
||||
engine_core.scheduler.requests[req_id].num_tokens
|
||||
== expected_num_tokens[req_id]
|
||||
)
|
||||
expected_num_tokens[req_id] += 1
|
||||
req_id = (req_id + 1) % 2
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
def test_engine_core_tp():
|
||||
"""
|
||||
Test engine can initialize worker in tp properly
|
||||
"""
|
||||
|
||||
"""Setup the EngineCore."""
|
||||
engine_args = EngineArgs(
|
||||
model=MODEL_NAME,
|
||||
tensor_parallel_size=2,
|
||||
# Reduce startup time.
|
||||
enforce_eager=True,
|
||||
)
|
||||
vllm_config = engine_args.create_engine_config()
|
||||
executor_class = Executor.get_class(vllm_config)
|
||||
|
||||
with set_default_torch_num_threads(1):
|
||||
engine_core = EngineCore(
|
||||
vllm_config=vllm_config, executor_class=executor_class, log_stats=True
|
||||
)
|
||||
|
||||
def get_worker_cache_config_field(worker, key: str):
|
||||
return getattr(worker.cache_config, key)
|
||||
|
||||
num_gpu_blocks = engine_core.collective_rpc(
|
||||
get_worker_cache_config_field, args=("num_gpu_blocks",)
|
||||
)
|
||||
num_cpu_blocks = engine_core.collective_rpc(
|
||||
get_worker_cache_config_field, args=("num_cpu_blocks",)
|
||||
)
|
||||
assert all(x is not None for x in num_gpu_blocks)
|
||||
assert all(x is not None for x in num_cpu_blocks)
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_engine_core_invalid_request_id_type():
|
||||
"""Test that engine raises TypeError for non-string request_id."""
|
||||
engine_args = EngineArgs(model=MODEL_NAME)
|
||||
vllm_config = engine_args.create_engine_config()
|
||||
executor_class = Executor.get_class(vllm_config)
|
||||
|
||||
with set_default_torch_num_threads(1):
|
||||
engine_core = EngineCore(
|
||||
vllm_config=vllm_config, executor_class=executor_class, log_stats=True
|
||||
)
|
||||
|
||||
# Test with UUID object (common mistake)
|
||||
uuid_request = make_request()
|
||||
uuid_request.request_id = uuid.uuid4() # UUID object instead of string
|
||||
|
||||
with pytest.raises(TypeError, match="request_id must be a string, got.*UUID"):
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(uuid_request))
|
||||
|
||||
# Test with integer
|
||||
int_request = make_request()
|
||||
int_request.request_id = 12345
|
||||
|
||||
with pytest.raises(TypeError, match="request_id must be a string, got.*int"):
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(int_request))
|
||||
|
||||
# Test with None
|
||||
none_request = make_request()
|
||||
none_request.request_id = None
|
||||
|
||||
with pytest.raises(TypeError, match="request_id must be a string, got.*NoneType"):
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(none_request))
|
||||
|
||||
# Verify engine is still functional after errors
|
||||
valid_request = make_request()
|
||||
engine_core.add_request(*engine_core.preprocess_add_request(valid_request))
|
||||
assert len(engine_core.scheduler.waiting) == 1
|
||||
assert len(engine_core.scheduler.running) == 0
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
@pytest.mark.parametrize(
|
||||
("ec_role", "gpu_memory_utilization", "enable_prefix_caching"),
|
||||
[
|
||||
("ec_producer", 0.01, False),
|
||||
# NOTE: ec_producer never allows prefix caching
|
||||
("ec_consumer", 0.7, True),
|
||||
("ec_consumer", 0.7, False),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("use_kv_connector", [False, True])
|
||||
def test_encoder_instance_zero_kv_cache(
|
||||
ec_role: str,
|
||||
gpu_memory_utilization: float,
|
||||
enable_prefix_caching: bool,
|
||||
use_kv_connector: bool,
|
||||
):
|
||||
"""EPD (Encoder-Prefill-Decode) Encoder-cache-specific tests
|
||||
|
||||
This test verifies encoder-only instance initializes with 0 KV cache blocks.
|
||||
Under EPD disagg mode, Encoder instances (EC producer role) only execute
|
||||
vision encoder, so they don't need KV cache for text generation.
|
||||
"""
|
||||
# Form vllm config
|
||||
model_config = ModelConfig(
|
||||
model="llava-hf/llava-1.5-7b-hf", # Multimodal model
|
||||
enforce_eager=True,
|
||||
trust_remote_code=True,
|
||||
dtype="float16",
|
||||
seed=42,
|
||||
)
|
||||
scheduler_config = SchedulerConfig(
|
||||
max_num_seqs=10,
|
||||
max_num_batched_tokens=512,
|
||||
max_model_len=512,
|
||||
disable_hybrid_kv_cache_manager=True,
|
||||
is_encoder_decoder=model_config.is_encoder_decoder,
|
||||
)
|
||||
cache_config = CacheConfig(
|
||||
block_size=16,
|
||||
gpu_memory_utilization=gpu_memory_utilization,
|
||||
cache_dtype="auto",
|
||||
enable_prefix_caching=enable_prefix_caching,
|
||||
)
|
||||
kv_transfer_config = (
|
||||
KVTransferConfig(
|
||||
kv_connector="ExampleConnector",
|
||||
kv_role="kv_both",
|
||||
kv_connector_extra_config={"shared_storage_path": "local_storage"},
|
||||
)
|
||||
if use_kv_connector
|
||||
else None
|
||||
)
|
||||
ec_transfer_config = ECTransferConfig(
|
||||
ec_connector="ECExampleConnector",
|
||||
ec_role=ec_role,
|
||||
ec_connector_extra_config={"shared_storage_path": "/tmp/ec_test_encoder"},
|
||||
)
|
||||
|
||||
vllm_config = VllmConfig(
|
||||
model_config=model_config,
|
||||
cache_config=cache_config,
|
||||
scheduler_config=scheduler_config,
|
||||
kv_transfer_config=kv_transfer_config,
|
||||
ec_transfer_config=ec_transfer_config,
|
||||
)
|
||||
|
||||
executor_class = Executor.get_class(vllm_config)
|
||||
print(f"executor_class: {executor_class}")
|
||||
|
||||
with set_default_torch_num_threads(1):
|
||||
engine_core = EngineCore(
|
||||
vllm_config=vllm_config, executor_class=executor_class, log_stats=True
|
||||
)
|
||||
|
||||
# Check encoder cache manager exists
|
||||
assert engine_core.scheduler.encoder_cache_manager is not None, (
|
||||
"encoder_cache_manager should exist"
|
||||
)
|
||||
|
||||
if ec_role == "ec_producer":
|
||||
# Check 1: num_blocks should be 0
|
||||
# NOTE: num_blocks=1 as BlockPool always needs a null_block.
|
||||
kv_cache_config = engine_core.scheduler.kv_cache_manager.kv_cache_config
|
||||
print(f"kv_cache_config: {kv_cache_config}")
|
||||
assert kv_cache_config.num_blocks == 1, (
|
||||
f"ec_producer should only have 1 KV blocks, "
|
||||
f"got {kv_cache_config.num_blocks}"
|
||||
)
|
||||
|
||||
# Check 2: kv_cache_groups should be empty
|
||||
assert len(kv_cache_config.kv_cache_groups) == 0, (
|
||||
f"ec_producer should have 0 KV cache groups, "
|
||||
f"got {len(kv_cache_config.kv_cache_groups)}"
|
||||
)
|
||||
|
||||
# Check 3: kv_cache_tensors should be empty
|
||||
assert len(kv_cache_config.kv_cache_tensors) == 0, (
|
||||
f"Encoder instance should have 0 KV cache tensors, "
|
||||
f"got {len(kv_cache_config.kv_cache_tensors)}"
|
||||
)
|
||||
|
||||
# Check 4: Verify EC connector is initialized and is producer
|
||||
assert engine_core.scheduler.ec_connector is not None, (
|
||||
"Encoder instance should have EC connector"
|
||||
)
|
||||
assert engine_core.scheduler.ec_connector.is_producer, (
|
||||
"Encoder instance EC connector should be producer"
|
||||
)
|
||||
|
||||
# Check 5: Verify chunked prefill is disabled
|
||||
assert not vllm_config.scheduler_config.enable_chunked_prefill, (
|
||||
"Encoder instance should disable chunked prefill (no KV cache)"
|
||||
)
|
||||
|
||||
elif ec_role == "ec_consumer":
|
||||
# Check 1: num_blocks should be > 1
|
||||
kv_cache_config = engine_core.scheduler.kv_cache_manager.kv_cache_config
|
||||
print(f"kv_cache_config: {kv_cache_config}")
|
||||
assert kv_cache_config.num_blocks > 1, (
|
||||
f"ec_consumer should have >1 KV blocks, got {kv_cache_config.num_blocks}"
|
||||
)
|
||||
|
||||
# Check 2: kv_cache_groups should NOT be empty
|
||||
assert len(kv_cache_config.kv_cache_groups) > 0, (
|
||||
f"ec_consumer should have KV cache groups, "
|
||||
f"got {len(kv_cache_config.kv_cache_groups)}"
|
||||
)
|
||||
|
||||
# Check 3: Verify EC connector is consumer
|
||||
assert engine_core.scheduler.ec_connector is not None, (
|
||||
"Consumer instance should have EC connector"
|
||||
)
|
||||
assert not engine_core.scheduler.ec_connector.is_producer, (
|
||||
"Consumer instance EC connector should be consumer"
|
||||
)
|
||||
1214
third_party/vllm/tests/v1/engine/test_engine_core_client.py
vendored
Normal file
1214
third_party/vllm/tests/v1/engine/test_engine_core_client.py
vendored
Normal file
File diff suppressed because it is too large
Load Diff
196
third_party/vllm/tests/v1/engine/test_fast_incdec_prefix_err.py
vendored
Normal file
196
third_party/vllm/tests/v1/engine/test_fast_incdec_prefix_err.py
vendored
Normal file
@@ -0,0 +1,196 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
from vllm.v1.engine.detokenizer import IncrementalDetokenizer
|
||||
|
||||
# ruff: noqa: E501
|
||||
|
||||
|
||||
def test_fast_inc_detok_invalid_utf8_err_case():
|
||||
"""
|
||||
Test edge case where tokenizer can produce non-monotonic,
|
||||
invalid UTF-8 output, which breaks the internal state of
|
||||
tokenizers' DecodeStream.
|
||||
See https://github.com/vllm-project/vllm/issues/17448.
|
||||
|
||||
Thanks to reproducer from @fpaupier:
|
||||
https://gist.github.com/fpaupier/0ed1375bd7633c5be6c894b1c7ac1be3.
|
||||
"""
|
||||
tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-1b-it")
|
||||
|
||||
# Create a test request
|
||||
prompt_token_ids = [107, 4606, 236787, 107]
|
||||
params = SamplingParams(skip_special_tokens=True)
|
||||
request = EngineCoreRequest(
|
||||
request_id="test",
|
||||
external_req_id="test-ext",
|
||||
prompt_token_ids=prompt_token_ids,
|
||||
mm_features=None,
|
||||
sampling_params=params,
|
||||
pooling_params=None,
|
||||
arrival_time=0.0,
|
||||
lora_request=None,
|
||||
cache_salt=None,
|
||||
data_parallel_rank=None,
|
||||
)
|
||||
|
||||
detokenizer = IncrementalDetokenizer.from_new_request(tokenizer, request)
|
||||
|
||||
assert detokenizer.__class__.__name__ == "FastIncrementalDetokenizer", (
|
||||
"Should use FastIncrementalDetokenizer by default"
|
||||
)
|
||||
|
||||
# Process tokens incrementally
|
||||
test_tokens = [
|
||||
236840,
|
||||
107,
|
||||
138,
|
||||
236782,
|
||||
107,
|
||||
140,
|
||||
236775,
|
||||
6265,
|
||||
1083,
|
||||
623,
|
||||
121908,
|
||||
147418,
|
||||
827,
|
||||
107,
|
||||
140,
|
||||
236775,
|
||||
6265,
|
||||
236779,
|
||||
2084,
|
||||
1083,
|
||||
623,
|
||||
203292,
|
||||
827,
|
||||
107,
|
||||
140,
|
||||
236775,
|
||||
6265,
|
||||
236779,
|
||||
7777,
|
||||
1083,
|
||||
623,
|
||||
121908,
|
||||
147418,
|
||||
569,
|
||||
537,
|
||||
236789,
|
||||
65880,
|
||||
569,
|
||||
537,
|
||||
236789,
|
||||
62580,
|
||||
853,
|
||||
115693,
|
||||
210118,
|
||||
35178,
|
||||
16055,
|
||||
1270,
|
||||
759,
|
||||
215817,
|
||||
4758,
|
||||
1925,
|
||||
1117,
|
||||
827,
|
||||
107,
|
||||
140,
|
||||
236775,
|
||||
5654,
|
||||
1083,
|
||||
623,
|
||||
110733,
|
||||
46291,
|
||||
827,
|
||||
107,
|
||||
140,
|
||||
236775,
|
||||
5654,
|
||||
236779,
|
||||
2084,
|
||||
1083,
|
||||
623,
|
||||
136955,
|
||||
56731,
|
||||
827,
|
||||
107,
|
||||
140,
|
||||
236775,
|
||||
5654,
|
||||
236779,
|
||||
7777,
|
||||
1083,
|
||||
623,
|
||||
194776,
|
||||
2947,
|
||||
496,
|
||||
109811,
|
||||
1608,
|
||||
890,
|
||||
215817,
|
||||
4758,
|
||||
1925,
|
||||
1117,
|
||||
2789,
|
||||
432,
|
||||
398,
|
||||
602,
|
||||
31118,
|
||||
569,
|
||||
124866,
|
||||
134772,
|
||||
509,
|
||||
19478,
|
||||
1640,
|
||||
33779,
|
||||
236743,
|
||||
236770,
|
||||
236819,
|
||||
236825,
|
||||
236771,
|
||||
432,
|
||||
398,
|
||||
432,
|
||||
237167,
|
||||
827,
|
||||
107,
|
||||
140,
|
||||
236775,
|
||||
77984,
|
||||
1083,
|
||||
623,
|
||||
2709,
|
||||
236745,
|
||||
2555,
|
||||
513,
|
||||
236789,
|
||||
602,
|
||||
31118,
|
||||
569,
|
||||
]
|
||||
|
||||
output = ""
|
||||
for i, token_id in enumerate(test_tokens):
|
||||
detokenizer.update([token_id], False)
|
||||
|
||||
finished = i == len(test_tokens) - 1
|
||||
output += detokenizer.get_next_output_text(finished, delta=True)
|
||||
|
||||
assert (
|
||||
output
|
||||
== r"""[
|
||||
{
|
||||
"source": "Résultats",
|
||||
"source_type": "CONCEPT",
|
||||
"source_description": "Résultats de l'analyse de l'impact des opérations israéliennes sur la frontière libanaise",
|
||||
"target": "Israël",
|
||||
"target_type": "ORGANIZATION",
|
||||
"target_description": "Pays qui a obtenu à sa frontière libanaise « un niveau de calme inédit depuis les années 1960 »",
|
||||
"relationship": "Obtention d'un niveau de"""
|
||||
)
|
||||
54
third_party/vllm/tests/v1/engine/test_init_error_messaging.py
vendored
Normal file
54
third_party/vllm/tests/v1/engine/test_init_error_messaging.py
vendored
Normal file
@@ -0,0 +1,54 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.v1.core.kv_cache_utils import check_enough_kv_cache_memory
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec
|
||||
|
||||
|
||||
def test_kv_cache_oom_no_memory():
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
config = MagicMock()
|
||||
config.model_config.max_model_len = 2048
|
||||
|
||||
spec = {
|
||||
"layer_0": FullAttentionSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=8,
|
||||
head_size=128,
|
||||
dtype="float16",
|
||||
)
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
check_enough_kv_cache_memory(config, spec, 0)
|
||||
|
||||
|
||||
def test_kv_cache_oom_insufficient_memory(monkeypatch):
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
config = MagicMock()
|
||||
config.model_config.max_model_len = 2048
|
||||
config.cache_config.block_size = 16
|
||||
config.parallel_config.tensor_parallel_size = 1
|
||||
config.parallel_config.pipeline_parallel_size = 1
|
||||
config.parallel_config.decode_context_parallel_size = 1
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm.v1.core.kv_cache_utils.max_memory_usage_bytes",
|
||||
lambda c, s: 100 * 1024**3, # 100 GiB
|
||||
)
|
||||
|
||||
spec = {
|
||||
"layer_0": FullAttentionSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=8,
|
||||
head_size=128,
|
||||
dtype="float16",
|
||||
)
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
check_enough_kv_cache_memory(config, spec, 1024**3) # 1 GiB
|
||||
237
third_party/vllm/tests/v1/engine/test_llm_engine.py
vendored
Normal file
237
third_party/vllm/tests/v1/engine/test_llm_engine.py
vendored
Normal file
@@ -0,0 +1,237 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import random
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm import LLM
|
||||
from vllm.sampling_params import SamplingParams, StructuredOutputsParams
|
||||
from vllm.v1.metrics.reader import Counter, Gauge, Histogram, Metric, Vector
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tests.conftest import VllmRunner
|
||||
else:
|
||||
VllmRunner = object
|
||||
|
||||
MODEL = "facebook/opt-125m"
|
||||
DTYPE = "half"
|
||||
|
||||
|
||||
def _vllm_model(
|
||||
apc: bool,
|
||||
vllm_runner: type[VllmRunner],
|
||||
*,
|
||||
skip_tokenizer_init: bool = False,
|
||||
):
|
||||
"""Set up VllmRunner instance."""
|
||||
return vllm_runner(
|
||||
MODEL,
|
||||
dtype=DTYPE,
|
||||
max_model_len=128,
|
||||
enforce_eager=True,
|
||||
enable_prefix_caching=apc,
|
||||
gpu_memory_utilization=0.5,
|
||||
skip_tokenizer_init=skip_tokenizer_init,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
# Function scope decouples tests & allows
|
||||
# env var adjustment via monkeypatch
|
||||
scope="function",
|
||||
# Prefix caching
|
||||
params=[False, True],
|
||||
)
|
||||
def vllm_model(vllm_runner, request):
|
||||
"""VllmRunner test fixture parameterized by APC True/False."""
|
||||
with _vllm_model(request.param, vllm_runner) as vllm_model:
|
||||
yield vllm_model
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def vllm_model_apc(vllm_runner):
|
||||
"""VllmRunner test fixture with APC."""
|
||||
with _vllm_model(True, vllm_runner) as vllm_model:
|
||||
yield vllm_model
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
# Function scope decouples tests & allows
|
||||
# env var adjustment via monkeypatch
|
||||
scope="function",
|
||||
# Prefix caching
|
||||
params=[False, True],
|
||||
)
|
||||
def vllm_model_skip_tokenizer_init(vllm_runner, request):
|
||||
"""VllmRunner test fixture with APC."""
|
||||
with _vllm_model(
|
||||
request.param,
|
||||
vllm_runner,
|
||||
skip_tokenizer_init=True,
|
||||
) as vllm_model:
|
||||
yield vllm_model
|
||||
|
||||
|
||||
def _get_test_sampling_params(
|
||||
prompt_list: list[str],
|
||||
seed: int | None = 42,
|
||||
structured_outputs: bool = False,
|
||||
) -> tuple[list[SamplingParams], list[int]]:
|
||||
"""Generate random sampling params for a batch."""
|
||||
|
||||
def get_mostly_n_gt1() -> int:
|
||||
r"""Mostly n \in [2,20], ~1/3 n=1"""
|
||||
x = random.randint(0, 28)
|
||||
if x < 10:
|
||||
return 1
|
||||
else:
|
||||
return x - 8
|
||||
|
||||
n_list = [get_mostly_n_gt1() for _ in range(len(prompt_list))]
|
||||
# High temperature to maximize the chance of unique completions
|
||||
return [
|
||||
SamplingParams(
|
||||
temperature=0.95,
|
||||
top_p=0.95,
|
||||
n=n,
|
||||
seed=seed,
|
||||
structured_outputs=StructuredOutputsParams(regex="[0-9]+")
|
||||
if structured_outputs
|
||||
else None,
|
||||
)
|
||||
for n in n_list
|
||||
], n_list
|
||||
|
||||
|
||||
def test_compatibility_with_skip_tokenizer_init(
|
||||
vllm_model_skip_tokenizer_init: VllmRunner,
|
||||
example_prompts: list[str],
|
||||
):
|
||||
# Case 1: Structured output request should raise an error.
|
||||
sampling_params_list, _ = _get_test_sampling_params(
|
||||
example_prompts,
|
||||
structured_outputs=True,
|
||||
)
|
||||
llm: LLM = vllm_model_skip_tokenizer_init.llm
|
||||
with pytest.raises(ValueError):
|
||||
_ = llm.generate(example_prompts, sampling_params_list)
|
||||
|
||||
|
||||
def test_parallel_sampling(vllm_model, example_prompts) -> None:
|
||||
"""Test passes if parallel sampling `n>1` yields `n` unique completions.
|
||||
|
||||
Args:
|
||||
vllm_model: VllmRunner instance under test.
|
||||
example_prompt: test fixture providing prompts for testing.
|
||||
"""
|
||||
sampling_params_list, n_list = _get_test_sampling_params(example_prompts)
|
||||
llm: LLM = vllm_model.llm
|
||||
outputs = llm.generate(example_prompts, sampling_params_list)
|
||||
|
||||
# Validate each request response
|
||||
for out, n in zip(outputs, n_list):
|
||||
completion_counts: dict[str, int] = {}
|
||||
# Assert correct number of completions
|
||||
assert len(out.outputs) == n, f"{len(out.outputs)} completions; {n} expected."
|
||||
for idx in range(n):
|
||||
comp = out.outputs[idx]
|
||||
# Assert correct completion indices
|
||||
assert comp.index == idx, f"Index {comp.index}; expected {idx}."
|
||||
text = comp.text
|
||||
completion_counts[text] = completion_counts.get(text, 0) + 1
|
||||
# Assert unique completions
|
||||
if len(completion_counts) != n:
|
||||
repeats = {txt: num for (txt, num) in completion_counts.items() if num > 1}
|
||||
raise AssertionError(
|
||||
f"{len(completion_counts)} unique completions; expected"
|
||||
f" {n}. Repeats: {repeats}"
|
||||
)
|
||||
|
||||
|
||||
def test_engine_metrics(vllm_runner, example_prompts):
|
||||
max_tokens = 100
|
||||
# Use spec decoding to test num_accepted_tokens_per_pos
|
||||
speculative_config = {
|
||||
"method": "ngram",
|
||||
"prompt_lookup_max": 5,
|
||||
"prompt_lookup_min": 3,
|
||||
"num_speculative_tokens": 5,
|
||||
}
|
||||
|
||||
with vllm_runner(
|
||||
MODEL,
|
||||
speculative_config=speculative_config,
|
||||
disable_log_stats=False,
|
||||
) as vllm_model:
|
||||
llm: LLM = vllm_model.llm
|
||||
sampling_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
|
||||
outputs = llm.generate(example_prompts, sampling_params)
|
||||
|
||||
n_prompts = len(example_prompts)
|
||||
assert len(outputs) == n_prompts
|
||||
|
||||
total_tokens = 0
|
||||
for out in outputs:
|
||||
assert len(out.outputs) == 1
|
||||
total_tokens += len(out.outputs[0].token_ids)
|
||||
assert total_tokens == max_tokens * n_prompts
|
||||
|
||||
metrics = llm.get_metrics()
|
||||
|
||||
def find_metric(name) -> list[Metric]:
|
||||
found = []
|
||||
for metric in metrics:
|
||||
if metric.name == name:
|
||||
found.append(metric)
|
||||
return found
|
||||
|
||||
num_requests_running = find_metric("vllm:num_requests_running")
|
||||
assert len(num_requests_running) == 1
|
||||
assert isinstance(num_requests_running[0], Gauge)
|
||||
assert num_requests_running[0].value == 0.0
|
||||
|
||||
generation_tokens = find_metric("vllm:generation_tokens")
|
||||
assert len(generation_tokens) == 1
|
||||
assert isinstance(generation_tokens[0], Counter)
|
||||
assert generation_tokens[0].value == total_tokens
|
||||
|
||||
request_generation_tokens = find_metric("vllm:request_generation_tokens")
|
||||
assert len(request_generation_tokens) == 1
|
||||
assert isinstance(request_generation_tokens[0], Histogram)
|
||||
assert "+Inf" in request_generation_tokens[0].buckets
|
||||
assert request_generation_tokens[0].buckets["+Inf"] == n_prompts
|
||||
assert request_generation_tokens[0].count == n_prompts
|
||||
assert request_generation_tokens[0].sum == total_tokens
|
||||
|
||||
num_accepted_tokens_per_pos = find_metric(
|
||||
"vllm:spec_decode_num_accepted_tokens_per_pos"
|
||||
)
|
||||
assert len(num_accepted_tokens_per_pos) == 1
|
||||
assert isinstance(num_accepted_tokens_per_pos[0], Vector)
|
||||
assert len(num_accepted_tokens_per_pos[0].values) == 5
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["meta-llama/Llama-3.2-1B-Instruct"])
|
||||
def test_skip_tokenizer_initialization(model: str):
|
||||
# This test checks if the flag skip_tokenizer_init skips the initialization
|
||||
# of tokenizer and detokenizer. The generated output is expected to contain
|
||||
# token ids.
|
||||
llm = LLM(
|
||||
model=model,
|
||||
skip_tokenizer_init=True,
|
||||
enforce_eager=True,
|
||||
)
|
||||
sampling_params = SamplingParams(prompt_logprobs=True, detokenize=True)
|
||||
|
||||
with pytest.raises(ValueError, match="`skip_tokenizer_init=True`"):
|
||||
llm.generate("abc", sampling_params)
|
||||
|
||||
outputs = llm.generate(
|
||||
{"prompt_token_ids": [1, 2, 3]}, sampling_params=sampling_params
|
||||
)
|
||||
assert len(outputs) > 0
|
||||
completions = outputs[0].outputs
|
||||
assert len(completions) > 0
|
||||
assert completions[0].text == ""
|
||||
assert completions[0].token_ids
|
||||
1338
third_party/vllm/tests/v1/engine/test_output_processor.py
vendored
Normal file
1338
third_party/vllm/tests/v1/engine/test_output_processor.py
vendored
Normal file
File diff suppressed because it is too large
Load Diff
83
third_party/vllm/tests/v1/engine/test_parallel_sampling.py
vendored
Normal file
83
third_party/vllm/tests/v1/engine/test_parallel_sampling.py
vendored
Normal file
@@ -0,0 +1,83 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from vllm import SamplingParams
|
||||
from vllm.outputs import CompletionOutput
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
from vllm.v1.engine.parallel_sampling import ParentRequest
|
||||
|
||||
|
||||
def test_parent_request_to_output_stream() -> None:
|
||||
parent_request = ParentRequest(make_request(SamplingParams(n=2)))
|
||||
parent_request.child_requests = {"child_id_0", "child_id_1"}
|
||||
output_0 = CompletionOutput(
|
||||
index=0, text="child 0", token_ids=[], cumulative_logprob=None, logprobs=None
|
||||
)
|
||||
output_1 = CompletionOutput(
|
||||
index=1, text="child 1", token_ids=[], cumulative_logprob=None, logprobs=None
|
||||
)
|
||||
# Request not finished
|
||||
assert ([output_0], False) == parent_request.get_outputs("child_id_0", output_0)
|
||||
assert ([output_1], False) == parent_request.get_outputs("child_id_1", output_1)
|
||||
assert ([output_0], False) == parent_request.get_outputs("child_id_0", output_0)
|
||||
assert ([output_1], False) == parent_request.get_outputs("child_id_1", output_1)
|
||||
|
||||
# output_1 finished
|
||||
output_1.finish_reason = "ended"
|
||||
assert ([output_0], False) == parent_request.get_outputs("child_id_0", output_0)
|
||||
assert ([output_1], False) == parent_request.get_outputs("child_id_1", output_1)
|
||||
# Finished output_1 had already returned, DO NOT returned again
|
||||
assert ([output_0], False) == parent_request.get_outputs("child_id_0", output_0)
|
||||
assert parent_request.get_outputs("child_id_1", output_1) == ([], False)
|
||||
|
||||
# output_0 finished
|
||||
output_0.finish_reason = "ended"
|
||||
assert ([output_0], True) == parent_request.get_outputs("child_id_0", output_0)
|
||||
assert parent_request.get_outputs("child_id_1", output_1) == ([], True)
|
||||
# Finished output_0 had already returned, DO NOT returned again
|
||||
assert parent_request.get_outputs("child_id_0", output_0) == ([], True)
|
||||
assert parent_request.get_outputs("child_id_1", output_1) == ([], True)
|
||||
|
||||
|
||||
def test_parent_request_to_output_final_only() -> None:
|
||||
parent_request = ParentRequest(
|
||||
make_request(SamplingParams(n=2, output_kind=RequestOutputKind.FINAL_ONLY))
|
||||
)
|
||||
parent_request.child_requests = {"child_id_0", "child_id_1"}
|
||||
output_0 = CompletionOutput(
|
||||
index=0, text="child 0", token_ids=[], cumulative_logprob=None, logprobs=None
|
||||
)
|
||||
output_1 = CompletionOutput(
|
||||
index=1, text="child 1", token_ids=[], cumulative_logprob=None, logprobs=None
|
||||
)
|
||||
# Request not finished, return nothing
|
||||
assert parent_request.get_outputs("child_id_0", output_0) == ([], False)
|
||||
assert parent_request.get_outputs("child_id_1", output_1) == ([], False)
|
||||
# output_1 finished, but outputs won't be returned until all child requests finished
|
||||
output_1.finish_reason = "ended"
|
||||
assert parent_request.get_outputs("child_id_0", output_0) == ([], False)
|
||||
assert parent_request.get_outputs("child_id_1", output_1) == ([], False)
|
||||
# output_0 finished, as all child requests finished, the output would be returned
|
||||
output_0.finish_reason = "ended"
|
||||
assert ([output_0, output_1], True) == parent_request.get_outputs(
|
||||
"child_id_0", output_0
|
||||
)
|
||||
assert ([output_0, output_1], True) == parent_request.get_outputs(
|
||||
"child_id_1", output_1
|
||||
)
|
||||
|
||||
|
||||
def make_request(sampling_params: SamplingParams) -> EngineCoreRequest:
|
||||
return EngineCoreRequest(
|
||||
request_id="parent_id",
|
||||
external_req_id="ext_parent_id",
|
||||
prompt_token_ids=None,
|
||||
mm_features=None,
|
||||
sampling_params=sampling_params,
|
||||
pooling_params=None,
|
||||
arrival_time=0.0,
|
||||
lora_request=None,
|
||||
cache_salt=None,
|
||||
data_parallel_rank=None,
|
||||
)
|
||||
63
third_party/vllm/tests/v1/engine/test_preprocess_error_handling.py
vendored
Normal file
63
third_party/vllm/tests/v1/engine/test_preprocess_error_handling.py
vendored
Normal file
@@ -0,0 +1,63 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch.cuda
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
from vllm.v1.engine.core import EngineCore
|
||||
|
||||
MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
|
||||
|
||||
|
||||
def test_preprocess_error_handling(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that preprocessing errors are handled gracefully."""
|
||||
|
||||
if current_platform.is_rocm() or current_platform.is_xpu():
|
||||
pytest.skip(
|
||||
"Skipped on ROCm/XPU: this test only works with 'fork', "
|
||||
"but ROCm/XPU uses 'spawn'."
|
||||
)
|
||||
|
||||
assert not torch.cuda.is_initialized(), (
|
||||
"fork needs to be used for the engine "
|
||||
"core process and this isn't possible if cuda is already initialized"
|
||||
)
|
||||
|
||||
# Store original method to call for non-failing requests
|
||||
original_preprocess = EngineCore.preprocess_add_request
|
||||
|
||||
# Monkeypatch to make preprocess_add_request raise an exception
|
||||
# only for requests with "FAIL" in the first token
|
||||
def conditional_failing_preprocess(self, request: EngineCoreRequest):
|
||||
# Fail if the first token id is 333
|
||||
if request.prompt_token_ids and request.prompt_token_ids[0] == 333:
|
||||
raise ValueError("Simulated preprocessing error!")
|
||||
return original_preprocess(self, request)
|
||||
|
||||
monkeypatch.setattr(
|
||||
EngineCore, "preprocess_add_request", conditional_failing_preprocess
|
||||
)
|
||||
|
||||
llm = LLM(model=MODEL_NAME)
|
||||
|
||||
# Create a failing request by crafting a request with an invalid token
|
||||
# We need to use a direct approach since LLM.generate tokenizes for us
|
||||
from vllm.inputs import TokensPrompt
|
||||
|
||||
# This should raise an exception due to the preprocessing failure
|
||||
# Special token id to trigger the failure
|
||||
failing_prompt = TokensPrompt(prompt_token_ids=[333])
|
||||
outputs = llm.generate(failing_prompt, SamplingParams(max_tokens=10)) # type: ignore
|
||||
assert len(outputs) == 1
|
||||
assert len(outputs[0].outputs[0].token_ids) == 0
|
||||
assert outputs[0].finished
|
||||
assert outputs[0].outputs[0].finish_reason == "error"
|
||||
|
||||
# Verify the engine is still functional with a normal request
|
||||
outputs = llm.generate("Hello, my name is", SamplingParams(max_tokens=10))
|
||||
assert len(outputs) == 1
|
||||
assert len(outputs[0].outputs[0].token_ids) > 0
|
||||
assert outputs[0].outputs[0].finish_reason in ("stop", "length")
|
||||
411
third_party/vllm/tests/v1/engine/utils.py
vendored
Normal file
411
third_party/vllm/tests/v1/engine/utils.py
vendored
Normal file
@@ -0,0 +1,411 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from typing import TypeAlias
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
|
||||
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
from vllm.v1.engine import EngineCoreOutput, FinishReason
|
||||
from vllm.v1.outputs import LogprobsLists, LogprobsTensors
|
||||
|
||||
GeneralTokenizerType: TypeAlias = PreTrainedTokenizer | PreTrainedTokenizerFast
|
||||
|
||||
# Number of sample logprobs to request when testing sample logprobs
|
||||
NUM_SAMPLE_LOGPROBS_UNDER_TEST = 5
|
||||
# Number of prompt logprobs to request when testing prompt logprobs
|
||||
NUM_PROMPT_LOGPROBS_UNDER_TEST = 7
|
||||
|
||||
TOKENIZER_NAME = "meta-llama/Llama-3.2-1B"
|
||||
|
||||
FULL_STRINGS = [
|
||||
"My name is Robert from Neural Magic and I love working on vLLM so much!",
|
||||
"Red Hat is the best open source company by far across Linux, K8s, and AI.",
|
||||
"Nick is the name of my brother in addition to my colleague from Red Hat.",
|
||||
]
|
||||
STOP_STRINGS = ["I love working on", "company by far", "brother in"]
|
||||
PROMPT_LEN = 5
|
||||
|
||||
random.seed(42)
|
||||
|
||||
|
||||
def _create_random_top_logprob_test_vector(
|
||||
num_logprobs: int,
|
||||
lower: float,
|
||||
upper: float,
|
||||
) -> torch.Tensor:
|
||||
"""Create a random vector of top logprob float values.
|
||||
|
||||
Use to create fake sample logprobs for testing.
|
||||
|
||||
Note that a real production scenario would require
|
||||
logprobs to be sorted in descending order, something
|
||||
which is omitted in this function.
|
||||
|
||||
Args:
|
||||
num_logprobs: number of top logprobs
|
||||
lower: lower range of logprob float values
|
||||
upper: upper range of logprob float values
|
||||
|
||||
Returns:
|
||||
1D length-`num_logprobs` torch Tensor of float logprob values
|
||||
"""
|
||||
return torch.rand(num_logprobs) * (upper - lower) + lower
|
||||
|
||||
|
||||
def _create_random_top_logprob_test_matrix(
|
||||
shape: tuple,
|
||||
lower: float,
|
||||
upper: float,
|
||||
) -> torch.Tensor:
|
||||
"""Create a random matrix of top logprob float values.
|
||||
|
||||
Use to create fake prompt logprobs for testing.
|
||||
|
||||
Note that a real production scenario would require
|
||||
logprobs to be sorted in descending order along rows,
|
||||
something which is omitted in this function.
|
||||
|
||||
Args:
|
||||
shape: (num_tokens,num_logprobs) tuple representing
|
||||
matrix shape
|
||||
lower: lower range of logprob float values
|
||||
upper: upper range of logprob float values
|
||||
|
||||
Returns:
|
||||
2D num_tokens x num_logprobs torch Tensor of float logprob values
|
||||
"""
|
||||
return torch.rand(*shape) * (upper - lower) + lower
|
||||
|
||||
|
||||
def _create_random_top_token_test_vector(
|
||||
num_logprobs: int,
|
||||
lower: int,
|
||||
upper: int,
|
||||
sampled_token_id: int,
|
||||
adjust_num_logprobs: bool = True,
|
||||
) -> tuple[torch.Tensor, int]:
|
||||
"""Create a random vector of top logprob token indices
|
||||
|
||||
Use to create fake sample logprobs for testing. The sampled token
|
||||
ID must always be one of the top logprobs, which this dummy test
|
||||
vector generator enforces. OpenAI API
|
||||
compatible engines must be able to return an additional sample
|
||||
logprob for the sampled token if the sampled token was not
|
||||
among the top sample logprobs; `adjust_num_logprobs` emulates
|
||||
this behavior by increasing the vector length by 1 if
|
||||
`adjust_num_logprobs` is set.
|
||||
|
||||
Args:
|
||||
num_logprobs: number of top logprobs
|
||||
lower: lower range of token ids
|
||||
upper: upper range of token ids
|
||||
sampled_token_id: the token actually sampled
|
||||
adjust_num_logprobs: if True, emulate situation where sampled
|
||||
token logprob must be injected into top
|
||||
logprobs
|
||||
|
||||
Returns:
|
||||
1D length-x torch Tensor of token ids where x is
|
||||
`num_logprobs+1` if `adjust_num_logprobs` and
|
||||
`num_logprobs` otherwise
|
||||
sampled_token_rank: the rank of sampled_token_id in the vocab
|
||||
vector when sorted in descending order by
|
||||
logprob
|
||||
"""
|
||||
|
||||
# Calculate the final number of logprobs required
|
||||
total_logprobs = num_logprobs + 1 if adjust_num_logprobs else num_logprobs
|
||||
|
||||
# Generate random indices using torch
|
||||
choice_tensor = torch.randperm(upper - lower)[:total_logprobs] + lower
|
||||
|
||||
# Ensure the sampled token ID is included in the tensor
|
||||
choice_tensor[0] = sampled_token_id
|
||||
|
||||
# Check if the sampled_token_id occurs in choice_tensor[1:]
|
||||
if sampled_token_id in choice_tensor[1:]:
|
||||
sampled_token_rank = (
|
||||
(choice_tensor[1:] == sampled_token_id).nonzero(as_tuple=True)[0].item()
|
||||
)
|
||||
else:
|
||||
# If not found, assign a random int between num_logprobs and 50700
|
||||
sampled_token_rank = random.randint(num_logprobs, 50700)
|
||||
|
||||
return choice_tensor, sampled_token_rank
|
||||
|
||||
|
||||
def _create_random_top_token_test_matrix(
|
||||
shape: tuple[int, int],
|
||||
lower: int,
|
||||
upper: int,
|
||||
tokens_list: list[int],
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Create a random matrix of top logprob token indices
|
||||
|
||||
Use to create fake prompt logprobs for testing.
|
||||
|
||||
Token ids are generated randomly and sampled without
|
||||
replacement.
|
||||
|
||||
Args:
|
||||
shape: (num_tokens, num_logprobs) tuple representing
|
||||
matrix shape
|
||||
lower: lower range of token ids
|
||||
upper: upper range of token ids
|
||||
|
||||
Returns:
|
||||
tuple containing:
|
||||
- 2D num_tokens x num_logprobs+1 torch Tensor of token ids
|
||||
- 1D tensor of ranks of prompt tokens in their respective
|
||||
rows, or random values
|
||||
"""
|
||||
num_elements = shape[0] * shape[1]
|
||||
choice_tensor = torch.randperm(upper - lower)[:num_elements] + lower
|
||||
matrix = torch.cat(
|
||||
(
|
||||
torch.tensor(tokens_list, dtype=torch.int).unsqueeze(-1),
|
||||
choice_tensor.view(shape),
|
||||
),
|
||||
dim=1,
|
||||
)
|
||||
|
||||
# Initialize the tensor for storing the ranks
|
||||
prompt_token_ranks = torch.empty(shape[0], dtype=torch.int)
|
||||
|
||||
# Iterate over each row to check presence of
|
||||
# tokens_list[rdx] and determine its index
|
||||
for rdx in range(shape[0]):
|
||||
row = matrix[rdx, 1:] # Skip the first column as it contains the token list
|
||||
token_index = (row == tokens_list[rdx]).nonzero(as_tuple=True)[0]
|
||||
if token_index.numel() > 0:
|
||||
prompt_token_ranks[rdx] = token_index.item()
|
||||
else:
|
||||
prompt_token_ranks[rdx] = random.randint(shape[1], 50700)
|
||||
|
||||
return matrix, prompt_token_ranks
|
||||
|
||||
|
||||
def decode_token(
|
||||
tok_id: int,
|
||||
tokenizer: PreTrainedTokenizer,
|
||||
) -> 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,
|
||||
tokenizer: PreTrainedTokenizer,
|
||||
) -> list[tuple[list[int], list[float], int]]:
|
||||
"""Generate dummy sample logprobs
|
||||
|
||||
Generate a test data structure which imitates the list of sample logprobs
|
||||
which would be assembled in the engine core during decode phase.
|
||||
|
||||
Args:
|
||||
sampled_tokens_list: list of sampled tokens
|
||||
num_logprobs: return `num_logprobs` or `num_logprobs+1` logprobs per token
|
||||
tokenizer: model tokenizer to use for detokenization
|
||||
|
||||
Returns
|
||||
list of (top token ids vector, logprobs vector, sampled token rank)
|
||||
Python lists tuples; in each tuple the logprobs and top token ids
|
||||
vectors have the same length which is either `num_logprobs` or
|
||||
`num_logprobs+1`. Sampled token rank is the rank (index+1) of the
|
||||
sampled token within the vocab vector when sorted by logprob in
|
||||
descending order.
|
||||
"""
|
||||
res = []
|
||||
for sampled_token_id in sampled_tokens_list:
|
||||
(
|
||||
token_vector,
|
||||
sampled_token_rank,
|
||||
) = _create_random_top_token_test_vector(
|
||||
num_logprobs, 0, len(tokenizer.vocab) - 1, sampled_token_id
|
||||
)
|
||||
|
||||
res.append(
|
||||
(
|
||||
token_vector,
|
||||
_create_random_top_logprob_test_vector(num_logprobs + 1, -100, 0),
|
||||
sampled_token_rank,
|
||||
)
|
||||
)
|
||||
|
||||
# Convert tensors in the list tuples to Python lists
|
||||
res_list_format = [
|
||||
(log_probs_tensor.tolist(), token_ids_tensor.tolist(), sampled_token_rank)
|
||||
for log_probs_tensor, token_ids_tensor, sampled_token_rank in res
|
||||
]
|
||||
|
||||
return res_list_format
|
||||
|
||||
|
||||
def generate_dummy_prompt_logprobs_tensors(
|
||||
prompt_tokens_list: list,
|
||||
num_logprobs: int,
|
||||
tokenizer: PreTrainedTokenizer,
|
||||
) -> LogprobsTensors:
|
||||
"""Generate dummy prompt logprobs tensors
|
||||
|
||||
Generate a test data structure which imitates the torch Tensors of prompt
|
||||
logprobs which would be assembled in the engine core during chunked
|
||||
prefill.
|
||||
|
||||
Args:
|
||||
prompt_tokens_list: list of prompt tokens
|
||||
num_logprobs: return `num_logprobs` logprobs per token
|
||||
tokenizer: model tokenizer to use for detokenization
|
||||
|
||||
Returns
|
||||
Single tuple of (logprobs matrix, top token ids matrix) torch Tensor,
|
||||
where both matrices have dimensions
|
||||
num_prompt_tokens x num_logprobs
|
||||
"""
|
||||
# For now, assume the whole prompt is processed in one chunk; thus,
|
||||
# the number of non-`None` prompt logprobs is `len(prompt_tokens_list)-1`.
|
||||
# Prior to injecting `None` at the beginning of prompt logprobs (which
|
||||
# happens later in the detokenizer, not here), the prompt logprobs in
|
||||
# the ith position are predicting the probability distribution of the
|
||||
# prompt token in (i+1)st position. Thus, we concat
|
||||
# `prompt_tokens_list[1:]` to the dummy token ids, just as the engine
|
||||
# would.
|
||||
num_prompt_logprobs = len(prompt_tokens_list) - 1
|
||||
(
|
||||
token_vector,
|
||||
prompt_token_ranks,
|
||||
) = _create_random_top_token_test_matrix(
|
||||
(num_prompt_logprobs, num_logprobs),
|
||||
0,
|
||||
len(tokenizer.vocab) - 1,
|
||||
prompt_tokens_list[1:],
|
||||
)
|
||||
return LogprobsTensors(
|
||||
token_vector,
|
||||
_create_random_top_logprob_test_matrix(
|
||||
(num_prompt_logprobs, num_logprobs + 1), -100, 0
|
||||
),
|
||||
prompt_token_ranks,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DummyOutputProcessorTestVectors:
|
||||
"""Dummy test vectors for output processor tests"""
|
||||
|
||||
tokenizer: GeneralTokenizerType
|
||||
vllm_config: EngineArgs
|
||||
full_tokens: list[list[int]] # Prompt + generated tokens
|
||||
prompt_tokens: list[list[int]]
|
||||
generation_tokens: list[list[int]]
|
||||
# Each request is associated with a tuple of
|
||||
# (top tokens, top logprobs, ranks) prompt logprobs tensors
|
||||
prompt_logprobs: list[LogprobsTensors]
|
||||
# Each request is associated with a sample logprobs; a request's
|
||||
# sample logprobs are a list of (top tokens, top logprobs, ranks)
|
||||
# sample logprobs tensors at each sequence position
|
||||
generation_logprobs: list[list[tuple[list[int], list[float], int]]]
|
||||
prompt_strings: list[str]
|
||||
prompt_strings_len: list[int]
|
||||
generation_strings: list[str]
|
||||
|
||||
|
||||
class MockEngineCore:
|
||||
"""Mock engine core outputs form premade tokens lists."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokens_list: list[list[int]],
|
||||
# For each request, for each sampled token offset,
|
||||
# a tuple of
|
||||
# (list of topk token ids, list of sample logprob vals, rank)
|
||||
generated_logprobs_raw: list[list[tuple[list[int], list[float], int]]]
|
||||
| None = None,
|
||||
# For each request, a tuple of
|
||||
# (prompt logprob val matrix, prompt logprob tok id matrix);
|
||||
# each matrix has dimensions
|
||||
# (num prompt toks) x (num prompt logprobs+1)
|
||||
prompt_logprobs_raw: list[LogprobsTensors] | None = None,
|
||||
eos_token_id: int | None = None,
|
||||
stop_token_ids: list[int] | None = None,
|
||||
request_ids: list[str] | None = None,
|
||||
) -> None:
|
||||
self.num_requests = len(tokens_list)
|
||||
self.tokens_list = tokens_list
|
||||
self.current_idx = 0
|
||||
self.generated_logprobs_raw = generated_logprobs_raw
|
||||
self.do_logprobs = generated_logprobs_raw is not None
|
||||
self.prompt_logprobs_raw = prompt_logprobs_raw
|
||||
self.do_prompt_logprobs = prompt_logprobs_raw is not None
|
||||
self.request_finished = [False for _ in range(self.num_requests)]
|
||||
self.eos_token_id = eos_token_id
|
||||
self.stop_token_ids = stop_token_ids
|
||||
self.request_ids = (
|
||||
request_ids
|
||||
if request_ids is not None
|
||||
else [f"request-{i}" for i in range(self.num_requests)]
|
||||
)
|
||||
|
||||
def get_outputs(self) -> list[EngineCoreOutput]:
|
||||
do_logprobs = self.do_logprobs
|
||||
do_prompt_logprobs = self.do_prompt_logprobs
|
||||
token_idx = self.current_idx
|
||||
|
||||
outputs = []
|
||||
for req_idx, token_ids in enumerate(self.tokens_list):
|
||||
if not self.request_finished[req_idx]:
|
||||
if do_logprobs:
|
||||
assert self.generated_logprobs_raw is not None
|
||||
(logprobs_token_ids_, logprobs_, sampled_token_ranks_) = (
|
||||
self.generated_logprobs_raw[req_idx][token_idx]
|
||||
)
|
||||
logprobs = LogprobsLists(
|
||||
np.array([logprobs_token_ids_]),
|
||||
np.array([logprobs_]),
|
||||
np.array([sampled_token_ranks_]),
|
||||
)
|
||||
else:
|
||||
logprobs = None
|
||||
if do_prompt_logprobs:
|
||||
if self.current_idx == 0:
|
||||
assert self.prompt_logprobs_raw is not None
|
||||
prompt_logprobs = self.prompt_logprobs_raw[req_idx]
|
||||
else:
|
||||
prompt_logprobs = None
|
||||
else:
|
||||
prompt_logprobs = None
|
||||
new_token_id = token_ids[token_idx]
|
||||
output = EngineCoreOutput(
|
||||
request_id=self.request_ids[req_idx],
|
||||
new_token_ids=[new_token_id],
|
||||
new_logprobs=logprobs,
|
||||
new_prompt_logprobs_tensors=prompt_logprobs,
|
||||
)
|
||||
if token_idx == len(token_ids) - 1:
|
||||
output.finish_reason = FinishReason.LENGTH
|
||||
self.request_finished[req_idx] = True
|
||||
if new_token_id == self.eos_token_id:
|
||||
output.finish_reason = FinishReason.STOP
|
||||
self.request_finished[req_idx] = True
|
||||
if new_token_id in (self.stop_token_ids or ()):
|
||||
output.finish_reason = FinishReason.STOP
|
||||
output.stop_reason = new_token_id
|
||||
self.request_finished[req_idx] = True
|
||||
outputs.append(output)
|
||||
|
||||
self.current_idx += 1
|
||||
return outputs
|
||||
0
third_party/vllm/tests/v1/entrypoints/__init__.py
vendored
Normal file
0
third_party/vllm/tests/v1/entrypoints/__init__.py
vendored
Normal file
173
third_party/vllm/tests/v1/entrypoints/conftest.py
vendored
Normal file
173
third_party/vllm/tests/v1/entrypoints/conftest.py
vendored
Normal file
@@ -0,0 +1,173 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_prompts():
|
||||
return [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_token_ids():
|
||||
return [
|
||||
[0],
|
||||
[0, 1],
|
||||
[0, 2, 1],
|
||||
[0, 3, 1, 2],
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_regex():
|
||||
return (
|
||||
r"((25[0-5]|(2[0-4]|1\d|[1-9]|)\d)\.){3}"
|
||||
r"(25[0-5]|(2[0-4]|1\d|[1-9]|)\d)"
|
||||
)
|
||||
|
||||
|
||||
# Note: Ensure this only uses attributes compatible with xgrammar
|
||||
@pytest.fixture
|
||||
def sample_json_schema():
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"age": {"type": "integer"},
|
||||
"skills": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
"grade": {
|
||||
"type": "string",
|
||||
"pattern": "^[A-D]$", # Regex pattern
|
||||
},
|
||||
"email": {
|
||||
"type": "string",
|
||||
"pattern": "^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}$",
|
||||
},
|
||||
"work_history": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"company": {"type": "string"},
|
||||
"duration": {
|
||||
"type": "number",
|
||||
"minimum": 0.0,
|
||||
"maximum": 100.0, # Numeric range
|
||||
},
|
||||
"position": {"type": "string"},
|
||||
},
|
||||
"required": ["company", "duration", "position"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"minItems": 0,
|
||||
"maxItems": 3,
|
||||
},
|
||||
},
|
||||
"required": ["name", "age", "skills", "grade", "email", "work_history"],
|
||||
"additionalProperties": False,
|
||||
"minProperties": 1,
|
||||
"maxProperties": 10,
|
||||
}
|
||||
|
||||
|
||||
# A schema unsupported by xgrammar
|
||||
@pytest.fixture
|
||||
def unsupported_json_schema():
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"score": {
|
||||
"type": "integer",
|
||||
"multipleOf": 5, # Numeric multiple
|
||||
},
|
||||
"tags": {
|
||||
"type": "array",
|
||||
"items": {"type": "string", "minLength": 10, "maxLength": 20},
|
||||
},
|
||||
},
|
||||
"required": ["score", "tags"],
|
||||
"additionalProperties": False,
|
||||
"patternProperties": {
|
||||
"^score$": {"type": "integer"},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_definition_json_schema():
|
||||
return {
|
||||
"$defs": {
|
||||
"Step": {
|
||||
"properties": {
|
||||
"explanation": {"title": "Explanation", "type": "string"},
|
||||
"output": {"title": "Output", "type": "string"},
|
||||
},
|
||||
"required": ["explanation", "output"],
|
||||
"title": "Step",
|
||||
"type": "object",
|
||||
}
|
||||
},
|
||||
"properties": {
|
||||
"steps": {
|
||||
"items": {"$ref": "#/$defs/Step"},
|
||||
"title": "Steps",
|
||||
"type": "array",
|
||||
},
|
||||
"final_answer": {"title": "Final Answer", "type": "string"},
|
||||
},
|
||||
"required": ["steps", "final_answer"],
|
||||
"title": "MathReasoning",
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_structured_outputs_choices():
|
||||
return [
|
||||
"Python",
|
||||
"Java",
|
||||
"JavaScript",
|
||||
"C++",
|
||||
"C#",
|
||||
"PHP",
|
||||
"TypeScript",
|
||||
"Ruby",
|
||||
"Swift",
|
||||
"Kotlin",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_sql_ebnf():
|
||||
return """
|
||||
root ::= select_statement
|
||||
select_statement ::= "SELECT" column "from" table "where" condition
|
||||
column ::= "col_1" | "col_2"
|
||||
table ::= "table_1" | "table_2"
|
||||
condition ::= column "=" number
|
||||
number ::= "1" | "2"
|
||||
"""
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_sql_lark():
|
||||
return """
|
||||
start: select_statement
|
||||
select_statement: "SELECT" column "from" table "where" condition
|
||||
column: "col_1" | "col_2"
|
||||
table: "table_1" | "table_2"
|
||||
condition: column "=" number
|
||||
number: "1" | "2"
|
||||
"""
|
||||
0
third_party/vllm/tests/v1/entrypoints/llm/__init__.py
vendored
Normal file
0
third_party/vllm/tests/v1/entrypoints/llm/__init__.py
vendored
Normal file
925
third_party/vllm/tests/v1/entrypoints/llm/test_struct_output_generate.py
vendored
Normal file
925
third_party/vllm/tests/v1/entrypoints/llm/test_struct_output_generate.py
vendored
Normal file
@@ -0,0 +1,925 @@
|
||||
# ruff: noqa: E501
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
import jsonschema
|
||||
import pytest
|
||||
import regex as re
|
||||
import torch
|
||||
from pydantic import BaseModel
|
||||
|
||||
from tests.reasoning.utils import run_reasoning_extraction
|
||||
from vllm.config import StructuredOutputsConfig
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.entrypoints.llm import LLM
|
||||
from vllm.outputs import RequestOutput
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.reasoning.abs_reasoning_parsers import ReasoningParserManager
|
||||
from vllm.sampling_params import (
|
||||
SamplingParams,
|
||||
StructuredOutputsParams,
|
||||
)
|
||||
|
||||
NGRAM_SPEC_CONFIG = {
|
||||
"model": "[ngram]",
|
||||
"num_speculative_tokens": 5,
|
||||
"prompt_lookup_max": 5,
|
||||
"prompt_lookup_min": 1,
|
||||
}
|
||||
|
||||
EAGLE_SPEC_CONFIG = {
|
||||
"method": "eagle",
|
||||
"model": "yuhuili/EAGLE-LLaMA3.1-Instruct-8B",
|
||||
"num_speculative_tokens": 5,
|
||||
}
|
||||
|
||||
PARAMS_MODELS_BACKENDS_TOKENIZER_MODE = [
|
||||
("mistralai/Ministral-8B-Instruct-2410", "xgrammar", "auto", None),
|
||||
# FIXME: Since "auto" will use Mistral tokenizer and these backends do not support
|
||||
# it, we skip these tests for now.
|
||||
# ("mistralai/Ministral-8B-Instruct-2410", "guidance", "auto", None),
|
||||
# ("mistralai/Ministral-8B-Instruct-2410", "lm-format-enforcer", "auto", None),
|
||||
("mistralai/Ministral-8B-Instruct-2410", "guidance", "hf", None),
|
||||
pytest.param(
|
||||
"mistralai/Ministral-8B-Instruct-2410",
|
||||
"lm-format-enforcer",
|
||||
"hf",
|
||||
None,
|
||||
marks=pytest.mark.skip(
|
||||
reason=(
|
||||
"Flaky: lm-format-enforcer intermittently returns"
|
||||
"incomplete JSON."
|
||||
"See https://github.com/noamgat/lm-format-enforcer/issues/169"
|
||||
)
|
||||
),
|
||||
),
|
||||
("mistralai/Ministral-8B-Instruct-2410", "xgrammar", "mistral", None),
|
||||
("Qwen/Qwen2.5-1.5B-Instruct", "xgrammar", "auto", None),
|
||||
pytest.param(
|
||||
"Qwen/Qwen2.5-1.5B-Instruct",
|
||||
"lm-format-enforcer",
|
||||
"auto",
|
||||
None,
|
||||
marks=pytest.mark.skip(
|
||||
reason=(
|
||||
"Flaky: lm-format-enforcer intermittently returns"
|
||||
"incomplete JSON."
|
||||
"See https://github.com/noamgat/lm-format-enforcer/issues/169"
|
||||
)
|
||||
),
|
||||
),
|
||||
# FIXME: This tests are flaky on CI thus disabled. Tracking in Issue #24402
|
||||
# ("mistralai/Ministral-8B-Instruct-2410", "outlines", "auto", None),
|
||||
# ("mistralai/Ministral-8B-Instruct-2410", "outlines", "mistral", None),
|
||||
# ("Qwen/Qwen2.5-1.5B-Instruct", "guidance", "auto"),
|
||||
("mistralai/Ministral-8B-Instruct-2410", "outlines", "auto", NGRAM_SPEC_CONFIG),
|
||||
("mistralai/Ministral-8B-Instruct-2410", "guidance", "hf", NGRAM_SPEC_CONFIG),
|
||||
("Qwen/Qwen2.5-1.5B-Instruct", "xgrammar", "auto", NGRAM_SPEC_CONFIG),
|
||||
("meta-llama/Meta-Llama-3.1-8B-Instruct", "xgrammar", "auto", EAGLE_SPEC_CONFIG),
|
||||
]
|
||||
|
||||
PARAMS_MODELS_TOKENIZER_MODE = [
|
||||
("mistralai/Ministral-8B-Instruct-2410", "auto"),
|
||||
("Qwen/Qwen2.5-1.5B-Instruct", "auto"),
|
||||
]
|
||||
|
||||
platform_args = {}
|
||||
if current_platform.is_rocm():
|
||||
platform_args["async_scheduling"] = False
|
||||
|
||||
|
||||
class CarType(str, Enum):
|
||||
sedan = "sedan"
|
||||
suv = "SUV"
|
||||
truck = "Truck"
|
||||
coupe = "Coupe"
|
||||
|
||||
|
||||
class CarDescription(BaseModel):
|
||||
brand: str
|
||||
model: str
|
||||
car_type: CarType
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name, backend, tokenizer_mode, speculative_config",
|
||||
PARAMS_MODELS_BACKENDS_TOKENIZER_MODE,
|
||||
)
|
||||
def test_structured_output(
|
||||
sample_json_schema: dict[str, Any],
|
||||
unsupported_json_schema: dict[str, Any],
|
||||
sample_sql_ebnf: str,
|
||||
sample_sql_lark: str,
|
||||
sample_regex: str,
|
||||
sample_structured_outputs_choices: str,
|
||||
backend: str,
|
||||
tokenizer_mode: str,
|
||||
model_name: str,
|
||||
speculative_config: dict[str, Any],
|
||||
):
|
||||
if current_platform.is_tpu() and speculative_config:
|
||||
pytest.skip("TPU does not support speculative decoding")
|
||||
|
||||
# Use a single LLM instance for several scenarios to
|
||||
# speed up the test suite.
|
||||
llm = LLM(
|
||||
model=model_name,
|
||||
enforce_eager=True,
|
||||
max_model_len=1024,
|
||||
structured_outputs_config=dict(
|
||||
backend=backend, disable_any_whitespace=backend in {"xgrammar", "guidance"}
|
||||
),
|
||||
seed=120,
|
||||
tokenizer_mode=tokenizer_mode,
|
||||
load_format="auto" if not model_name.startswith("mistralai/") else "hf",
|
||||
config_format="auto" if not model_name.startswith("mistralai/") else "hf",
|
||||
speculative_config=speculative_config,
|
||||
**platform_args,
|
||||
)
|
||||
|
||||
#
|
||||
# Test 1: Generate JSON output based on a provided schema
|
||||
#
|
||||
sampling_params = SamplingParams(
|
||||
temperature=1.0,
|
||||
max_tokens=4096,
|
||||
structured_outputs=StructuredOutputsParams(json=sample_json_schema),
|
||||
)
|
||||
|
||||
prompt = (
|
||||
"Give an example JSON for an employee profile that fits this "
|
||||
"schema. Make the response as short as possible. Schema: "
|
||||
f"{sample_json_schema}"
|
||||
)
|
||||
outputs = llm.generate(
|
||||
[prompt] * 2,
|
||||
sampling_params=sampling_params,
|
||||
use_tqdm=True,
|
||||
)
|
||||
|
||||
assert outputs is not None
|
||||
|
||||
for output in outputs:
|
||||
assert output is not None
|
||||
assert isinstance(output, RequestOutput)
|
||||
prompt = output.prompt
|
||||
|
||||
generated_text = output.outputs[0].text
|
||||
assert generated_text is not None
|
||||
if backend != "lm-format-enforcer":
|
||||
assert "\n" not in generated_text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
try:
|
||||
output_json = json.loads(generated_text)
|
||||
except json.JSONDecodeError as e:
|
||||
pytest.fail(
|
||||
f"Invalid JSON from backend={backend}: {generated_text!r}\n"
|
||||
f"Schema: {sample_json_schema}\nError: {e}"
|
||||
)
|
||||
jsonschema.validate(instance=output_json, schema=sample_json_schema)
|
||||
|
||||
#
|
||||
# Test 2: Generate JSON object without a schema
|
||||
#
|
||||
if backend != "outlines":
|
||||
sampling_params = SamplingParams(
|
||||
temperature=1.0,
|
||||
max_tokens=4096,
|
||||
n=2,
|
||||
structured_outputs=StructuredOutputsParams(json_object=True),
|
||||
)
|
||||
|
||||
outputs = llm.generate(
|
||||
prompts=(
|
||||
"Generate a JSON object with curly braces for a person with "
|
||||
"name and age fields for John Smith who is 31 years old. "
|
||||
"Make the response as short as possible."
|
||||
),
|
||||
sampling_params=sampling_params,
|
||||
use_tqdm=True,
|
||||
)
|
||||
|
||||
assert outputs is not None
|
||||
for output in outputs:
|
||||
assert output is not None
|
||||
assert isinstance(output, RequestOutput)
|
||||
|
||||
for i in range(2):
|
||||
generated_text = output.outputs[i].text
|
||||
print(generated_text)
|
||||
assert generated_text is not None
|
||||
|
||||
# Parse to verify it is a valid JSON object
|
||||
parsed_json = json.loads(generated_text)
|
||||
assert isinstance(parsed_json, dict)
|
||||
|
||||
#
|
||||
# Test 3: test a jsonschema incompatible with xgrammar
|
||||
#
|
||||
sampling_params = SamplingParams(
|
||||
temperature=1.0,
|
||||
max_tokens=4096,
|
||||
structured_outputs=StructuredOutputsParams(json=unsupported_json_schema),
|
||||
)
|
||||
if backend.startswith("xgrammar"):
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="The provided JSON schema contains features "
|
||||
"not supported by xgrammar.",
|
||||
):
|
||||
prompt = (
|
||||
f"Give an example JSON for an employee profile that "
|
||||
f"fits this schema: {unsupported_json_schema}. "
|
||||
f"Make the response as short as possible."
|
||||
)
|
||||
llm.generate(
|
||||
[prompt] * 2,
|
||||
sampling_params=sampling_params,
|
||||
use_tqdm=True,
|
||||
)
|
||||
else:
|
||||
prompt = (
|
||||
f"Give an example JSON object for a grade that "
|
||||
f"fits this schema: {unsupported_json_schema}. "
|
||||
f"Make the response as short as possible."
|
||||
)
|
||||
outputs = llm.generate(
|
||||
prompt,
|
||||
sampling_params=sampling_params,
|
||||
use_tqdm=True,
|
||||
)
|
||||
assert outputs is not None
|
||||
for output in outputs:
|
||||
assert output is not None
|
||||
assert isinstance(output, RequestOutput)
|
||||
generated_text = output.outputs[0].text
|
||||
assert generated_text is not None
|
||||
print(generated_text)
|
||||
|
||||
# Parse to verify it is valid JSON
|
||||
parsed_json = json.loads(generated_text)
|
||||
assert isinstance(parsed_json, dict)
|
||||
|
||||
if backend not in ["outlines", "lm-format-enforcer"]:
|
||||
#
|
||||
# Test 4: Generate SQL statement using EBNF grammar
|
||||
#
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.8,
|
||||
top_p=0.95,
|
||||
max_tokens=1000,
|
||||
structured_outputs=StructuredOutputsParams(grammar=sample_sql_ebnf),
|
||||
)
|
||||
outputs = llm.generate(
|
||||
(
|
||||
"Generate a sql statement that selects col_1 from "
|
||||
"table_1 where it is equal to 1. Make the response as short as "
|
||||
"possible."
|
||||
),
|
||||
sampling_params=sampling_params,
|
||||
use_tqdm=True,
|
||||
)
|
||||
|
||||
assert outputs is not None
|
||||
for output in outputs:
|
||||
assert output is not None
|
||||
assert isinstance(output, RequestOutput)
|
||||
prompt = output.prompt
|
||||
|
||||
generated_text = output.outputs[0].text
|
||||
assert generated_text is not None
|
||||
|
||||
# remove spaces for comparison b/c we removed them in the grammar
|
||||
ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(" ", "")
|
||||
|
||||
assert generated_text.strip() == ground_truth
|
||||
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
|
||||
#
|
||||
# Test 5: Generate SQL statement using Lark grammar
|
||||
#
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.8,
|
||||
top_p=0.95,
|
||||
max_tokens=1000,
|
||||
structured_outputs=StructuredOutputsParams(grammar=sample_sql_lark),
|
||||
)
|
||||
outputs = llm.generate(
|
||||
(
|
||||
"Generate a sql statement that selects col_1 from "
|
||||
"table_1 where it is equal to 1. Make the response as short as "
|
||||
"possible."
|
||||
),
|
||||
sampling_params=sampling_params,
|
||||
use_tqdm=True,
|
||||
)
|
||||
|
||||
assert outputs is not None
|
||||
for output in outputs:
|
||||
assert output is not None
|
||||
assert isinstance(output, RequestOutput)
|
||||
prompt = output.prompt
|
||||
|
||||
generated_text = output.outputs[0].text
|
||||
assert generated_text is not None
|
||||
|
||||
# use Lark to parse the output, and make sure it's a valid parse tree
|
||||
from lark import Lark
|
||||
|
||||
parser = Lark(sample_sql_lark)
|
||||
parser.parse(generated_text)
|
||||
|
||||
# remove spaces for comparison b/c we removed them in the grammar
|
||||
ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(" ", "")
|
||||
|
||||
assert generated_text.strip() == ground_truth
|
||||
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
|
||||
#
|
||||
# Test 6: Test invalid grammar input
|
||||
#
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.8,
|
||||
top_p=0.95,
|
||||
max_tokens=1000,
|
||||
structured_outputs=StructuredOutputsParams(grammar="not a grammar"),
|
||||
)
|
||||
with pytest.raises(ValueError, match="Failed to convert the grammar "):
|
||||
llm.generate(
|
||||
(
|
||||
"Generate a sql statement that selects col_1 from "
|
||||
"table_1 where it is equal to 1. Make the response as short "
|
||||
"as possible."
|
||||
),
|
||||
sampling_params=sampling_params,
|
||||
use_tqdm=True,
|
||||
)
|
||||
|
||||
#
|
||||
# Test 7: Generate text based on a regex pattern
|
||||
#
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.8,
|
||||
top_p=0.95,
|
||||
structured_outputs=StructuredOutputsParams(regex=sample_regex),
|
||||
)
|
||||
|
||||
prompt = (
|
||||
f"Give an example IPv4 address with this regex: {sample_regex}. "
|
||||
f"Make the response as short as possible."
|
||||
)
|
||||
outputs = llm.generate(
|
||||
[prompt] * 2,
|
||||
sampling_params=sampling_params,
|
||||
use_tqdm=True,
|
||||
)
|
||||
|
||||
assert outputs is not None
|
||||
for output in outputs:
|
||||
assert output is not None
|
||||
assert isinstance(output, RequestOutput)
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(generated_text)
|
||||
assert generated_text is not None
|
||||
assert re.fullmatch(sample_regex, generated_text) is not None
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
|
||||
#
|
||||
# Test 8: Generate text based on a choices
|
||||
#
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.8,
|
||||
top_p=0.95,
|
||||
structured_outputs=StructuredOutputsParams(
|
||||
choice=sample_structured_outputs_choices
|
||||
),
|
||||
)
|
||||
|
||||
outputs = llm.generate(
|
||||
(
|
||||
"The best language for type-safe systems programming is "
|
||||
"(Make the response as short as possible.) "
|
||||
),
|
||||
sampling_params=sampling_params,
|
||||
use_tqdm=True,
|
||||
)
|
||||
assert outputs is not None
|
||||
for output in outputs:
|
||||
assert output is not None
|
||||
assert isinstance(output, RequestOutput)
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(generated_text)
|
||||
assert generated_text is not None
|
||||
assert generated_text in sample_structured_outputs_choices
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
|
||||
#
|
||||
# Test 9: Generate structured output using a Pydantic model with an enum
|
||||
#
|
||||
json_schema = CarDescription.model_json_schema()
|
||||
sampling_params = SamplingParams(
|
||||
temperature=1.0,
|
||||
max_tokens=1000,
|
||||
structured_outputs=StructuredOutputsParams(json=json_schema),
|
||||
)
|
||||
|
||||
outputs = llm.generate(
|
||||
(
|
||||
"Generate a JSON with the brand, model and car_type of the most "
|
||||
"iconic car from the 90's. Make the response as short as "
|
||||
"possible."
|
||||
),
|
||||
sampling_params=sampling_params,
|
||||
use_tqdm=True,
|
||||
)
|
||||
|
||||
assert outputs is not None
|
||||
|
||||
for output in outputs:
|
||||
assert output is not None
|
||||
assert isinstance(output, RequestOutput)
|
||||
prompt = output.prompt
|
||||
|
||||
generated_text = output.outputs[0].text
|
||||
assert generated_text is not None
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
try:
|
||||
output_json = json.loads(generated_text)
|
||||
except json.JSONDecodeError as e:
|
||||
pytest.fail(
|
||||
f"Invalid JSON from backend={backend}: {generated_text!r}\n"
|
||||
f"Schema: {json_schema}\nError: {e}"
|
||||
)
|
||||
jsonschema.validate(instance=output_json, schema=json_schema)
|
||||
|
||||
#
|
||||
# Test 10: Generate structured with minLength and maxLength
|
||||
#
|
||||
min_length = 50
|
||||
max_length = 50
|
||||
json_schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"description": {
|
||||
"type": "string",
|
||||
"maxLength": max_length,
|
||||
"minLength": min_length,
|
||||
}
|
||||
},
|
||||
"required": ["description"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
temperature=1.0,
|
||||
max_tokens=4096,
|
||||
structured_outputs=StructuredOutputsParams(json=json_schema),
|
||||
)
|
||||
|
||||
outputs = llm.generate(
|
||||
(
|
||||
"Generate a description of a frog using 50 characters. "
|
||||
"Make the response as short as possible."
|
||||
),
|
||||
sampling_params=sampling_params,
|
||||
use_tqdm=True,
|
||||
)
|
||||
|
||||
assert outputs is not None
|
||||
|
||||
for output in outputs:
|
||||
assert output is not None
|
||||
assert isinstance(output, RequestOutput)
|
||||
prompt = output.prompt
|
||||
|
||||
generated_text = output.outputs[0].text
|
||||
assert generated_text is not None
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
try:
|
||||
output_json = json.loads(generated_text)
|
||||
except json.JSONDecodeError as e:
|
||||
pytest.fail(
|
||||
f"Invalid JSON from backend={backend}: {generated_text!r}\n"
|
||||
f"Schema: {json_schema}\nError: {e}"
|
||||
)
|
||||
jsonschema.validate(instance=output_json, schema=json_schema)
|
||||
|
||||
if backend not in ["outlines", "lm-format-enforcer"]:
|
||||
#
|
||||
# Test 11: Generate structured output using structural_tag format
|
||||
#
|
||||
structural_tag_config = {
|
||||
"type": "structural_tag",
|
||||
"structures": [
|
||||
{
|
||||
"begin": "<function=get_weather>",
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"end": "</function>",
|
||||
}
|
||||
],
|
||||
"triggers": ["<function="],
|
||||
}
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.0,
|
||||
max_tokens=4096,
|
||||
structured_outputs=StructuredOutputsParams(
|
||||
structural_tag=json.dumps(structural_tag_config)
|
||||
),
|
||||
)
|
||||
|
||||
prompt = """
|
||||
You have access to the following function to retrieve the weather in a city:
|
||||
|
||||
{
|
||||
"name": "get_weather",
|
||||
"parameters": {
|
||||
"city": {
|
||||
"param_type": "string",
|
||||
"description": "The city to get the weather for",
|
||||
"required": True
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
If a you choose to call a function ONLY reply in the following format:
|
||||
<{start_tag}={function_name}>{parameters}{end_tag}
|
||||
where
|
||||
|
||||
start_tag => `<function`
|
||||
parameters => a JSON dict with the function argument name
|
||||
as key and function argument value as value.
|
||||
end_tag => `</function>`
|
||||
|
||||
Here is an example,
|
||||
<function=example_function_name>{"example_name": "example_value"}</function>
|
||||
|
||||
Reminder:
|
||||
- Function calls MUST follow the specified format
|
||||
- Required parameters MUST be specified
|
||||
- Only call one function at a time
|
||||
- Put the entire function call reply on one line
|
||||
- Always add your sources when using search results to answer the user query
|
||||
|
||||
You are a helpful assistant.
|
||||
|
||||
Given the previous instructions, what is the weather in New York City? \
|
||||
Make the response as short as possible.
|
||||
"""
|
||||
|
||||
# Change this once other backends support structural_tag
|
||||
outputs = llm.generate(prompt, sampling_params=sampling_params, use_tqdm=True)
|
||||
assert outputs is not None
|
||||
|
||||
for output in outputs:
|
||||
assert output is not None
|
||||
assert isinstance(output, RequestOutput)
|
||||
generated_text = output.outputs[0].text
|
||||
assert generated_text is not None
|
||||
|
||||
# Search for function call pattern in the response
|
||||
function_call_pattern = r"<function=get_weather>(.*?)</function>"
|
||||
matches = re.findall(function_call_pattern, generated_text)
|
||||
|
||||
if not matches:
|
||||
print(
|
||||
f"Warning: No function calls found in response: {generated_text!r}"
|
||||
)
|
||||
continue
|
||||
|
||||
# Take the first function call if multiple are found
|
||||
json_str = matches[0]
|
||||
try:
|
||||
json_content = json.loads(json_str)
|
||||
assert "city" in json_content
|
||||
assert isinstance(json_content["city"], str)
|
||||
print(f"Found valid function call: {generated_text!r}")
|
||||
except (json.JSONDecodeError, AssertionError) as e:
|
||||
pytest.fail(
|
||||
f"Invalid function call format: {generated_text!r}\nError: {str(e)}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name, backend, tokenizer_mode, reasoning_parser, speculative_config, async_scheduling", # noqa: E501
|
||||
[
|
||||
(
|
||||
"deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
|
||||
"xgrammar",
|
||||
"auto",
|
||||
"deepseek_r1",
|
||||
NGRAM_SPEC_CONFIG,
|
||||
False,
|
||||
),
|
||||
("Qwen/Qwen3-1.7B", "xgrammar", "auto", "deepseek_r1", None, False),
|
||||
("Qwen/Qwen3-1.7B", "xgrammar", "auto", "deepseek_r1", None, True),
|
||||
],
|
||||
)
|
||||
def test_structured_output_with_reasoning_matrices(
|
||||
backend: str,
|
||||
tokenizer_mode: str,
|
||||
reasoning_parser: str,
|
||||
model_name: str,
|
||||
speculative_config: dict[str, Any] | None,
|
||||
async_scheduling: bool,
|
||||
):
|
||||
if current_platform.is_tpu() and speculative_config:
|
||||
pytest.skip("TPU does not support speculative decoding")
|
||||
|
||||
# Use a single LLM instance for several scenarios to
|
||||
# speed up the test suite.
|
||||
llm = LLM(
|
||||
model=model_name,
|
||||
# Don't use eager execution on TPUs because we want to test for no
|
||||
# recompilation at runtime
|
||||
enforce_eager=bool(not current_platform.is_tpu()),
|
||||
max_model_len=1024,
|
||||
max_num_seqs=16,
|
||||
structured_outputs_config=dict(
|
||||
backend=backend,
|
||||
disable_any_whitespace=backend in {"xgrammar", "guidance"},
|
||||
reasoning_parser=reasoning_parser,
|
||||
),
|
||||
tokenizer_mode=tokenizer_mode,
|
||||
speculative_config=speculative_config,
|
||||
async_scheduling=async_scheduling,
|
||||
)
|
||||
tokenizer = llm.get_tokenizer()
|
||||
reasoner = ReasoningParserManager.get_reasoning_parser(reasoning_parser)(
|
||||
tokenizer=tokenizer
|
||||
)
|
||||
|
||||
reasoning_prompt = "Solve the following math problem step-by-step, then provide the final answer as JSON object with a single key 'result'. Make sure to correct your reasoning if there are any issue should it arise.\nProblem: What is 5 * 8 + 2?" # noqa: E501
|
||||
reasoning_schema = {
|
||||
"type": "object",
|
||||
"properties": {"result": {"type": "integer"}},
|
||||
"required": ["result"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
if "Qwen3" in model_name:
|
||||
reasoning_prompt += "<think>\n"
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.1,
|
||||
max_tokens=8192,
|
||||
structured_outputs=StructuredOutputsParams(json=reasoning_schema),
|
||||
)
|
||||
outputs = llm.generate(
|
||||
[reasoning_prompt],
|
||||
sampling_params=sampling_params,
|
||||
use_tqdm=True,
|
||||
)
|
||||
|
||||
assert outputs is not None
|
||||
output = outputs[0]
|
||||
assert output is not None and isinstance(output, RequestOutput)
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
reasoning, content = run_reasoning_extraction(reasoner, [generated_text])
|
||||
print(f"Prompt: {prompt!r}\nReasoning: {reasoning!r}\nContent: {content!r}")
|
||||
|
||||
if "Qwen3" in model_name:
|
||||
assert content is not None
|
||||
|
||||
assert reasoning is not None
|
||||
|
||||
if content is not None:
|
||||
output_json = json.loads(content)
|
||||
jsonschema.validate(instance=output_json, schema=reasoning_schema)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name, tokenizer_mode", PARAMS_MODELS_TOKENIZER_MODE)
|
||||
def test_structured_output_auto_mode(
|
||||
unsupported_json_schema: dict[str, Any],
|
||||
model_name: str,
|
||||
tokenizer_mode: str,
|
||||
):
|
||||
llm = LLM(
|
||||
model=model_name,
|
||||
max_model_len=1024,
|
||||
structured_outputs_config=dict(backend="auto"),
|
||||
tokenizer_mode=tokenizer_mode,
|
||||
load_format="auto",
|
||||
config_format="auto",
|
||||
)
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
temperature=1.0,
|
||||
max_tokens=1000,
|
||||
structured_outputs=StructuredOutputsParams(json=unsupported_json_schema),
|
||||
)
|
||||
|
||||
prompts = (
|
||||
"Give an example JSON object for a grade "
|
||||
"that fits this schema: "
|
||||
f"{unsupported_json_schema}. Make the response as short as possible."
|
||||
)
|
||||
# This would fail with the default of "xgrammar", but in "auto"
|
||||
# we will handle fallback automatically.
|
||||
outputs = llm.generate(prompts, sampling_params=sampling_params, use_tqdm=True)
|
||||
# Make sure `auto` backend handling doesn't mess up sampling_params
|
||||
# and that we can reuse it without error.
|
||||
outputs.extend(
|
||||
llm.generate(prompts, sampling_params=sampling_params, use_tqdm=True)
|
||||
)
|
||||
|
||||
assert outputs is not None
|
||||
for output in outputs:
|
||||
assert output is not None
|
||||
assert isinstance(output, RequestOutput)
|
||||
generated_text = output.outputs[0].text
|
||||
assert generated_text is not None
|
||||
print(generated_text)
|
||||
|
||||
# Parse to verify it is valid JSON
|
||||
parsed_json = json.loads(generated_text)
|
||||
assert isinstance(parsed_json, dict)
|
||||
|
||||
|
||||
def test_guidance_no_additional_properties():
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen2.5-1.5B-Instruct",
|
||||
max_model_len=1024,
|
||||
structured_outputs_config=dict(
|
||||
backend="guidance",
|
||||
disable_any_whitespace=True,
|
||||
disable_additional_properties=True,
|
||||
),
|
||||
)
|
||||
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"a1": {"type": "string"},
|
||||
"a2": {"type": "string"},
|
||||
"a3": {"type": "string"},
|
||||
},
|
||||
"required": ["a1", "a2", "a3"],
|
||||
}
|
||||
|
||||
prompt = (
|
||||
"<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a "
|
||||
"helpful assistant.<|im_end|>\n<|im_start|>user\nPlease generate a "
|
||||
"large JSON object with key-value pairs a1=b1, a2=b2, ..., a20=b20. "
|
||||
"Make the response as short as possible."
|
||||
"<|im_end|>\n<|im_start|>assistant\n"
|
||||
)
|
||||
|
||||
def generate_with_backend(backend):
|
||||
structured_outputs_params = StructuredOutputsParams(
|
||||
json=schema,
|
||||
backend=backend,
|
||||
disable_any_whitespace=True,
|
||||
disable_additional_properties=True,
|
||||
)
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0, max_tokens=256, structured_outputs=structured_outputs_params
|
||||
)
|
||||
|
||||
outputs = llm.generate(prompt, sampling_params=sampling_params)
|
||||
assert outputs is not None
|
||||
generated_text = outputs[0].outputs[0].text
|
||||
assert generated_text is not None
|
||||
parsed_json = json.loads(generated_text)
|
||||
assert isinstance(parsed_json, dict)
|
||||
jsonschema.validate(instance=parsed_json, schema=schema)
|
||||
return parsed_json
|
||||
|
||||
generated = generate_with_backend("guidance")
|
||||
assert "a1" in generated
|
||||
assert "a2" in generated
|
||||
assert "a3" in generated
|
||||
assert "a4" not in generated
|
||||
assert "a5" not in generated
|
||||
assert "a6" not in generated
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["guidance", "xgrammar", "outlines"])
|
||||
def test_structured_output_batched_with_non_structured_outputs_requests(
|
||||
sample_json_schema: dict[str, Any],
|
||||
backend: str,
|
||||
):
|
||||
# Don't use eager execution on TPUs because we want to test for no
|
||||
# recompilation at runtime
|
||||
enforce_eager = bool(not current_platform.is_tpu())
|
||||
|
||||
llm = LLM(
|
||||
model="meta-llama/Meta-Llama-3.1-8B-Instruct",
|
||||
enforce_eager=enforce_eager,
|
||||
max_model_len=1024,
|
||||
structured_outputs_config=StructuredOutputsConfig(
|
||||
backend=backend,
|
||||
disable_any_whitespace=backend in {"xgrammar", "guidance"},
|
||||
),
|
||||
)
|
||||
|
||||
structured_outputs_prompt = (
|
||||
"Give an example JSON for an employee profile that fits this "
|
||||
"schema. Make the response as short as possible. Schema: "
|
||||
f"{sample_json_schema}"
|
||||
)
|
||||
|
||||
non_structured_outputs_prompt = "The diameter of the Earth in kilometers is "
|
||||
|
||||
prompts = [structured_outputs_prompt, non_structured_outputs_prompt]
|
||||
sampling_params = [
|
||||
SamplingParams(
|
||||
temperature=1.0,
|
||||
max_tokens=400,
|
||||
structured_outputs=StructuredOutputsParams(json=sample_json_schema),
|
||||
),
|
||||
# No max tokens, temp=0 to assert on contents
|
||||
SamplingParams(
|
||||
seed=42,
|
||||
temperature=0,
|
||||
top_p=1.0,
|
||||
),
|
||||
]
|
||||
|
||||
outputs = llm.generate(
|
||||
prompts=prompts, sampling_params=sampling_params, use_tqdm=True
|
||||
)
|
||||
|
||||
assert outputs is not None
|
||||
|
||||
# Free memory as soon as possible as failed assertions
|
||||
# will short circuit and not free up memory
|
||||
del llm
|
||||
torch.accelerator.empty_cache()
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
for index, output in enumerate(outputs):
|
||||
assert output is not None
|
||||
assert isinstance(output, RequestOutput)
|
||||
prompt = output.prompt
|
||||
|
||||
generated_text = output.outputs[0].text
|
||||
assert generated_text is not None
|
||||
print(f"Prompt:\n{prompt!r}\nGenerated text:\n{generated_text!r}")
|
||||
|
||||
if index == 0:
|
||||
# First prompt is structured outputs, expect valid JSON
|
||||
assert "\n" not in generated_text
|
||||
output_json = json.loads(generated_text)
|
||||
jsonschema.validate(instance=output_json, schema=sample_json_schema)
|
||||
else:
|
||||
# Second prompt is not structured outputs, expect valid output
|
||||
# Cannot assert on exact output, but we can expect it to be factual
|
||||
assert "12,742" in generated_text
|
||||
|
||||
# non-structured outputs requests should not return a valid JSON here
|
||||
with pytest.raises(ValueError):
|
||||
output_json = json.loads(generated_text)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["xgrammar"])
|
||||
def test_structured_output_with_structural_tag(backend: str):
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen2.5-1.5B-Instruct",
|
||||
structured_outputs_config=StructuredOutputsConfig(backend=backend),
|
||||
)
|
||||
|
||||
structural_tag_config = {
|
||||
"type": "structural_tag",
|
||||
"format": {
|
||||
"type": "triggered_tags",
|
||||
"tags": [
|
||||
{"begin": "hello_flag", "content": {"type": "any_text"}, "end": "hello"}
|
||||
],
|
||||
"triggers": ["hello"],
|
||||
"stop_after_first": False,
|
||||
},
|
||||
}
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.0,
|
||||
max_tokens=500,
|
||||
structured_outputs=StructuredOutputsParams(
|
||||
structural_tag=json.dumps(structural_tag_config)
|
||||
),
|
||||
)
|
||||
|
||||
prompt = "Hello and repeat hello 10 times, do not say anything else. Only say hello hello hello, now start"
|
||||
outputs = llm.generate(prompt, sampling_params=sampling_params, use_tqdm=True)
|
||||
assert outputs is not None
|
||||
for output in outputs:
|
||||
assert output is not None
|
||||
assert isinstance(output, RequestOutput)
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
assert generated_text is not None
|
||||
assert "hello_flag" in generated_text, (
|
||||
f"Expected 'hello_flag' to be in generated text, but got: {generated_text}"
|
||||
)
|
||||
0
third_party/vllm/tests/v1/entrypoints/openai/serving_responses/__init__.py
vendored
Normal file
0
third_party/vllm/tests/v1/entrypoints/openai/serving_responses/__init__.py
vendored
Normal file
44
third_party/vllm/tests/v1/entrypoints/openai/serving_responses/conftest.py
vendored
Normal file
44
third_party/vllm/tests/v1/entrypoints/openai/serving_responses/conftest.py
vendored
Normal file
@@ -0,0 +1,44 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
# Use a small reasoning model to test the responses API.
|
||||
MODEL_NAME = "Qwen/Qwen3-1.7B"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def default_server_args():
|
||||
return [
|
||||
"--max-model-len",
|
||||
"8192",
|
||||
"--enforce-eager", # For faster startup.
|
||||
"--enable-auto-tool-choice",
|
||||
"--structured-outputs-config.backend",
|
||||
"xgrammar",
|
||||
"--tool-call-parser",
|
||||
"hermes",
|
||||
"--reasoning-parser",
|
||||
"qwen3",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server_with_store(default_server_args):
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME,
|
||||
default_server_args,
|
||||
env_dict={
|
||||
"VLLM_ENABLE_RESPONSES_API_STORE": "1",
|
||||
"VLLM_SERVER_DEV_MODE": "1",
|
||||
},
|
||||
) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(server_with_store):
|
||||
async with server_with_store.get_async_client() as async_client:
|
||||
yield async_client
|
||||
93
third_party/vllm/tests/v1/entrypoints/openai/serving_responses/test_basic.py
vendored
Normal file
93
third_party/vllm/tests/v1/entrypoints/openai/serving_responses/test_basic.py
vendored
Normal file
@@ -0,0 +1,93 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import openai # use the official client for correctness check
|
||||
import openai.types.responses as openai_responses_types
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_simple_input(client: openai.AsyncOpenAI):
|
||||
response = await client.responses.create(input="What is 13 * 24?")
|
||||
print(response)
|
||||
|
||||
outputs = response.output
|
||||
# Whether the output contains the answer.
|
||||
assert outputs[-1].type == "message"
|
||||
assert "312" in outputs[-1].content[0].text
|
||||
|
||||
# Whether the output contains the reasoning.
|
||||
assert outputs[0].type == "reasoning"
|
||||
assert outputs[0].content[0].text != ""
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_instructions(client: openai.AsyncOpenAI):
|
||||
response = await client.responses.create(
|
||||
instructions="Finish the answer with QED.",
|
||||
input="What is 13 * 24?",
|
||||
)
|
||||
print(response)
|
||||
|
||||
output_text = response.output[-1].content[0].text
|
||||
assert "312" in output_text
|
||||
assert "QED" in output_text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat(client: openai.AsyncOpenAI):
|
||||
response = await client.responses.create(
|
||||
input=[
|
||||
{"role": "system", "content": "Finish the answer with QED."},
|
||||
{"role": "user", "content": "What is 5 * 3?"},
|
||||
{"role": "assistant", "content": "15. QED."},
|
||||
{"role": "user", "content": "Multiply the result by 2."},
|
||||
],
|
||||
)
|
||||
print(response)
|
||||
|
||||
output_text = response.output[-1].content[0].text
|
||||
assert "30" in output_text
|
||||
assert "QED" in output_text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_with_input_type(client: openai.AsyncOpenAI):
|
||||
response = await client.responses.create(
|
||||
input=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": "Hello!"}],
|
||||
},
|
||||
],
|
||||
)
|
||||
print(response)
|
||||
assert response.status == "completed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logprobs(client: openai.AsyncOpenAI):
|
||||
response = await client.responses.create(
|
||||
include=["message.output_text.logprobs"],
|
||||
input="What is 13 * 24?",
|
||||
top_logprobs=5,
|
||||
)
|
||||
print(response)
|
||||
outputs = response.output
|
||||
assert outputs[-1].content[-1].logprobs
|
||||
assert len(outputs[-1].content[-1].logprobs[0].top_logprobs) == 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming(client: openai.AsyncOpenAI):
|
||||
stream = await client.responses.create(
|
||||
input="What is 13 * 24?",
|
||||
stream=True,
|
||||
)
|
||||
events = [event async for event in stream]
|
||||
assert isinstance(events[0], openai_responses_types.ResponseCreatedEvent)
|
||||
assert any(
|
||||
isinstance(event, openai_responses_types.ResponseTextDeltaEvent)
|
||||
for event in events
|
||||
)
|
||||
assert isinstance(events[-1], openai_responses_types.ResponseCompletedEvent)
|
||||
304
third_party/vllm/tests/v1/entrypoints/openai/serving_responses/test_function_call.py
vendored
Normal file
304
third_party/vllm/tests/v1/entrypoints/openai/serving_responses/test_function_call.py
vendored
Normal file
@@ -0,0 +1,304 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
|
||||
import openai # use the official client for correctness check
|
||||
import pytest
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen3-1.7B"
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "get_current_weather",
|
||||
"description": "Get the current weather in a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "The city to find the weather for, e.g. 'Vienna'",
|
||||
"default": "Vienna",
|
||||
},
|
||||
"country": {
|
||||
"type": "string",
|
||||
"description": "The country that the city is in, e.g. 'Austria'",
|
||||
},
|
||||
"unit": {
|
||||
"type": "string",
|
||||
"description": "The unit to fetch the temperature in",
|
||||
"enum": ["celsius", "fahrenheit"],
|
||||
},
|
||||
"options": {
|
||||
"$ref": "#/$defs/WeatherOptions",
|
||||
"description": "Optional parameters for weather query",
|
||||
},
|
||||
},
|
||||
"required": ["country", "unit"],
|
||||
"$defs": {
|
||||
"WeatherOptions": {
|
||||
"title": "WeatherOptions",
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"unit": {
|
||||
"type": "string",
|
||||
"enum": ["celsius", "fahrenheit"],
|
||||
"default": "celsius",
|
||||
"description": "Temperature unit",
|
||||
"title": "Temperature Unit",
|
||||
},
|
||||
"include_forecast": {
|
||||
"type": "boolean",
|
||||
"default": False,
|
||||
"description": "Whether to include a 24-hour forecast",
|
||||
"title": "Include Forecast",
|
||||
},
|
||||
"language": {
|
||||
"type": "string",
|
||||
"default": "zh-CN",
|
||||
"description": "Language of the response",
|
||||
"title": "Language",
|
||||
"enum": ["zh-CN", "en-US", "ja-JP"],
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"name": "get_forecast",
|
||||
"description": "Get the weather forecast for a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"description": "The city to get the forecast for, e.g. 'Vienna'",
|
||||
"default": "Vienna",
|
||||
},
|
||||
"country": {
|
||||
"type": "string",
|
||||
"description": "The country that the city is in, e.g. 'Austria'",
|
||||
},
|
||||
"days": {
|
||||
"type": "integer",
|
||||
"description": "Number of days to get the forecast for (1-7)",
|
||||
},
|
||||
"unit": {
|
||||
"type": "string",
|
||||
"description": "The unit to fetch the temperature in",
|
||||
"enum": ["celsius", "fahrenheit"],
|
||||
},
|
||||
},
|
||||
"required": ["country", "days", "unit"],
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("tool_choice", ["auto", "required"])
|
||||
async def test_function_tool_use(
|
||||
client: openai.AsyncOpenAI, model_name: str, tool_choice: str
|
||||
):
|
||||
prompt = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Can you tell me what the current weather is in Berlin and the "
|
||||
"forecast for the next 5 days, in fahrenheit?",
|
||||
},
|
||||
]
|
||||
response = await client.responses.create(
|
||||
model=model_name,
|
||||
input=prompt,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
temperature=0.0,
|
||||
)
|
||||
|
||||
assert len(response.output) >= 1
|
||||
tool_call = None
|
||||
reasoning = None
|
||||
for out in response.output:
|
||||
if out.type == "function_call":
|
||||
tool_call = out
|
||||
if out.type == "reasoning":
|
||||
reasoning = out
|
||||
assert tool_call is not None
|
||||
assert tool_call.type == "function_call"
|
||||
assert json.loads(tool_call.arguments) is not None
|
||||
assert reasoning is not None
|
||||
assert reasoning.type == "reasoning"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_named_tool_use(client: openai.AsyncOpenAI):
|
||||
def get_weather(latitude: float, longitude: float) -> str:
|
||||
"""
|
||||
Mock function to simulate getting weather data.
|
||||
In a real application, this would call an external weather API.
|
||||
"""
|
||||
return f"Current temperature at ({latitude}, {longitude}) is 20°C."
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "get_weather",
|
||||
"description": (
|
||||
"Get current temperature for provided coordinates in celsius."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"latitude": {"type": "number"},
|
||||
"longitude": {"type": "number"},
|
||||
},
|
||||
"required": ["latitude", "longitude"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"strict": True,
|
||||
}
|
||||
]
|
||||
|
||||
input_messages = [
|
||||
{"role": "user", "content": "What's the weather like in Paris today?"}
|
||||
]
|
||||
|
||||
response = await client.responses.create(
|
||||
model=MODEL_NAME,
|
||||
input=input_messages,
|
||||
tools=tools,
|
||||
tool_choice={"type": "function", "name": "get_weather"},
|
||||
)
|
||||
assert len(response.output) >= 1
|
||||
for out in response.output:
|
||||
if out.type == "function_call":
|
||||
tool_call = out
|
||||
assert tool_call is not None
|
||||
assert tool_call.type == "function_call"
|
||||
assert tool_call.name == "get_weather"
|
||||
args = json.loads(tool_call.arguments)
|
||||
assert args["latitude"] is not None
|
||||
assert args["longitude"] is not None
|
||||
# call the tool
|
||||
result = get_weather(args["latitude"], args["longitude"])
|
||||
input_messages.append(tool_call) # append model's function call message
|
||||
input_messages.append(
|
||||
{ # append result message
|
||||
"type": "function_call_output",
|
||||
"call_id": tool_call.call_id,
|
||||
"output": str(result),
|
||||
}
|
||||
)
|
||||
# create a new response with the tool call result
|
||||
response_2 = await client.responses.create(model=MODEL_NAME, input=input_messages)
|
||||
# check the output
|
||||
assert len(response_2.output_text) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_function_calling_with_streaming_expected_arguments(
|
||||
client: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "get_weather",
|
||||
"description": "Get current temperature for provided location in celsius.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {"type": "string"},
|
||||
},
|
||||
"required": ["location"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"strict": True,
|
||||
}
|
||||
]
|
||||
|
||||
stream_response = await client.responses.create(
|
||||
model=model_name,
|
||||
input="Can you tell me what the current weather is in Berlin?",
|
||||
tools=tools,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
tool_call_item = None
|
||||
completed_event = None
|
||||
async for event in stream_response:
|
||||
if (
|
||||
event.type == "response.output_item.added"
|
||||
and event.item.type == "function_call"
|
||||
):
|
||||
tool_call_item = event.item
|
||||
elif event.type == "response.function_call_arguments.delta" and tool_call_item:
|
||||
tool_call_item.arguments += event.delta
|
||||
elif (
|
||||
event.type == "response.output_item.done"
|
||||
and event.item.type == "function_call"
|
||||
):
|
||||
completed_event = event
|
||||
assert tool_call_item is not None
|
||||
assert tool_call_item.type == "function_call"
|
||||
assert tool_call_item.name == "get_weather"
|
||||
assert completed_event is not None
|
||||
assert tool_call_item.arguments == completed_event.item.arguments
|
||||
assert tool_call_item.name == completed_event.item.name
|
||||
args = json.loads(tool_call_item.arguments)
|
||||
assert "location" in args
|
||||
assert args["location"] is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_function_calling_with_streaming_types(
|
||||
client: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
# this links the "done" type with the "start" type
|
||||
# so every "done" type should have a corresponding "start" type
|
||||
# and every open block should be closed by the end of the stream
|
||||
pairs_of_event_types = {
|
||||
"response.completed": "response.created",
|
||||
"response.output_item.done": "response.output_item.added",
|
||||
"response.output_text.done": "response.output_text.delta",
|
||||
"response.content_part.done": "response.content_part.added",
|
||||
"response.reasoning_text.done": "response.reasoning_text.delta",
|
||||
"response.reasoning_part.done": "response.reasoning_part.added",
|
||||
"response.function_call_arguments.done": "response.function_call_arguments.delta", # noqa
|
||||
}
|
||||
|
||||
input_list = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Can you tell me what the current weather is in Berlin?",
|
||||
}
|
||||
]
|
||||
stream_response = await client.responses.create(
|
||||
model=model_name,
|
||||
input=input_list,
|
||||
tools=tools,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
stack_of_event_types = []
|
||||
async for event in stream_response:
|
||||
if event.type == "response.created":
|
||||
stack_of_event_types.append(event.type)
|
||||
elif event.type == "response.completed":
|
||||
assert stack_of_event_types[-1] == pairs_of_event_types[event.type]
|
||||
stack_of_event_types.pop()
|
||||
if event.type.endswith("added"):
|
||||
stack_of_event_types.append(event.type)
|
||||
elif event.type.endswith("delta"):
|
||||
if stack_of_event_types[-1] == event.type:
|
||||
continue
|
||||
stack_of_event_types.append(event.type)
|
||||
elif event.type.endswith("done"):
|
||||
assert stack_of_event_types[-1] == pairs_of_event_types[event.type]
|
||||
stack_of_event_types.pop()
|
||||
assert len(stack_of_event_types) == 0
|
||||
171
third_party/vllm/tests/v1/entrypoints/openai/serving_responses/test_image.py
vendored
Normal file
171
third_party/vllm/tests/v1/entrypoints/openai/serving_responses/test_image.py
vendored
Normal file
@@ -0,0 +1,171 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.multimodal.utils import encode_image_url
|
||||
|
||||
# Use a small vision model for testing
|
||||
MODEL_NAME = "Qwen/Qwen2.5-VL-3B-Instruct"
|
||||
MAXIMUM_IMAGES = 2
|
||||
# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
|
||||
TEST_IMAGE_ASSETS = [
|
||||
"2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
|
||||
"Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
|
||||
"1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png",
|
||||
"RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def default_image_server_args():
|
||||
return [
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"6000",
|
||||
"--max-num-seqs",
|
||||
"128",
|
||||
"--limit-mm-per-prompt",
|
||||
json.dumps({"image": MAXIMUM_IMAGES}),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def image_server(default_image_server_args):
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME,
|
||||
default_image_server_args,
|
||||
env_dict={"VLLM_ENABLE_RESPONSES_API_STORE": "1"},
|
||||
) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(image_server):
|
||||
async with image_server.get_async_client() as async_client:
|
||||
yield async_client
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def url_encoded_image(local_asset_server) -> dict[str, str]:
|
||||
return {
|
||||
image_url: encode_image_url(local_asset_server.get_image_asset(image_url))
|
||||
for image_url in TEST_IMAGE_ASSETS
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
|
||||
async def test_single_chat_session_image(
|
||||
client: openai.AsyncOpenAI, model_name: str, image_url: str
|
||||
):
|
||||
content_text = "What's in this image?"
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "input_image",
|
||||
"image_url": image_url,
|
||||
"detail": "auto",
|
||||
},
|
||||
{"type": "input_text", "text": content_text},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
# test image url
|
||||
response = await client.responses.create(
|
||||
model=model_name,
|
||||
input=messages,
|
||||
)
|
||||
assert len(response.output_text) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("raw_image_url", TEST_IMAGE_ASSETS)
|
||||
async def test_single_chat_session_image_base64encoded(
|
||||
client: openai.AsyncOpenAI,
|
||||
model_name: str,
|
||||
raw_image_url: str,
|
||||
url_encoded_image: dict[str, str],
|
||||
):
|
||||
content_text = "What's in this image?"
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "input_image",
|
||||
"image_url": url_encoded_image[raw_image_url],
|
||||
"detail": "auto",
|
||||
},
|
||||
{"type": "input_text", "text": content_text},
|
||||
],
|
||||
}
|
||||
]
|
||||
# test image base64
|
||||
response = await client.responses.create(
|
||||
model=model_name,
|
||||
input=messages,
|
||||
)
|
||||
assert len(response.output_text) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize(
|
||||
"image_urls",
|
||||
[TEST_IMAGE_ASSETS[:i] for i in range(2, len(TEST_IMAGE_ASSETS))],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_multi_image_input(
|
||||
client: openai.AsyncOpenAI, model_name: str, image_urls: list[str]
|
||||
):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
*(
|
||||
{
|
||||
"type": "input_image",
|
||||
"image_url": image_url,
|
||||
"detail": "auto",
|
||||
}
|
||||
for image_url in image_urls
|
||||
),
|
||||
{"type": "input_text", "text": "What's in this image?"},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
if len(image_urls) > MAXIMUM_IMAGES:
|
||||
with pytest.raises(openai.BadRequestError): # test multi-image input
|
||||
await client.responses.create(
|
||||
model=model_name,
|
||||
input=messages,
|
||||
)
|
||||
# the server should still work afterwards
|
||||
response = await client.responses.create(
|
||||
model=model_name,
|
||||
input=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What's the weather like in Paris today?",
|
||||
}
|
||||
],
|
||||
)
|
||||
assert len(response.output_text) > 0
|
||||
else:
|
||||
response = await client.responses.create(
|
||||
model=model_name,
|
||||
input=messages,
|
||||
)
|
||||
assert len(response.output_text) > 0
|
||||
152
third_party/vllm/tests/v1/entrypoints/openai/serving_responses/test_stateful.py
vendored
Normal file
152
third_party/vllm/tests/v1/entrypoints/openai/serving_responses/test_stateful.py
vendored
Normal file
@@ -0,0 +1,152 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import asyncio
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store(client: openai.AsyncOpenAI):
|
||||
# By default, store is True.
|
||||
response = await client.responses.create(input="Hello!")
|
||||
assert response.status == "completed"
|
||||
|
||||
# Retrieve the response.
|
||||
response = await client.responses.retrieve(response.id)
|
||||
assert response.status == "completed"
|
||||
|
||||
# Test store=False.
|
||||
response = await client.responses.create(
|
||||
input="Hello!",
|
||||
store=False,
|
||||
)
|
||||
assert response.status == "completed"
|
||||
|
||||
# The response should not be found.
|
||||
with pytest.raises(openai.NotFoundError, match="Response with id .* not found."):
|
||||
await client.responses.retrieve(response.id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_background(client: openai.AsyncOpenAI):
|
||||
# NOTE: This query should be easy enough for the model to answer
|
||||
# within the 10 seconds.
|
||||
response = await client.responses.create(
|
||||
input="Hello!",
|
||||
background=True,
|
||||
)
|
||||
assert response.status == "queued"
|
||||
|
||||
max_retries = 10
|
||||
for _ in range(max_retries):
|
||||
await asyncio.sleep(1)
|
||||
response = await client.responses.retrieve(response.id)
|
||||
if response.status != "queued":
|
||||
break
|
||||
print(response)
|
||||
|
||||
assert response.status == "completed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_background_error(client: openai.AsyncOpenAI):
|
||||
with pytest.raises(
|
||||
openai.BadRequestError, match="background can only be used when `store` is true"
|
||||
):
|
||||
_ = await client.responses.create(
|
||||
input="What is 13 * 24?",
|
||||
background=True,
|
||||
store=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_background_cancel(client: openai.AsyncOpenAI):
|
||||
response = await client.responses.create(
|
||||
input="Write a long story about a cat.",
|
||||
background=True,
|
||||
)
|
||||
assert response.status == "queued"
|
||||
|
||||
# Cancel the response before it is completed.
|
||||
# Poll until the response is no longer queued (started processing) or timeout
|
||||
loop = asyncio.get_running_loop()
|
||||
start_time = loop.time()
|
||||
max_wait_seconds = 5.0
|
||||
poll_interval = 0.1
|
||||
while loop.time() - start_time < max_wait_seconds:
|
||||
response = await client.responses.retrieve(response.id)
|
||||
if response.status != "queued":
|
||||
# Started processing or completed - try to cancel
|
||||
break
|
||||
await asyncio.sleep(poll_interval)
|
||||
|
||||
response = await client.responses.cancel(response.id)
|
||||
assert response.status == "cancelled"
|
||||
|
||||
# Make sure the response status remains unchanged after some time.
|
||||
max_retries = 10
|
||||
for _ in range(max_retries):
|
||||
await asyncio.sleep(0.5)
|
||||
response = await client.responses.retrieve(response.id)
|
||||
# Verify status is still cancelled
|
||||
assert response.status == "cancelled"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_completed(client: openai.AsyncOpenAI):
|
||||
response = await client.responses.create(input="Hello")
|
||||
assert response.status == "completed"
|
||||
|
||||
with pytest.raises(
|
||||
openai.BadRequestError, match="Cannot cancel a synchronous response."
|
||||
):
|
||||
await client.responses.cancel(response.id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_previous_response_id(client: openai.AsyncOpenAI):
|
||||
response1 = await client.responses.create(
|
||||
instructions="You are tested on your ability to retrieve the correct "
|
||||
"information from the previous response.",
|
||||
input="Hello, my name is John.",
|
||||
)
|
||||
|
||||
response2 = await client.responses.create(
|
||||
input="Actually, my name is not John. My real name is Mark.",
|
||||
previous_response_id=response1.id,
|
||||
)
|
||||
|
||||
response3 = await client.responses.create(
|
||||
input="What is my real name again? Answer in one word.",
|
||||
previous_response_id=response2.id,
|
||||
)
|
||||
print(response3)
|
||||
assert "Mark" in response3.output[-1].content[0].text
|
||||
assert "John" not in response3.output[-1].content[0].text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_two_responses_with_same_prev_id(client: openai.AsyncOpenAI):
|
||||
response1 = await client.responses.create(
|
||||
instructions="You are tested on your ability to retrieve the correct "
|
||||
"information from the previous response.",
|
||||
input="Hello, my name is John.",
|
||||
)
|
||||
|
||||
# Both response 2 and 3 use response 1 as the previous response.
|
||||
response2 = client.responses.create(
|
||||
input="Actually, my name is not John. My name is Mark.",
|
||||
previous_response_id=response1.id,
|
||||
)
|
||||
response3 = client.responses.create(
|
||||
input="What is my name again? Answer in one word.",
|
||||
previous_response_id=response1.id,
|
||||
)
|
||||
|
||||
_ = await response2
|
||||
response3_result = await response3
|
||||
print(response3_result)
|
||||
assert "John" in response3_result.output[-1].content[0].text
|
||||
assert "Mark" not in response3_result.output[-1].content[0].text
|
||||
78
third_party/vllm/tests/v1/entrypoints/openai/serving_responses/test_structured_output.py
vendored
Normal file
78
third_party/vllm/tests/v1/entrypoints/openai/serving_responses/test_structured_output.py
vendored
Normal file
@@ -0,0 +1,78 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import json
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_structured_output(client: openai.AsyncOpenAI):
|
||||
response = await client.responses.create(
|
||||
input=[
|
||||
{"role": "system", "content": "Extract the event information."},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Alice and Bob are going to a science fair on Friday.",
|
||||
},
|
||||
],
|
||||
text={
|
||||
"format": {
|
||||
"type": "json_schema",
|
||||
"name": "calendar_event",
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"event_name": {"type": "string"},
|
||||
"date": {"type": "string"},
|
||||
"participants": {"type": "array", "items": {"type": "string"}},
|
||||
},
|
||||
"required": ["event_name", "date", "participants"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"description": "A calendar event.",
|
||||
"strict": True,
|
||||
}
|
||||
},
|
||||
)
|
||||
print(response)
|
||||
|
||||
# NOTE: The JSON schema is applied to the output text, not reasoning.
|
||||
output_text = response.output[-1].content[0].text
|
||||
event = json.loads(output_text)
|
||||
|
||||
assert event["event_name"].lower() == "science fair"
|
||||
assert event["date"] == "Friday"
|
||||
participants = event["participants"]
|
||||
assert len(participants) == 2
|
||||
assert participants[0] == "Alice"
|
||||
assert participants[1] == "Bob"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_structured_output_with_parse(client: openai.AsyncOpenAI):
|
||||
class CalendarEvent(BaseModel):
|
||||
event_name: str
|
||||
date: str
|
||||
participants: list[str]
|
||||
|
||||
response = await client.responses.parse(
|
||||
model=None,
|
||||
instructions="Extract the event information.",
|
||||
input="Alice and Bob are going to a science fair on Friday.",
|
||||
text_format=CalendarEvent,
|
||||
)
|
||||
print(response)
|
||||
|
||||
# The output is successfully parsed.
|
||||
event = response.output_parsed
|
||||
assert event is not None
|
||||
|
||||
# The output is correct.
|
||||
assert event.event_name.lower() == "science fair"
|
||||
assert event.date == "Friday"
|
||||
participants = event.participants
|
||||
assert len(participants) == 2
|
||||
assert participants[0] == "Alice"
|
||||
assert participants[1] == "Bob"
|
||||
160
third_party/vllm/tests/v1/entrypoints/openai/test_chat_completion.py
vendored
Normal file
160
third_party/vllm/tests/v1/entrypoints/openai/test_chat_completion.py
vendored
Normal file
@@ -0,0 +1,160 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import openai # use the official client for correctness check
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
# any model with a chat template defined in tokenizer_config should work here
|
||||
MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def default_server_args():
|
||||
return [
|
||||
# use half precision for speed and memory savings in CI environment
|
||||
"--max-model-len",
|
||||
"2048",
|
||||
"--max-num-seqs",
|
||||
"128",
|
||||
"--enforce-eager",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server(default_server_args):
|
||||
with RemoteOpenAIServer(MODEL_NAME, default_server_args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(server):
|
||||
async with server.get_async_client() as async_client:
|
||||
yield async_client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_invalid_json_schema(client: openai.AsyncOpenAI, model_name: str) -> None:
|
||||
invalid_json_schema = {
|
||||
"$defs": {
|
||||
"CarType": {
|
||||
"enum": ["sedan", "SUV", "Truck", "Coupe"],
|
||||
"title": "CarType",
|
||||
"type": "string",
|
||||
}
|
||||
},
|
||||
"properties": {
|
||||
"brand": {"title": "Brand", "type": "string"},
|
||||
"model": {"title": "Model", "type": "string"},
|
||||
"car_type": {"$ref": "#/$defs/CarType"},
|
||||
"foo": "bar",
|
||||
},
|
||||
"required": ["brand", "model", "car_type"],
|
||||
"title": "CarDescription",
|
||||
"type": "object",
|
||||
}
|
||||
prompt = (
|
||||
"Generate a JSON with the brand, model and car_type of"
|
||||
"the most iconic car from the 90's"
|
||||
)
|
||||
with pytest.raises((openai.BadRequestError, openai.APIError)):
|
||||
await client.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt,
|
||||
}
|
||||
],
|
||||
extra_body={"structured_outputs": {"json": invalid_json_schema}},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_invalid_regex(client: openai.AsyncOpenAI, model_name: str):
|
||||
prompt = (
|
||||
"Generate an email address for Alan Turing, who works in Enigma."
|
||||
"End in .com and new line. Example result:"
|
||||
"alan.turing@enigma.com\n"
|
||||
)
|
||||
|
||||
with pytest.raises((openai.BadRequestError, openai.APIError)):
|
||||
await client.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt,
|
||||
}
|
||||
],
|
||||
extra_body={"structured_outputs": {"regex": r"[.*"}, "stop": ["\n"]},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_invalid_grammar(client: openai.AsyncOpenAI, model_name: str):
|
||||
invalid_simplified_sql_grammar = """
|
||||
root ::= select_statementinvalidsyntax
|
||||
|
||||
select_statement ::= "SELECT " column " from " table " where " condition
|
||||
|
||||
column ::= "col_1 " | "col_2 "
|
||||
|
||||
table ::= "table_1 " | "table_2 "
|
||||
|
||||
condition ::= column "= " number
|
||||
|
||||
number ::= "1 " | "2 "
|
||||
"""
|
||||
|
||||
prompt = (
|
||||
"Generate an SQL query to show the 'username' and 'email'"
|
||||
"from the 'users' table."
|
||||
)
|
||||
with pytest.raises((openai.BadRequestError, openai.APIError)):
|
||||
await client.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt,
|
||||
}
|
||||
],
|
||||
extra_body={
|
||||
"structured_outputs": {"grammar": invalid_simplified_sql_grammar}
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_empty_grammar(client: openai.AsyncOpenAI, model_name: str) -> None:
|
||||
prompt = "Say hello"
|
||||
with pytest.raises((openai.BadRequestError, openai.APIError)):
|
||||
await client.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt,
|
||||
}
|
||||
],
|
||||
extra_body={"structured_outputs": {"grammar": ""}},
|
||||
)
|
||||
699
third_party/vllm/tests/v1/entrypoints/openai/test_completion.py
vendored
Normal file
699
third_party/vllm/tests/v1/entrypoints/openai/test_completion.py
vendored
Normal file
@@ -0,0 +1,699 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import openai # use the official client for correctness check
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import regex as re
|
||||
from openai import BadRequestError
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.tokenizers import get_tokenizer
|
||||
|
||||
# any model with a chat template should work here
|
||||
MODEL_NAME = "facebook/opt-125m"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def default_server_args():
|
||||
return [
|
||||
"--dtype",
|
||||
"float32",
|
||||
"--max-model-len",
|
||||
"2048",
|
||||
"--max-num-seqs",
|
||||
"128",
|
||||
"--enforce-eager",
|
||||
"--enable-prompt-tokens-details",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
scope="module",
|
||||
params=[
|
||||
["--no-enable-prefix-caching"],
|
||||
["--no-enable-prefix-caching", "--disable-frontend-multiprocessing"],
|
||||
],
|
||||
)
|
||||
def server(default_server_args, request):
|
||||
if request.param:
|
||||
default_server_args = default_server_args + request.param
|
||||
with RemoteOpenAIServer(MODEL_NAME, default_server_args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(server):
|
||||
async with server.get_async_client() as async_client:
|
||||
yield async_client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_single_completion(client: openai.AsyncOpenAI, model_name: str) -> None:
|
||||
completion = await client.completions.create(
|
||||
model=model_name, prompt="Hello, my name is", max_tokens=5, temperature=0.0
|
||||
)
|
||||
|
||||
assert completion.id is not None
|
||||
assert completion.choices is not None and len(completion.choices) == 1
|
||||
|
||||
choice = completion.choices[0]
|
||||
assert len(choice.text) >= 5
|
||||
assert choice.finish_reason == "length"
|
||||
assert completion.usage == openai.types.CompletionUsage(
|
||||
completion_tokens=5, prompt_tokens=6, total_tokens=11
|
||||
)
|
||||
|
||||
# test using token IDs
|
||||
completion = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=[0, 0, 0, 0, 0],
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
)
|
||||
assert len(completion.choices[0].text) >= 1
|
||||
assert completion.choices[0].prompt_logprobs is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_no_logprobs(client: openai.AsyncOpenAI, model_name: str):
|
||||
# test using token IDs
|
||||
completion = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=[0, 0, 0, 0, 0],
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
logprobs=None,
|
||||
)
|
||||
choice = completion.choices[0]
|
||||
assert choice.logprobs is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_zero_logprobs(client: openai.AsyncOpenAI, model_name: str):
|
||||
# test using token IDs
|
||||
completion = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=[0, 0, 0, 0, 0],
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
logprobs=0,
|
||||
)
|
||||
choice = completion.choices[0]
|
||||
assert choice.logprobs is not None
|
||||
assert choice.logprobs.token_logprobs is not None
|
||||
assert choice.logprobs.top_logprobs is not None
|
||||
assert len(choice.logprobs.top_logprobs[0]) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_some_logprobs(client: openai.AsyncOpenAI, model_name: str):
|
||||
# test using token IDs
|
||||
completion = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=[0, 0, 0, 0, 0],
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
logprobs=5,
|
||||
)
|
||||
choice = completion.choices[0]
|
||||
assert choice.logprobs is not None
|
||||
assert choice.logprobs.token_logprobs is not None
|
||||
assert choice.logprobs.top_logprobs is not None
|
||||
assert 5 <= len(choice.logprobs.top_logprobs[0]) <= 6
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_too_many_completion_logprobs(
|
||||
client: openai.AsyncOpenAI, model_name: str
|
||||
) -> None:
|
||||
with pytest.raises(
|
||||
(openai.BadRequestError, openai.APIError)
|
||||
): # test using token IDs
|
||||
await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=[0, 0, 0, 0, 0],
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
# vLLM has higher default max_logprobs (20 instead of 5) to support
|
||||
# both Completion API and Chat Completion API
|
||||
logprobs=21,
|
||||
)
|
||||
...
|
||||
with pytest.raises(
|
||||
(openai.BadRequestError, openai.APIError)
|
||||
): # test using token IDs
|
||||
stream = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=[0, 0, 0, 0, 0],
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
# vLLM has higher default max_logprobs (20 instead of 5) to support
|
||||
# both Completion API and Chat Completion API
|
||||
logprobs=30,
|
||||
stream=True,
|
||||
)
|
||||
async for chunk in stream:
|
||||
...
|
||||
|
||||
# the server should still work afterwards
|
||||
completion = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=[0, 0, 0, 0, 0],
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
)
|
||||
assert len(completion.choices[0].text) >= 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name, prompt_logprobs",
|
||||
[(MODEL_NAME, -1), (MODEL_NAME, 0), (MODEL_NAME, 1), (MODEL_NAME, None)],
|
||||
)
|
||||
async def test_prompt_logprobs_completion(
|
||||
client: openai.AsyncOpenAI, model_name: str, prompt_logprobs: int | None
|
||||
):
|
||||
params: dict = {
|
||||
"prompt": ["A robot may not injure another robot", "My name is"],
|
||||
"model": model_name,
|
||||
}
|
||||
if prompt_logprobs is not None:
|
||||
params["extra_body"] = {"prompt_logprobs": prompt_logprobs}
|
||||
|
||||
if prompt_logprobs is not None and prompt_logprobs < 0:
|
||||
with pytest.raises(BadRequestError):
|
||||
await client.completions.create(**params)
|
||||
else:
|
||||
completion = await client.completions.create(**params)
|
||||
if prompt_logprobs is not None:
|
||||
assert completion.choices[0].prompt_logprobs is not None
|
||||
assert len(completion.choices[0].prompt_logprobs) > 0
|
||||
|
||||
assert completion.choices[1].prompt_logprobs is not None
|
||||
assert len(completion.choices[1].prompt_logprobs) > 0
|
||||
|
||||
else:
|
||||
assert completion.choices[0].prompt_logprobs is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_completion_streaming(
|
||||
client: openai.AsyncOpenAI, model_name: str
|
||||
) -> None:
|
||||
prompt = "What is an LLM?"
|
||||
|
||||
single_completion = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
)
|
||||
single_output = single_completion.choices[0].text
|
||||
stream = await client.completions.create(
|
||||
model=model_name, prompt=prompt, max_tokens=5, temperature=0.0, stream=True
|
||||
)
|
||||
chunks: list[str] = []
|
||||
finish_reason_count = 0
|
||||
async for chunk in stream:
|
||||
chunks.append(chunk.choices[0].text)
|
||||
if chunk.choices[0].finish_reason is not None:
|
||||
finish_reason_count += 1
|
||||
# finish reason should only return in last block
|
||||
assert finish_reason_count == 1
|
||||
assert chunk.choices[0].finish_reason == "length"
|
||||
assert chunk.choices[0].text
|
||||
assert "".join(chunks) == single_output
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_parallel_no_streaming(client: openai.AsyncOpenAI, model_name: str):
|
||||
"""Parallel sampling without streaming.
|
||||
A single request output contains a list of completions.
|
||||
"""
|
||||
|
||||
prompt = "What is an LLM?"
|
||||
n = 3
|
||||
max_tokens = 50 # we want some to finish earlier than others
|
||||
|
||||
# High temperature to maximize chance of unique completions.
|
||||
completion = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=max_tokens,
|
||||
n=n,
|
||||
temperature=1.0,
|
||||
stream=False,
|
||||
logprobs=0,
|
||||
seed=42,
|
||||
)
|
||||
|
||||
# Assert `n` completions
|
||||
num_completions = len(completion.choices)
|
||||
assert num_completions == n, f"Num completions {num_completions} but expected {n}."
|
||||
completion_repeats: dict[str, int] = {}
|
||||
output_token_lengths = set()
|
||||
for idx, choice in enumerate(completion.choices):
|
||||
# Assert correct completion index & some finish reason.
|
||||
assert choice.index == idx, f"Index {choice.index} but expected {idx}."
|
||||
assert choice.finish_reason is not None, "None finish_reason is invalid."
|
||||
text = choice.text
|
||||
completion_repeats[text] = completion_repeats.get(text, 0) + 1
|
||||
output_token_lengths.add(len(choice.logprobs.tokens))
|
||||
# Assert subrequests finished at different times
|
||||
assert len(output_token_lengths) > 1
|
||||
# Assert `n` unique completions
|
||||
num_unique = len(completion_repeats)
|
||||
if num_unique != n:
|
||||
repeats = {txt: num for (txt, num) in completion_repeats.items() if num > 1}
|
||||
raise AssertionError(
|
||||
f"Expected {n} unique completions, got {num_unique}; repeats: {repeats}."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str):
|
||||
"""Streaming for parallel sampling.
|
||||
The tokens from multiple samples, are flattened into a single stream,
|
||||
with an index to indicate which sample the token belongs to.
|
||||
"""
|
||||
|
||||
prompt = "What is an LLM?"
|
||||
n = 3
|
||||
max_tokens = 50 # we want some to finish earlier than others
|
||||
|
||||
stream = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=max_tokens,
|
||||
n=n,
|
||||
temperature=1.0,
|
||||
stream=True,
|
||||
seed=42,
|
||||
)
|
||||
chunks: list[list[str]] = [[] for _ in range(n)]
|
||||
finish_reason_count = 0
|
||||
async for chunk in stream:
|
||||
index = chunk.choices[0].index
|
||||
text = chunk.choices[0].text
|
||||
chunks[index].append(text)
|
||||
if chunk.choices[0].finish_reason is not None:
|
||||
finish_reason_count += 1
|
||||
# Assert `n` completions with correct finish reasons
|
||||
assert finish_reason_count == n, (
|
||||
f"Expected {n} completions with valid indices and finish_reason."
|
||||
)
|
||||
completion_repeats: dict[str, int] = {}
|
||||
chunk_lengths = set()
|
||||
for chunk in chunks:
|
||||
chunk_len = len(chunk)
|
||||
# Assert correct number of completion tokens
|
||||
chunk_lengths.add(chunk_len)
|
||||
assert chunk_len <= max_tokens, (
|
||||
f"max_tokens={max_tokens} but chunk len is {chunk_len}."
|
||||
)
|
||||
text = "".join(chunk)
|
||||
completion_repeats[text] = completion_repeats.get(text, 0) + 1
|
||||
print(text)
|
||||
# Assert subrequests finished at different times
|
||||
assert len(chunk_lengths) > 1
|
||||
# Assert `n` unique completions
|
||||
num_unique = len(completion_repeats)
|
||||
if num_unique != n:
|
||||
repeats = {txt: num for (txt, num) in completion_repeats.items() if num > 1}
|
||||
raise AssertionError(
|
||||
f"{num_unique} unique completions, expected {n}; repeats: {repeats}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_completion_stream_options(client: openai.AsyncOpenAI, model_name: str):
|
||||
prompt = "What is the capital of France?"
|
||||
|
||||
# Test stream=True, stream_options=
|
||||
# {"include_usage": False, "continuous_usage_stats": False}
|
||||
stream = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
stream=True,
|
||||
stream_options={
|
||||
"include_usage": False,
|
||||
"continuous_usage_stats": False,
|
||||
},
|
||||
)
|
||||
|
||||
async for chunk in stream:
|
||||
assert chunk.usage is None
|
||||
|
||||
# Test stream=True, stream_options=
|
||||
# {"include_usage": False, "continuous_usage_stats": True}
|
||||
stream = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
stream=True,
|
||||
stream_options={
|
||||
"include_usage": False,
|
||||
"continuous_usage_stats": True,
|
||||
},
|
||||
)
|
||||
async for chunk in stream:
|
||||
assert chunk.usage is None
|
||||
|
||||
# Test stream=True, stream_options=
|
||||
# {"include_usage": True, "continuous_usage_stats": False}
|
||||
stream = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
stream=True,
|
||||
stream_options={
|
||||
"include_usage": True,
|
||||
"continuous_usage_stats": False,
|
||||
},
|
||||
)
|
||||
async for chunk in stream:
|
||||
if chunk.choices[0].finish_reason is None:
|
||||
assert chunk.usage is None
|
||||
else:
|
||||
assert chunk.usage is None
|
||||
final_chunk = await anext(stream)
|
||||
assert final_chunk.usage is not None
|
||||
assert final_chunk.usage.prompt_tokens > 0
|
||||
assert final_chunk.usage.completion_tokens > 0
|
||||
assert final_chunk.usage.total_tokens == (
|
||||
final_chunk.usage.prompt_tokens + final_chunk.usage.completion_tokens
|
||||
)
|
||||
assert final_chunk.choices == []
|
||||
|
||||
# Test stream=True, stream_options=
|
||||
# {"include_usage": True, "continuous_usage_stats": True}
|
||||
stream = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
stream=True,
|
||||
stream_options={
|
||||
"include_usage": True,
|
||||
"continuous_usage_stats": True,
|
||||
},
|
||||
)
|
||||
async for chunk in stream:
|
||||
assert chunk.usage is not None
|
||||
assert chunk.usage.prompt_tokens > 0
|
||||
assert chunk.usage.completion_tokens > 0
|
||||
assert chunk.usage.total_tokens == (
|
||||
chunk.usage.prompt_tokens + chunk.usage.completion_tokens
|
||||
)
|
||||
if chunk.choices[0].finish_reason is not None:
|
||||
final_chunk = await anext(stream)
|
||||
assert final_chunk.usage is not None
|
||||
assert final_chunk.usage.prompt_tokens > 0
|
||||
assert final_chunk.usage.completion_tokens > 0
|
||||
assert final_chunk.usage.total_tokens == (
|
||||
final_chunk.usage.prompt_tokens + final_chunk.usage.completion_tokens
|
||||
)
|
||||
assert final_chunk.choices == []
|
||||
|
||||
# Test stream=True, stream_options={}
|
||||
stream = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
stream=True,
|
||||
stream_options={},
|
||||
)
|
||||
async for chunk in stream:
|
||||
assert chunk.usage is None
|
||||
|
||||
# Test stream=False, stream_options=
|
||||
# {"include_usage": None}
|
||||
with pytest.raises(BadRequestError):
|
||||
await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
stream=False,
|
||||
stream_options={"include_usage": None},
|
||||
)
|
||||
|
||||
# Test stream=False, stream_options=
|
||||
# {"include_usage": True}
|
||||
with pytest.raises(BadRequestError):
|
||||
await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
stream=False,
|
||||
stream_options={"include_usage": True},
|
||||
)
|
||||
|
||||
# Test stream=False, stream_options=
|
||||
# {"continuous_usage_stats": None}
|
||||
with pytest.raises(BadRequestError):
|
||||
await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
stream=False,
|
||||
stream_options={"continuous_usage_stats": None},
|
||||
)
|
||||
|
||||
# Test stream=False, stream_options=
|
||||
# {"continuous_usage_stats": True}
|
||||
with pytest.raises(BadRequestError):
|
||||
await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
stream=False,
|
||||
stream_options={"continuous_usage_stats": True},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_batch_completions(client: openai.AsyncOpenAI, model_name: str):
|
||||
# test both text and token IDs
|
||||
for prompts in (["Hello, my name is"] * 2, [[0, 0, 0, 0, 0]] * 2):
|
||||
# test simple list
|
||||
batch = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompts,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
)
|
||||
assert len(batch.choices) == 2
|
||||
assert batch.choices[0].text == batch.choices[1].text
|
||||
|
||||
# test n = 2
|
||||
batch = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompts,
|
||||
n=2,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
extra_body=dict(
|
||||
# NOTE: this has to be true for n > 1 in vLLM, but
|
||||
# not necessary for official client.
|
||||
use_beam_search=True
|
||||
),
|
||||
)
|
||||
assert len(batch.choices) == 4
|
||||
assert batch.choices[0].text != batch.choices[1].text, (
|
||||
"beam search should be different"
|
||||
)
|
||||
assert batch.choices[0].text == batch.choices[2].text, (
|
||||
"two copies of the same prompt should be the same"
|
||||
)
|
||||
assert batch.choices[1].text == batch.choices[3].text, (
|
||||
"two copies of the same prompt should be the same"
|
||||
)
|
||||
|
||||
# test streaming
|
||||
batch = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompts,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
stream=True,
|
||||
)
|
||||
texts = [""] * 2
|
||||
async for chunk in batch:
|
||||
assert len(chunk.choices) == 1
|
||||
choice = chunk.choices[0]
|
||||
texts[choice.index] += choice.text
|
||||
assert texts[0] == texts[1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
@pytest.mark.parametrize("logprobs_arg", [1, 0])
|
||||
async def test_echo_logprob_completion(
|
||||
client: openai.AsyncOpenAI, model_name: str, logprobs_arg: int
|
||||
):
|
||||
tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME)
|
||||
# test using text and token IDs
|
||||
for prompt in ("Hello, my name is", [0, 0, 0, 0, 0]):
|
||||
completion = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
echo=True,
|
||||
logprobs=logprobs_arg,
|
||||
)
|
||||
|
||||
prompt_text = tokenizer.decode(prompt) if isinstance(prompt, list) else prompt
|
||||
assert re.search(r"^" + prompt_text, completion.choices[0].text)
|
||||
logprobs = completion.choices[0].logprobs
|
||||
assert logprobs is not None
|
||||
assert len(logprobs.text_offset) > 5
|
||||
assert len(logprobs.token_logprobs) > 5 and logprobs.token_logprobs[0] is None
|
||||
assert len(logprobs.top_logprobs) > 5 and logprobs.top_logprobs[0] is None
|
||||
for top_logprobs in logprobs.top_logprobs[1:]:
|
||||
assert max(logprobs_arg, 1) <= len(top_logprobs) <= logprobs_arg + 1
|
||||
assert len(logprobs.tokens) > 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_invalid_json_schema(client: openai.AsyncOpenAI, model_name: str) -> None:
|
||||
invalid_json_schema = {
|
||||
"$defs": {
|
||||
"CarType": {
|
||||
"enum": ["sedan", "SUV", "Truck", "Coupe"],
|
||||
"title": "CarType",
|
||||
"type": "string",
|
||||
}
|
||||
},
|
||||
"properties": {
|
||||
"brand": {"title": "Brand", "type": "string"},
|
||||
"model": {"title": "Model", "type": "string"},
|
||||
"car_type": {"$ref": "#/$defs/CarType"},
|
||||
"foo": "bar",
|
||||
},
|
||||
"required": ["brand", "model", "car_type"],
|
||||
"title": "CarDescription",
|
||||
"type": "object",
|
||||
}
|
||||
prompt = (
|
||||
"Generate a JSON with the brand, model and car_type of"
|
||||
"the most iconic car from the 90's"
|
||||
)
|
||||
with pytest.raises((openai.BadRequestError, openai.APIError)):
|
||||
await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
extra_body={"structured_outputs": {"json": invalid_json_schema}},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_invalid_regex(client: openai.AsyncOpenAI, model_name: str):
|
||||
prompt = (
|
||||
"Generate an email address for Alan Turing, who works in Enigma."
|
||||
"End in .com and new line. Example result:"
|
||||
"alan.turing@enigma.com\n"
|
||||
)
|
||||
|
||||
with pytest.raises((openai.BadRequestError, openai.APIError)):
|
||||
await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
extra_body={"structured_outputs": {"regex": r"[.*"}, "stop": ["\n"]},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_invalid_grammar(client: openai.AsyncOpenAI, model_name: str):
|
||||
invalid_simplified_sql_grammar = """
|
||||
root ::= select_statementinvalidsyntax
|
||||
|
||||
select_statement ::= "SELECT " column " from " table " where " condition
|
||||
|
||||
column ::= "col_1 " | "col_2 "
|
||||
|
||||
table ::= "table_1 " | "table_2 "
|
||||
|
||||
condition ::= column "= " number
|
||||
|
||||
number ::= "1 " | "2 "
|
||||
"""
|
||||
|
||||
prompt = (
|
||||
"Generate an SQL query to show the 'username' and 'email'"
|
||||
"from the 'users' table."
|
||||
)
|
||||
with pytest.raises((openai.BadRequestError, openai.APIError)):
|
||||
await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
extra_body={
|
||||
"structured_outputs": {"grammar": invalid_simplified_sql_grammar}
|
||||
},
|
||||
)
|
||||
86
third_party/vllm/tests/v1/entrypoints/openai/test_completion_with_image_embeds.py
vendored
Normal file
86
third_party/vllm/tests/v1/entrypoints/openai/test_completion_with_image_embeds.py
vendored
Normal file
@@ -0,0 +1,86 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
|
||||
import openai # use the official client for correctness check
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import torch
|
||||
from transformers import AutoConfig
|
||||
|
||||
from tests.conftest import ImageTestAssets
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.utils.serial_utils import tensor2base64
|
||||
|
||||
# any model with a chat template should work here
|
||||
MODEL_NAME = "llava-hf/llava-1.5-7b-hf"
|
||||
CONFIG = AutoConfig.from_pretrained(MODEL_NAME)
|
||||
MAXIMUM_IMAGES = 2
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def default_image_embeds_server_args() -> list[str]:
|
||||
return [
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-model-len",
|
||||
"2048",
|
||||
"--max-num-seqs",
|
||||
"4",
|
||||
"--enforce-eager",
|
||||
"--limit-mm-per-prompt",
|
||||
json.dumps({"image": MAXIMUM_IMAGES}),
|
||||
"--enable-mm-embeds",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server_with_image_embeds(default_image_embeds_server_args):
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME, default_image_embeds_server_args, max_wait_seconds=600
|
||||
) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client_with_image_embeds(server_with_image_embeds):
|
||||
async with server_with_image_embeds.get_async_client() as async_client:
|
||||
yield async_client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("dtype", [torch.half, torch.float16, torch.float32])
|
||||
async def test_completions_with_image_embeds(
|
||||
client_with_image_embeds: openai.AsyncOpenAI,
|
||||
model_name: str,
|
||||
image_assets: ImageTestAssets,
|
||||
dtype: torch.dtype,
|
||||
):
|
||||
# Test case: Single image embeds input
|
||||
image_embeds = image_assets[0].image_embeds.to(dtype=dtype)
|
||||
base64_image_embedding = tensor2base64(image_embeds)
|
||||
chat_completion = await client_with_image_embeds.chat.completions.create(
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Describe these images separately. For each image,"
|
||||
"reply with a short sentence (no more than 10 words).",
|
||||
},
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": base64_image_embedding,
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
model=model_name,
|
||||
)
|
||||
assert chat_completion.choices[0].message.content is not None
|
||||
assert isinstance(chat_completion.choices[0].message.content, str)
|
||||
assert len(chat_completion.choices[0].message.content) > 0
|
||||
175
third_party/vllm/tests/v1/entrypoints/openai/test_multi_api_servers.py
vendored
Normal file
175
third_party/vllm/tests/v1/entrypoints/openai/test_multi_api_servers.py
vendored
Normal file
@@ -0,0 +1,175 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
import openai # use the official client for correctness check
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from tests.v1.utils import check_request_balancing
|
||||
|
||||
MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
|
||||
|
||||
DP_SIZE = os.getenv("DP_SIZE", "1")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def default_server_args():
|
||||
return [
|
||||
# use half precision for speed and memory savings in CI environment
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-model-len",
|
||||
"2048",
|
||||
"--max-num-seqs",
|
||||
"128",
|
||||
"--enforce-eager",
|
||||
"--api-server-count",
|
||||
"4",
|
||||
"--data_parallel_size",
|
||||
DP_SIZE,
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server(default_server_args):
|
||||
with RemoteOpenAIServer(MODEL_NAME, default_server_args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(server):
|
||||
async with server.get_async_client() as async_client:
|
||||
yield async_client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_single_completion(
|
||||
client: openai.AsyncOpenAI, server: RemoteOpenAIServer, model_name: str
|
||||
) -> None:
|
||||
async def make_request():
|
||||
completion = await client.completions.create(
|
||||
model=model_name, prompt="Hello, my name is", max_tokens=10, temperature=1.0
|
||||
)
|
||||
|
||||
assert completion.id is not None
|
||||
assert completion.choices is not None and len(completion.choices) == 1
|
||||
|
||||
choice = completion.choices[0]
|
||||
# The exact number of tokens can vary slightly with temperature=1.0,
|
||||
# so we check for a reasonable minimum length.
|
||||
assert len(choice.text) >= 1
|
||||
# Finish reason might not always be 'length' if the model finishes early
|
||||
# or due to other reasons, especially with high temperature.
|
||||
# So, we'll accept 'length' or 'stop'.
|
||||
assert choice.finish_reason in ("length", "stop")
|
||||
|
||||
# Token counts can also vary, so we check they are positive.
|
||||
assert completion.usage.completion_tokens > 0
|
||||
assert completion.usage.prompt_tokens > 0
|
||||
assert completion.usage.total_tokens > 0
|
||||
return completion
|
||||
|
||||
# Test single request
|
||||
result = await make_request()
|
||||
assert result is not None
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Send two bursts of requests
|
||||
num_requests = 100
|
||||
tasks = [make_request() for _ in range(num_requests)]
|
||||
results = await asyncio.gather(*tasks)
|
||||
assert len(results) == num_requests
|
||||
assert all(completion is not None for completion in results)
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
tasks = [make_request() for _ in range(num_requests)]
|
||||
results = await asyncio.gather(*tasks)
|
||||
assert len(results) == num_requests
|
||||
assert all(completion is not None for completion in results)
|
||||
|
||||
# Check request balancing via Prometheus metrics if DP_SIZE > 1
|
||||
check_request_balancing(server, int(DP_SIZE))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[MODEL_NAME],
|
||||
)
|
||||
async def test_completion_streaming(
|
||||
client: openai.AsyncOpenAI, server: RemoteOpenAIServer, model_name: str
|
||||
) -> None:
|
||||
prompt = "What is an LLM?"
|
||||
|
||||
async def make_streaming_request():
|
||||
# Perform a non-streaming request to get the expected full output
|
||||
single_completion = await client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt,
|
||||
max_tokens=5,
|
||||
temperature=0.0,
|
||||
)
|
||||
single_output = single_completion.choices[0].text
|
||||
|
||||
# Perform the streaming request
|
||||
stream = await client.completions.create(
|
||||
model=model_name, prompt=prompt, max_tokens=5, temperature=0.0, stream=True
|
||||
)
|
||||
chunks: list[str] = []
|
||||
finish_reason_count = 0
|
||||
last_chunk = None
|
||||
async for chunk in stream:
|
||||
chunks.append(chunk.choices[0].text)
|
||||
if chunk.choices[0].finish_reason is not None:
|
||||
finish_reason_count += 1
|
||||
last_chunk = chunk # Keep track of the last chunk
|
||||
|
||||
# finish reason should only return in the last block for OpenAI API
|
||||
assert finish_reason_count == 1, "Finish reason should appear exactly once."
|
||||
assert last_chunk is not None, "Stream should have yielded at least one chunk."
|
||||
assert last_chunk.choices[0].finish_reason == "length", (
|
||||
"Finish reason should be 'length'."
|
||||
)
|
||||
# Check that the combined text matches the non-streamed version.
|
||||
assert "".join(chunks) == single_output, (
|
||||
"Streamed output should match non-streamed output."
|
||||
)
|
||||
return True # Indicate success for this request
|
||||
|
||||
# Test single request
|
||||
result = await make_streaming_request()
|
||||
assert result is not None
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Send two bursts of requests
|
||||
num_requests = 100
|
||||
tasks = [make_streaming_request() for _ in range(num_requests)]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
assert len(results) == num_requests, (
|
||||
f"Expected {num_requests} results, got {len(results)}"
|
||||
)
|
||||
assert all(results), "Not all streaming requests completed successfully."
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
tasks = [make_streaming_request() for _ in range(num_requests)]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
assert len(results) == num_requests, (
|
||||
f"Expected {num_requests} results, got {len(results)}"
|
||||
)
|
||||
assert all(results), "Not all streaming requests completed successfully."
|
||||
|
||||
# Check request balancing via Prometheus metrics if DP_SIZE > 1
|
||||
check_request_balancing(server, int(DP_SIZE))
|
||||
0
third_party/vllm/tests/v1/executor/__init__.py
vendored
Normal file
0
third_party/vllm/tests/v1/executor/__init__.py
vendored
Normal file
136
third_party/vllm/tests/v1/executor/test_executor.py
vendored
Normal file
136
third_party/vllm/tests/v1/executor/test_executor.py
vendored
Normal file
@@ -0,0 +1,136 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
from concurrent.futures import Future
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
from vllm.v1.engine.llm_engine import LLMEngine
|
||||
from vllm.v1.executor.multiproc_executor import MultiprocExecutor
|
||||
|
||||
|
||||
class Mock: ...
|
||||
|
||||
|
||||
class CustomMultiprocExecutor(MultiprocExecutor):
|
||||
def collective_rpc(
|
||||
self,
|
||||
method: str | Callable,
|
||||
timeout: float | None = None,
|
||||
args: tuple = (),
|
||||
kwargs: dict | None = None,
|
||||
non_block: bool = False,
|
||||
unique_reply_rank: int | None = None,
|
||||
kv_output_aggregator: KVOutputAggregator = None,
|
||||
) -> Any | list[Any] | Future[Any | list[Any]]:
|
||||
# Drop marker to show that this was run
|
||||
with open(".marker", "w"):
|
||||
...
|
||||
return super().collective_rpc(
|
||||
method,
|
||||
timeout,
|
||||
args,
|
||||
kwargs,
|
||||
non_block,
|
||||
unique_reply_rank,
|
||||
kv_output_aggregator,
|
||||
)
|
||||
|
||||
|
||||
CustomMultiprocExecutorAsync = CustomMultiprocExecutor
|
||||
MODEL = "Qwen/Qwen3-0.6B"
|
||||
|
||||
|
||||
def test_custom_executor_type_checking():
|
||||
with pytest.raises(ValueError):
|
||||
engine_args = EngineArgs(
|
||||
model=MODEL,
|
||||
gpu_memory_utilization=0.2,
|
||||
max_model_len=8192,
|
||||
distributed_executor_backend=Mock,
|
||||
)
|
||||
LLMEngine.from_engine_args(engine_args)
|
||||
with pytest.raises(ValueError):
|
||||
engine_args = AsyncEngineArgs(
|
||||
model=MODEL,
|
||||
gpu_memory_utilization=0.2,
|
||||
max_model_len=8192,
|
||||
distributed_executor_backend=Mock,
|
||||
)
|
||||
AsyncLLM.from_engine_args(engine_args)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"distributed_executor_backend",
|
||||
[
|
||||
CustomMultiprocExecutor,
|
||||
"tests.v1.executor.test_executor.CustomMultiprocExecutor",
|
||||
],
|
||||
)
|
||||
def test_custom_executor(distributed_executor_backend, tmp_path):
|
||||
cwd = os.path.abspath(".")
|
||||
os.chdir(tmp_path)
|
||||
try:
|
||||
assert not os.path.exists(".marker")
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model=MODEL,
|
||||
gpu_memory_utilization=0.2,
|
||||
max_model_len=8192,
|
||||
distributed_executor_backend=distributed_executor_backend,
|
||||
enforce_eager=True, # reduce test time
|
||||
)
|
||||
engine = LLMEngine.from_engine_args(engine_args)
|
||||
sampling_params = SamplingParams(max_tokens=1)
|
||||
|
||||
engine.add_request("0", "foo", sampling_params)
|
||||
engine.step()
|
||||
|
||||
assert os.path.exists(".marker")
|
||||
finally:
|
||||
os.chdir(cwd)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"distributed_executor_backend",
|
||||
[
|
||||
CustomMultiprocExecutorAsync,
|
||||
"tests.v1.executor.test_executor.CustomMultiprocExecutorAsync",
|
||||
],
|
||||
)
|
||||
def test_custom_executor_async(distributed_executor_backend, tmp_path):
|
||||
cwd = os.path.abspath(".")
|
||||
os.chdir(tmp_path)
|
||||
try:
|
||||
assert not os.path.exists(".marker")
|
||||
|
||||
engine_args = AsyncEngineArgs(
|
||||
model=MODEL,
|
||||
gpu_memory_utilization=0.2,
|
||||
max_model_len=8192,
|
||||
distributed_executor_backend=distributed_executor_backend,
|
||||
enforce_eager=True, # reduce test time
|
||||
)
|
||||
engine = AsyncLLM.from_engine_args(engine_args)
|
||||
sampling_params = SamplingParams(max_tokens=1)
|
||||
|
||||
async def t():
|
||||
stream = engine.generate(
|
||||
request_id="0", prompt="foo", sampling_params=sampling_params
|
||||
)
|
||||
async for x in stream:
|
||||
...
|
||||
|
||||
asyncio.run(t())
|
||||
|
||||
assert os.path.exists(".marker")
|
||||
finally:
|
||||
os.chdir(cwd)
|
||||
0
third_party/vllm/tests/v1/kv_connector/__init__.py
vendored
Normal file
0
third_party/vllm/tests/v1/kv_connector/__init__.py
vendored
Normal file
0
third_party/vllm/tests/v1/kv_connector/extract_hidden_states_integration/__init__.py
vendored
Normal file
0
third_party/vllm/tests/v1/kv_connector/extract_hidden_states_integration/__init__.py
vendored
Normal file
120
third_party/vllm/tests/v1/kv_connector/extract_hidden_states_integration/predictable_llama.py
vendored
Normal file
120
third_party/vllm/tests/v1/kv_connector/extract_hidden_states_integration/predictable_llama.py
vendored
Normal file
@@ -0,0 +1,120 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Predictable dummy model for testing extract_hidden_states.
|
||||
|
||||
Subclasses LlamaForCausalLM but overrides the model to produce deterministic
|
||||
hidden states: layer i outputs values equal to (i).
|
||||
"""
|
||||
|
||||
from collections.abc import Iterable
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.model_executor.models.interfaces import EagleModelMixin
|
||||
from vllm.model_executor.models.llama import LlamaForCausalLM
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
|
||||
class PredictableLlamaModel(nn.Module, EagleModelMixin):
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
self.config = vllm_config.model_config.hf_config
|
||||
|
||||
# Create minimal embed_tokens for embedding
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
VocabParallelEmbedding,
|
||||
)
|
||||
|
||||
self.embed_tokens = VocabParallelEmbedding(
|
||||
self.config.vocab_size,
|
||||
self.config.hidden_size,
|
||||
)
|
||||
|
||||
# Required for pipeline parallelism
|
||||
from vllm.model_executor.models.utils import (
|
||||
make_empty_intermediate_tensors_factory,
|
||||
)
|
||||
|
||||
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
|
||||
["hidden_states", "residual"], self.config.hidden_size
|
||||
)
|
||||
|
||||
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
"""Embed input IDs."""
|
||||
return self.embed_tokens(input_ids)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**extra_layer_kwargs,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
|
||||
"""Forward pass that produces predictable outputs.
|
||||
|
||||
Returns:
|
||||
If aux_hidden_state_layers is set: (hidden_states, aux_hidden_states)
|
||||
Otherwise: hidden_states
|
||||
"""
|
||||
# Determine sequence length
|
||||
if inputs_embeds is not None:
|
||||
seq_len = inputs_embeds.shape[0]
|
||||
device = inputs_embeds.device
|
||||
elif input_ids is not None:
|
||||
seq_len = input_ids.shape[0] if input_ids.ndim == 1 else input_ids.shape[-1]
|
||||
device = input_ids.device
|
||||
else:
|
||||
raise ValueError("Either input_ids or inputs_embeds must be provided")
|
||||
|
||||
# Final hidden states (last layer value)
|
||||
hidden_states = torch.full(
|
||||
(seq_len, self.config.hidden_size),
|
||||
fill_value=float(self.config.num_hidden_layers),
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
# Check if we need auxiliary hidden states
|
||||
if len(self.aux_hidden_state_layers) > 0:
|
||||
aux_hidden_states = []
|
||||
for layer_idx in self.aux_hidden_state_layers:
|
||||
# Fill with (layer_idx) for predictability
|
||||
layer_hidden = torch.full(
|
||||
(seq_len, self.config.hidden_size),
|
||||
fill_value=float(layer_idx),
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
aux_hidden_states.append(layer_hidden)
|
||||
|
||||
return hidden_states, aux_hidden_states
|
||||
|
||||
return hidden_states
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
"""Skip weight loading."""
|
||||
return set()
|
||||
|
||||
|
||||
class PredictableLlamaForCausalLM(LlamaForCausalLM):
|
||||
"""Predictable Llama model for testing.
|
||||
|
||||
Overrides _init_model to use PredictableLlamaModel instead of LlamaModel.
|
||||
"""
|
||||
|
||||
def _init_model(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
prefix: str = "",
|
||||
layer_type: type[nn.Module] | None = None,
|
||||
):
|
||||
"""Initialize with predictable model."""
|
||||
return PredictableLlamaModel(vllm_config=vllm_config, prefix=prefix)
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
"""Skip weight loading for dummy model."""
|
||||
return set()
|
||||
155
third_party/vllm/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py
vendored
Normal file
155
third_party/vllm/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py
vendored
Normal file
@@ -0,0 +1,155 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import gc
|
||||
import os
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
|
||||
from vllm import LLM, ModelRegistry, SamplingParams
|
||||
|
||||
|
||||
def get_and_check_output(output, expected_shape):
|
||||
assert output.kv_transfer_params is not None
|
||||
hidden_states_path = output.kv_transfer_params.get("hidden_states_path")
|
||||
assert hidden_states_path is not None
|
||||
assert os.path.exists(hidden_states_path)
|
||||
|
||||
# Load and verify the saved tensors
|
||||
with safe_open(hidden_states_path, "pt") as f:
|
||||
# Check that token_ids and hidden_states are present
|
||||
tensor_names = f.keys()
|
||||
assert "token_ids" in tensor_names
|
||||
assert "hidden_states" in tensor_names
|
||||
|
||||
token_ids = f.get_tensor("token_ids")
|
||||
hidden_states = f.get_tensor("hidden_states")
|
||||
|
||||
prompt_token_ids = output.prompt_token_ids
|
||||
assert torch.equal(token_ids, torch.tensor(prompt_token_ids))
|
||||
|
||||
assert hidden_states.shape == expected_shape
|
||||
|
||||
# Verify hidden_states are not all zeros (i.e., they were actually computed)
|
||||
assert not torch.allclose(hidden_states, torch.zeros_like(hidden_states))
|
||||
|
||||
return token_ids, hidden_states
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def predictable_llama_config_path(tmp_path_factory):
|
||||
"""Create a minimal LlamaConfig for PredictableLlamaForCausalLM."""
|
||||
from transformers import LlamaConfig, LlamaTokenizerFast
|
||||
|
||||
config_dir = tmp_path_factory.mktemp("predictable_llama")
|
||||
|
||||
# Create a minimal Llama config with small dimensions
|
||||
config = LlamaConfig(
|
||||
vocab_size=1000,
|
||||
hidden_size=256,
|
||||
intermediate_size=512,
|
||||
num_hidden_layers=24, # Enough layers to test various layer_ids
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=4,
|
||||
max_position_embeddings=128,
|
||||
architectures=["PredictableLlamaForCausalLM"],
|
||||
)
|
||||
|
||||
# Save config
|
||||
config.save_pretrained(config_dir)
|
||||
|
||||
# Create a simple tokenizer
|
||||
tokenizer = LlamaTokenizerFast.from_pretrained(
|
||||
"TinyLlama/TinyLlama-1.1B-Chat-v1.0",
|
||||
cache_dir=os.path.expanduser("~/.cache/huggingface"),
|
||||
)
|
||||
tokenizer.save_pretrained(config_dir)
|
||||
|
||||
return str(config_dir)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", autouse=True)
|
||||
def register_predictable_model():
|
||||
"""Register the PredictableLlamaForCausalLM model."""
|
||||
from .predictable_llama import PredictableLlamaForCausalLM
|
||||
|
||||
if "PredictableLlamaForCausalLM" not in ModelRegistry.get_supported_archs():
|
||||
ModelRegistry.register_model(
|
||||
"PredictableLlamaForCausalLM", PredictableLlamaForCausalLM
|
||||
)
|
||||
yield
|
||||
|
||||
|
||||
def test_extract_hidden_states_with_predictable_dummy_model(
|
||||
predictable_llama_config_path, tmp_path
|
||||
):
|
||||
"""Comprehensive test using a predictable dummy model with synthetic weights.
|
||||
|
||||
The PredictableLlamaForCausalLM outputs deterministic hidden states where
|
||||
each layer produces values equal to (layer_index). This test verifies:
|
||||
1. Hidden states are correctly extracted from requested layers
|
||||
2. Values match the expected predictable pattern
|
||||
3. Layer ordering is preserved correctly (non-sequential layer IDs)
|
||||
4. Multiple prompts of different lengths produce consistent layer values
|
||||
"""
|
||||
# Test with non-sequential layer ordering to verify correct association
|
||||
layer_ids = [5, 2, 10]
|
||||
num_layers = len(layer_ids)
|
||||
|
||||
llm = LLM(
|
||||
model=predictable_llama_config_path,
|
||||
speculative_config={
|
||||
"method": "extract_hidden_states",
|
||||
"num_speculative_tokens": 1,
|
||||
"draft_model_config": {
|
||||
"hf_config": {"eagle_aux_hidden_state_layer_ids": layer_ids}
|
||||
},
|
||||
},
|
||||
kv_transfer_config={
|
||||
"kv_connector": "ExampleHiddenStatesConnector",
|
||||
"kv_role": "kv_producer",
|
||||
"kv_connector_extra_config": {"shared_storage_path": tmp_path},
|
||||
},
|
||||
max_model_len=128,
|
||||
enforce_eager=True,
|
||||
trust_remote_code=True,
|
||||
load_format="dummy", # Don't try to load real weights
|
||||
)
|
||||
|
||||
# Test with multiple prompts of different lengths
|
||||
prompts = [
|
||||
"Short",
|
||||
"Medium length",
|
||||
"Much longer prompt with many tokens",
|
||||
"Much longer prompt with many tokens", # repeated prompt
|
||||
]
|
||||
sampling_params = SamplingParams(max_tokens=1, temperature=0.0)
|
||||
hidden_size = llm.llm_engine.model_config.get_hidden_size()
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
del llm
|
||||
gc.collect()
|
||||
|
||||
assert len(outputs) == len(prompts)
|
||||
|
||||
for output in outputs:
|
||||
# hidden_states shape is [prompt_len, num_hidden_layers, hidden_size]
|
||||
expected_shape = (
|
||||
len(output.prompt_token_ids),
|
||||
num_layers,
|
||||
hidden_size,
|
||||
)
|
||||
_token_ids, hidden_states = get_and_check_output(output, expected_shape)
|
||||
|
||||
for idx, layer_id in enumerate(layer_ids):
|
||||
layer_hidden = hidden_states[:, idx, :]
|
||||
assert torch.allclose(
|
||||
layer_hidden,
|
||||
torch.full_like(layer_hidden, layer_id),
|
||||
atol=1e-5,
|
||||
), (
|
||||
f"Layer {layer_id} at position {idx} should output {float(layer_id)}, "
|
||||
f"but got mean={layer_hidden.mean():.3f}, "
|
||||
f"min={layer_hidden.min():.3f}, max={layer_hidden.max():.3f}"
|
||||
)
|
||||
90
third_party/vllm/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
vendored
Executable file
90
third_party/vllm/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
vendored
Executable file
@@ -0,0 +1,90 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
# Utility to run integration tests sequentially with varying TP configurations.
|
||||
SCRIPT="v1/kv_connector/nixl_integration/run_accuracy_test.sh"
|
||||
|
||||
# Define test configurations
|
||||
tp_configs=(
|
||||
"GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2"
|
||||
"GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2"
|
||||
"GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1"
|
||||
"GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA case
|
||||
"GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny"
|
||||
"GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny"
|
||||
"GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192" # SW model
|
||||
)
|
||||
dp_ep_configs=(
|
||||
"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP1, D-DPEP=2 (TP=1)
|
||||
"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1)
|
||||
)
|
||||
hybrid_ssm_configs=(
|
||||
"ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code"
|
||||
# TODO: (NickLucche) Address async scheduling issue with TP>1 separately as this may impact other models.
|
||||
"ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling"
|
||||
)
|
||||
|
||||
# Select config array based on DP_EP env var
|
||||
if [[ -n "${DP_EP:-}" ]]; then
|
||||
configs=("${dp_ep_configs[@]}")
|
||||
echo "DP_EP is set, using dp_ep_configs"
|
||||
elif [[ -n "${HYBRID_SSM:-}" ]]; then
|
||||
configs=("${hybrid_ssm_configs[@]}")
|
||||
echo "HYBRID_SSM is set, using hybrid_ssm_configs."
|
||||
else
|
||||
configs=("${tp_configs[@]}")
|
||||
fi
|
||||
|
||||
if [[ -n "${ENABLE_HMA_FLAG:-}" ]]; then
|
||||
# Append ENABLE_HMA_FLAG=1 to each config in the selected array
|
||||
echo "ENABLE_HMA_FLAG is set, appending ENABLE_HMA_FLAG=1 to each config"
|
||||
for i in "${!configs[@]}"; do
|
||||
configs[$i]="ENABLE_HMA_FLAG=1 ${configs[$i]}"
|
||||
done
|
||||
fi
|
||||
|
||||
run_tests() {
|
||||
local label=$1
|
||||
local extra_args=$2
|
||||
|
||||
echo "=== Running tests (${label}) ==="
|
||||
for cfg in "${configs[@]}"; do
|
||||
local -a cfg_parts extra_args_parts
|
||||
read -r -a cfg_parts <<< "$cfg"
|
||||
read -r -a extra_args_parts <<< "$extra_args"
|
||||
|
||||
echo "-> Running with ${cfg} ${extra_args:+and ${extra_args}}"
|
||||
# Use 'env' to safely set variables without eval
|
||||
# keep argv splitting safe and SC2086-clean via arrays.
|
||||
if ! env "${cfg_parts[@]}" bash "${SCRIPT}" "${extra_args_parts[@]}"; then
|
||||
echo "❌ Test failed for config: ${cfg} ${extra_args:+(${extra_args})}"
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
echo "✅ All ${label} tests passed!"
|
||||
}
|
||||
|
||||
# Set backend
|
||||
label="default backend"
|
||||
cmdline_args=""
|
||||
if [[ -n "${ROCM_ATTN:-}" ]]; then
|
||||
echo "ROCM_ATTN is set, running with --attention-backend ROCM_ATTN"
|
||||
label="ROCM_ATTN backend"
|
||||
cmdline_args=" --attention-backend ROCM_ATTN "
|
||||
elif [[ -n "${FLASHINFER:-}" ]]; then
|
||||
echo "FLASHINFER is set, running with --attention-backend FLASHINFER"
|
||||
label="FLASHINFER backend"
|
||||
cmdline_args=" --attention-backend FLASHINFER "
|
||||
else
|
||||
echo "running with default attention backend"
|
||||
fi
|
||||
|
||||
# Check if cross-layers is enabled (non-empty)
|
||||
if [[ -n "${CROSS_LAYERS_BLOCKS:-}" ]]; then
|
||||
echo "CROSS_LAYERS_BLOCKS is set, running with --enable-cross-layers"
|
||||
label+=" - CROSS_LAYERS_BLOCKS enabled"
|
||||
cmdline_args+=" --enable-cross-layers "
|
||||
fi
|
||||
|
||||
# Run tests
|
||||
run_tests "${label}" "${cmdline_args}"
|
||||
298
third_party/vllm/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
vendored
Executable file
298
third_party/vllm/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
vendored
Executable file
@@ -0,0 +1,298 @@
|
||||
#!/bin/bash
|
||||
set -xe
|
||||
|
||||
# Parse command line arguments
|
||||
KV_BUFFER_DEVICE="cuda" # Default to cuda
|
||||
ATTENTION_BACKEND="" # Default to empty (use vllm default)
|
||||
CROSS_LAYERS_BLOCKS="False"
|
||||
ENABLE_HMA_VAR="" # Default to empty (HMA disabled by default for kv connector)
|
||||
# Check for ENABLE_HMA_FLAG environment variable
|
||||
if [[ -n "${ENABLE_HMA_FLAG:-}" ]]; then
|
||||
ENABLE_HMA_VAR="--no-disable-hybrid-kv-cache-manager"
|
||||
fi
|
||||
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case $1 in
|
||||
--kv_buffer_device)
|
||||
KV_BUFFER_DEVICE="$2"
|
||||
shift 2
|
||||
;;
|
||||
--attention-backend)
|
||||
ATTENTION_BACKEND="$2"
|
||||
shift 2
|
||||
;;
|
||||
--enable-cross-layers)
|
||||
CROSS_LAYERS_BLOCKS="True"
|
||||
shift 1
|
||||
;;
|
||||
*)
|
||||
echo "Unknown option $1"
|
||||
echo "Usage: $0 [--kv_buffer_device <cuda|cpu>] [--attention-backend <backend>]"
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
echo "Running accuracy tests with kv_buffer_device=$KV_BUFFER_DEVICE"
|
||||
if [[ -n "$ATTENTION_BACKEND" ]]; then
|
||||
echo "Using attention backend: $ATTENTION_BACKEND"
|
||||
fi
|
||||
if [[ -n "$ENABLE_HMA_VAR" ]]; then
|
||||
echo "HMA (Hybrid KV Cache Manager) enabled"
|
||||
fi
|
||||
if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
|
||||
echo "vLLM serve extra args: $VLLM_SERVE_EXTRA_ARGS"
|
||||
fi
|
||||
|
||||
DECODER_KV_LAYOUT=${DECODER_KV_LAYOUT:-"HND"} # Default to HND, optional NHD
|
||||
if [[ "$DECODER_KV_LAYOUT" == "NHD" ]]; then
|
||||
KV_CONFIG_HETERO_LAYOUT=',"enable_permute_local_kv":"True"'
|
||||
else
|
||||
KV_CONFIG_HETERO_LAYOUT=''
|
||||
fi
|
||||
|
||||
if [[ "$CROSS_LAYERS_BLOCKS" == "True" ]]; then
|
||||
KV_EXTRA_CONFIG=',"kv_connector_extra_config":{"enable_cross_layers_blocks": "True"}'
|
||||
else
|
||||
KV_EXTRA_CONFIG=''
|
||||
fi
|
||||
|
||||
# Build the kv-transfer-config once
|
||||
if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then
|
||||
KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'
|
||||
else
|
||||
KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"
|
||||
fi
|
||||
|
||||
# Models to run
|
||||
MODEL_NAMES=${MODEL_NAMES:-}
|
||||
if [[ -n "$MODEL_NAMES" ]]; then
|
||||
MODELS=("$MODEL_NAMES")
|
||||
else
|
||||
MODELS=(
|
||||
"Qwen/Qwen3-0.6B"
|
||||
)
|
||||
fi
|
||||
|
||||
# Number of prefill and decode instances to create
|
||||
NUM_PREFILL_INSTANCES=${NUM_PREFILL_INSTANCES:-1} # Default to 1
|
||||
NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-1} # Default to 1
|
||||
PREFILLER_TP_SIZE=${PREFILLER_TP_SIZE:-1}
|
||||
DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}
|
||||
GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.2}
|
||||
PREFILL_BLOCK_SIZE=${PREFILL_BLOCK_SIZE:-128}
|
||||
DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128}
|
||||
# Comma-separated extra args for vllm serve (e.g. --max-model-len,2048)
|
||||
VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
|
||||
|
||||
# Find the git repository root directory
|
||||
GIT_ROOT=$(git rev-parse --show-toplevel)
|
||||
|
||||
SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "")
|
||||
|
||||
# Trap the SIGINT signal (triggered by Ctrl+C)
|
||||
trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT
|
||||
|
||||
# Waits for vLLM to start.
|
||||
wait_for_server() {
|
||||
local port=$1
|
||||
timeout 1200 bash -c "
|
||||
until curl -s localhost:${port}/v1/completions > /dev/null; do
|
||||
sleep 1
|
||||
done" && return 0 || return 1
|
||||
}
|
||||
|
||||
# Function to clean up previous instances
|
||||
cleanup_instances() {
|
||||
echo "Cleaning up any running vLLM instances..."
|
||||
pkill -f "vllm serve" || true
|
||||
sleep 2
|
||||
}
|
||||
|
||||
get_num_gpus() {
|
||||
if [[ "$SMI_BIN" == *"nvidia"* ]]; then
|
||||
$SMI_BIN --query-gpu=name --format=csv,noheader | wc -l
|
||||
elif [[ "$SMI_BIN" == *"rocm"* ]]; then
|
||||
$SMI_BIN -l | grep -c GPU
|
||||
else
|
||||
# works for non-cuda platforms,
|
||||
# assuming at least 1 device and
|
||||
# let system to decide which card to use
|
||||
echo "1"
|
||||
fi
|
||||
}
|
||||
|
||||
# Function to run tests for a specific model
|
||||
run_tests_for_model() {
|
||||
local model_name=$1
|
||||
echo "================================"
|
||||
echo "Testing model: $model_name"
|
||||
echo "================================"
|
||||
|
||||
# Arrays to store all hosts and ports
|
||||
PREFILL_HOSTS=()
|
||||
PREFILL_PORTS=()
|
||||
DECODE_HOSTS=()
|
||||
DECODE_PORTS=()
|
||||
|
||||
# Start prefill instances
|
||||
for i in $(seq 0 $((NUM_PREFILL_INSTANCES-1))); do
|
||||
# Calculate GPU ID - we'll distribute across available GPUs
|
||||
GPU_ID=$((i % $(get_num_gpus)))
|
||||
NEXT_GPU=${GPU_ID}
|
||||
# If PREFILLER_TP_SIZE is more than 1
|
||||
for (( j=1; j < PREFILLER_TP_SIZE; j++ )); do
|
||||
NEXT_GPU=$(((GPU_ID + j) % $(get_num_gpus)))
|
||||
GPU_ID="${GPU_ID},${NEXT_GPU}"
|
||||
done
|
||||
|
||||
# Calculate port number (base port + instance number)
|
||||
PORT=$((8100 + i))
|
||||
# Calculate side channel port. Avoid clash with with TP workers.
|
||||
SIDE_CHANNEL_PORT=$((5559 + i))
|
||||
|
||||
echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"
|
||||
|
||||
# Build the command with or without model-specific args
|
||||
BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
|
||||
VLLM_KV_CACHE_LAYOUT='HND' \
|
||||
UCX_NET_DEVICES=all \
|
||||
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
|
||||
vllm serve $model_name \
|
||||
--port $PORT \
|
||||
--enforce-eager \
|
||||
--block-size ${PREFILL_BLOCK_SIZE} \
|
||||
--gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
|
||||
--tensor-parallel-size $PREFILLER_TP_SIZE \
|
||||
--kv-transfer-config '$KV_CONFIG'"
|
||||
if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
|
||||
IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
|
||||
for arg in "${extra_args[@]}"; do
|
||||
BASE_CMD="${BASE_CMD} $arg"
|
||||
done
|
||||
fi
|
||||
|
||||
# Add attention backend config if specified
|
||||
if [[ -n "$ATTENTION_BACKEND" ]]; then
|
||||
BASE_CMD="${BASE_CMD} --attention-backend=$ATTENTION_BACKEND"
|
||||
fi
|
||||
|
||||
# Add HMA flag if specified
|
||||
if [[ -n "$ENABLE_HMA_VAR" ]]; then
|
||||
BASE_CMD="${BASE_CMD} $ENABLE_HMA_VAR"
|
||||
fi
|
||||
|
||||
FULL_CMD="$BASE_CMD"
|
||||
eval "$FULL_CMD &"
|
||||
|
||||
# Store host and port for proxy configuration
|
||||
PREFILL_HOSTS+=("localhost")
|
||||
PREFILL_PORTS+=("$PORT")
|
||||
done
|
||||
|
||||
# Start decode instances
|
||||
for i in $(seq 0 $((NUM_DECODE_INSTANCES-1))); do
|
||||
# Calculate GPU ID - we'll distribute across available GPUs, starting from after prefill GPUs
|
||||
GPU_ID=$(((i + NEXT_GPU + 1) % $(get_num_gpus)))
|
||||
# If DECODER_TP_SIZE is more than 1
|
||||
for (( j=1; j < DECODER_TP_SIZE; j++ )); do
|
||||
NEXT_GPU=$(((GPU_ID + j) % $(get_num_gpus)))
|
||||
GPU_ID="${GPU_ID},${NEXT_GPU}"
|
||||
done
|
||||
# Calculate port number (base port + instance number)
|
||||
PORT=$((8200 + i))
|
||||
# Calculate side channel port
|
||||
SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE))
|
||||
|
||||
echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"
|
||||
|
||||
# Build the command with or without model-specific args
|
||||
BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
|
||||
VLLM_KV_CACHE_LAYOUT=$DECODER_KV_LAYOUT \
|
||||
UCX_NET_DEVICES=all \
|
||||
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
|
||||
vllm serve $model_name \
|
||||
--port $PORT \
|
||||
--enforce-eager \
|
||||
--block-size ${DECODE_BLOCK_SIZE} \
|
||||
--gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
|
||||
--kv-transfer-config '$KV_CONFIG'"
|
||||
if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
|
||||
IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
|
||||
for arg in "${extra_args[@]}"; do
|
||||
BASE_CMD="${BASE_CMD} $arg"
|
||||
done
|
||||
fi
|
||||
|
||||
# Add attention backend config if specified
|
||||
if [[ -n "$ATTENTION_BACKEND" ]]; then
|
||||
BASE_CMD="${BASE_CMD} --attention-backend=$ATTENTION_BACKEND"
|
||||
fi
|
||||
|
||||
# Add HMA flag if specified
|
||||
if [[ -n "$ENABLE_HMA_VAR" ]]; then
|
||||
BASE_CMD="${BASE_CMD} $ENABLE_HMA_VAR"
|
||||
fi
|
||||
|
||||
# DP-EP attention mode
|
||||
if [[ -z "$DP_EP" ]]; then
|
||||
BASE_CMD="${BASE_CMD} --tensor-parallel-size $DECODER_TP_SIZE"
|
||||
else
|
||||
echo "DP-EP Attention enabled, deploying with dp=DECODER_TP_SIZE and tp=1"
|
||||
BASE_CMD="${BASE_CMD} --data-parallel-size $DECODER_TP_SIZE \
|
||||
--tensor-parallel-size 1 --enable-expert-parallel"
|
||||
fi
|
||||
|
||||
FULL_CMD="$BASE_CMD"
|
||||
|
||||
eval "$FULL_CMD &"
|
||||
|
||||
# Store host and port for proxy configuration
|
||||
DECODE_HOSTS+=("localhost")
|
||||
DECODE_PORTS+=("$PORT")
|
||||
done
|
||||
|
||||
# Wait for all instances to start
|
||||
for PORT in "${PREFILL_PORTS[@]}"; do
|
||||
echo "Waiting for prefill instance on port $PORT to start..."
|
||||
wait_for_server "$PORT"
|
||||
done
|
||||
|
||||
for PORT in "${DECODE_PORTS[@]}"; do
|
||||
echo "Waiting for decode instance on port $PORT to start..."
|
||||
wait_for_server "$PORT"
|
||||
done
|
||||
|
||||
# Build the command for the proxy server with all the hosts and ports
|
||||
PROXY_CMD="python3 ${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py --port 8192"
|
||||
|
||||
# Add all prefill hosts and ports
|
||||
PROXY_CMD+=" --prefiller-hosts ${PREFILL_HOSTS[*]}"
|
||||
PROXY_CMD+=" --prefiller-ports ${PREFILL_PORTS[*]}"
|
||||
|
||||
# Add all decode hosts and ports
|
||||
PROXY_CMD+=" --decoder-hosts ${DECODE_HOSTS[*]}"
|
||||
PROXY_CMD+=" --decoder-ports ${DECODE_PORTS[*]}"
|
||||
|
||||
# Start the proxy server
|
||||
echo "Starting proxy server with command: $PROXY_CMD"
|
||||
$PROXY_CMD &
|
||||
|
||||
# Wait for the proxy to start
|
||||
sleep 5
|
||||
|
||||
# Run lm eval for this model
|
||||
echo "Running tests for $model_name"
|
||||
TEST_MODEL=$model_name python3 -m pytest -s -x "${GIT_ROOT}"/tests/v1/kv_connector/nixl_integration/test_accuracy.py
|
||||
|
||||
# Clean up before running next model
|
||||
cleanup_instances
|
||||
sleep 3
|
||||
}
|
||||
|
||||
# Run tests for each model
|
||||
for model in "${MODELS[@]}"; do
|
||||
run_tests_for_model "$model"
|
||||
done
|
||||
|
||||
echo "All tests completed!"
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user