[CI] Ignore ruff D209, rejoin the docstrings it split, and silence incompatible-rule warnings (#57212)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
1085b64425
commit
651a88c09a
@@ -136,8 +136,7 @@ async def _client_loop(
|
||||
client_id: int,
|
||||
):
|
||||
"""A single async client: pulls prompts, submits to engine, consumes
|
||||
hidden states as soon as each request finishes.
|
||||
"""
|
||||
hidden states as soon as each request finishes."""
|
||||
loop = asyncio.get_event_loop()
|
||||
while True:
|
||||
item = await prompt_queue.get()
|
||||
|
||||
@@ -84,8 +84,7 @@ def make_inputs(
|
||||
|
||||
def build_configs(head_configs, num_tokens_list, input_dtypes, tp_sizes):
|
||||
"""Build (num_tokens, num_heads, head_size, dtype_str) config tuples,
|
||||
applying TP division to num_heads and skipping invalid combos.
|
||||
"""
|
||||
applying TP division to num_heads and skipping invalid combos."""
|
||||
configs = []
|
||||
for (_, nh, hs), nt, dtype, tp in itertools.product(
|
||||
head_configs, num_tokens_list, input_dtypes, tp_sizes
|
||||
|
||||
@@ -32,8 +32,7 @@ FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
def old_default_config(M, E, N, K, topk, dtype=None, block_shape=None):
|
||||
"""The original defaults before https://github.com/vllm-project/vllm/pull/34846,
|
||||
for comparison.
|
||||
"""
|
||||
for comparison."""
|
||||
if dtype == "fp8_w8a8" and block_shape is not None:
|
||||
return {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
|
||||
@@ -97,8 +97,7 @@ def graph_replay(fn):
|
||||
|
||||
def calculate_diff(vocab_size, hidden_size, tp_size, num_tokens):
|
||||
"""The fused kernel must be bit-exact with the path it replaces, on every
|
||||
rank of the shard.
|
||||
"""
|
||||
rank of the shard."""
|
||||
dtype = torch.bfloat16
|
||||
for tp_rank in range(tp_size):
|
||||
input_ids, weight, shard_indices = make_inputs(
|
||||
|
||||
@@ -75,8 +75,7 @@ class Example:
|
||||
|
||||
If path is a file, it returns the path itself. If path is a directory, it
|
||||
searches for Markdown files (*.md) in the directory and returns the first one
|
||||
found. If no Markdown files are found, it returns None.
|
||||
"""
|
||||
found. If no Markdown files are found, it returns None."""
|
||||
# Single file example
|
||||
if self.path.is_file():
|
||||
return self.path
|
||||
@@ -91,8 +90,7 @@ class Example:
|
||||
"""Determine other files in the directory excluding the main file.
|
||||
|
||||
If path is a file, it returns an empty list. Otherwise, it returns every file
|
||||
in the directory except the main file in a list.
|
||||
"""
|
||||
in the directory except the main file in a list."""
|
||||
# Single file example
|
||||
if self.path.is_file():
|
||||
return []
|
||||
|
||||
@@ -55,8 +55,7 @@ logger = init_logger(__name__)
|
||||
|
||||
class DummyPerReqLogitsProcessor:
|
||||
"""The request-level logits processor masks out all logits except the
|
||||
token id identified by `target_token`
|
||||
"""
|
||||
token id identified by `target_token`"""
|
||||
|
||||
def __init__(self, target_token: int) -> None:
|
||||
"""Specify `target_token`."""
|
||||
@@ -75,8 +74,7 @@ class DummyPerReqLogitsProcessor:
|
||||
|
||||
class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):
|
||||
"""Example of wrapping a fake request-level logit processor to create a
|
||||
batch-level logits processor
|
||||
"""
|
||||
batch-level logits processor"""
|
||||
|
||||
@classmethod
|
||||
def validate_params(cls, params: SamplingParams):
|
||||
|
||||
@@ -56,8 +56,7 @@ logger = init_logger(__name__)
|
||||
|
||||
class DummyPerReqLogitsProcessor:
|
||||
"""The request-level logits processor masks out all logits except the
|
||||
token id identified by `target_token`
|
||||
"""
|
||||
token id identified by `target_token`"""
|
||||
|
||||
def __init__(self, target_token: int) -> None:
|
||||
"""Specify `target_token`."""
|
||||
@@ -76,8 +75,7 @@ class DummyPerReqLogitsProcessor:
|
||||
|
||||
class WrappedPerReqLogitsProcessor(AdapterLogitsProcessor):
|
||||
"""Example of overriding the wrapper class `__init__()` in order to utilize
|
||||
info about the device type
|
||||
"""
|
||||
info about the device type"""
|
||||
|
||||
@classmethod
|
||||
def validate_params(cls, params: SamplingParams):
|
||||
|
||||
@@ -101,8 +101,7 @@ PROMPTS = ["The capital of France is", "The future of AI is"]
|
||||
class FSDPTrainWorker:
|
||||
"""One FSDP2 rank per GPU. Identical to the 8-GPU example's worker: the
|
||||
sharded-RDT engine owns the NIXL serve surface, so this stays a plain Ray
|
||||
actor with no producer mixin.
|
||||
"""
|
||||
actor with no producer mixin."""
|
||||
|
||||
def __init__(self, model_name, rank, world_size, master_addr, master_port):
|
||||
self.rank = rank
|
||||
|
||||
+3
-1
@@ -90,7 +90,9 @@ ignore = [
|
||||
# don't enforce that everything has a docstring
|
||||
"D100", "D101", "D102", "D103", "D104", "D105", "D106", "D107",
|
||||
# unwanted pydocstyle rules
|
||||
"D205", "D301", "D400", "D401", "D404", "D415",
|
||||
"D205", "D209", "D301", "D400", "D401", "D404", "D415",
|
||||
# incompatible pydocstyle pairs; pin the half `ruff` already picks
|
||||
"D203", "D213",
|
||||
# star imports
|
||||
"F405", "F403",
|
||||
# lambda expression assignment
|
||||
|
||||
@@ -724,8 +724,7 @@ class precompiled_wheel_utils:
|
||||
wheels: list[dict], repo_url: str, arch: str
|
||||
) -> None:
|
||||
"""Warn if installed torch differs from the custom ROCm build on
|
||||
wheels.vllm.ai and suggest the correct install command.
|
||||
"""
|
||||
wheels.vllm.ai and suggest the correct install command."""
|
||||
try:
|
||||
installed = torch.__version__
|
||||
except Exception:
|
||||
|
||||
@@ -112,8 +112,7 @@ def test_basic_cumem():
|
||||
@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
|
||||
def test_discard_tags():
|
||||
"""Test that discard(tags) selectively frees GPU memory for specific
|
||||
tags while keeping other tags mapped and usable.
|
||||
"""
|
||||
tags while keeping other tags mapped and usable."""
|
||||
allocator = get_mem_allocator_instance()
|
||||
|
||||
with allocator.use_memory_pool("weights"):
|
||||
|
||||
@@ -13,8 +13,7 @@ from vllm.benchmarks.datasets import BFCLDataset, get_samples
|
||||
|
||||
def _patch_hf_api(side_effect):
|
||||
"""Return a patch context that swaps `hf_api()` to a stub whose
|
||||
`.hf_hub_download` attribute uses `side_effect`.
|
||||
"""
|
||||
`.hf_hub_download` attribute uses `side_effect`."""
|
||||
fake_api = MagicMock()
|
||||
fake_api.hf_hub_download.side_effect = side_effect
|
||||
return patch("vllm.benchmarks.datasets.datasets.hf_api", return_value=fake_api)
|
||||
@@ -116,8 +115,7 @@ def test_bfcl_dataset_translates_schema_and_attaches_tools(
|
||||
) -> None:
|
||||
"""BFCLDataset should translate schemas to OpenAI tool format, set
|
||||
`messages` directly on SampleRequest, and attach tools/tool_choice via
|
||||
request_overrides.
|
||||
"""
|
||||
request_overrides."""
|
||||
paths = _write_fake_files(tmp_path)
|
||||
|
||||
def fake_download(_repo, filename, **_kwargs):
|
||||
@@ -168,8 +166,7 @@ def test_bfcl_dataset_missing_category_raises_clear_error(
|
||||
hf_tokenizer: PreTrainedTokenizerBase,
|
||||
) -> None:
|
||||
"""A typo'd category should produce an actionable ValueError, not an
|
||||
opaque huggingface_hub exception.
|
||||
"""
|
||||
opaque huggingface_hub exception."""
|
||||
from huggingface_hub.errors import EntryNotFoundError
|
||||
|
||||
args = _args_for_bfcl(categories=["simpl"]) # typo
|
||||
@@ -187,8 +184,7 @@ def test_bfcl_dataset_missing_category_raises_clear_error(
|
||||
@pytest.mark.benchmark
|
||||
def test_chat_backend_uses_messages_field_when_set() -> None:
|
||||
"""When RequestFuncInput.chat_messages is set, the chat backend must use
|
||||
it verbatim and skip default content construction from `prompt`.
|
||||
"""
|
||||
it verbatim and skip default content construction from `prompt`."""
|
||||
import asyncio
|
||||
|
||||
from vllm.benchmarks.lib.endpoint_request_func import (
|
||||
@@ -248,8 +244,7 @@ def test_chat_backend_uses_messages_field_when_set() -> None:
|
||||
def test_bfcl_prompt_len_includes_tools(tmp_path: Path) -> None:
|
||||
"""prompt_len must reflect tokens from both messages *and* tool schemas,
|
||||
so percentile buckets and input-distribution summaries aren't biased
|
||||
low for tool-heavy traffic.
|
||||
"""
|
||||
low for tool-heavy traffic."""
|
||||
paths = _write_fake_files(tmp_path)
|
||||
|
||||
def fake_download(_repo, filename, **_kwargs):
|
||||
@@ -296,8 +291,7 @@ def test_bfcl_prompt_len_falls_back_when_tokenizer_rejects_tools(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Older tokenizers don't accept tools=; fallback must still produce a
|
||||
non-zero prompt_len without crashing.
|
||||
"""
|
||||
non-zero prompt_len without crashing."""
|
||||
paths = _write_fake_files(tmp_path)
|
||||
|
||||
def fake_download(_repo, filename, **_kwargs):
|
||||
|
||||
@@ -12,8 +12,7 @@ from vllm.benchmarks.datasets import get_samples
|
||||
class _RecordingTokenizer:
|
||||
"""Minimal tokenizer stub that records the kwargs forwarded to
|
||||
apply_chat_template, so we can assert chat_template_kwargs propagation
|
||||
without loading a real model/template.
|
||||
"""
|
||||
without loading a real model/template."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.captured_kwargs: dict | None = None
|
||||
|
||||
@@ -10,8 +10,7 @@ from vllm.tokenizers import TokenizerLike
|
||||
|
||||
class _FakeTokenizer(TokenizerLike):
|
||||
"""Minimal tokenizer implementing the TokenizerLike protocol
|
||||
for testing get_sampling_params.
|
||||
"""
|
||||
for testing get_sampling_params."""
|
||||
|
||||
def __init__(self, vocab_size: int = 1000, num_special_tokens: int = 0) -> None:
|
||||
self._vocab_size = vocab_size
|
||||
@@ -172,8 +171,7 @@ class TestGetSamplingParams:
|
||||
|
||||
def test_output_low_clamped_to_one(self):
|
||||
"""Even with a high ratio that would push output_low to 0,
|
||||
the function clamps it to 1.
|
||||
"""
|
||||
the function clamps it to 1."""
|
||||
rng = np.random.default_rng(0)
|
||||
# output_len=1, ratio=0.99 → floor(1*0.01)=0, should clamp to 1
|
||||
_, output_lens, _ = get_sampling_params(
|
||||
@@ -235,8 +233,7 @@ class TestGetSamplingParams:
|
||||
|
||||
def test_input_len_zero_with_special_tokens(self):
|
||||
"""input_len < num_special_tokens → real_input_len = 0, which is fine
|
||||
(range [0, 0]).
|
||||
"""
|
||||
(range [0, 0])."""
|
||||
rng = np.random.default_rng(0)
|
||||
input_lens, _, _ = get_sampling_params(
|
||||
rng, 5, 0.0, 5, 50, self._tok(num_special=10)
|
||||
|
||||
@@ -52,8 +52,7 @@ def _write_dataset(path: Path) -> None:
|
||||
|
||||
def _args(dataset_path: str) -> argparse.Namespace:
|
||||
"""Reproduce the argparse.Namespace that serve.py builds from the
|
||||
failing command, including skip_tokenizer_init=True.
|
||||
"""
|
||||
failing command, including skip_tokenizer_init=True."""
|
||||
return argparse.Namespace(
|
||||
# dataset
|
||||
dataset_name="custom",
|
||||
|
||||
@@ -732,8 +732,7 @@ def test_aiter_fusion_rmsnorm_gated_quant_no_gdn_layers(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Verify that without GDN layers in static_forward_context,
|
||||
the gated pattern is not registered and no matches occur.
|
||||
"""
|
||||
the gated pattern is not registered and no matches occur."""
|
||||
group_shape = GroupShape(1, 128)
|
||||
vllm_config = VllmConfig(
|
||||
model_config=ModelConfig(dtype=dtype),
|
||||
|
||||
@@ -15,8 +15,7 @@ DEVICE_TYPE = current_platform.device_type
|
||||
|
||||
class SplitCoalescingModel(torch.nn.Module):
|
||||
"""Model with 3 separate split_with_sizes calls on the same input,
|
||||
simulating the B200+FP8 graph where CSE fails to merge them.
|
||||
"""
|
||||
simulating the B200+FP8 graph where CSE fails to merge them."""
|
||||
|
||||
def __init__(self, q_size: int, kv_size: int) -> None:
|
||||
super().__init__()
|
||||
|
||||
@@ -77,8 +77,7 @@ def test_non_primitive_kwargs_lifted_to_consts(
|
||||
into the namespace, so ``device`` is unbound. The fix collects such
|
||||
objects into ``__vllm_consts__`` and references them by index. The
|
||||
unqualified ``device(type=...)`` form must never appear in the
|
||||
generated source.
|
||||
"""
|
||||
generated source."""
|
||||
split_gm = _trace_and_split(model_fn, (x,), split_ops)
|
||||
code, submod_names, consts = generate_execution_code(split_gm)
|
||||
|
||||
@@ -99,8 +98,7 @@ def test_dtype_singleton_deduped(x: torch.Tensor) -> None:
|
||||
"""``torch.float16`` is a process-wide singleton, so two ops referring
|
||||
to it in the traced graph share a single consts slot via ``id()``-based
|
||||
dedup. Distinct expressions (``x.to(...)`` vs ``(x*2).to(...)``) ensure
|
||||
the tracer can't CSE the two ops into a single node.
|
||||
"""
|
||||
the tracer can't CSE the two ops into a single node."""
|
||||
|
||||
def model_fn(x: torch.Tensor) -> torch.Tensor:
|
||||
return x.to(dtype=torch.float16) + (x * 2).to(dtype=torch.float16)
|
||||
@@ -149,8 +147,7 @@ def test_distinct_dtypes_get_distinct_slots(x: torch.Tensor) -> None:
|
||||
def test_consts_ordering_deterministic(x: torch.Tensor) -> None:
|
||||
"""Two independent traces of the same model must produce equal consts
|
||||
lists *in the same order*. Cache artifacts identify const slots by
|
||||
index, so a non-deterministic order would invalidate cached code.
|
||||
"""
|
||||
index, so a non-deterministic order would invalidate cached code."""
|
||||
|
||||
def model_fn(x: torch.Tensor) -> torch.Tensor:
|
||||
# Multiple distinct non-primitives encountered in a fixed graph order.
|
||||
@@ -188,8 +185,7 @@ def test_consts_shared_across_split_submods(x: torch.Tensor) -> None:
|
||||
The function below splits into three inlined submods, two of which
|
||||
independently reference ``torch.float16``. The shared ``const_index``
|
||||
threaded through recursive ``generate_execution_code_with_name`` calls
|
||||
must collapse the dtype to a single slot used from both submods.
|
||||
"""
|
||||
must collapse the dtype to a single slot used from both submods."""
|
||||
|
||||
def model_fn(x: torch.Tensor) -> torch.Tensor:
|
||||
a = x.to(dtype=torch.float16) # submod_0: _to_copy(fp16)
|
||||
@@ -234,8 +230,7 @@ def test_non_graphmodule_submod_uses_indexed_callable(x: torch.Tensor) -> None:
|
||||
"""When a child of split_gm is *not* a ``torch.fx.GraphModule`` — as
|
||||
happens in production once ``PiecewiseBackend`` replaces submods —
|
||||
codegen emits ``__vllm_submods__[idx](...)`` instead of inlining, and
|
||||
the runtime callable is bound from ``submod_callables``.
|
||||
"""
|
||||
the runtime callable is bound from ``submod_callables``."""
|
||||
|
||||
def model_fn(x: torch.Tensor) -> torch.Tensor:
|
||||
return x.relu().sigmoid()
|
||||
@@ -285,8 +280,7 @@ def test_non_graphmodule_submod_uses_indexed_callable(x: torch.Tensor) -> None:
|
||||
def test_getitem_in_stitching_graph(x: torch.Tensor) -> None:
|
||||
"""``operator.getitem`` on submod tuple returns is the ``call_function``
|
||||
special case at codegen.py — emitted as ``name = source[index]``
|
||||
rather than a function call.
|
||||
"""
|
||||
rather than a function call."""
|
||||
|
||||
def model_fn(x: torch.Tensor) -> torch.Tensor:
|
||||
return x.relu().sigmoid()
|
||||
@@ -304,8 +298,7 @@ def test_getitem_in_stitching_graph(x: torch.Tensor) -> None:
|
||||
def test_del_emitted_for_intermediate_values(x: torch.Tensor) -> None:
|
||||
"""The codegen schedules ``del`` after a value's last use to free
|
||||
memory early. Multi-submod splits naturally have intermediates whose
|
||||
last use is not the output node.
|
||||
"""
|
||||
last use is not the output node."""
|
||||
|
||||
def model_fn(x: torch.Tensor) -> torch.Tensor:
|
||||
return x.relu().sigmoid().tanh()
|
||||
@@ -325,8 +318,7 @@ def test_with_submod_false_rejects_call_module() -> None:
|
||||
"""``generate_execution_code_with_name(with_submod=False)`` is the
|
||||
recursive entry for inlining a GraphModule into its parent. It must
|
||||
refuse a graph that itself contains ``call_module`` nodes — the parent
|
||||
is responsible for handling those.
|
||||
"""
|
||||
is responsible for handling those."""
|
||||
g = fx.Graph()
|
||||
x_node = g.placeholder("x")
|
||||
root = torch.nn.Module()
|
||||
@@ -343,8 +335,7 @@ def test_node_ref_recurses_through_containers() -> None:
|
||||
"""``_node_ref`` is the recursive walker that lifts non-primitives
|
||||
nested inside list/tuple/dict args. Real aten ops rarely produce such
|
||||
structures, but the path is needed for DTensor placement lists and
|
||||
other future cases — unit-test the walker directly.
|
||||
"""
|
||||
other future cases — unit-test the walker directly."""
|
||||
consts: list = []
|
||||
const_index: dict[int, int] = {}
|
||||
cpu = torch.device("cpu")
|
||||
@@ -368,8 +359,7 @@ def test_node_ref_recurses_through_containers() -> None:
|
||||
def test_legacy_code_without_consts() -> None:
|
||||
"""``compile_execution_fn(consts=None)`` must still load code that has
|
||||
no ``__vllm_consts__`` reference, so older serialized cache artifacts
|
||||
keep working.
|
||||
"""
|
||||
keep working."""
|
||||
# Pre-consts codegen: no __vllm_consts__ reference, only torch/operator.
|
||||
legacy_code = (
|
||||
"import torch\n"
|
||||
|
||||
@@ -129,8 +129,7 @@ def test_compile_config_get_compile_ranges():
|
||||
|
||||
class PostGradStaticShapeChecker(InductorPass):
|
||||
"""Asserts that compile_sizes entries produce graphs with fully concrete
|
||||
(non-symbolic) shapes, and compile_ranges entries have symbolic shapes.
|
||||
"""
|
||||
(non-symbolic) shapes, and compile_ranges entries have symbolic shapes."""
|
||||
|
||||
def __init__(self):
|
||||
self.num_static_calls = 0
|
||||
@@ -171,8 +170,7 @@ class PostGradStaticShapeChecker(InductorPass):
|
||||
|
||||
def test_compile_sizes_produce_static_shapes(disable_vllm_compile_cache):
|
||||
"""Verify that compile_sizes entries are compiled with fully concrete
|
||||
shapes (no SymInts), while compile_ranges entries retain dynamic shapes.
|
||||
"""
|
||||
shapes (no SymInts), while compile_ranges entries retain dynamic shapes."""
|
||||
checker = PostGradStaticShapeChecker()
|
||||
torch.set_default_device("cuda")
|
||||
vllm_config = VllmConfig(
|
||||
|
||||
@@ -1229,8 +1229,7 @@ def test_compile_sizes_padding_validation():
|
||||
|
||||
def test_inductor_asserts_default_disabled(monkeypatch):
|
||||
"""Test that inductor runtime asserts are disabled by default
|
||||
(INFO logging level) on torch < 2.12.
|
||||
"""
|
||||
(INFO logging level) on torch < 2.12."""
|
||||
monkeypatch.setenv("VLLM_LOGGING_LEVEL", "INFO")
|
||||
|
||||
import importlib
|
||||
@@ -1248,8 +1247,7 @@ def test_inductor_asserts_default_disabled(monkeypatch):
|
||||
|
||||
def test_inductor_asserts_enabled_in_debug(monkeypatch):
|
||||
"""Test that VLLM_LOGGING_LEVEL=DEBUG enables inductor runtime asserts
|
||||
on torch < 2.12.
|
||||
"""
|
||||
on torch < 2.12."""
|
||||
monkeypatch.setenv("VLLM_LOGGING_LEVEL", "DEBUG")
|
||||
|
||||
import importlib
|
||||
@@ -1285,8 +1283,7 @@ def test_get_inductor_factors_includes_configs():
|
||||
|
||||
def test_inductor_asserts_user_override(monkeypatch):
|
||||
"""Test that explicit inductor_compile_config overrides the
|
||||
debug-logging default.
|
||||
"""
|
||||
debug-logging default."""
|
||||
monkeypatch.setenv("VLLM_LOGGING_LEVEL", "INFO")
|
||||
|
||||
import importlib
|
||||
|
||||
@@ -236,8 +236,7 @@ class TestInductorFallbackFastPath:
|
||||
|
||||
def test_inner_set_membership_still_works_for_standard_ops(self):
|
||||
"""The patch must not break Inductor's existing fallback decisions
|
||||
for non-vllm ops such as ``torchvision::roi_align``.
|
||||
"""
|
||||
for non-vllm ops such as ``torchvision::roi_align``."""
|
||||
import torch._inductor.lowering as _lowering
|
||||
|
||||
allow_list = _lowering.FALLBACK_ALLOW_LIST
|
||||
|
||||
@@ -372,8 +372,7 @@ class _HeterogeneousConfig(PretrainedConfig):
|
||||
"""A heterogeneous config with no convertor of its own.
|
||||
|
||||
Mirrors the parts vLLM uses: per-layer configs are shallow copies with the
|
||||
varying attributes applied and heterogeneity stripped, so they do not recurse.
|
||||
"""
|
||||
varying attributes applied and heterogeneity stripped, so they do not recurse."""
|
||||
|
||||
is_heterogeneous = True
|
||||
|
||||
@@ -472,8 +471,7 @@ def test_deepseek_v4_convertor_splits_vision_architecture():
|
||||
"""The vision checkpoint shares model_type/architectures with the
|
||||
text-only DeepSeek-V4; the convertor routes it to the VL wrapper class
|
||||
by rewriting hf_config.architectures (model-class resolution reads the
|
||||
raw attribute).
|
||||
"""
|
||||
raw attribute)."""
|
||||
from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config
|
||||
from vllm.transformers_utils.model_arch_config_convertor import (
|
||||
DeepseekV4ModelArchConfigConvertor,
|
||||
|
||||
@@ -42,8 +42,7 @@ def test_mm_encoder_attn_backend_hash_updates():
|
||||
|
||||
def test_language_model_only_does_not_affect_mm_hash():
|
||||
"""language_model_only does not affect the ViT computation graph,
|
||||
so it should not change the multimodal config hash.
|
||||
"""
|
||||
so it should not change the multimodal config hash."""
|
||||
base_hash = MultiModalConfig().compute_hash()
|
||||
lm_only_hash = MultiModalConfig(language_model_only=True).compute_hash()
|
||||
assert base_hash == lm_only_hash
|
||||
@@ -51,8 +50,7 @@ def test_language_model_only_does_not_affect_mm_hash():
|
||||
|
||||
def test_language_model_only_affects_model_hash():
|
||||
"""language_model_only affects the LM computation graph,
|
||||
so it should change the model config hash.
|
||||
"""
|
||||
so it should change the model config hash."""
|
||||
model = "llava-hf/llava-1.5-7b-hf"
|
||||
base_hash = ModelConfig(model).compute_hash()
|
||||
lm_only_hash = ModelConfig(model, language_model_only=True).compute_hash()
|
||||
@@ -388,8 +386,7 @@ def _resolve_mm_video_decode_device(
|
||||
torchcodec_available: bool = True,
|
||||
) -> dict:
|
||||
"""Run processor-device then video-decode-device resolution and report
|
||||
the resulting video media IO kwargs.
|
||||
"""
|
||||
the resulting video media IO kwargs."""
|
||||
mm_config = MultiModalConfig(
|
||||
mm_processor_kwargs={} if device is None else {"device": device},
|
||||
mm_tensor_ipc=mm_tensor_ipc, # type: ignore[arg-type]
|
||||
@@ -419,8 +416,7 @@ def _resolve_mm_video_decode_device(
|
||||
|
||||
def test_auto_video_decode_uses_nvdec_on_encoder_instance():
|
||||
"""The processor runs on the accelerator there, so decoded frames should
|
||||
stay on-device too.
|
||||
"""
|
||||
stay on-device too."""
|
||||
assert _resolve_mm_video_decode_device(ec_role="ec_producer") == {
|
||||
"backend": "torchcodec",
|
||||
"device": "cuda",
|
||||
@@ -432,8 +428,7 @@ def test_auto_video_decode_stays_on_cpu_off_encoder_instance(
|
||||
ec_role: ECRole | None,
|
||||
):
|
||||
"""The processor stays on CPU off encode-only instances, so video
|
||||
decoding should too.
|
||||
"""
|
||||
decoding should too."""
|
||||
assert _resolve_mm_video_decode_device(ec_role=ec_role) == {}
|
||||
|
||||
|
||||
|
||||
@@ -34,8 +34,7 @@ def _make_hf_config(**kwargs) -> PretrainedConfig:
|
||||
@pytest.mark.cpu_test
|
||||
def test_dict_overrides_are_not_forwarded_to_draft():
|
||||
"""Dict overrides are target-specific key patches; the draft must get
|
||||
only the architecture-mapping override.
|
||||
"""
|
||||
only the architecture-mapping override."""
|
||||
composed = SpeculativeConfig.compose_draft_hf_overrides(
|
||||
{"max_position_embeddings": 1234}
|
||||
)
|
||||
@@ -51,8 +50,7 @@ def test_none_overrides_fall_back_to_arch_mapping():
|
||||
@pytest.mark.cpu_test
|
||||
def test_callable_overrides_reach_the_draft_config():
|
||||
"""A callable override (config-to-config transform) composes with the
|
||||
architecture-mapping override and is applied to the draft config.
|
||||
"""
|
||||
architecture-mapping override and is applied to the draft config."""
|
||||
|
||||
def shrink(hf_config: PretrainedConfig) -> PretrainedConfig:
|
||||
hf_config.num_hidden_layers = 1
|
||||
@@ -69,8 +67,7 @@ def test_callable_overrides_reach_the_draft_config():
|
||||
@pytest.mark.cpu_test
|
||||
def test_arch_mapping_applies_before_callable_override():
|
||||
"""The static arch-mapping override runs first, so the user callable
|
||||
observes (and may adjust) the post-mapping config.
|
||||
"""
|
||||
observes (and may adjust) the post-mapping config."""
|
||||
seen_architectures: list[str] = []
|
||||
|
||||
def record(hf_config: PretrainedConfig) -> PretrainedConfig:
|
||||
@@ -131,8 +128,7 @@ def test_composed_override_is_picklable():
|
||||
the composed override must be picklable. A nested local closure is not
|
||||
(it raised ``Can't get local object`` on DFlashDraftModel); a
|
||||
``functools.partial`` over a module-referenceable static method is.
|
||||
Guard against regressing to a closure.
|
||||
"""
|
||||
Guard against regressing to a closure."""
|
||||
composed = SpeculativeConfig.compose_draft_hf_overrides(_module_level_shrink)
|
||||
|
||||
assert isinstance(composed, functools.partial)
|
||||
|
||||
@@ -99,8 +99,7 @@ def test_independent_draft_model_keeps_its_own_limit(
|
||||
vllm_caplog: pytest.LogCaptureFixture,
|
||||
):
|
||||
"""An independent AR draft may genuinely have a smaller context than the
|
||||
target; its max_position_embeddings must not be resized.
|
||||
"""
|
||||
target; its max_position_embeddings must not be resized."""
|
||||
target_model_config = ModelConfig(
|
||||
AR_MODEL, hf_overrides={"max_position_embeddings": 8192}
|
||||
)
|
||||
|
||||
@@ -179,8 +179,7 @@ class MockSubscriber:
|
||||
@pytest.fixture
|
||||
def enable_ray_v2_backend():
|
||||
"""Set env vars for the Ray V2 executor backend and shut down Ray
|
||||
between tests.
|
||||
"""
|
||||
between tests."""
|
||||
import ray
|
||||
|
||||
saved = {
|
||||
|
||||
@@ -124,8 +124,7 @@ def _build_processed_layer(
|
||||
backend: str, raw: dict[str, torch.Tensor], device: torch.device
|
||||
) -> tuple[CompressedTensorsW4A4Nvfp4MoEMethod, _RoutedExpertsStub]:
|
||||
"""Create the method + layer, load ``raw`` and run the real
|
||||
process_weights_after_loading (which builds the flashinfer kernel).
|
||||
"""
|
||||
process_weights_after_loading (which builds the flashinfer kernel)."""
|
||||
moe_config = _make_moe_config(backend)
|
||||
method = CompressedTensorsW4A4Nvfp4MoEMethod(moe_config, "layer.0", use_a16=False)
|
||||
layer = _RoutedExpertsStub(moe_config).to(device)
|
||||
@@ -158,8 +157,7 @@ def _assert_tensors_equal(
|
||||
|
||||
def _simulate_eplb_rearrangement(layer: torch.nn.Module, perm: torch.Tensor) -> None:
|
||||
"""Permute experts of every registered Parameter in place, the way
|
||||
EPLB's in-place rearrangement moves expert slices along dim 0.
|
||||
"""
|
||||
EPLB's in-place rearrangement moves expert slices along dim 0."""
|
||||
with torch.no_grad():
|
||||
for _, param in layer.named_parameters():
|
||||
param.copy_(param[perm])
|
||||
@@ -167,8 +165,7 @@ def _simulate_eplb_rearrangement(layer: torch.nn.Module, perm: torch.Tensor) ->
|
||||
|
||||
def _ensure_world1_distributed() -> None:
|
||||
"""The enable_eplb config makes weight processing all-reduce activation
|
||||
scale amaxes over the EP group, so a (single-rank) EP group must exist.
|
||||
"""
|
||||
scale amaxes over the EP group, so a (single-rank) EP group must exist."""
|
||||
from vllm.distributed.parallel_state import (
|
||||
ensure_model_parallel_initialized,
|
||||
init_distributed_environment,
|
||||
@@ -256,8 +253,7 @@ def test_nvfp4_eplb_rearrangement_matches_reload(backend: str) -> None:
|
||||
def test_fp8_per_tensor_alphas_registered_and_aliased() -> None:
|
||||
"""CPU-only: the fp8 per-tensor oracle must register the fused
|
||||
(w_scale * a_scale) products as layer Parameters aliased by the quant
|
||||
config, so EPLB rearrangement keeps kernels consistent.
|
||||
"""
|
||||
config, so EPLB rearrangement keeps kernels consistent."""
|
||||
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
|
||||
Fp8MoeBackend,
|
||||
make_fp8_moe_quant_config,
|
||||
|
||||
@@ -160,8 +160,7 @@ def _make_forward_context(rank, world_size, num_tokens_per_rank):
|
||||
class _AttnMeta:
|
||||
"""Minimal placeholder so set_forward_context's
|
||||
``attn_metadata is not None`` guard (forward_context.py:334)
|
||||
is satisfied. The real DPMetadata is built from num_tokens_across_dp.
|
||||
"""
|
||||
is satisfied. The real DPMetadata is built from num_tokens_across_dp."""
|
||||
|
||||
dp_metadata = None
|
||||
|
||||
@@ -460,8 +459,7 @@ def _one_sided_workspace_grow_worker(rank, world_size):
|
||||
def test_one_sided_manager_workspace_grow(world_size):
|
||||
"""A later initialize() with a larger per-token payload must grow the
|
||||
workspace and rebuild MoeAlltoAll; a later initialize() with a smaller
|
||||
payload must no-op.
|
||||
"""
|
||||
payload must no-op."""
|
||||
_spawn_workers(
|
||||
_one_sided_workspace_grow_worker,
|
||||
world_size,
|
||||
|
||||
@@ -21,8 +21,7 @@ from vllm.utils.network_utils import get_ip
|
||||
|
||||
def test_mq_bind_with_local_ip():
|
||||
"""MessageQueue with remote readers should successfully bind
|
||||
when connect_ip is the local node's IP.
|
||||
"""
|
||||
when connect_ip is the local node's IP."""
|
||||
# n_reader=2, n_local_reader=1 means 1 remote reader,
|
||||
# which triggers the remote ZMQ socket bind.
|
||||
mq = MessageQueue(
|
||||
@@ -44,8 +43,7 @@ def test_mq_bind_with_local_ip():
|
||||
def test_mq_bind_with_non_local_ip_fails():
|
||||
"""MessageQueue should fail to bind when connect_ip is a
|
||||
non-local IP address (simulating the bug where master_addr
|
||||
from a different node was used).
|
||||
"""
|
||||
from a different node was used)."""
|
||||
# Use a non-local IP that we definitely can't bind to.
|
||||
# 198.51.100.1 is from TEST-NET-2 (RFC 5737), never locally assigned.
|
||||
non_local_ip = "198.51.100.1"
|
||||
@@ -59,8 +57,7 @@ def test_mq_bind_with_non_local_ip_fails():
|
||||
|
||||
def test_mq_bind_defaults_to_local_ip():
|
||||
"""When connect_ip is None, MessageQueue should auto-detect
|
||||
the local IP and bind successfully.
|
||||
"""
|
||||
the local IP and bind successfully."""
|
||||
mq = MessageQueue(
|
||||
n_reader=2,
|
||||
n_local_reader=1,
|
||||
|
||||
@@ -495,8 +495,7 @@ def test_pynccl_suspend_resume_idempotent():
|
||||
|
||||
def test_pynccl_suspend_noop_without_symbol():
|
||||
"""RCCL / NCCL < 2.29.7 lack ncclCommSuspend; suspend() must no-op, not
|
||||
crash, and resume() must then no-op too since nothing was suspended.
|
||||
"""
|
||||
crash, and resume() must then no-op too since nothing was suspended."""
|
||||
from unittest.mock import Mock
|
||||
|
||||
comm = object.__new__(PyNcclCommunicator)
|
||||
|
||||
@@ -23,8 +23,7 @@ def _ray_init():
|
||||
|
||||
Without this, workers cannot unpickle actor classes defined in the
|
||||
``tests`` package, causing FunctionActorManager to fall back to
|
||||
TemporaryActor which drops async method signatures.
|
||||
"""
|
||||
TemporaryActor which drops async method signatures."""
|
||||
project_root = str(pathlib.Path(__file__).resolve().parents[2])
|
||||
ray.init(
|
||||
ignore_reinit_error=True,
|
||||
|
||||
@@ -53,8 +53,7 @@ META = torch.device("meta")
|
||||
|
||||
class _FakeLayer:
|
||||
"""Stand-in for a leaf module. The planner only stores it and keys on
|
||||
``id()``; nothing in the plan path touches its attributes.
|
||||
"""
|
||||
``id()``; nothing in the plan path touches its attributes."""
|
||||
|
||||
def __init__(self, tag: str) -> None:
|
||||
self.tag = tag
|
||||
@@ -65,8 +64,7 @@ class _FakeLayer:
|
||||
|
||||
def _module(layer, copies):
|
||||
"""What the bake records for one leaf module: a scatter list whose entries
|
||||
all carry that module as their ``layer``.
|
||||
"""
|
||||
all carry that module as their ``layer``."""
|
||||
for c in copies:
|
||||
c.layer = layer
|
||||
return copies
|
||||
@@ -83,8 +81,7 @@ def _copy(
|
||||
dtype=torch.bfloat16,
|
||||
):
|
||||
"""A `_Scatter` as the bake would record it: source key, owning layer, and
|
||||
the meta destination region.
|
||||
"""
|
||||
the meta destination region."""
|
||||
stride = (
|
||||
(1,)
|
||||
if len(shape) == 1
|
||||
@@ -143,8 +140,7 @@ def _planner(baked, *, name_meta, live=None, held_by=None, group_lens=None):
|
||||
|
||||
def _one_module_per_layer(n_layers, *, dtype="bfloat16", numel=4):
|
||||
"""`(baked, name_meta, names, group_lens)` for a pre / N-layer / post model
|
||||
where every gather group is one fully-baked leaf module.
|
||||
"""
|
||||
where every gather group is one fully-baked leaf module."""
|
||||
names = ["embed.weight"]
|
||||
names += [f"model.layers.{i}.w" for i in range(n_layers)]
|
||||
names += ["norm.weight"]
|
||||
@@ -165,8 +161,7 @@ def _one_module_per_layer(n_layers, *, dtype="bfloat16", numel=4):
|
||||
class TestFakeOpChains:
|
||||
"""Every allowlisted op must append itself to the chain and hand back a
|
||||
child whose shape/dtype PyTorch itself computed. The chain is the wire
|
||||
format the producer replays, so its exact contents are load-bearing.
|
||||
"""
|
||||
format the producer replays, so its exact contents are load-bearing."""
|
||||
|
||||
def _fake(self, shape=(4, 6), dtype=torch.bfloat16, sink=None):
|
||||
return FakeRDTTensor(
|
||||
@@ -227,8 +222,7 @@ class TestFakeOpChains:
|
||||
@pytest.mark.parametrize("op,n", [("chunk", 2), ("unbind", 4)])
|
||||
def test_multi_return_op_emits_one_child_per_output(self, op, n):
|
||||
"""chunk/unbind hand back a tuple; each child carries the base op plus a
|
||||
trailing __getitem__(i) so the producer can index the replayed result.
|
||||
"""
|
||||
trailing __getitem__(i) so the producer can index the replayed result."""
|
||||
parts = getattr(self._fake(), op)(*((n,) if op == "chunk" else ()), 0)
|
||||
assert isinstance(parts, tuple)
|
||||
assert len(parts) == n
|
||||
@@ -244,8 +238,7 @@ class TestFakeOpChains:
|
||||
|
||||
class TestFakeUnsupportedOps:
|
||||
"""Anything that needs real data must fail loudly at bake time rather than
|
||||
silently transferring the wrong bytes.
|
||||
"""
|
||||
silently transferring the wrong bytes."""
|
||||
|
||||
def _fake(self):
|
||||
return FakeRDTTensor(
|
||||
@@ -285,8 +278,7 @@ class TestFakeUnsupportedOps:
|
||||
|
||||
class TestBakeRecording:
|
||||
"""During the dry run the fake's ``copy_`` is the data sink: it records the
|
||||
source chain plus the meta destination's strided region and moves nothing.
|
||||
"""
|
||||
source chain plus the meta destination's strided region and moves nothing."""
|
||||
|
||||
def _recorder_and_fake(self, shape=(4, 6)):
|
||||
rec = BakeSink()
|
||||
@@ -322,8 +314,7 @@ class TestBakeRecording:
|
||||
|
||||
def test_unattributed_copy_is_live_but_unrecorded(self):
|
||||
"""A copy_ with no loader stamp cannot be attributed to a param, so its
|
||||
module must fall back to the plain load — but the name still moved data.
|
||||
"""
|
||||
module must fall back to the plain load — but the name still moved data."""
|
||||
rec, fake = self._recorder_and_fake()
|
||||
rec.current = None
|
||||
torch.empty((4, 6), dtype=torch.bfloat16, device=META).copy_(fake)
|
||||
@@ -363,8 +354,7 @@ class TestBakeRecording:
|
||||
producer packs what the replay yields, so the record must carry the
|
||||
POST-chain dtype: taking it from the source name's metadata instead sizes
|
||||
the slice with the wrong itemsize and shifts every later slice in the
|
||||
chunk, carving the packed blob differently on the two sides.
|
||||
"""
|
||||
chunk, carving the packed blob differently on the two sides."""
|
||||
rec, fake = self._recorder_and_fake(shape=(4,))
|
||||
layer = _FakeLayer("reinterpreted")
|
||||
rec.current = (layer, "weight")
|
||||
@@ -384,8 +374,7 @@ class TestBakeRecording:
|
||||
the producer packs what the chain yields, and nothing downstream can
|
||||
catch a mismatch — the consumer's buffer view is exactly prod(dest.shape)
|
||||
elements, so it reshapes cleanly over bytes laid out at other offsets.
|
||||
So the bake refuses rather than recording a slice it cannot carve.
|
||||
"""
|
||||
So the bake refuses rather than recording a slice it cannot carve."""
|
||||
rec, fake = self._recorder_and_fake(shape=(1, 6))
|
||||
rec.current = (_FakeLayer("broadcast"), "weight")
|
||||
param = torch.empty((4, 6), dtype=torch.bfloat16, device=META)
|
||||
@@ -401,8 +390,7 @@ class TestBakeRecording:
|
||||
class TestConsumerIdentity:
|
||||
"""Every worker in the fleet needs a DISTINCT id in 0..C-1: it selects the
|
||||
worker's producer block and keys the producer's per-consumer serve ring, so
|
||||
a collision silently serves two workers out of one ring.
|
||||
"""
|
||||
a collision silently serves two workers out of one ring."""
|
||||
|
||||
def _engine(self, *, dp_index, rank, world_size):
|
||||
eng = object.__new__(ShardedRDTWeightTransferEngine)
|
||||
@@ -439,8 +427,7 @@ class TestConsumerIdentity:
|
||||
def test_independent_engines_offset_into_distinct_ranges(self):
|
||||
"""A fleet of separate engines restarts _global_worker_index at 0 in
|
||||
each, so without the replica offset every engine would claim 0..w-1.
|
||||
The driver sets replica_rank on the payload; these are the ids it buys.
|
||||
"""
|
||||
The driver sets replica_rank on the payload; these are the ids it buys."""
|
||||
ids = [
|
||||
self._ids(num_consumers=8, workers=range(4), replica_rank=r, num_replicas=2)
|
||||
for r in (0, 1)
|
||||
@@ -466,8 +453,7 @@ def _router(
|
||||
group_lens=None,
|
||||
):
|
||||
"""A router over ``n`` single-name groups unless told otherwise, so a test
|
||||
can state expectations in group indices.
|
||||
"""
|
||||
can state expectations in group indices."""
|
||||
if names is None:
|
||||
n = len(name_owner_class or []) or (len(group_lens or []) or 1)
|
||||
names = [f"g{i}" for i in range(n)]
|
||||
@@ -496,8 +482,7 @@ class TestRdtRouter:
|
||||
def test_every_producer_nic_carries_traffic(self):
|
||||
"""16 producers / 8 consumers: the block rule spreads each consumer's
|
||||
pulls over its block, alternating by group, so no producer sits idle.
|
||||
Plain `consumer_id % len(owners)` would idle half of them.
|
||||
"""
|
||||
Plain `consumer_id % len(owners)` would idle half of them."""
|
||||
r = _router(16, 8, group_lens=[1] * 95)
|
||||
names = [f"g{i}" for i in range(95)]
|
||||
for c in range(8):
|
||||
@@ -515,8 +500,7 @@ class TestRdtRouter:
|
||||
|
||||
def test_pipeline_stages_route_to_the_owning_stage(self):
|
||||
"""2 stages x 8 ranks: groups 0-2 on stage 0, 3-5 on stage 1. Two owner
|
||||
sets, one per stage; each consumer reaches both.
|
||||
"""
|
||||
sets, one per stage; each consumer reaches both."""
|
||||
r = _router(
|
||||
16,
|
||||
8,
|
||||
@@ -532,8 +516,7 @@ class TestRdtRouter:
|
||||
|
||||
def test_expert_names_route_to_the_holding_ranks(self):
|
||||
"""A name held by only some ranks routes to those; a replicated name
|
||||
keeps the full set.
|
||||
"""
|
||||
keeps the full set."""
|
||||
r = RdtRouter(
|
||||
8,
|
||||
4,
|
||||
@@ -552,8 +535,7 @@ class TestRdtRouter:
|
||||
|
||||
def test_consumers_spread_over_an_owner_set(self):
|
||||
"""Several ranks hold a name (its TP peers); the block rule must spread
|
||||
consumers across them, not funnel through one.
|
||||
"""
|
||||
consumers across them, not funnel through one."""
|
||||
r = _router(
|
||||
4, 4, owner_sets=[[0, 1]], name_owner_class=[0] * 6, group_lens=[1] * 6
|
||||
)
|
||||
@@ -562,8 +544,7 @@ class TestRdtRouter:
|
||||
|
||||
def test_a_group_frees_at_every_owner_of_any_of_its_names(self):
|
||||
"""The free barrier is per group, so its fan-out is the UNION over the
|
||||
group's names: a rank holding none of them is not signalled.
|
||||
"""
|
||||
group's names: a rank holding none of them is not signalled."""
|
||||
r = RdtRouter(4, 2, [[0, 1], [2, 3]], [0, 1, 0], ["a", "b", "c"], [2, 1])
|
||||
assert r.group_owners(0) == [0, 1, 2, 3]
|
||||
assert r.group_owners(1) == [0, 1]
|
||||
@@ -633,8 +614,7 @@ class TestReplicaOverlay:
|
||||
|
||||
def test_adding_a_deployment_does_not_move_the_first_one(self):
|
||||
"""The strongest form: an existing consumer's routes are byte-identical
|
||||
before and after another deployment is provisioned.
|
||||
"""
|
||||
before and after another deployment is provisioned."""
|
||||
one, names = self._fleet(16, 8, 1)
|
||||
two, _ = self._fleet(16, 8, 2)
|
||||
for consumer_id in range(8):
|
||||
@@ -644,8 +624,7 @@ class TestReplicaOverlay:
|
||||
|
||||
def test_one_deployment_is_the_historical_carve(self):
|
||||
"""Default (no width) and width == the whole fleet must agree with the
|
||||
plain block rule, so a single-deployment fleet is untouched.
|
||||
"""
|
||||
plain block rule, so a single-deployment fleet is untouched."""
|
||||
plain = _router(16, 8, group_lens=[1] * 6)
|
||||
explicit, names = self._fleet(16, 8, 1)
|
||||
for consumer_id in range(8):
|
||||
@@ -659,8 +638,7 @@ class TestReplicaOverlay:
|
||||
def test_the_overlay_holds_for_partially_owned_names(self):
|
||||
"""Pipeline stages and experts too: two owner sets, one narrow. The
|
||||
narrow one already forced agreement (one owner, no block to carve); the
|
||||
wide one is what the carve decides.
|
||||
"""
|
||||
wide one is what the carve decides."""
|
||||
r, names = self._fleet(
|
||||
8,
|
||||
4,
|
||||
@@ -676,16 +654,14 @@ class TestReplicaOverlay:
|
||||
def test_a_width_that_does_not_divide_the_fleet_raises(self):
|
||||
"""A uniform fleet is asserted upstream in ``get_world_size``; if it ever
|
||||
is not, the overlay would map two workers of one deployment onto the same
|
||||
block index, so refuse rather than serve the wrong bytes.
|
||||
"""
|
||||
block index, so refuse rather than serve the wrong bytes."""
|
||||
with pytest.raises(ValueError, match="does not divide"):
|
||||
RdtRouter(4, 6, None, None, ["a"], [1], workers_per_replica=4)
|
||||
|
||||
def test_pulls_still_carry_the_global_consumer_id(self):
|
||||
"""The carve uses the index WITHIN a deployment, but the wire must keep
|
||||
the fleet-global id: it is what lets the producer tell the sharers of a
|
||||
slot apart, and count their arrivals separately.
|
||||
"""
|
||||
slot apart, and count their arrivals separately."""
|
||||
|
||||
class _Method:
|
||||
def __init__(self):
|
||||
@@ -728,8 +704,7 @@ def _producer_pack_offsets(slices):
|
||||
class TestPackedLayout:
|
||||
def test_consumer_layout_matches_the_producer_rule(self):
|
||||
"""Mixed dtypes, sizes that do not land on 16B — i.e. the Kimi-style
|
||||
group (fp8 weights + fp32 scales + bf16 norms) that exposes alignment.
|
||||
"""
|
||||
group (fp8 weights + fp32 scales + bf16 norms) that exposes alignment."""
|
||||
specs = [
|
||||
("w.fp8", "float8_e4m3fn", [17]),
|
||||
("w.scale", "float32", [3]),
|
||||
@@ -804,8 +779,7 @@ class TestPackedLayout:
|
||||
class TestChunkModuleScatters:
|
||||
"""The chunk cut: one chunk per distinct owner class present in the copies,
|
||||
ascending by class index. Derived purely from the bake plus the ownership
|
||||
table, so vLLM's expert placement needs no cases.
|
||||
"""
|
||||
table, so vLLM's expert placement needs no cases."""
|
||||
|
||||
def test_unstamped_copies_form_one_chunk(self):
|
||||
layer = _FakeLayer("l")
|
||||
@@ -819,8 +793,7 @@ class TestChunkModuleScatters:
|
||||
"""Classes are numbered by first appearance in metadata order, so the
|
||||
replicated class (0) leads and the rest follow — the ordering the old
|
||||
"-1 first, then ascending stamp" rule produced, now falling out of the
|
||||
table instead of a special case.
|
||||
"""
|
||||
table instead of a special case."""
|
||||
layer = _FakeLayer("moe")
|
||||
names = ["norm", "e5", "e0", "e2"]
|
||||
copies = [_copy(n, shape=(4,)) for n in names]
|
||||
@@ -840,8 +813,7 @@ class TestChunkModuleScatters:
|
||||
|
||||
def test_copy_order_within_a_chunk_is_bake_order(self):
|
||||
"""Scatter order within a chunk must follow the bake, not the stamp
|
||||
sort: the packed layout and the producer's replay agree on keys order.
|
||||
"""
|
||||
sort: the packed layout and the producer's replay agree on keys order."""
|
||||
layer = _FakeLayer("moe")
|
||||
names = ["e3", "e1", "e2"] # one stamp, deliberately unsorted names
|
||||
copies = [_copy(n, shape=(4,)) for n in names]
|
||||
@@ -855,8 +827,7 @@ class TestChunkModuleScatters:
|
||||
|
||||
def test_stamps_interleave_across_modules_without_reordering_within(self):
|
||||
"""Two modules' copies bucket by stamp; within each bucket the order is
|
||||
module-then-bake order.
|
||||
"""
|
||||
module-then-bake order."""
|
||||
a, b = _FakeLayer("a"), _FakeLayer("b")
|
||||
mods = [
|
||||
_module(a, [_copy("a0", shape=(4,)), _copy("a1", shape=(4,))]),
|
||||
@@ -876,8 +847,7 @@ class TestChunkModuleScatters:
|
||||
def test_scatters_carry_their_own_dtype(self):
|
||||
"""Dtype rides the record from the bake, where it is the fake's dtype
|
||||
AFTER its op chain — not a plan-time lookup of the source name, which
|
||||
would be wrong for any chain that reinterprets dtype.
|
||||
"""
|
||||
would be wrong for any chain that reinterprets dtype."""
|
||||
layer = _FakeLayer("l")
|
||||
eng = _planner({}, name_meta={"w": ("float32", [10])})
|
||||
((_er, scatters),) = eng._chunk_module_scatters(
|
||||
@@ -916,8 +886,7 @@ class TestBuildCallPlan:
|
||||
|
||||
def test_materialize_fires_on_a_modules_first_chunk_only(self):
|
||||
"""Empty HF params are allocated once per module, by construction —
|
||||
including a FusedMoE-like module whose copies span owner-class chunks.
|
||||
"""
|
||||
including a FusedMoE-like module whose copies span owner-class chunks."""
|
||||
layer = _FakeLayer("spanning")
|
||||
names = [f"w{i}" for i in range(4)]
|
||||
copies = [_copy(n, shape=(4,)) for n in names]
|
||||
@@ -952,8 +921,7 @@ class TestBuildCallPlan:
|
||||
def test_free_signal_waits_for_a_groups_last_chunk_when_it_is_split(self):
|
||||
"""The load-bearing case: with the group cut across owner-class chunks the
|
||||
signal must hang off the LAST one. Signaling earlier lets the producers
|
||||
drop the gather buffers while a later chunk's RDMA is still reading.
|
||||
"""
|
||||
drop the gather buffers while a later chunk's RDMA is still reading."""
|
||||
layer = _FakeLayer("l")
|
||||
names = [f"w{i}" for i in range(4)]
|
||||
copies = [_copy(n, shape=(4,)) for n in names]
|
||||
@@ -969,8 +937,7 @@ class TestBuildCallPlan:
|
||||
|
||||
def test_free_signal_timing_across_two_split_groups(self):
|
||||
"""Two groups x 2 owner-class chunks each: signals land on each group's
|
||||
last chunk only.
|
||||
"""
|
||||
last chunk only."""
|
||||
names_a = [f"model.layers.0.w{i}" for i in range(4)]
|
||||
names_b = [f"model.layers.1.w{i}" for i in range(4)]
|
||||
baked, name_meta, stamps = {}, {}, {}
|
||||
@@ -988,8 +955,7 @@ class TestBuildCallPlan:
|
||||
def test_a_group_with_nothing_local_is_signaled_at_sync_start(self):
|
||||
"""The owners still published it, so this consumer must signal it —
|
||||
immediately, before the pipeline (signal-before-publish is tolerated),
|
||||
never hung off another group's chunk.
|
||||
"""
|
||||
never hung off another group's chunk."""
|
||||
baked, name_meta, names, group_lens = _one_module_per_layer(2)
|
||||
del baked["model.layers.1.w"] # nothing baked, and not live => no pull
|
||||
eng = _planner(
|
||||
@@ -1010,8 +976,7 @@ class TestBuildCallPlan:
|
||||
|
||||
def test_modules_are_deduped_within_a_group(self):
|
||||
"""Several names of one fused module map to the same scatter list; the
|
||||
plan must not chunk it twice.
|
||||
"""
|
||||
plan must not chunk it twice."""
|
||||
layer = _FakeLayer("fused")
|
||||
names = ["qkv.q", "qkv.k", "qkv.v"]
|
||||
copies = [_copy(n, shape=(4,)) for n in names]
|
||||
@@ -1027,8 +992,7 @@ class TestBuildCallPlan:
|
||||
|
||||
def test_plan_is_a_pure_function_of_its_inputs(self):
|
||||
"""Called twice with the same arguments it must produce the same plan —
|
||||
that is what lets the engine build it once and reuse it every sync.
|
||||
"""
|
||||
that is what lets the engine build it once and reuse it every sync."""
|
||||
baked, name_meta, names, group_lens = _one_module_per_layer(2)
|
||||
eng = _planner(baked, name_meta=name_meta)
|
||||
a = eng._build_call_plan(names, group_lens)
|
||||
@@ -1048,8 +1012,7 @@ class TestBuildCallPlan:
|
||||
class TestUnbakedNamesGuard:
|
||||
"""A live name with no baked plan has no way to load — there is no fallback,
|
||||
and its pull would target groups the pipeline already freed — so the plan
|
||||
build fails at init, naming the names.
|
||||
"""
|
||||
build fails at init, naming the names."""
|
||||
|
||||
def test_a_live_unbaked_name_fails_the_plan_build(self):
|
||||
baked, name_meta, names, group_lens = _one_module_per_layer(2)
|
||||
@@ -1060,8 +1023,7 @@ class TestUnbakedNamesGuard:
|
||||
|
||||
def test_never_copied_names_are_dropped_entirely(self):
|
||||
"""Experts owned by another EP rank no-op in their loader: not an error,
|
||||
and not a chunk — their group is pre-freed if nothing else fills it.
|
||||
"""
|
||||
and not a chunk — their group is pre-freed if nothing else fills it."""
|
||||
baked, name_meta, names, group_lens = _one_module_per_layer(2)
|
||||
del baked["norm.weight"]
|
||||
eng = _planner(baked, name_meta=name_meta, live=set(names) - {"norm.weight"})
|
||||
@@ -1073,15 +1035,13 @@ class TestUnbakedNamesGuard:
|
||||
class TestCallPlanRouting:
|
||||
"""Under partial ownership each chunk must be pulled from a producer that
|
||||
actually holds every name in it. Consumers bind EVERY producer (signals fan
|
||||
out to all owners), so local index == trainer rank.
|
||||
"""
|
||||
out to all owners), so local index == trainer rank."""
|
||||
|
||||
def _routed_planner(
|
||||
self, held_by_group, num_consumers, consumer_id, *, held_by=None
|
||||
):
|
||||
"""``held_by_group[g]`` = the ranks holding group g's names, so a test
|
||||
can still speak in groups; ``held_by`` overrides individual names.
|
||||
"""
|
||||
can still speak in groups; ``held_by`` overrides individual names."""
|
||||
baked, name_meta, names, group_lens = _one_module_per_layer(
|
||||
len(held_by_group) - 2
|
||||
)
|
||||
@@ -1108,16 +1068,14 @@ class TestCallPlanRouting:
|
||||
|
||||
def test_two_stage_ownership_splits_chunks_between_owners(self):
|
||||
"""Two PP stages, 4 groups: stage 0 holds the first half, stage 1 the
|
||||
second. Every chunk goes to a holder of its names.
|
||||
"""
|
||||
second. Every chunk goes to a holder of its names."""
|
||||
eng, names, group_lens = self._routed_planner([[0], [0], [1], [1]], 1, 0)
|
||||
plan = eng._build_call_plan(names, group_lens)
|
||||
assert [c.owner for c in plan.chunks] == [0, 0, 1, 1]
|
||||
|
||||
def test_a_multi_owner_group_still_resolves_to_one_producer(self):
|
||||
"""Every chunk is served by exactly ONE producer: splitting a pull only
|
||||
multiplies produce calls, since the consumer's NIC bounds it.
|
||||
"""
|
||||
multiplies produce calls, since the consumer's NIC bounds it."""
|
||||
eng, names, group_lens = self._routed_planner([[0, 1]] * 4, 1, 0)
|
||||
plan = eng._build_call_plan(names, group_lens)
|
||||
assert all(0 <= c.owner < 2 for c in plan.chunks)
|
||||
@@ -1125,8 +1083,7 @@ class TestCallPlanRouting:
|
||||
def test_an_expert_name_routes_to_the_rank_holding_it(self):
|
||||
"""PP ∩ EP: stage 0 = ranks 0-1, stage 1 = ranks 2-3, and one expert
|
||||
name inside a stage-1 group is held by rank 3 alone. The name's own owner
|
||||
set is what routes it.
|
||||
"""
|
||||
set is what routes it."""
|
||||
eng, names, group_lens = self._routed_planner(
|
||||
[[0, 1], [0, 1], [2, 3], [2, 3]],
|
||||
1,
|
||||
@@ -1141,8 +1098,7 @@ class TestCallPlanRouting:
|
||||
def test_an_unheld_name_is_rejected_at_init(self):
|
||||
"""A name no rank holds can never be served. The trainer raises when it
|
||||
transposes the holdings; the router's validate() is the consumer-side
|
||||
backstop, and either way it fails at init rather than at first pull.
|
||||
"""
|
||||
backstop, and either way it fails at init rather than at first pull."""
|
||||
with pytest.raises(ValueError, match="empty"):
|
||||
RdtRouter(2, 1, [[]], [0], ["w"], [1]).validate()
|
||||
|
||||
@@ -1151,16 +1107,14 @@ class TestSignalCompleteness:
|
||||
"""Every gather group is signaled exactly once — after its last chunk when
|
||||
the worker pulls from it, at sync start otherwise — for ANY expert
|
||||
placement. A group signaled twice over-credits the barrier; a group never
|
||||
signaled parks a producer credit and hangs end_sync.
|
||||
"""
|
||||
signaled parks a producer credit and hangs end_sync."""
|
||||
|
||||
def _moe_planner(self, worker_experts, *, ep_size=4, n_experts=8, n_layers=2):
|
||||
"""Pre / n_layers MoE layers / post. Each layer: one norm (its own
|
||||
module, stamp -1) + n_experts expert names stamped ``e // n_local``.
|
||||
The worker's bake covers only ``worker_experts`` (its placement);
|
||||
foreign expert names never copied => dropped entirely, as in the real
|
||||
bake.
|
||||
"""
|
||||
bake."""
|
||||
n_local = n_experts // ep_size
|
||||
names = ["embed.weight"]
|
||||
group_lens = [1]
|
||||
@@ -1220,8 +1174,7 @@ class TestSignalCompleteness:
|
||||
|
||||
def test_placement_changes_only_the_chunk_count_never_the_signals(self):
|
||||
"""Linear experts 0-3 hit 2 owner classes; round_robin 0,2,4,6
|
||||
hits all 4. More chunks per group, identical signal set.
|
||||
"""
|
||||
hits all 4. More chunks per group, identical signal set."""
|
||||
plans = {}
|
||||
for label, experts in (("linear", [0, 1, 2, 3]), ("round_robin", [0, 2, 4, 6])):
|
||||
eng, names, group_lens = self._moe_planner(worker_experts=experts)
|
||||
@@ -1235,8 +1188,7 @@ class TestSignalCompleteness:
|
||||
|
||||
def test_the_fused_module_materializes_first_and_quants_last_across_chunks(self):
|
||||
"""The FusedMoE shape: one module's copies span every owner-class chunk of
|
||||
its group; materialize on its first, quant on its last.
|
||||
"""
|
||||
its group; materialize on its first, quant on its last."""
|
||||
eng, names, group_lens = self._moe_planner(
|
||||
worker_experts=[0, 2, 4, 6], n_layers=1
|
||||
)
|
||||
@@ -1294,8 +1246,7 @@ class TestLayerwiseGroups:
|
||||
def test_flattening_a_partition_returns_the_input_order(self):
|
||||
"""Un-indexed names group by POSITION, not by name class — so one that
|
||||
looks like a pre-block name still lands after the layers it follows.
|
||||
Group index therefore means the same thing on every rank.
|
||||
"""
|
||||
Group index therefore means the same thing on every rank."""
|
||||
names = [
|
||||
"embed.weight",
|
||||
"model.layers.0.a",
|
||||
@@ -1313,8 +1264,7 @@ class TestLayerwiseGroups:
|
||||
|
||||
def test_stacks_keep_prefix_appearance_order_and_sort_within(self):
|
||||
"""Two stacks stay in the order they appear — the vision tower before the
|
||||
text stack — while each sorts by index internally.
|
||||
"""
|
||||
text stack — while each sorts by index internally."""
|
||||
names = [
|
||||
"visual.blocks.1.w",
|
||||
"visual.blocks.0.w",
|
||||
@@ -1331,8 +1281,7 @@ class TestLayerwiseGroups:
|
||||
def test_the_post_block_lands_last_however_early_it_arrives(self):
|
||||
"""Megatron-Bridge streams the last pipeline stage's output block before
|
||||
its layers. Sweeping un-indexed names after the first layer into a
|
||||
trailing group is what lets the gather loop walk groups ascending.
|
||||
"""
|
||||
trailing group is what lets the gather loop walk groups ascending."""
|
||||
names = ["model.layers.0.w", "model.norm.weight", "model.layers.1.w"]
|
||||
assert layerwise_groups(names) == [
|
||||
["model.layers.0.w"],
|
||||
@@ -1342,8 +1291,7 @@ class TestLayerwiseGroups:
|
||||
|
||||
def test_names_sharing_a_layer_coalesce_when_the_source_interleaves(self):
|
||||
"""A raw-checkpoint source may yield in shard-packing order. Coalescing on
|
||||
the key keeps that from shattering one layer into many groups.
|
||||
"""
|
||||
the key keeps that from shattering one layer into many groups."""
|
||||
names = [
|
||||
"model.layers.0.a",
|
||||
"model.layers.1.a",
|
||||
@@ -1385,15 +1333,13 @@ class TestLayerwiseGroups:
|
||||
def test_every_supported_naming_convention_partitions_per_layer(self, prefix):
|
||||
"""One decoder layer per group under any prefix, with no per-architecture
|
||||
table. A literal `model.layers.` match silently yields one whole-model
|
||||
group for most of these.
|
||||
"""
|
||||
group for most of these."""
|
||||
names = [f"{prefix}{i}.self_attn.q_proj.weight" for i in range(3)]
|
||||
assert layerwise_groups(names) == [[n] for n in names]
|
||||
|
||||
def test_a_moe_layer_stays_whole(self):
|
||||
"""The OUTERMOST index wins, so per-expert names group by layer. Splitting
|
||||
on the expert would cut FusedMoE — one leaf module — across groups.
|
||||
"""
|
||||
on the expert would cut FusedMoE — one leaf module — across groups."""
|
||||
names = [
|
||||
"model.layers.0.mlp.experts.0.w1.weight",
|
||||
"model.layers.0.mlp.experts.1.w1.weight",
|
||||
@@ -1410,8 +1356,7 @@ class TestLayerwiseGroups:
|
||||
def test_a_vlm_partitions_the_vision_tower_and_the_text_stack(self):
|
||||
"""Qwen-VL layout. Today's literal `model.layers.` match puts every one of
|
||||
these in ONE group — the whole model — because the text stack ships as
|
||||
`model.language_model.layers.`.
|
||||
"""
|
||||
`model.language_model.layers.`."""
|
||||
names = [
|
||||
"visual.patch_embed.proj.weight",
|
||||
"visual.blocks.0.attn.qkv.weight",
|
||||
@@ -1430,8 +1375,7 @@ class TestLayerwiseGroups:
|
||||
|
||||
def test_two_stacks_sharing_an_index_do_not_merge(self):
|
||||
"""The key carries the prefix, so vision block 0 and text layer 0 are
|
||||
different groups even though both are index 0.
|
||||
"""
|
||||
different groups even though both are index 0."""
|
||||
names = [
|
||||
"visual.blocks.0.attn.qkv.weight",
|
||||
"model.language_model.layers.0.self_attn.q_proj.weight",
|
||||
@@ -1468,23 +1412,20 @@ class TestOpAllowlistAgreement:
|
||||
|
||||
def test_every_bake_recordable_op_is_replayable(self):
|
||||
"""Anything the bake can record must be serveable, or the failure lands a
|
||||
whole sync later than the mistake.
|
||||
"""
|
||||
whole sync later than the mistake."""
|
||||
assert not set(SUPPORTED_OPS.values()) - set(ALLOWED_OPS)
|
||||
|
||||
def test_the_producer_allows_nothing_the_consumer_cannot_emit(self):
|
||||
"""The allowlist is also a guard against a misbehaving or spoofed
|
||||
consumer invoking arbitrary methods on trainer tensors, so it must not be
|
||||
wider than the ops the bake can actually produce. ``to`` in particular
|
||||
would let a replay change dtype or device.
|
||||
"""
|
||||
would let a replay change dtype or device."""
|
||||
assert not set(ALLOWED_OPS) - set(SUPPORTED_OPS.values())
|
||||
assert "to" not in ALLOWED_OPS
|
||||
|
||||
def test_transpose_via_t_is_serveable(self):
|
||||
"""The specific regression: a chain recorded from ``.t()`` must pass the
|
||||
producer's guard.
|
||||
"""
|
||||
producer's guard."""
|
||||
fake = FakeRDTTensor(
|
||||
name="w",
|
||||
shape=torch.Size((4, 6)),
|
||||
@@ -1529,8 +1470,7 @@ BAKE_MODELS = [
|
||||
class TestRequiresTheRayExecutor:
|
||||
"""The engine's data plane is Ray's, so a non-Ray executor cannot work. The
|
||||
check is at construction because the alternative is an opaque failure during
|
||||
the first handshake, long after the misconfiguration.
|
||||
"""
|
||||
the first handshake, long after the misconfiguration."""
|
||||
|
||||
def _construct(self, backend):
|
||||
cfg = SimpleNamespace(
|
||||
@@ -1557,8 +1497,7 @@ class TestRequiresTheRayExecutor:
|
||||
|
||||
def test_a_custom_executor_class_is_left_alone(self):
|
||||
"""A `type[Executor]` override is deliberate and unjudgeable here, so it
|
||||
must not be rejected for merely not being the string "ray".
|
||||
"""
|
||||
must not be rejected for merely not being the string "ray"."""
|
||||
|
||||
class _CustomExecutor:
|
||||
pass
|
||||
|
||||
@@ -236,8 +236,7 @@ class TestPublishAndRebuild:
|
||||
def test_views_of_one_storage_do_not_overlap(self, server_factory):
|
||||
"""One IPC export per storage, one as_strided view per name — the whole
|
||||
point of the storage/view split. Overlapping views would serve the same
|
||||
bytes for two different weights.
|
||||
"""
|
||||
bytes for two different weights."""
|
||||
server = server_factory()
|
||||
server.begin_sync(1)
|
||||
_publish(server, GI_A, GROUP_A, nbytes=64)
|
||||
@@ -260,8 +259,7 @@ class TestPublishAndRebuild:
|
||||
def test_the_default_lookahead_is_one(self):
|
||||
"""1 is the sweet spot: group N+1 is gathered AND pullable while N is
|
||||
being pulled, with resident memory at its floor of 2 groups — the bound
|
||||
the larger-model runs size against.
|
||||
"""
|
||||
the larger-model runs size against."""
|
||||
assert DEFAULT_GATHER_LOOKAHEAD == 1
|
||||
|
||||
|
||||
@@ -276,8 +274,7 @@ class TestFreeBarrier:
|
||||
|
||||
def test_group_is_held_until_the_last_live_consumer_signals(self, server_factory):
|
||||
"""Every live consumer signals every owner; releasing on the first
|
||||
signal would drop storage other consumers are still reading.
|
||||
"""
|
||||
signal would drop storage other consumers are still reading."""
|
||||
server = server_factory()
|
||||
server.begin_sync(3)
|
||||
_publish(server, GI_A, GROUP_A)
|
||||
@@ -304,8 +301,7 @@ class TestFreeBarrier:
|
||||
):
|
||||
"""A consumer with nothing to pull from a group signals it at sync
|
||||
start, which can precede the publish. The publish must then release the
|
||||
group rather than wait for a signal that will never come again.
|
||||
"""
|
||||
group rather than wait for a signal that will never come again."""
|
||||
server = server_factory()
|
||||
server.begin_sync(1)
|
||||
server.free_group(GI_A)
|
||||
@@ -328,8 +324,7 @@ class TestFreeBarrier:
|
||||
the server says the group is gone — and it hears that ONLY through
|
||||
``wait_freed`` / ``end_sync``, never a publish return (a freed notice
|
||||
riding an unharvested async publish while the engine blocks in
|
||||
``wait_freed`` would wedge the loop).
|
||||
"""
|
||||
``wait_freed`` would wedge the loop)."""
|
||||
server = server_factory()
|
||||
server.begin_sync(1)
|
||||
_publish(server, GI_A, GROUP_A)
|
||||
@@ -358,8 +353,7 @@ class TestFreeBarrier:
|
||||
|
||||
class _CountingSource:
|
||||
"""`iter_groups` yielding one CPU-tensor group per call, built lazily so a
|
||||
group's tensors exist only once the engine's credit gate lets it gather.
|
||||
"""
|
||||
group's tensors exist only once the engine's credit gate lets it gather."""
|
||||
|
||||
def __init__(self, groups):
|
||||
self._groups = groups
|
||||
@@ -372,8 +366,7 @@ class _CountingSource:
|
||||
class _ResidencyDict(dict):
|
||||
"""The engine's `_inflight`, instrumented: `max_len` is the high-water mark
|
||||
of groups whose CUDA-IPC export refs were alive at once — the trainer's
|
||||
resident-memory bound in groups.
|
||||
"""
|
||||
resident-memory bound in groups."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
@@ -387,8 +380,7 @@ class _ResidencyDict(dict):
|
||||
def _loop_engine(server, n_groups, *, lookahead):
|
||||
"""A real engine wired straight to a real (non-Ray) server, with only the
|
||||
state `_run_gather_loop` reads. The `_rpc` seam dispatches inline;
|
||||
`_publish_async` sees no `.remote` and runs inline too.
|
||||
"""
|
||||
`_publish_async` sees no `.remote` and runs inline too."""
|
||||
groups = [(f"model.layers.{gi}.w",) for gi in range(n_groups)]
|
||||
e = ShardedRDTTrainerWeightTransferEngine.__new__(
|
||||
ShardedRDTTrainerWeightTransferEngine
|
||||
@@ -411,16 +403,14 @@ class TestGatherCredit:
|
||||
serveable immediately, and the ENGINE's loop stops gathering while more than
|
||||
`gather_lookahead` groups are unfreed. So at most `lookahead + 1` groups are
|
||||
resident — two at the default — while group N+1 is pulled the instant N's
|
||||
pulls finish.
|
||||
"""
|
||||
pulls finish."""
|
||||
|
||||
@pytest.fixture
|
||||
def gather_engine(self, monkeypatch):
|
||||
"""Neutralize the two CUDA touches in the export path so the REAL
|
||||
gather loop runs on CPU: `.cuda()` becomes identity, and the storage
|
||||
export ships the `(nbytes, ..., device, ...)` tuple the fixture's
|
||||
stubbed `rebuild_cuda_tensor` expects.
|
||||
"""
|
||||
stubbed `rebuild_cuda_tensor` expects."""
|
||||
monkeypatch.setattr(torch.Tensor, "cuda", lambda self: self)
|
||||
monkeypatch.setattr(
|
||||
trainer_mod,
|
||||
@@ -432,8 +422,7 @@ class TestGatherCredit:
|
||||
def test_publish_never_blocks(self, server_factory):
|
||||
"""Publishing must not park on a credit: the whole point is that a
|
||||
gathered group becomes serveable immediately. Three publishes with no
|
||||
frees, inline on this thread — a block here hangs the test.
|
||||
"""
|
||||
frees, inline on this thread — a block here hangs the test."""
|
||||
server = server_factory(gather_lookahead=1)
|
||||
server.begin_sync(1)
|
||||
for gi in range(3):
|
||||
@@ -447,8 +436,7 @@ class TestGatherCredit:
|
||||
than 2 groups of gathered weights, no matter how the consumer paces —
|
||||
and the pipeline still overlaps, because the consumer here refuses to
|
||||
free group N until N+1 is already published (pullable). A gate an
|
||||
off-by-one too tight deadlocks this test; too loose fails max_len.
|
||||
"""
|
||||
off-by-one too tight deadlocks this test; too loose fails max_len."""
|
||||
server = server_factory(gather_lookahead=1)
|
||||
n_groups = 8
|
||||
engine = gather_engine(server, n_groups, lookahead=1)
|
||||
@@ -490,8 +478,7 @@ class TestGatherCredit:
|
||||
):
|
||||
"""0 is the serialized baseline: the loop waits for group N to be freed
|
||||
before gathering N+1, so exactly ONE group is ever resident and nothing
|
||||
overlaps. A floor of 1 on the bound would silently pipeline instead.
|
||||
"""
|
||||
overlaps. A floor of 1 on the bound would silently pipeline instead."""
|
||||
server = server_factory(gather_lookahead=0)
|
||||
engine = gather_engine(server, 4, lookahead=0)
|
||||
|
||||
@@ -517,8 +504,7 @@ class TestGatherCredit:
|
||||
self, server_factory, gather_engine
|
||||
):
|
||||
"""Lookahead + 1, not a hardcoded 2: at lookahead=2 a free-nothing
|
||||
consumer sees exactly 3 groups gathered before the loop parks.
|
||||
"""
|
||||
consumer sees exactly 3 groups gathered before the loop parks."""
|
||||
server = server_factory(gather_lookahead=2)
|
||||
engine = gather_engine(server, 8, lookahead=2)
|
||||
|
||||
@@ -554,8 +540,7 @@ class TestGatherCredit:
|
||||
def test_a_gather_error_releases_a_blocked_wait_freed(self, server_factory):
|
||||
"""Otherwise a trainer-side failure on another rank deadlocks this one
|
||||
inside its credit gate. It must RAISE, not return empty — an empty
|
||||
return would spin the engine straight back into the wait.
|
||||
"""
|
||||
return would spin the engine straight back into the wait."""
|
||||
server = server_factory()
|
||||
server.begin_sync(1)
|
||||
_publish(server, GI_A, GROUP_A)
|
||||
@@ -636,8 +621,7 @@ class TestBeginSync:
|
||||
|
||||
def test_the_target_is_floored_at_one(self, server_factory):
|
||||
"""A bare/zero call must not make every group free instantly (or divide
|
||||
the barrier by zero).
|
||||
"""
|
||||
the barrier by zero)."""
|
||||
server = server_factory()
|
||||
server.begin_sync(0)
|
||||
assert server._live_count == 1
|
||||
@@ -645,8 +629,7 @@ class TestBeginSync:
|
||||
def test_a_degraded_sync_lowers_the_target(self, server_factory):
|
||||
"""FT degraded sync: the live count is the WHOLE degraded-sync
|
||||
mechanism on this side — the group releases after the live consumers
|
||||
alone.
|
||||
"""
|
||||
alone."""
|
||||
server = server_factory()
|
||||
server.begin_sync(2) # 4 provisioned, 2 alive
|
||||
_publish(server, GI_A, GROUP_A)
|
||||
@@ -670,8 +653,7 @@ class TestBeginSync:
|
||||
def test_a_straggler_signal_cannot_credit_the_next_sync(self, server_factory):
|
||||
"""begin_sync resets the counts, so a signal that leaked past the sync
|
||||
boundary would over-credit — the reason the consumer drains its fired
|
||||
signals before finishing. Pin the reset side of that contract.
|
||||
"""
|
||||
signals before finishing. Pin the reset side of that contract."""
|
||||
server = server_factory()
|
||||
server.begin_sync(2)
|
||||
_publish(server, GI_A, GROUP_A)
|
||||
@@ -693,8 +675,7 @@ class TestBeginSync:
|
||||
self, server_factory
|
||||
):
|
||||
"""The packed layout repeats every sync — that is what makes caching the
|
||||
destination views worth ~5ms per 384-spec group.
|
||||
"""
|
||||
destination views worth ~5ms per 384-spec group."""
|
||||
server = server_factory()
|
||||
server._pack_dsts[("sentinel",)] = (0, [])
|
||||
server.begin_sync(1)
|
||||
@@ -705,8 +686,7 @@ class TestServedNamesGuard:
|
||||
"""A producer serves only the names it holds: its stage's groups, and
|
||||
within them its own EP coordinate's experts plus the replicated names.
|
||||
Without this guard a misrouted pull would block forever in the cache wait
|
||||
for a name this rank never publishes.
|
||||
"""
|
||||
for a name this rank never publishes."""
|
||||
|
||||
def test_a_pull_for_an_unserved_name_fails_loudly(self, server_factory):
|
||||
server = server_factory(served_names=list(GROUP_A))
|
||||
@@ -722,8 +702,7 @@ class TestServedNamesGuard:
|
||||
|
||||
def test_served_names_none_accepts_anything(self, server_factory):
|
||||
"""Gather-to-all: every producer holds every group, so there is nothing
|
||||
to guard. The call proceeds to the cache wait (not exercised here).
|
||||
"""
|
||||
to guard. The call proceeds to the cache wait (not exercised here)."""
|
||||
server = server_factory(served_names=None)
|
||||
assert server._served_names is None
|
||||
|
||||
@@ -770,8 +749,7 @@ class TestStallWatchdog:
|
||||
|
||||
def test_end_sync_fails_when_a_consumer_never_signals(self, server_factory):
|
||||
"""The exact mid-sync death signature: the group is published, one of
|
||||
the two live consumers that owed a signal is gone.
|
||||
"""
|
||||
the two live consumers that owed a signal is gone."""
|
||||
server = server_factory(stall_timeout_s=0.3)
|
||||
server.begin_sync(2)
|
||||
_publish(server, GI_A, GROUP_A)
|
||||
@@ -785,8 +763,7 @@ class TestStallWatchdog:
|
||||
"""The watchdog fires on the existing `set_gather_error` channel, so one
|
||||
stall unwinds the whole rank through one path rather than each waiter
|
||||
timing out separately — here both a consumer parked in the serve cache
|
||||
wait and the engine parked in its credit gate.
|
||||
"""
|
||||
wait and the engine parked in its credit gate."""
|
||||
server = server_factory(gather_lookahead=1, stall_timeout_s=0.3)
|
||||
server.begin_sync(1)
|
||||
_publish(server, GI_A, GROUP_A)
|
||||
@@ -819,8 +796,7 @@ class TestStallWatchdog:
|
||||
"""A consumer that is slow but signaling must never trip it: the stamp
|
||||
is global to the producer, so a steady trickle of signals holds the
|
||||
whole rank open. Nine sequential publishes with a 0.3s timeout each take
|
||||
longer in total than the timeout, and none may fire.
|
||||
"""
|
||||
longer in total than the timeout, and none may fire."""
|
||||
server = server_factory(gather_lookahead=1, stall_timeout_s=0.3)
|
||||
server.begin_sync(1)
|
||||
for gi in range(9):
|
||||
@@ -833,8 +809,7 @@ class TestStallWatchdog:
|
||||
|
||||
def test_begin_sync_resets_the_progress_stamp(self, server_factory):
|
||||
"""Syncs are minutes apart. Without the reset the first publish of sync N+1
|
||||
would measure its stall from somewhere inside sync N and fire immediately.
|
||||
"""
|
||||
would measure its stall from somewhere inside sync N and fire immediately."""
|
||||
server = server_factory(gather_lookahead=1, stall_timeout_s=0.3)
|
||||
server.begin_sync(1)
|
||||
_publish(server, GI_A, GROUP_A)
|
||||
@@ -856,8 +831,7 @@ class TestConcurrentPublishAndFree:
|
||||
self, server_factory
|
||||
):
|
||||
"""Every group published must end up freed and out of the cache, with
|
||||
end_sync returning cleanly — the property whose violation is a hang.
|
||||
"""
|
||||
end_sync returning cleanly — the property whose violation is a hang."""
|
||||
server = server_factory(gather_lookahead=2)
|
||||
server.begin_sync(1)
|
||||
groups = [(gi, (f"model.layers.{gi}.w",)) for gi in range(12)]
|
||||
@@ -892,8 +866,7 @@ class TestConcurrentPublishAndFree:
|
||||
):
|
||||
"""All live consumers signaling at once: the release must happen on the
|
||||
last signal and only once, or the engine drops IPC refs twice (or
|
||||
never).
|
||||
"""
|
||||
never)."""
|
||||
server = server_factory()
|
||||
server.begin_sync(8)
|
||||
_publish(server, GI_A, GROUP_A)
|
||||
@@ -920,8 +893,7 @@ class TestFakeServerAgreesWithTheRealOne:
|
||||
"""`test_sharded_rdt_trainer.py`'s `_FakeProducerServer` is a second,
|
||||
independent model of this protocol that the engine-side tests assert on. Pin
|
||||
the two against each other so the fake cannot silently drift from the
|
||||
semantics the real server enforces.
|
||||
"""
|
||||
semantics the real server enforces."""
|
||||
|
||||
@staticmethod
|
||||
def _fake():
|
||||
@@ -952,8 +924,7 @@ class TestExportRing:
|
||||
"""The packed export path. Its one hard invariant is that the trainer packs
|
||||
a group exactly the way ``publish_group`` rebuilds it — a layout
|
||||
disagreement would serve the wrong bytes with nothing downstream to catch
|
||||
it.
|
||||
"""
|
||||
it."""
|
||||
|
||||
@staticmethod
|
||||
def _engine():
|
||||
@@ -1045,8 +1016,7 @@ class TestExportRingSlotSafety:
|
||||
):
|
||||
"""Free the later group of each pair first and hold the earlier one live
|
||||
until the NEXT group is gathered -- the window a modular slot counter got
|
||||
wrong (at lookahead=1 there are 2 slots, so group 2 took group 0's).
|
||||
"""
|
||||
wrong (at lookahead=1 there are 2 slots, so group 2 took group 0's)."""
|
||||
server = server_factory(gather_lookahead=1)
|
||||
n_groups = 8
|
||||
engine = gather_engine(server, n_groups, lookahead=1)
|
||||
@@ -1170,8 +1140,7 @@ class TestSharedServeSlots:
|
||||
|
||||
def test_one_deployment_never_shares(self, sharing_server):
|
||||
"""With the width equal to the live count every group is a singleton, so
|
||||
each call packs for itself out of its own ring.
|
||||
"""
|
||||
each call packs for itself out of its own ring."""
|
||||
server = sharing_server(workers_per_replica=2)
|
||||
server.begin_sync(2, [0, 1])
|
||||
_publish(server, GI_A, GROUP_A)
|
||||
@@ -1251,8 +1220,7 @@ class TestSharedServeSlots:
|
||||
|
||||
def test_a_seqless_call_is_refused(self, sharing_server):
|
||||
"""The slot comes from ``seq``, so a caller that sends none cannot be
|
||||
served safely.
|
||||
"""
|
||||
served safely."""
|
||||
server = sharing_server(workers_per_replica=2)
|
||||
server.begin_sync(1, [0])
|
||||
_publish(server, GI_A, GROUP_A)
|
||||
@@ -1262,8 +1230,7 @@ class TestSharedServeSlots:
|
||||
|
||||
def test_a_failed_pack_reaches_every_sharer(self, sharing_server):
|
||||
"""Only the packer touches the buffer, so its failure has to be
|
||||
republished: a waiter must raise rather than read a half-written slot.
|
||||
"""
|
||||
republished: a waiter must raise rather than read a half-written slot."""
|
||||
server = sharing_server(workers_per_replica=2)
|
||||
server.begin_sync(4, [0, 1, 2, 3])
|
||||
_publish(server, GI_A, GROUP_A)
|
||||
@@ -1287,8 +1254,7 @@ class TestSharedServeSlots:
|
||||
|
||||
def test_a_degraded_sync_narrows_the_rendezvous(self, sharing_server):
|
||||
"""A dead deployment must not be waited for: a survivor left alone in
|
||||
its group packs for itself instead of blocking on the corpse.
|
||||
"""
|
||||
its group packs for itself instead of blocking on the corpse."""
|
||||
server = sharing_server(workers_per_replica=2, stall_timeout_s=30.0)
|
||||
server.begin_sync(3, [0, 1, 3]) # consumer 2 died
|
||||
_publish(server, GI_A, GROUP_A)
|
||||
@@ -1302,8 +1268,7 @@ class TestSharedServeSlots:
|
||||
|
||||
def test_a_sharer_that_never_arrives_trips_the_watchdog(self, sharing_server):
|
||||
"""A consumer that dies INSIDE the sync window has no detector, exactly
|
||||
as for the free barrier; the rendezvous waits on the same channel.
|
||||
"""
|
||||
as for the free barrier; the rendezvous waits on the same channel."""
|
||||
server = sharing_server(workers_per_replica=2, stall_timeout_s=0.3)
|
||||
server.begin_sync(4, [0, 1, 2, 3])
|
||||
_publish(server, GI_A, GROUP_A)
|
||||
@@ -1315,8 +1280,7 @@ class TestSharedServeSlots:
|
||||
|
||||
class TestSharedReservation:
|
||||
"""``reserve_serve_buffer`` under sharing: one ring per GROUP, and the
|
||||
init-time check that the group's members pull the same chunks.
|
||||
"""
|
||||
init-time check that the group's members pull the same chunks."""
|
||||
|
||||
def test_matching_plans_reserve_one_ring_between_them(self, sharing_server):
|
||||
server = sharing_server(workers_per_replica=2)
|
||||
@@ -1338,8 +1302,7 @@ class TestSharedReservation:
|
||||
def test_mismatched_plans_fail_at_reservation(self, sharing_server):
|
||||
"""Sharing rests on the sharers pulling the same chunks. Checked here,
|
||||
where it names both consumers, rather than at serve time, where it is a
|
||||
rendezvous nobody completes.
|
||||
"""
|
||||
rendezvous nobody completes."""
|
||||
server = sharing_server(workers_per_replica=2)
|
||||
server.reserve_serve_buffer(0, 1024, "plan-a")
|
||||
|
||||
@@ -1348,8 +1311,7 @@ class TestSharedReservation:
|
||||
|
||||
def test_different_groups_may_differ(self, sharing_server):
|
||||
"""Only the members of ONE group have to agree; worker 0 and worker 1 of
|
||||
the same deployment hold different slices and never share a slot.
|
||||
"""
|
||||
the same deployment hold different slices and never share a slot."""
|
||||
server = sharing_server(workers_per_replica=2)
|
||||
server.reserve_serve_buffer(0, 1024, "plan-a")
|
||||
server.reserve_serve_buffer(1, 1024, "plan-b")
|
||||
@@ -1358,8 +1320,7 @@ class TestSharedReservation:
|
||||
|
||||
def test_a_missing_digest_skips_the_check(self, sharing_server):
|
||||
"""The digest is optional on the wire, so an engine that does not send
|
||||
one still reserves.
|
||||
"""
|
||||
one still reserves."""
|
||||
server = sharing_server(workers_per_replica=2)
|
||||
server.reserve_serve_buffer(0, 1024, None)
|
||||
server.reserve_serve_buffer(2, 1024, None)
|
||||
|
||||
@@ -34,8 +34,7 @@ from vllm.distributed.weight_transfer.sharded_rdt_trainer import (
|
||||
|
||||
class _ListSource(WeightSource):
|
||||
"""A WeightSource over an explicit ordered (name, cpu-tensor) list, so the
|
||||
sharded-RDT group/order logic can be tested without a real model.
|
||||
"""
|
||||
sharded-RDT group/order logic can be tested without a real model."""
|
||||
|
||||
def __init__(self, pairs):
|
||||
self._pairs = list(pairs)
|
||||
@@ -55,8 +54,7 @@ class _FakeProducerServer:
|
||||
per-group barrier: publish/free are keyed by GROUP INDEX, signals count
|
||||
to the ``begin_sync`` live total, publish_group returns nothing, and freed
|
||||
groups flow back only through wait_freed / end_sync (see
|
||||
test_sharded_rdt_producer.TestFakeServerAgreesWithTheRealOne).
|
||||
"""
|
||||
test_sharded_rdt_producer.TestFakeServerAgreesWithTheRealOne)."""
|
||||
|
||||
def __init__(self, auto_free=True):
|
||||
self.order: list[str] = []
|
||||
@@ -100,8 +98,7 @@ class _FakeProducerServer:
|
||||
must already have a banked credit when the gate asks (auto_free, or a
|
||||
test's own free_group) — anything else is the deadlock the real
|
||||
watchdog would kill, so fail loudly. Not appended to ``order``: it is
|
||||
pacing, not a lifecycle milestone.
|
||||
"""
|
||||
pacing, not a lifecycle milestone."""
|
||||
assert self._pending_freed, (
|
||||
"wait_freed with nothing freed: the gather loop would deadlock"
|
||||
)
|
||||
@@ -124,8 +121,7 @@ def _rdt_engine_with_fake_server(
|
||||
):
|
||||
"""Build a ShardedRDTTrainerWeightTransferEngine wired to an in-process fake
|
||||
server (no Ray, no CUDA IPC): bypass trainer_init's spawn, set the
|
||||
group-major metadata, and route _rpc to the fake.
|
||||
"""
|
||||
group-major metadata, and route _rpc to the fake."""
|
||||
import vllm.distributed.weight_transfer.sharded_rdt_trainer as mod
|
||||
|
||||
# reduce_tensor needs CUDA; the fake server never rebuilds, so stub it.
|
||||
@@ -325,8 +321,7 @@ def test_sharded_rdt_send_weights_surfaces_update_error(monkeypatch):
|
||||
|
||||
class _OwnedSource(_ListSource):
|
||||
"""A source holding only some groups' names, like a pipeline-parallel rank.
|
||||
Takes group indices for convenience and declares the names inside them.
|
||||
"""
|
||||
Takes group indices for convenience and declares the names inside them."""
|
||||
|
||||
def __init__(self, pairs, owned_group_idx):
|
||||
super().__init__(pairs)
|
||||
@@ -392,8 +387,7 @@ def test_sharded_rdt_owned_group_order_mismatch_raises(monkeypatch):
|
||||
)
|
||||
def test_sharded_rdt_begin_sync_carries_the_live_count(monkeypatch):
|
||||
"""The free barrier's target is one integer per sync: the live consumer
|
||||
count, defaulting to the whole provisioned fleet.
|
||||
"""
|
||||
count, defaulting to the whole provisioned fleet."""
|
||||
server = _FakeProducerServer(auto_free=True)
|
||||
engine = _rdt_engine_with_fake_server(
|
||||
_rdt_source_two_layers(),
|
||||
@@ -411,8 +405,7 @@ def test_sharded_rdt_begin_sync_carries_the_live_count(monkeypatch):
|
||||
class TestLiveCountPlumbing:
|
||||
"""``send_weights(live_consumer_ids)`` -> the barrier target. The
|
||||
provisioned geometry is frozen; a degraded sync only lowers the target and
|
||||
narrows each slot-sharing group's rendezvous.
|
||||
"""
|
||||
narrows each slot-sharing group's rendezvous."""
|
||||
|
||||
@staticmethod
|
||||
def _engine(num_consumers, world=2, rank=0):
|
||||
@@ -445,8 +438,7 @@ class TestLiveCountPlumbing:
|
||||
|
||||
def test_which_consumers_died_does_not_matter_only_how_many(self):
|
||||
"""The whole point of the barrier: no routed per-producer targets, so
|
||||
the identity of the dead consumer is irrelevant to the producers.
|
||||
"""
|
||||
the identity of the dead consumer is irrelevant to the producers."""
|
||||
counts = []
|
||||
for live in ([0, 1, 2, 3], [4, 5, 6, 7], [0, 2, 4, 6]):
|
||||
engine, got = self._engine(num_consumers=8)
|
||||
@@ -457,8 +449,7 @@ class TestLiveCountPlumbing:
|
||||
def test_the_live_ids_travel_with_the_count(self):
|
||||
"""The count sizes the free barrier, the ids size the slot-sharing
|
||||
rendezvous, so a producer that shares slots can tell WHICH consumers it
|
||||
is still waiting for. They must describe the same set.
|
||||
"""
|
||||
is still waiting for. They must describe the same set."""
|
||||
engine, got = self._engine(num_consumers=8)
|
||||
engine.send_weights([6, 0, 2, 2])
|
||||
((count, ids),) = got
|
||||
@@ -526,8 +517,7 @@ def test_sharded_rdt_rejects_a_rank_holding_nothing(monkeypatch):
|
||||
|
||||
def test_sharded_rdt_rejects_a_name_no_rank_holds(monkeypatch):
|
||||
"""Every name must be held somewhere or it can never be served — caught
|
||||
when the holdings are transposed, naming the orphan.
|
||||
"""
|
||||
when the holdings are transposed, naming the orphan."""
|
||||
with pytest.raises(ValueError, match="no trainer rank holds"):
|
||||
_rdt_engine_with_fake_server(
|
||||
_OwnedSource(_rdt_source_two_layers()._pairs, [0, 1]),
|
||||
@@ -541,8 +531,7 @@ def test_sharded_rdt_rejects_a_name_no_rank_holds(monkeypatch):
|
||||
|
||||
def test_sharded_rdt_rejects_metadata_disagreement_across_ranks(monkeypatch):
|
||||
"""Only the sender's metadata reaches the consumers, so a rank describing
|
||||
just its own share must fail loudly rather than silently drop the rest.
|
||||
"""
|
||||
just its own share must fail loudly rather than silently drop the rest."""
|
||||
engine = _rdt_engine_with_fake_server(
|
||||
_OwnedSource(_rdt_source_two_layers()._pairs, [1, 2]),
|
||||
is_sender=False,
|
||||
@@ -565,8 +554,7 @@ def _serve_ring_server(src_name, src):
|
||||
"""A producer server with one cached tensor and a pre-seeded serve ring, so a
|
||||
pull needs no Ray and no NIXL registration. Returns (server, serve) where
|
||||
``serve(chain)`` packs one spec into the SAME ring slot every time — which is
|
||||
what puts the destination-view cache, and only it, under test.
|
||||
"""
|
||||
what puts the destination-view cache, and only it, under test."""
|
||||
from vllm.distributed.weight_transfer.sharded_rdt_trainer import (
|
||||
_RDTProducerServer,
|
||||
)
|
||||
|
||||
@@ -419,8 +419,7 @@ def test_tensor_broadcast():
|
||||
|
||||
def _dumps_oob(obj) -> tuple[bytes, list]:
|
||||
"""Pickle `obj` the same way `MessageQueue.enqueue` does: tensor
|
||||
dispatch table + out-of-band buffers >= 1MiB.
|
||||
"""
|
||||
dispatch table + out-of-band buffers >= 1MiB."""
|
||||
buffers = []
|
||||
|
||||
def callback(buf: pickle.PickleBuffer) -> bool:
|
||||
@@ -491,8 +490,7 @@ def test_tensor_pickle_roundtrip(case: str):
|
||||
@pytest.mark.parametrize("case", ["cuda", "requires_grad", "conj"])
|
||||
def test_reduce_tensor_fallback(case: str):
|
||||
"""Tensors the zero-copy reducer can't safely alias must fall back to
|
||||
torch's default reduction.
|
||||
"""
|
||||
torch's default reduction."""
|
||||
if case == "cuda":
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("requires CUDA")
|
||||
@@ -755,8 +753,7 @@ def test_remote_subscribe_addr_unique_concurrent_writers(
|
||||
Pre-fix, the writer probed a port with get_open_port() and bound it
|
||||
afterwards; pinning the probe to one free port makes every writer
|
||||
bind the same port and fail deterministically on that code path,
|
||||
while the late-binding implementation never consults the probe.
|
||||
"""
|
||||
while the late-binding implementation never consults the probe."""
|
||||
from vllm.distributed.device_communicators import shm_broadcast
|
||||
|
||||
colliding_port = get_open_port()
|
||||
|
||||
@@ -31,8 +31,7 @@ class _CUDADeviceCountStatelessTestActor:
|
||||
|
||||
def test_cuda_device_count_stateless():
|
||||
"""Test that cuda_device_count_stateless changes return value if
|
||||
CUDA_VISIBLE_DEVICES is changed.
|
||||
"""
|
||||
CUDA_VISIBLE_DEVICES is changed."""
|
||||
if current_platform.is_rocm():
|
||||
pytest.skip("Skip for ROCm because Ray uses HIP_VISIBLE_DEVICES.")
|
||||
actor = _CUDADeviceCountStatelessTestActor.options( # type: ignore
|
||||
|
||||
@@ -1166,8 +1166,7 @@ class RecordingClient:
|
||||
|
||||
def _module_with(*pairs):
|
||||
"""A tiny nn.Module exposing the given (name, tensor) pairs as parameters,
|
||||
so trainer tests can build a ModuleSource without a real model.
|
||||
"""
|
||||
so trainer tests can build a ModuleSource without a real model."""
|
||||
module = torch.nn.Module()
|
||||
for name, tensor in pairs:
|
||||
module.register_parameter(name, torch.nn.Parameter(tensor, requires_grad=False))
|
||||
@@ -1288,8 +1287,7 @@ class TestModuleSource:
|
||||
|
||||
def test_metadata_agrees_with_iteration(self):
|
||||
"""The two channels must line up element-for-element: engines declare
|
||||
the round from `metadata()` and then send what iteration yields.
|
||||
"""
|
||||
the round from `metadata()` and then send what iteration yields."""
|
||||
source = ModuleSource(
|
||||
_module_with(("w", torch.zeros(2, 3)), ("b", torch.zeros(3)))
|
||||
)
|
||||
@@ -1303,13 +1301,11 @@ class TestModuleSource:
|
||||
class TestWeightSourceGroupContract:
|
||||
"""`groups()` / `iter_groups()` on the WeightSource ABC. Groups are what
|
||||
backends gather and free by, so the default must agree with
|
||||
`layerwise_groups` over `metadata()`, restricted to what this rank holds.
|
||||
"""
|
||||
`layerwise_groups` over `metadata()`, restricted to what this rank holds."""
|
||||
|
||||
class _Source(WeightSource):
|
||||
"""Minimal source over an ordered (name, tensor) list, optionally holding
|
||||
only some names (in which case it iterates only their groups).
|
||||
"""
|
||||
only some names (in which case it iterates only their groups)."""
|
||||
|
||||
def __init__(self, names, held=None, reverse=False):
|
||||
self._pairs = [(n, torch.full((2,), float(i))) for i, n in enumerate(names)]
|
||||
@@ -1364,8 +1360,7 @@ class TestWeightSourceGroupContract:
|
||||
def test_groups_order_follows_metadata_not_the_declaration(self):
|
||||
"""The declaration is a SET of names; the groups it selects still come
|
||||
out in metadata order, so each pairs with the right ``iter_groups()``
|
||||
batch however the source listed them.
|
||||
"""
|
||||
batch however the source listed them."""
|
||||
names = ["embed.w", "model.layers.0.a", "model.layers.1.a", "norm.w"]
|
||||
source = self._source(
|
||||
names, held=["model.layers.1.a", "model.layers.0.a", "model.layers.1.a"]
|
||||
@@ -1399,16 +1394,14 @@ class TestWeightSourceGroupContract:
|
||||
|
||||
def test_out_of_order_iteration_raises(self):
|
||||
"""Materializing is usually a collective, so a rank that iterates out of
|
||||
order deadlocks its peers -- fail loudly instead.
|
||||
"""
|
||||
order deadlocks its peers -- fail loudly instead."""
|
||||
source = self._source(["model.layers.0.a", "model.layers.0.b"], reverse=True)
|
||||
with pytest.raises(RuntimeError, match="iteration order must match"):
|
||||
list(source.iter_groups())
|
||||
|
||||
def test_a_source_may_override_iter_groups(self):
|
||||
"""The extension point: materialize a whole group in one step instead of
|
||||
one generator resume per tensor.
|
||||
"""
|
||||
one generator resume per tensor."""
|
||||
calls = []
|
||||
|
||||
class _Batched(TestWeightSourceGroupContract._Source):
|
||||
@@ -1429,8 +1422,7 @@ class TestDeferredProcessingContract:
|
||||
caller that takes over the update tail (running its own
|
||||
`finalize_layerwise_reload` instead of going through `finish_weight_update`)
|
||||
reads the flag and calls the method. Both must be answerable on any engine, or
|
||||
that caller ends up reaching through a getattr.
|
||||
"""
|
||||
that caller ends up reaching through a getattr."""
|
||||
|
||||
def _engines(self):
|
||||
return {
|
||||
@@ -1444,8 +1436,7 @@ class TestDeferredProcessingContract:
|
||||
|
||||
def test_every_engine_can_be_drained(self):
|
||||
"""The default is a no-op, so a caller never has to check whether the
|
||||
method exists before calling it.
|
||||
"""
|
||||
method exists before calling it."""
|
||||
for name, cls in self._engines().items():
|
||||
assert callable(cls.drain_pending), name
|
||||
|
||||
@@ -1539,8 +1530,7 @@ class TestTrainerEngineBase:
|
||||
)
|
||||
def test_ipc_trainer_send_weights_drives_client_in_order():
|
||||
"""send_weights issues start -> update -> finish and ships per-round metadata;
|
||||
the packed wire param rides the init info, not the per-round update_info.
|
||||
"""
|
||||
the packed wire param rides the init info, not the per-round update_info."""
|
||||
client = RecordingClient()
|
||||
engine = IPCTrainerWeightTransferEngine(
|
||||
client=client,
|
||||
@@ -1559,8 +1549,7 @@ def test_ipc_trainer_send_weights_drives_client_in_order():
|
||||
|
||||
def test_ipc_trainer_init_ships_packed_to_worker():
|
||||
"""trainer_init drives the inference-side init handshake and propagates the
|
||||
must-agree `packed` flag to the worker.
|
||||
"""
|
||||
must-agree `packed` flag to the worker."""
|
||||
if torch.accelerator.device_count() < 1:
|
||||
pytest.skip("Need at least 1 GPU (CUDA IPC handles).")
|
||||
|
||||
@@ -1581,8 +1570,7 @@ def test_ipc_trainer_init_ships_packed_to_worker():
|
||||
def test_nccl_trainer_init_ships_worker_init_info(monkeypatch):
|
||||
"""The sender's trainer_init drives the inference-side init handshake with
|
||||
the worker-shaped init info (rank_offset=1) while opening its own endpoint,
|
||||
and propagates the must-agree wire params to the worker.
|
||||
"""
|
||||
and propagates the must-agree wire params to the worker."""
|
||||
import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod
|
||||
|
||||
# Bypass the real NCCL rendezvous.
|
||||
@@ -1623,8 +1611,7 @@ def test_nccl_trainer_init_ships_worker_init_info(monkeypatch):
|
||||
def test_nccl_worker_learns_wire_params_from_init_handshake(monkeypatch):
|
||||
"""The worker engine reads packed + buffer geometry from the
|
||||
trainer-supplied init info at the handshake, not from the config or the
|
||||
per-round update info.
|
||||
"""
|
||||
per-round update info."""
|
||||
import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod
|
||||
|
||||
monkeypatch.setattr(
|
||||
@@ -1657,8 +1644,7 @@ def test_nccl_worker_learns_wire_params_from_init_handshake(monkeypatch):
|
||||
|
||||
def test_nccl_trainer_init_non_sender_skips_rendezvous_and_client():
|
||||
"""Non-sender trainer ranks build an engine without opening an endpoint or
|
||||
touching the client; they only join the collectives in send_weights.
|
||||
"""
|
||||
touching the client; they only join the collectives in send_weights."""
|
||||
client = RecordingClient()
|
||||
engine = WeightTransferTrainerFactory.trainer_init(
|
||||
init_info=NCCLTrainerInitInfo(
|
||||
@@ -1688,8 +1674,7 @@ def test_nccl_trainer_init_non_sender_skips_rendezvous_and_client():
|
||||
def test_nccl_trainer_send_weights_drives_client_in_order():
|
||||
"""send_weights issues start -> update -> finish and ships per-round
|
||||
metadata; the packed wire params ride the init handshake, not the
|
||||
per-round update_info.
|
||||
"""
|
||||
per-round update_info."""
|
||||
client = RecordingClient()
|
||||
engine = NCCLTrainerWeightTransferEngine(
|
||||
client=client,
|
||||
@@ -1710,8 +1695,7 @@ def test_nccl_trainer_send_weights_drives_client_in_order():
|
||||
|
||||
class _ScriptedSource(WeightSource):
|
||||
"""Declares `meta` but yields whatever `pairs` says — used to drive the
|
||||
metadata/iteration agreement checks.
|
||||
"""
|
||||
metadata/iteration agreement checks."""
|
||||
|
||||
def __init__(self, meta, pairs):
|
||||
self._meta = meta
|
||||
@@ -1747,8 +1731,7 @@ def test_nccl_trainer_init_requires_source():
|
||||
|
||||
def test_nccl_trainer_send_weights_rejects_reordered_source(monkeypatch):
|
||||
"""The worker sizes its buffers (and cuts packed chunks) from metadata(), so
|
||||
iteration disagreeing with it must raise rather than corrupt the stream.
|
||||
"""
|
||||
iteration disagreeing with it must raise rather than corrupt the stream."""
|
||||
meta = [
|
||||
ParamMeta("w", torch.float32, (4,)),
|
||||
ParamMeta("b", torch.float32, (2,)),
|
||||
@@ -1762,8 +1745,7 @@ def test_nccl_trainer_send_weights_rejects_reordered_source(monkeypatch):
|
||||
|
||||
def test_nccl_trainer_send_weights_rejects_dtype_disagreement(monkeypatch):
|
||||
"""A source that declares one wire dtype and materializes another would make
|
||||
the two sides disagree on every byte offset.
|
||||
"""
|
||||
the two sides disagree on every byte offset."""
|
||||
meta = [ParamMeta("w", torch.float32, (4,))]
|
||||
engine = _mock_group_engine(
|
||||
_ScriptedSource(meta, [("w", torch.zeros(4, dtype=torch.bfloat16))]),
|
||||
@@ -1790,8 +1772,7 @@ def test_nccl_trainer_send_weights_rejects_truncated_source(monkeypatch):
|
||||
|
||||
def test_nccl_trainer_send_weights_broadcasts_contiguous(monkeypatch):
|
||||
"""NCCL sends numel elements from data_ptr(), so a non-contiguous view must
|
||||
be linearized first or the worker receives unrelated memory.
|
||||
"""
|
||||
be linearized first or the worker receives unrelated memory."""
|
||||
base = torch.arange(6, dtype=torch.float32).reshape(2, 3)
|
||||
view = base.t() # non-contiguous
|
||||
meta = [ParamMeta("w", torch.float32, tuple(view.shape))]
|
||||
@@ -1861,8 +1842,7 @@ def _sparse_patch(device: str = "cpu") -> SparseWeightPatch:
|
||||
def test_sparse_nccl_trainer_init_ships_worker_init_info(monkeypatch):
|
||||
"""The sender's trainer_init drives the init handshake with the
|
||||
worker-shaped init info; sparse ships no packed wire params, so the worker
|
||||
keeps its unpacked defaults. Sparse takes no `source`.
|
||||
"""
|
||||
keeps its unpacked defaults. Sparse takes no `source`."""
|
||||
import vllm.distributed.weight_transfer.sparse_nccl_engine as sparse_mod
|
||||
|
||||
monkeypatch.setattr(sparse_mod, "open_trainer_endpoint", lambda info: MagicMock())
|
||||
@@ -1915,8 +1895,7 @@ def test_sparse_nccl_trainer_send_weights_drives_client_in_order(monkeypatch):
|
||||
|
||||
def test_sparse_nccl_trainer_send_weights_empty_round_is_noop():
|
||||
"""A round with no patches must not touch the client (an empty sparse
|
||||
update info is invalid by construction).
|
||||
"""
|
||||
update info is invalid by construction)."""
|
||||
client = RecordingClient()
|
||||
engine = SparseNCCLTrainerWeightTransferEngine(client=client)
|
||||
engine.model_update_group = MagicMock()
|
||||
@@ -1952,8 +1931,7 @@ def test_sparse_nccl_trainer_rejects_source():
|
||||
def test_sparse_nccl_trainer_validates_patch_before_any_rpc():
|
||||
"""Malformed patches must fail on the trainer, before start_weight_update:
|
||||
the worker's own checks only run once the broadcasts are already in flight,
|
||||
where a size mismatch wedges both sides instead of raising.
|
||||
"""
|
||||
where a size mismatch wedges both sides instead of raising."""
|
||||
client = RecordingClient()
|
||||
engine = SparseNCCLTrainerWeightTransferEngine(client=client)
|
||||
engine.model_update_group = MagicMock()
|
||||
|
||||
@@ -403,8 +403,7 @@ def trainer_broadcast_tensor_torch_free(
|
||||
def _run_uid_transfer(trainer_task) -> None:
|
||||
"""Mint a unique id, run the worker and the given trainer peer concurrently,
|
||||
and assert the tensor round-trips. Both ranks must enter init together --
|
||||
there is no store barrier on the UID path.
|
||||
"""
|
||||
there is no store barrier on the UID path."""
|
||||
from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary
|
||||
|
||||
_init_ray_for_weight_transfer()
|
||||
@@ -451,8 +450,7 @@ def _run_uid_transfer(trainer_task) -> None:
|
||||
)
|
||||
def test_nccl_weight_transfer_between_processes_uid():
|
||||
"""Weight transfer over a pre-shared ncclUniqueId (no TCPStore), with both
|
||||
ranks on vLLM's ``PyNcclCommunicator``.
|
||||
"""
|
||||
ranks on vLLM's ``PyNcclCommunicator``."""
|
||||
_run_uid_transfer(trainer_broadcast_tensor_uid)
|
||||
|
||||
|
||||
@@ -462,7 +460,6 @@ def test_nccl_weight_transfer_between_processes_uid():
|
||||
)
|
||||
def test_nccl_weight_transfer_torch_free_trainer():
|
||||
"""UID weight transfer against a foreign peer that uses no torch and no vLLM
|
||||
communicator, only a cupy buffer and raw ``ncclBroadcast``.
|
||||
"""
|
||||
communicator, only a cupy buffer and raw ``ncclBroadcast``."""
|
||||
pytest.importorskip("cupy")
|
||||
_run_uid_transfer(trainer_broadcast_tensor_torch_free)
|
||||
|
||||
@@ -570,8 +570,7 @@ def test_attention_config():
|
||||
def test_multi_node_world_size_includes_pcp(monkeypatch):
|
||||
"""PCP expands the process world size, so the --nnodes divisibility check
|
||||
must include it. Without this, TP=1/PCP=2 over 2 nodes computes a world
|
||||
size of 1 and the launch is rejected before the engine starts.
|
||||
"""
|
||||
size of 1 and the launch is rejected before the engine starts."""
|
||||
import vllm.config.vllm
|
||||
|
||||
# PCP requires the V2 model runner, which is gated on Triton.
|
||||
@@ -797,8 +796,7 @@ def test_expand_json_human_readable_numbers(input_json, expected_json):
|
||||
)
|
||||
def test_cloud_storage_uri_skips_get_model_path(uri, monkeypatch):
|
||||
"""Cloud storage URIs should not be passed to get_model_path()
|
||||
when HF_HUB_OFFLINE=1, as they are not valid HF repo IDs.
|
||||
"""
|
||||
when HF_HUB_OFFLINE=1, as they are not valid HF repo IDs."""
|
||||
import huggingface_hub
|
||||
|
||||
monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True)
|
||||
@@ -810,8 +808,7 @@ def test_cloud_storage_uri_skips_get_model_path(uri, monkeypatch):
|
||||
|
||||
def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch):
|
||||
"""Cloud storage tokenizer URI should not be passed to
|
||||
get_model_path() when HF_HUB_OFFLINE=1.
|
||||
"""
|
||||
get_model_path() when HF_HUB_OFFLINE=1."""
|
||||
import huggingface_hub
|
||||
|
||||
monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True)
|
||||
@@ -973,8 +970,7 @@ class TestDpDeviceIdSharding:
|
||||
def test_dp_supervisor_device_ids_stay_env_relative(self):
|
||||
"""Regression test: the DP supervisor must pass env-relative indices,
|
||||
not physical IDs, because each child re-resolves --device-ids
|
||||
against its inherited device-control env var.
|
||||
"""
|
||||
against its inherited device-control env var."""
|
||||
import argparse
|
||||
|
||||
from vllm.entrypoints.launchers.dp_supervisor import _build_device_ids
|
||||
@@ -1001,8 +997,7 @@ class TestDpDeviceIdSharding:
|
||||
|
||||
def test_dp_rank_shards_user_assigned_gpu_ids(self):
|
||||
"""get_physical_gpu_ids_for_local_dp_rank slices the user-provided
|
||||
--device-ids list instead of recomputing from the env var.
|
||||
"""
|
||||
--device-ids list instead of recomputing from the env var."""
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.engine.utils import get_physical_gpu_ids_for_local_dp_rank
|
||||
|
||||
|
||||
@@ -241,8 +241,7 @@ class TestToolResultContent:
|
||||
self, tool_result_content
|
||||
) -> AnthropicMessagesRequest:
|
||||
"""Build a request with assistant tool_use followed by user
|
||||
tool_result.
|
||||
"""
|
||||
tool_result."""
|
||||
return _make_request(
|
||||
[
|
||||
{
|
||||
@@ -327,8 +326,7 @@ class TestToolResultContent:
|
||||
|
||||
def test_tool_result_with_text_and_image(self):
|
||||
"""Mixed text+image tool_result: text in tool msg, image in user
|
||||
msg.
|
||||
"""
|
||||
msg."""
|
||||
request = self._make_tool_result_request(
|
||||
[
|
||||
{"type": "text", "text": "Here is the screenshot"},
|
||||
@@ -425,8 +423,7 @@ class TestToolResultContent:
|
||||
class TestAttributionHeaderStripping:
|
||||
def test_billing_header_stripped_from_system(self):
|
||||
"""Claude Code's x-anthropic-billing-header block should be
|
||||
stripped to preserve prefix caching.
|
||||
"""
|
||||
stripped to preserve prefix caching."""
|
||||
request = _make_request(
|
||||
[{"role": "user", "content": "Hello"}],
|
||||
system=[
|
||||
@@ -1150,8 +1147,7 @@ class TestMessageStreamConverterToolUseContentBuffering:
|
||||
async def test_buffered_content_flushed_on_done_without_usage_chunk(self):
|
||||
"""Content buffered during tool_use must be emitted even if the
|
||||
stream jumps straight from finish_reason to [DONE], skipping the
|
||||
empty-choices usage chunk.
|
||||
"""
|
||||
empty-choices usage chunk."""
|
||||
|
||||
async def sse_input():
|
||||
yield _make_stream_chunk(
|
||||
@@ -1255,8 +1251,7 @@ class TestStreamingCacheUsageSemantics:
|
||||
async def test_streaming_cache_fields_absent_then_populated(self):
|
||||
"""First chunk lacks prompt_tokens_details (vLLM contract);
|
||||
message_start omits cache fields. The final chunk carries
|
||||
prompt_tokens_details, so message_delta carries resolved values.
|
||||
"""
|
||||
prompt_tokens_details, so message_delta carries resolved values."""
|
||||
|
||||
async def sse_input():
|
||||
yield _make_stream_chunk(
|
||||
@@ -1301,8 +1296,7 @@ class TestStreamingCacheUsageSemantics:
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_no_cache_hit(self):
|
||||
"""When the final chunk reports cached_tokens=0, message_delta carries
|
||||
cache fields = 0 (cache miss); message_start still omits them.
|
||||
"""
|
||||
cache fields = 0 (cache miss); message_start still omits them."""
|
||||
|
||||
async def sse_input():
|
||||
yield _make_stream_chunk(
|
||||
@@ -1343,8 +1337,7 @@ class TestStreamingCacheUsageSemantics:
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_no_prompt_tokens_details_at_all(self):
|
||||
"""If --enable-prompt-tokens-details is off, no chunk carries cache
|
||||
info; both message_start and message_delta omit cache fields.
|
||||
"""
|
||||
info; both message_start and message_delta omit cache fields."""
|
||||
|
||||
async def sse_input():
|
||||
yield _make_stream_chunk(
|
||||
@@ -1465,8 +1458,7 @@ class TestMessagesFullConverter:
|
||||
class TestCacheSalt:
|
||||
def test_cache_salt_passed_through(self):
|
||||
"""cache_salt on the Anthropic request reaches the converted
|
||||
ChatCompletionRequest so prefix-cache isolation works via /v1/messages.
|
||||
"""
|
||||
ChatCompletionRequest so prefix-cache isolation works via /v1/messages."""
|
||||
request = _make_request(
|
||||
[{"role": "user", "content": "Hello"}],
|
||||
cache_salt="tenant-abc-secret-salt",
|
||||
@@ -1671,8 +1663,7 @@ class TestClientErrorResponses:
|
||||
@staticmethod
|
||||
def _conversion_error() -> ValidationError:
|
||||
"""A real pydantic ValidationError like the one ChatCompletionRequest
|
||||
construction raises when Anthropic input violates the OpenAI schema.
|
||||
"""
|
||||
construction raises when Anthropic input violates the OpenAI schema."""
|
||||
|
||||
class _StubRequest(BaseModel):
|
||||
stop: Annotated[list[str], Field(max_length=4)] | None = None
|
||||
@@ -1683,8 +1674,7 @@ class TestClientErrorResponses:
|
||||
|
||||
def test_validation_error_returns_bad_request(self):
|
||||
"""A pydantic ValidationError during Anthropic->OpenAI conversion is
|
||||
surfaced as a 400 BadRequestError, not a 500.
|
||||
"""
|
||||
surfaced as a 400 BadRequestError, not a 500."""
|
||||
handler = MagicMock(spec=AnthropicServingMessages)
|
||||
handler.create_messages.side_effect = self._conversion_error()
|
||||
|
||||
|
||||
@@ -87,8 +87,7 @@ def _make_chunk(
|
||||
citations: list[Any] | None = None,
|
||||
) -> str:
|
||||
"""Build the ``data: {...}\\n\\n`` SSE frame the production code
|
||||
consumes.
|
||||
"""
|
||||
consumes."""
|
||||
delta: dict[str, Any] = {}
|
||||
if role is not None:
|
||||
delta["role"] = role
|
||||
@@ -136,8 +135,7 @@ async def _stream_from(items: list[str]) -> AsyncGenerator[str, None]:
|
||||
|
||||
async def _drain(serving: CohereServingChatV2, items: list[str]) -> list[str]:
|
||||
"""Drive ``_chat_completion_stream_to_v2`` over ``items`` and collect
|
||||
the emitted SSE frames.
|
||||
"""
|
||||
the emitted SSE frames."""
|
||||
request = CohereChatV2Request(
|
||||
model="m", messages=[{"role": "user", "content": "hi"}], stream=True
|
||||
)
|
||||
|
||||
@@ -50,8 +50,7 @@ def update_admission_stats_worker(listen_address, sock, args, client_config):
|
||||
# importable by `multiprocessing.spawn` (no closures, no nesting).
|
||||
def defer_addresses_stub_worker(listen_address, sock, args, client_config):
|
||||
"""Bind ROUTER/PULL with a kernel-assigned port, report the actual
|
||||
endpoints back via the pipe, then exit.
|
||||
"""
|
||||
endpoints back via the pipe, then exit."""
|
||||
ctx = zmq.Context()
|
||||
try:
|
||||
in_sock = make_zmq_socket(
|
||||
@@ -358,8 +357,7 @@ def test_external_process_monitoring(api_server_args):
|
||||
def test_gather_actual_addresses_end_to_end():
|
||||
"""Each child binds ROUTER/PULL with a kernel-picked port and reports
|
||||
the bound endpoints back via its per-child pipe; the manager surfaces
|
||||
them via :py:meth:`gather_actual_addresses`.
|
||||
"""
|
||||
them via :py:meth:`gather_actual_addresses`."""
|
||||
host = "127.0.0.1"
|
||||
num_servers = 4
|
||||
|
||||
@@ -409,8 +407,7 @@ def test_gather_actual_addresses_end_to_end():
|
||||
@pytest.mark.timeout(30)
|
||||
def test_gather_actual_addresses_child_crash_before_report():
|
||||
"""A child that exits before sending its endpoints must surface a
|
||||
clear ``RuntimeError`` rather than hang or return ``None`` slots.
|
||||
"""
|
||||
clear ``RuntimeError`` rather than hang or return ``None`` slots."""
|
||||
host = "127.0.0.1"
|
||||
num_servers = 2
|
||||
placeholder_inputs = [
|
||||
@@ -447,8 +444,7 @@ def test_gather_actual_addresses_child_crash_before_report():
|
||||
|
||||
def test_rust_frontend_launch_log_redacts_credentials(monkeypatch, caplog):
|
||||
"""The Rust frontend command carries every non-default arg as JSON.
|
||||
Credentials must not reach the log line.
|
||||
"""
|
||||
Credentials must not reach the log line."""
|
||||
import subprocess as subprocess_mod
|
||||
|
||||
from vllm.entrypoints.launchers.cli_args import make_arg_parser
|
||||
|
||||
@@ -186,8 +186,7 @@ def test_enable_auto_choice_fails_with_enable_reasoning(serve_parser):
|
||||
|
||||
def test_passes_with_reasoning_parser(serve_parser):
|
||||
"""Ensure validation passes if reasoning is enabled
|
||||
with a reasoning parser
|
||||
"""
|
||||
with a reasoning parser"""
|
||||
args = serve_parser.parse_args(
|
||||
args=[
|
||||
"--reasoning-parser",
|
||||
@@ -382,8 +381,7 @@ def test_default_chat_template_kwargs_invalid_json(serve_parser):
|
||||
)
|
||||
def test_served_model_name_parsing(tmp_path, vllm_parser, args, raises):
|
||||
"""Ensure that users don't misuse --served-model-name and end up with the default
|
||||
model tag instead of the one they intended to serve.
|
||||
"""
|
||||
model tag instead of the one they intended to serve."""
|
||||
# Call the serve subparser
|
||||
args.insert(0, "serve")
|
||||
# Create a dummy config file if the test case includes it
|
||||
|
||||
+1
-2
@@ -2,8 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""E2E test for mixing `prompt_embeds` with `audio_embeds` in a single
|
||||
Chat Completions request.
|
||||
"""
|
||||
Chat Completions request."""
|
||||
|
||||
import json
|
||||
|
||||
|
||||
+1
-2
@@ -68,8 +68,7 @@ async def client(server):
|
||||
def image_url() -> str:
|
||||
"""Stable real image as a data URL, kept identical across both the
|
||||
text and prompt_embeds requests so any output difference must come from
|
||||
how the text content is delivered.
|
||||
"""
|
||||
how the text content is delivered."""
|
||||
return encode_image_url(ImageAsset("stop_sign").pil_image)
|
||||
|
||||
|
||||
|
||||
@@ -145,8 +145,7 @@ def describe_image_messages(
|
||||
) -> list[dict]:
|
||||
"""Build the system + user messages used by the completions-with-image
|
||||
family of tests. *extra_image_fields* is merged into the top-level
|
||||
image content block (for uuid / bad-key tests).
|
||||
"""
|
||||
image content block (for uuid / bad-key tests)."""
|
||||
image_block: dict = {
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_url},
|
||||
@@ -176,8 +175,7 @@ async def complete_and_check(
|
||||
temperature: float = 0.0,
|
||||
) -> str:
|
||||
"""Run a chat completion and assert the output is non-empty.
|
||||
Returns the content string.
|
||||
"""
|
||||
Returns the content string."""
|
||||
chat_completion = await client.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=messages,
|
||||
|
||||
@@ -465,8 +465,7 @@ def test_system_message_warns_on_video(video_content):
|
||||
def test_json_schema_response_format_missing_schema():
|
||||
"""When response_format type is 'json_schema' but the json_schema field
|
||||
is not provided, request construction should raise a validation error
|
||||
so the API returns 400 instead of 500.
|
||||
"""
|
||||
so the API returns 400 instead of 500."""
|
||||
with pytest.raises(Exception, match="json_schema.*must be provided"):
|
||||
ChatCompletionRequest(
|
||||
model=MODEL_NAME,
|
||||
@@ -520,8 +519,7 @@ def test_structured_outputs_structural_tag_invalid(structural_tag):
|
||||
@pytest.mark.parametrize("field_name", ["prompt_logprobs", "top_logprobs"])
|
||||
def test_non_numeric_logprobs_rejected(field_name):
|
||||
"""A non-numeric logprobs value must be a clean 400 validation error, not a
|
||||
TypeError from the mode='before' comparison (which surfaces as HTTP 500).
|
||||
"""
|
||||
TypeError from the mode='before' comparison (which surfaces as HTTP 500)."""
|
||||
with pytest.raises(VLLMValidationError, match=f"`{field_name}` must be an integer"):
|
||||
ChatCompletionRequest(
|
||||
model=MODEL_NAME,
|
||||
|
||||
@@ -81,8 +81,7 @@ async def test_chat_logit_bias_invalid(client):
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_logit_bias_non_integer_key(client):
|
||||
"""Test that a non-integer logit_bias key is rejected with a clean,
|
||||
informative error instead of a raw 'invalid literal for int()' message.
|
||||
"""
|
||||
informative error instead of a raw 'invalid literal for int()' message."""
|
||||
with pytest.raises(openai.BadRequestError) as excinfo:
|
||||
await client.chat.completions.create(
|
||||
model=MODEL_NAME,
|
||||
@@ -102,8 +101,7 @@ async def test_chat_logit_bias_non_integer_key(client):
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_logit_bias_non_numeric_value(client):
|
||||
"""Test that a non-numeric logit_bias value is rejected with a message
|
||||
that names the specific offending token, not just a generic TypeError.
|
||||
"""
|
||||
that names the specific offending token, not just a generic TypeError."""
|
||||
with pytest.raises(openai.BadRequestError) as excinfo:
|
||||
await client.chat.completions.create(
|
||||
model=MODEL_NAME,
|
||||
@@ -122,8 +120,7 @@ async def test_chat_logit_bias_non_numeric_value(client):
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_logit_bias_multiple_non_integer_keys(client):
|
||||
"""Test that ALL invalid logit_bias keys are reported together,
|
||||
not just the first one encountered.
|
||||
"""
|
||||
not just the first one encountered."""
|
||||
with pytest.raises(openai.BadRequestError) as excinfo:
|
||||
await client.chat.completions.create(
|
||||
model=MODEL_NAME,
|
||||
|
||||
@@ -54,8 +54,7 @@ def stop_sign_image_url():
|
||||
async def test_translategemma_extra_lang_code_fields(client: openai.AsyncOpenAI):
|
||||
"""En -> es translation through TranslateGemma's bundled chat template's
|
||||
``text`` content branch, which depends on ``source_lang_code`` /
|
||||
``target_lang_code`` extra fields being preserved on the content part.
|
||||
"""
|
||||
``target_lang_code`` extra fields being preserved on the content part."""
|
||||
completion = await client.chat.completions.create(
|
||||
model=MODEL_NAME,
|
||||
messages=[
|
||||
@@ -99,8 +98,7 @@ async def test_translategemma_image_extra_lang_code_fields(
|
||||
``source_lang_code`` / ``target_lang_code`` extras must survive parsing
|
||||
and end up alongside the ``{"type": "image"}`` placeholder so the
|
||||
template's image branch (``content["type"] == 'image'``) renders the
|
||||
right OCR-translation prompt.
|
||||
"""
|
||||
right OCR-translation prompt."""
|
||||
completion = await client.chat.completions.create(
|
||||
model=MODEL_NAME,
|
||||
messages=[
|
||||
|
||||
@@ -191,8 +191,7 @@ async def test_logprob_token_ids_stream_with_no_top_k(server):
|
||||
async def test_logprob_token_ids_default_behavior_unchanged(server):
|
||||
"""Without `logprob_token_ids`, the response carries the natural top-k
|
||||
most-likely tokens. This guards against the new field accidentally
|
||||
changing the default-path output.
|
||||
"""
|
||||
changing the default-path output."""
|
||||
async with server.get_async_client() as client:
|
||||
resp = await client.chat.completions.create(
|
||||
model=MODEL_NAME,
|
||||
|
||||
@@ -1406,8 +1406,7 @@ async def test_serving_chat_did_set_correct_cache_salt(model_type):
|
||||
@pytest.mark.asyncio
|
||||
async def test_serving_chat_data_parallel_rank_extraction():
|
||||
"""Test that data_parallel_rank is properly extracted from header and
|
||||
passed to engine.
|
||||
"""
|
||||
passed to engine."""
|
||||
mock_engine = MagicMock(spec=AsyncLLM)
|
||||
mock_engine.errored = False
|
||||
mock_engine.model_config = MockModelConfig()
|
||||
@@ -1729,8 +1728,7 @@ class TestServingChatWithHarmony:
|
||||
@pytest.mark.asyncio
|
||||
async def test_system_message_without_tools(self, serving_chat, stream):
|
||||
"""Leading system message produces a developer message with
|
||||
DeveloperContent (# Instructions header).
|
||||
"""
|
||||
DeveloperContent (# Instructions header)."""
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Hello"},
|
||||
@@ -1752,8 +1750,7 @@ class TestServingChatWithHarmony:
|
||||
@pytest.mark.asyncio
|
||||
async def test_system_message_with_tools(self, serving_chat, stream, weather_tools):
|
||||
"""Leading system message is folded into the developer message
|
||||
alongside tool definitions.
|
||||
"""
|
||||
alongside tool definitions."""
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "What's the weather?"},
|
||||
@@ -2168,8 +2165,7 @@ class TestServingChatWithHarmony:
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_choice_validation_without_parser():
|
||||
"""Test that tool_choice='required' or named tool without tool_parser
|
||||
returns an appropriate error message.
|
||||
"""
|
||||
returns an appropriate error message."""
|
||||
mock_engine = MagicMock(spec=AsyncLLM)
|
||||
mock_engine.errored = False
|
||||
mock_engine.model_config = MockModelConfig()
|
||||
|
||||
@@ -151,8 +151,7 @@ async def client(request, server, server_with_auto_reasoning_config):
|
||||
@pytest.mark.parametrize("client", ["default", "auto_config"], indirect=True)
|
||||
async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI):
|
||||
"""Test that mixed requests (some with thinking_token_budget, some without)
|
||||
complete successfully without errors.
|
||||
"""
|
||||
complete successfully without errors."""
|
||||
response_with_budget = await client.chat.completions.create(
|
||||
model=MODEL_NAME,
|
||||
messages=MESSAGES,
|
||||
|
||||
@@ -521,8 +521,7 @@ async def test_completion_error_stream():
|
||||
def test_json_schema_response_format_missing_schema():
|
||||
"""When response_format type is 'json_schema' but the json_schema field
|
||||
is not provided, request construction should raise a validation error
|
||||
so the API returns 400 instead of 500.
|
||||
"""
|
||||
so the API returns 400 instead of 500."""
|
||||
with pytest.raises(Exception, match="json_schema.*must be provided"):
|
||||
CompletionRequest(
|
||||
model=MODEL_NAME,
|
||||
@@ -585,8 +584,7 @@ def test_negative_prompt_token_ids_flat():
|
||||
def test_logprobs_minus_one_allowed():
|
||||
"""logprobs=-1 means "return all logprobs". The sampling layer and the chat
|
||||
top_logprobs / prompt_logprobs validators all accept -1, so the completion
|
||||
logprobs validator must accept it too instead of rejecting it as negative.
|
||||
"""
|
||||
logprobs validator must accept it too instead of rejecting it as negative."""
|
||||
request = CompletionRequest(
|
||||
model=MODEL_NAME,
|
||||
prompt="Test prompt",
|
||||
@@ -746,8 +744,7 @@ class TestCompletionPromptListLimit:
|
||||
@pytest.mark.parametrize("field_name", ["prompt_logprobs", "logprobs"])
|
||||
def test_non_numeric_logprobs_rejected(field_name):
|
||||
"""A non-numeric logprobs value must be a clean 400 validation error, not a
|
||||
TypeError from the mode='before' comparison (which surfaces as HTTP 500).
|
||||
"""
|
||||
TypeError from the mode='before' comparison (which surfaces as HTTP 500)."""
|
||||
with pytest.raises(VLLMValidationError, match=f"`{field_name}` must be an integer"):
|
||||
CompletionRequest(
|
||||
model=MODEL_NAME,
|
||||
|
||||
@@ -43,8 +43,7 @@ def _system():
|
||||
|
||||
class TestResponseInputToHarmonyRenderParity:
|
||||
"""Each test drives the same conversation through both APIs and asserts
|
||||
identical Harmony messages and rendered token sequences.
|
||||
"""
|
||||
identical Harmony messages and rendered token sequences."""
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# Single-message cases
|
||||
@@ -52,8 +51,7 @@ class TestResponseInputToHarmonyRenderParity:
|
||||
|
||||
def test_developer_message(self):
|
||||
"""Both APIs must render developer messages identically using
|
||||
DeveloperContent (with the '# Instructions' header).
|
||||
"""
|
||||
DeveloperContent (with the '# Instructions' header)."""
|
||||
chat_msgs = parse_chat_input_to_harmony_message(
|
||||
{"role": "developer", "content": "Be concise."}
|
||||
)
|
||||
@@ -248,8 +246,7 @@ class TestResponseInputToHarmonyRenderParity:
|
||||
def test_reasoning_combined_with_function_call(self):
|
||||
"""Chat API packs reasoning + tool_calls into one dict; responses API
|
||||
represents them as two separate items. Both must produce the same two
|
||||
Harmony messages in the same order: analysis then commentary.
|
||||
"""
|
||||
Harmony messages in the same order: analysis then commentary."""
|
||||
chat_msgs = parse_chat_input_to_harmony_message(
|
||||
{
|
||||
"role": "assistant",
|
||||
@@ -532,8 +529,7 @@ class TestResponseInputToHarmonyRenderParity:
|
||||
def test_completed_turns_drop_reasoning(self):
|
||||
"""Validates that reasoning from completed turns is dropped, while
|
||||
reasoning from the current in-progress tool-call turn is preserved
|
||||
in both chat completions and responses previous_input_messages.
|
||||
"""
|
||||
in both chat completions and responses previous_input_messages."""
|
||||
first_turn_reasoning = "FIRST_TURN_REASONING"
|
||||
second_turn_reasoning = "SECOND_TURN_REASONING"
|
||||
|
||||
|
||||
@@ -323,8 +323,7 @@ class TestCommonParseInputToHarmonyMessage:
|
||||
|
||||
def test_system_message(self, parse_function):
|
||||
"""Test parsing system messages, which are parsed into developer messages
|
||||
with DeveloperContent.
|
||||
"""
|
||||
with DeveloperContent."""
|
||||
chat_msg = {
|
||||
"role": "system",
|
||||
"content": "You are a helpful assistant",
|
||||
|
||||
@@ -12,8 +12,7 @@ from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
|
||||
|
||||
def test_function_call_dict_converted_to_object():
|
||||
"""Test that function_call dictionaries are correctly parsed into
|
||||
ResponseFunctionToolCall objects.
|
||||
"""
|
||||
ResponseFunctionToolCall objects."""
|
||||
# Create a request with function_call as dict
|
||||
request_data = {
|
||||
"model": "gpt-oss",
|
||||
@@ -330,8 +329,7 @@ def test_validator_handles_empty_iterator():
|
||||
|
||||
def test_assistant_string_content_stays_easyinput():
|
||||
"""EasyInput assistant message with plain string content is not
|
||||
coerced into a ResponseOutputMessage.
|
||||
"""
|
||||
coerced into a ResponseOutputMessage."""
|
||||
request_data = {
|
||||
"model": "test-model",
|
||||
"input": [
|
||||
@@ -353,8 +351,7 @@ def test_assistant_string_content_stays_easyinput():
|
||||
|
||||
def test_assistant_output_style_content_coerced():
|
||||
"""Assistant message whose content is output-message-shaped (list of
|
||||
output_text items) should be coerced to ResponseOutputMessage.
|
||||
"""
|
||||
output_text items) should be coerced to ResponseOutputMessage."""
|
||||
request_data = {
|
||||
"model": "test-model",
|
||||
"input": [
|
||||
|
||||
@@ -67,8 +67,7 @@ class TestResponseInputToHarmonyMessage:
|
||||
|
||||
def test_system_message(self):
|
||||
"""System messages carry developer instructions and must be rendered
|
||||
as developer messages with DeveloperContent.
|
||||
"""
|
||||
as developer messages with DeveloperContent."""
|
||||
msg = response_input_to_harmony(
|
||||
{"type": "message", "role": "system", "content": "Be helpful."},
|
||||
prev_responses=[],
|
||||
@@ -90,8 +89,7 @@ class TestResponseInputToHarmonyMessage:
|
||||
|
||||
def test_developer_message_gets_instructions_prefix(self):
|
||||
"""Developer messages must use DeveloperContent which adds the
|
||||
'# Instructions' header the model was trained on.
|
||||
"""
|
||||
'# Instructions' header the model was trained on."""
|
||||
msg = response_input_to_harmony(
|
||||
{"type": "message", "role": "developer", "content": "Be concise."},
|
||||
prev_responses=[],
|
||||
@@ -121,8 +119,7 @@ class TestResponseInputToHarmonyMessage:
|
||||
|
||||
def test_developer_message_array_content_concatenated(self):
|
||||
"""Array content in developer messages is flattened and rendered
|
||||
via DeveloperContent with the '# Instructions' header.
|
||||
"""
|
||||
via DeveloperContent with the '# Instructions' header."""
|
||||
msg = response_input_to_harmony(
|
||||
{
|
||||
"type": "message",
|
||||
@@ -213,8 +210,7 @@ class TestResponseInputToHarmonyMessage:
|
||||
|
||||
def test_function_call_output_uses_most_recent_matching_call(self):
|
||||
"""When multiple prev_responses share a call_id, the last one wins
|
||||
because the search is reversed.
|
||||
"""
|
||||
because the search is reversed."""
|
||||
earlier = ResponseFunctionToolCall(
|
||||
id="fc_old",
|
||||
call_id="call_test",
|
||||
|
||||
@@ -839,13 +839,11 @@ class TestConstructChatMessagesCombinePolicy:
|
||||
|
||||
class TestConstructInputMessagesInstructionsLeak:
|
||||
"""Regression tests for #37697: instructions from a prior response
|
||||
should NOT leak through previous_response_id.
|
||||
"""
|
||||
should NOT leak through previous_response_id."""
|
||||
|
||||
def test_old_instructions_stripped_from_prev_msg(self):
|
||||
"""System message in prev_msg must be dropped so the new request's
|
||||
instructions are the only system message in the conversation.
|
||||
"""
|
||||
instructions are the only system message in the conversation."""
|
||||
prev = [
|
||||
{"role": "system", "content": "old instructions"},
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
@@ -862,8 +860,7 @@ class TestConstructInputMessagesInstructionsLeak:
|
||||
|
||||
def test_no_instructions_in_new_request(self):
|
||||
"""If the new request has no instructions, old ones should still
|
||||
be stripped -- they must not carry over.
|
||||
"""
|
||||
be stripped -- they must not carry over."""
|
||||
prev = [
|
||||
{"role": "system", "content": "old instructions"},
|
||||
{"role": "user", "content": "Hi"},
|
||||
|
||||
@@ -1163,8 +1163,7 @@ class TestHarmonyPreambleStreaming:
|
||||
|
||||
def test_preamble_done_emits_text_done_events(self) -> None:
|
||||
"""Completed preamble should emit text done + content_part done +
|
||||
output_item done, same shape as final channel.
|
||||
"""
|
||||
output_item done, same shape as final channel."""
|
||||
from vllm.entrypoints.openai.responses.streaming_events import (
|
||||
emit_previous_item_done_events,
|
||||
)
|
||||
@@ -1185,8 +1184,7 @@ class TestHarmonyPreambleStreaming:
|
||||
|
||||
def test_commentary_with_recipient_no_preamble_done(self) -> None:
|
||||
"""Commentary + recipient='functions.X' should route to function call
|
||||
done, not preamble done.
|
||||
"""
|
||||
done, not preamble done."""
|
||||
from vllm.entrypoints.openai.responses.streaming_events import (
|
||||
emit_previous_item_done_events,
|
||||
)
|
||||
@@ -1332,8 +1330,7 @@ def _mock_parser_with_reasoning(serving, delta_sequence: list[DeltaMessage]):
|
||||
class TestStreamingReasoningToContentTransition:
|
||||
"""Tests for _process_simple_streaming_events reasoning-to-content
|
||||
transition, specifically the fix for mixed deltas that carry both
|
||||
reasoning and content simultaneously.
|
||||
"""
|
||||
reasoning and content simultaneously."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mixed_delta_reasoning_and_content_emits_reasoning_delta(
|
||||
@@ -1343,8 +1340,7 @@ class TestStreamingReasoningToContentTransition:
|
||||
and content set (e.g. reasoning end and content start in the same
|
||||
chunk), the trailing reasoning text must be emitted as a
|
||||
ResponseReasoningTextDeltaEvent and included in the
|
||||
ResponseReasoningTextDoneEvent text.
|
||||
"""
|
||||
ResponseReasoningTextDoneEvent text."""
|
||||
monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False)
|
||||
serving = _make_serving_instance_with_reasoning()
|
||||
|
||||
@@ -1413,8 +1409,7 @@ class TestStreamingReasoningToContentTransition:
|
||||
self, monkeypatch
|
||||
):
|
||||
"""When the transition from reasoning to content is clean (no mixed
|
||||
delta), no extra reasoning delta event should be emitted.
|
||||
"""
|
||||
delta), no extra reasoning delta event should be emitted."""
|
||||
monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False)
|
||||
serving = _make_serving_instance_with_reasoning()
|
||||
|
||||
@@ -1475,8 +1470,7 @@ class TestStreamingReasoningToContentTransition:
|
||||
async def test_reasoning_only_stream_no_content(self, monkeypatch):
|
||||
"""When the stream has only reasoning deltas and no content, the
|
||||
reasoning done event should be emitted at finalization with the
|
||||
full accumulated text, and no text delta events should appear.
|
||||
"""
|
||||
full accumulated text, and no text delta events should appear."""
|
||||
monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False)
|
||||
serving = _make_serving_instance_with_reasoning()
|
||||
|
||||
|
||||
@@ -143,8 +143,7 @@ async def test_streaming_output_consistency(client: OpenAI, model_name: str):
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_streaming_logprobs(client: OpenAI, model_name: str):
|
||||
"""Test that streaming with logprobs returns valid logprob data on
|
||||
output_text.delta events and that top_logprobs has the requested count.
|
||||
"""
|
||||
output_text.delta events and that top_logprobs has the requested count."""
|
||||
response = await client.responses.create(
|
||||
model=model_name,
|
||||
input="Say hello.",
|
||||
|
||||
@@ -124,8 +124,7 @@ class TestProcessorCompoundDeltas:
|
||||
|
||||
def test_split_name_and_args_same_index(self):
|
||||
"""Regression: parsers like KimiK2 emit name and args as separate
|
||||
DeltaToolCalls at the same index within one DeltaMessage.
|
||||
"""
|
||||
DeltaToolCalls at the same index within one DeltaMessage."""
|
||||
tc_name = _make_tool_call(0, name="get_weather")
|
||||
tc_args = _make_tool_call(0, arguments='{"city":"SF"}')
|
||||
delta = DeltaMessage(tool_calls=[tc_name, tc_args])
|
||||
@@ -141,8 +140,7 @@ class TestProcessorCompoundDeltas:
|
||||
|
||||
def test_reasoning_to_content_transition(self):
|
||||
"""Regression: the old special case in emit_delta handled this;
|
||||
now split_delta handles it generically.
|
||||
"""
|
||||
now split_delta handles it generically."""
|
||||
processor = SimpleStreamingEventProcessor()
|
||||
_run_through_processor(processor, DeltaMessage(reasoning="think"))
|
||||
assert processor.state.current_state == _StateType.REASONING
|
||||
|
||||
@@ -25,8 +25,7 @@ def _model_config() -> Mock:
|
||||
def test_completion_flag_forwarded_to_tok_params():
|
||||
"""build_tok_params must forward return_token_offsets, defaulting to
|
||||
False (zero behavioral change for existing callers) and coercing JSON
|
||||
null to False via the bool() guard.
|
||||
"""
|
||||
null to False via the bool() guard."""
|
||||
cfg = _model_config()
|
||||
|
||||
default = CompletionRequest(model="m", prompt="hi")
|
||||
@@ -41,8 +40,7 @@ def test_completion_flag_forwarded_to_tok_params():
|
||||
|
||||
def test_chat_flag_forwarded_to_tok_params():
|
||||
"""Chat build_tok_params has its own (max_completion_tokens) branch, so
|
||||
its return_token_offsets forwarding is verified independently.
|
||||
"""
|
||||
its return_token_offsets forwarding is verified independently."""
|
||||
cfg = _model_config()
|
||||
messages = [{"role": "user", "content": "hi"}]
|
||||
|
||||
@@ -66,8 +64,7 @@ def test_generate_request_token_offsets_default_none():
|
||||
|
||||
def test_generate_request_token_offsets_survive_json_round_trip():
|
||||
"""GenerateRequest crosses the disagg serialization boundary; the
|
||||
tuple[int, int] offsets must survive model_dump and re-validate.
|
||||
"""
|
||||
tuple[int, int] offsets must survive model_dump and re-validate."""
|
||||
req = GenerateRequest(
|
||||
token_ids=[10, 20],
|
||||
sampling_params=SamplingParams(),
|
||||
|
||||
@@ -111,8 +111,7 @@ def test_basic_embed(server: RemoteOpenAIServer, model_name: str):
|
||||
|
||||
def test_unsupported_input_type_rejected(server: RemoteOpenAIServer, model_name: str):
|
||||
"""An input_type not defined in the model's prompt config should be
|
||||
rejected with a 400 error.
|
||||
"""
|
||||
rejected with a 400 error."""
|
||||
body = {
|
||||
"model": model_name,
|
||||
"input_type": "nonexistent_type",
|
||||
@@ -139,8 +138,7 @@ def test_omitted_input_type_accepted(server: RemoteOpenAIServer, model_name: str
|
||||
|
||||
def test_v1_v2_parity(server: RemoteOpenAIServer, model_name: str):
|
||||
"""v1 (OpenAI) and v2 (Cohere) endpoints should produce the same
|
||||
float embeddings for a generic model.
|
||||
"""
|
||||
float embeddings for a generic model."""
|
||||
texts = ["hello world"]
|
||||
v2 = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"])
|
||||
v1 = _openai_embed(server, model_name, texts)
|
||||
|
||||
@@ -76,8 +76,7 @@ def test_single_text_parity(server: RemoteOpenAIServer):
|
||||
|
||||
def test_batch_parity(server: RemoteOpenAIServer):
|
||||
"""A batch of texts should produce equivalent embeddings via both APIs,
|
||||
in the same order.
|
||||
"""
|
||||
in the same order."""
|
||||
texts = [
|
||||
"machine learning",
|
||||
"deep learning",
|
||||
|
||||
@@ -241,8 +241,7 @@ class TestResolveTruncation:
|
||||
|
||||
def test_truncate_none_with_max_tokens(self):
|
||||
"""truncate=NONE should NOT set truncate_prompt_tokens; the
|
||||
max_tokens limit is enforced separately via _check_max_tokens.
|
||||
"""
|
||||
max_tokens limit is enforced separately via _check_max_tokens."""
|
||||
req = self._make_request(truncate="NONE", max_tokens=10)
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens is None
|
||||
|
||||
@@ -240,8 +240,7 @@ async def test_long_text_embedding_1500_chars(
|
||||
client_with_chunked_processing: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
"""Test embedding processing for ~1500 character long text
|
||||
(~1028 tokens, exceeding 512 token limit).
|
||||
"""
|
||||
(~1028 tokens, exceeding 512 token limit)."""
|
||||
# Verify text length
|
||||
# Verify text has sufficient word count (approximately 1500 words)
|
||||
word_count = len(LONG_TEXT_1500_WORDS.split())
|
||||
@@ -289,8 +288,7 @@ async def test_long_text_embedding_2500_chars(
|
||||
client_with_chunked_processing: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
"""Test embedding processing for ~2500 character long text
|
||||
(~2048 tokens, requiring multiple chunks).
|
||||
"""
|
||||
(~2048 tokens, requiring multiple chunks)."""
|
||||
# Verify text length
|
||||
# Verify text has sufficient word count (approximately 2500 words)
|
||||
word_count = len(LONG_TEXT_2500_WORDS.split())
|
||||
@@ -382,8 +380,7 @@ async def test_chunked_vs_normal_consistency(
|
||||
client_with_chunked_processing: openai.AsyncOpenAI, model_name: str
|
||||
):
|
||||
"""Test consistency between chunked and
|
||||
normal processing (using short text).
|
||||
"""
|
||||
normal processing (using short text)."""
|
||||
# Use a short text within the 512 token limit
|
||||
short_text = (
|
||||
"Artificial intelligence technology is changing our world, "
|
||||
|
||||
@@ -1,8 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for Cohere embed protocol: build_typed_embeddings and its
|
||||
underlying packing helpers, plus Cohere-specific serving helpers.
|
||||
"""
|
||||
underlying packing helpers, plus Cohere-specific serving helpers."""
|
||||
|
||||
import struct
|
||||
|
||||
|
||||
@@ -454,8 +454,7 @@ async def test_derender_completion_empty_generate_responses(client):
|
||||
@pytest.mark.asyncio
|
||||
async def test_derender_completion_logprobs(client):
|
||||
"""token_id:N placeholders in logprobs are resolved; CompletionLogProbs
|
||||
flat-list structure is returned with non-empty tokens and text_offsets.
|
||||
"""
|
||||
flat-list structure is returned with non-empty tokens and text_offsets."""
|
||||
gr1 = await _render_completion(client, "Hello world")
|
||||
ids1 = gr1["token_ids"][:3]
|
||||
token_id = ids1[0]
|
||||
|
||||
@@ -121,8 +121,7 @@ async def _disagg(
|
||||
|
||||
def _tool_sig(response_choice: dict) -> list[tuple[str, dict]]:
|
||||
"""[(name, json normalized args)] so key ordering / whitespace don't
|
||||
cause false negatives.
|
||||
"""
|
||||
cause false negatives."""
|
||||
return [
|
||||
(tc["function"]["name"], json.loads(tc["function"]["arguments"]))
|
||||
for tc in (response_choice["message"].get("tool_calls") or [])
|
||||
@@ -144,8 +143,7 @@ async def _run_parity_case(
|
||||
client: httpx.AsyncClient, messages: list[dict], **extra
|
||||
) -> tuple[dict, dict]:
|
||||
"""Run the coupled request then feed its generated tokens into the
|
||||
disaggregated derender endpoint. Returns (coupled, disagg).
|
||||
"""
|
||||
disaggregated derender endpoint. Returns (coupled, disagg)."""
|
||||
coupled = await _coupled(client, messages, **extra)
|
||||
ch = coupled["choices"][0]
|
||||
chat_request = {"model": MODEL, "messages": messages, **extra}
|
||||
@@ -205,8 +203,7 @@ async def test_parity_tool_call(client):
|
||||
@pytest.mark.asyncio
|
||||
async def test_parity_reasoning_and_tool_call(client):
|
||||
"""Combined reasoning + tool call parity means the highest drift risk
|
||||
since it exercises both parser branches on the same output.
|
||||
"""
|
||||
since it exercises both parser branches on the same output."""
|
||||
messages = [{"role": "user", "content": "What's the weather in Paris?"}]
|
||||
coupled, disagg = await _run_parity_case(
|
||||
client,
|
||||
|
||||
@@ -455,8 +455,7 @@ async def test_chat_completion_render_with_sampling_params(client):
|
||||
@pytest.mark.asyncio
|
||||
async def test_completion_render_emits_token_offsets(client):
|
||||
"""With return_token_offsets, /v1/completions/render returns per-token
|
||||
(start, end) char offsets aligned with token_ids.
|
||||
"""
|
||||
(start, end) char offsets aligned with token_ids."""
|
||||
prompt = "Hello, world."
|
||||
response = await client.post(
|
||||
"/v1/completions/render",
|
||||
@@ -481,8 +480,7 @@ async def test_completion_render_emits_token_offsets(client):
|
||||
@pytest.mark.asyncio
|
||||
async def test_completion_render_default_no_token_offsets(client):
|
||||
"""Without the flag, token_offsets must be null (existing responses
|
||||
unchanged).
|
||||
"""
|
||||
unchanged)."""
|
||||
response = await client.post(
|
||||
"/v1/completions/render",
|
||||
json={
|
||||
@@ -499,8 +497,7 @@ async def test_completion_render_default_no_token_offsets(client):
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_render_emits_token_offsets(client):
|
||||
"""With return_token_offsets, /v1/chat/completions/render returns
|
||||
per-token offsets relative to the templated prompt string.
|
||||
"""
|
||||
per-token offsets relative to the templated prompt string."""
|
||||
response = await client.post(
|
||||
"/v1/chat/completions/render",
|
||||
json={
|
||||
@@ -674,8 +671,7 @@ async def test_chat_completion_render_assistant_tokens_mask_with_generation_tags
|
||||
client,
|
||||
):
|
||||
"""With a ``{% generation %}``-enabled template, the mask marks assistant
|
||||
tokens and the masked tokens decode to the assistant content.
|
||||
"""
|
||||
tokens and the masked tokens decode to the assistant content."""
|
||||
response = await client.post(
|
||||
"/v1/chat/completions/render",
|
||||
json={
|
||||
|
||||
@@ -25,8 +25,7 @@ def _base_payload() -> dict:
|
||||
|
||||
def test_omitted_max_tokens_is_not_provided():
|
||||
"""Body without ``max_tokens`` must surface as 'not provided' so the
|
||||
server can apply its own default instead of the dataclass 16.
|
||||
"""
|
||||
server can apply its own default instead of the dataclass 16."""
|
||||
req = GenerateRequest.model_validate(_base_payload())
|
||||
# SamplingParams' dataclass default leaks through the parsed instance —
|
||||
# this is exactly the bug the server-side defaulting works around.
|
||||
@@ -36,8 +35,7 @@ def test_omitted_max_tokens_is_not_provided():
|
||||
|
||||
def test_explicit_max_tokens_is_provided():
|
||||
"""Even when the client picks the same value as the dataclass default,
|
||||
it must register as explicitly set so the server won't override it.
|
||||
"""
|
||||
it must register as explicitly set so the server won't override it."""
|
||||
payload = _base_payload()
|
||||
payload["sampling_params"] = {"max_tokens": 16}
|
||||
req = GenerateRequest.model_validate(payload)
|
||||
@@ -69,8 +67,7 @@ def test_json_roundtrip_preserves_provided_keys():
|
||||
def test_internal_instance_construction_treats_all_as_provided():
|
||||
"""When internal callers build ``GenerateRequest`` from a pre-resolved
|
||||
``SamplingParams`` instance, every field is considered explicitly set
|
||||
so server-side defaulting can't clobber values resolved upstream.
|
||||
"""
|
||||
so server-side defaulting can't clobber values resolved upstream."""
|
||||
sp = SamplingParams(max_tokens=500, temperature=0.0)
|
||||
req = GenerateRequest(token_ids=[1, 2, 3], sampling_params=sp)
|
||||
assert req.is_sampling_param_provided("max_tokens")
|
||||
|
||||
@@ -22,8 +22,7 @@ def test_sanitize_message():
|
||||
|
||||
class TestSanitizeMessageFilePaths:
|
||||
"""sanitize_message should also strip file paths and traceback
|
||||
frames, not just memory addresses - see #31683.
|
||||
"""
|
||||
frames, not just memory addresses - see #31683."""
|
||||
|
||||
def test_strips_traceback_style_frame(self):
|
||||
msg = (
|
||||
@@ -45,8 +44,7 @@ class TestSanitizeMessageFilePaths:
|
||||
|
||||
def test_strips_single_parent_container_path(self):
|
||||
"""Regression: /app/server.py and /workspace/server.py (common in
|
||||
container deployments) were missed by the original {2,} quantifier.
|
||||
"""
|
||||
container deployments) were missed by the original {2,} quantifier."""
|
||||
assert "/app/" not in sanitize_message("Error in /app/server.py")
|
||||
assert "/workspace/" not in sanitize_message("Error in /workspace/server.py")
|
||||
|
||||
|
||||
@@ -24,8 +24,7 @@ from vllm.entrypoints.serve.exception_handling.handlers.validation import (
|
||||
|
||||
def _fake_request(log_error_stack: bool = False) -> SimpleNamespace:
|
||||
"""Minimal stand-in for a FastAPI Request - just enough for the
|
||||
handler to read req.app.state.args.log_error_stack.
|
||||
"""
|
||||
handler to read req.app.state.args.log_error_stack."""
|
||||
return SimpleNamespace(
|
||||
app=SimpleNamespace(
|
||||
state=SimpleNamespace(args=SimpleNamespace(log_error_stack=log_error_stack))
|
||||
@@ -36,8 +35,7 @@ def _fake_request(log_error_stack: bool = False) -> SimpleNamespace:
|
||||
|
||||
class TestValidationErrorParamFallback:
|
||||
"""Ensure `param` falls back to the Pydantic error's `loc` when no
|
||||
custom VLLMValidationError context is present.
|
||||
"""
|
||||
custom VLLMValidationError context is present."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("error_type", "msg"),
|
||||
@@ -61,8 +59,7 @@ class TestValidationErrorParamFallback:
|
||||
async def test_param_fallback_does_not_crash_on_non_dict_error(self):
|
||||
"""Schemathesis fuzzing found that errors[0] isn't always a dict.
|
||||
The fallback must not crash in that case - it should just leave
|
||||
param as None instead of raising.
|
||||
"""
|
||||
param as None instead of raising."""
|
||||
exc = RequestValidationError(["some unexpected non-dict error"])
|
||||
|
||||
response = await validation_exception_handler(_fake_request(), exc)
|
||||
@@ -100,8 +97,7 @@ class TestCleanLocForParam:
|
||||
|
||||
def test_all_internal_falls_back_to_raw_join(self):
|
||||
"""If every loc segment looks internal, fall back to the raw
|
||||
dot-join rather than returning an empty string.
|
||||
"""
|
||||
dot-join rather than returning an empty string."""
|
||||
loc = ("function-wrap[__log_extra_fields__()]",)
|
||||
assert clean_loc_for_param(loc) == "function-wrap[__log_extra_fields__()]"
|
||||
|
||||
@@ -109,8 +105,7 @@ class TestCleanLocForParam:
|
||||
class TestValidationErrorDoesNotLeakServerPaths:
|
||||
"""FastAPI sets endpoint_file/line/function/path on the exception
|
||||
during routing, not in the constructor - so we set them manually
|
||||
here to reproduce the real leak. See #31683.
|
||||
"""
|
||||
here to reproduce the real leak. See #31683."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handler_strips_endpoint_file_context(self):
|
||||
|
||||
@@ -245,8 +245,7 @@ async def test_multiple_lora_adapters(
|
||||
client: openai.AsyncOpenAI, tmp_path, qwen3_lora_files
|
||||
):
|
||||
"""Validate that many loras can be dynamically registered and inferenced
|
||||
with concurrently
|
||||
"""
|
||||
with concurrently"""
|
||||
|
||||
# This test file configures the server with --max-cpu-loras=2 and this test
|
||||
# will concurrently load 10 adapters, so it should flex the LRU cache
|
||||
|
||||
@@ -72,8 +72,7 @@ class MockModelConfig:
|
||||
@pytest.fixture(params=get_args(SupportedTask))
|
||||
def task_routes(request, monkeypatch) -> tuple[str, list[tuple[str, list[str]]]]:
|
||||
"""For each supported task, build an app with only that task's routers,
|
||||
extract all routes, and return the task name and routes.
|
||||
"""
|
||||
extract all routes, and return the task name and routes."""
|
||||
task = request.param
|
||||
# Enable development mode to register all routes (including dev-only routes).
|
||||
monkeypatch.setenv("VLLM_SERVER_DEV_MODE", "1")
|
||||
@@ -99,8 +98,7 @@ def task_routes(request, monkeypatch) -> tuple[str, list[tuple[str, list[str]]]]
|
||||
|
||||
def test_auto_discovered_protected_routes_require_auth(task_routes):
|
||||
"""For every auto-discovered route that starts with a guarded prefix,
|
||||
verify that authentication is enforced.
|
||||
"""
|
||||
verify that authentication is enforced."""
|
||||
task, routes = task_routes
|
||||
app = _create_app_with_mock_routes(routes)
|
||||
client = TestClient(app)
|
||||
@@ -134,8 +132,7 @@ def test_auto_discovered_protected_routes_require_auth(task_routes):
|
||||
|
||||
def test_auto_discovered_unprotected_routes_no_auth(task_routes):
|
||||
"""For every auto-discovered route that does NOT start with a guarded
|
||||
prefix, verify that no authentication is required.
|
||||
"""
|
||||
prefix, verify that no authentication is required."""
|
||||
task, routes = task_routes
|
||||
app = _create_app_with_mock_routes(routes)
|
||||
client = TestClient(app)
|
||||
|
||||
@@ -114,8 +114,7 @@ class TestHandlerOverrideIntegration:
|
||||
|
||||
def test_customer_script_functions_auto_loaded(self, monkeypatch, tmp_path):
|
||||
"""Test customer scenario: script functions automatically override
|
||||
framework defaults.
|
||||
"""
|
||||
framework defaults."""
|
||||
try:
|
||||
from model_hosting_container_standards.sagemaker.config import (
|
||||
SageMakerEnvVars,
|
||||
@@ -176,8 +175,7 @@ async def custom_sagemaker_invocation_handler(request: Request):
|
||||
|
||||
def test_customer_decorator_usage(self, monkeypatch, tmp_path):
|
||||
"""Test customer scenario: using @custom_ping_handler and
|
||||
@custom_invocation_handler decorators.
|
||||
"""
|
||||
@custom_invocation_handler decorators."""
|
||||
try:
|
||||
from model_hosting_container_standards.sagemaker.config import (
|
||||
SageMakerEnvVars,
|
||||
@@ -233,8 +231,7 @@ async def my_invoke(request: Request):
|
||||
|
||||
def test_handler_priority_order(self, monkeypatch, tmp_path):
|
||||
"""Test priority: @custom_ping_handler/@custom_invocation_handler
|
||||
decorators vs script functions.
|
||||
"""
|
||||
decorators vs script functions."""
|
||||
try:
|
||||
from model_hosting_container_standards.sagemaker.config import (
|
||||
SageMakerEnvVars,
|
||||
@@ -307,8 +304,7 @@ async def custom_sagemaker_invocation_handler(request: Request):
|
||||
|
||||
def test_environment_variable_script_loading(self, monkeypatch, tmp_path):
|
||||
"""Test that environment variables correctly specify script location
|
||||
and loading.
|
||||
"""
|
||||
and loading."""
|
||||
try:
|
||||
from model_hosting_container_standards.sagemaker.config import (
|
||||
SageMakerEnvVars,
|
||||
@@ -420,8 +416,7 @@ async def custom_sagemaker_invocation_handler(request: Request):
|
||||
|
||||
def test_handler_env_var_override(self, monkeypatch, tmp_path):
|
||||
"""Test CUSTOM_FASTAPI_PING_HANDLER and CUSTOM_FASTAPI_INVOCATION_HANDLER
|
||||
environment variable overrides.
|
||||
"""
|
||||
environment variable overrides."""
|
||||
try:
|
||||
from model_hosting_container_standards.common.fastapi.config import (
|
||||
FastAPIEnvVars,
|
||||
@@ -515,8 +510,7 @@ async def custom_sagemaker_invocation_handler(request: Request):
|
||||
|
||||
def test_env_var_priority_over_decorator_and_script(self, monkeypatch, tmp_path):
|
||||
"""Test that environment variables have highest priority over decorators
|
||||
and script functions for both ping and invocation handlers.
|
||||
"""
|
||||
and script functions for both ping and invocation handlers."""
|
||||
try:
|
||||
from model_hosting_container_standards.common.fastapi.config import (
|
||||
FastAPIEnvVars,
|
||||
|
||||
@@ -51,8 +51,7 @@ async def test_tokenize_then_chat_completion_with_image(
|
||||
local_asset_server,
|
||||
):
|
||||
"""Tokenize a multimodal message, then send the same message to chat
|
||||
completions. The chat completion must succeed (not 500).
|
||||
"""
|
||||
completions. The chat completion must succeed (not 500)."""
|
||||
image_url = local_asset_server.url_for("stop_sign.jpg")
|
||||
messages = [
|
||||
{
|
||||
|
||||
@@ -36,13 +36,11 @@ def test_should_include_usage_force_enables_continuous_usage(stream_options, exp
|
||||
class TestGetMaxTokens:
|
||||
"""Tests for get_max_tokens() to ensure generation_config's max_tokens
|
||||
acts as a default when from model author, and as a ceiling when
|
||||
explicitly set by the user.
|
||||
"""
|
||||
explicitly set by the user."""
|
||||
|
||||
def test_default_sampling_params_used_when_no_request_max_tokens(self):
|
||||
"""When user doesn't specify max_tokens, generation_config default
|
||||
should apply.
|
||||
"""
|
||||
should apply."""
|
||||
result = get_max_tokens(
|
||||
max_model_len=24000,
|
||||
max_tokens=None,
|
||||
@@ -53,8 +51,7 @@ class TestGetMaxTokens:
|
||||
|
||||
def test_request_max_tokens_not_capped_by_default_sampling_params(self):
|
||||
"""When user specifies max_tokens in request, model author's
|
||||
generation_config max_tokens must NOT cap it (fixes #34005).
|
||||
"""
|
||||
generation_config max_tokens must NOT cap it (fixes #34005)."""
|
||||
result = get_max_tokens(
|
||||
max_model_len=24000,
|
||||
max_tokens=5000,
|
||||
@@ -97,8 +94,7 @@ class TestGetMaxTokens:
|
||||
|
||||
def test_request_max_tokens_smaller_than_default(self):
|
||||
"""When user explicitly requests fewer tokens than gen_config default,
|
||||
that should be respected.
|
||||
"""
|
||||
that should be respected."""
|
||||
result = get_max_tokens(
|
||||
max_model_len=24000,
|
||||
max_tokens=512,
|
||||
|
||||
@@ -239,8 +239,7 @@ def test_request_logger_log_outputs_integration():
|
||||
|
||||
def test_streaming_complete_logs_full_text_content():
|
||||
"""Test that streaming complete logging includes
|
||||
full accumulated text, not just token count.
|
||||
"""
|
||||
full accumulated text, not just token count."""
|
||||
mock_logger = MagicMock()
|
||||
|
||||
with patch("vllm.entrypoints.serve.utils.request_logger.logger", mock_logger):
|
||||
|
||||
@@ -17,8 +17,7 @@ from vllm.entrypoints.serve.utils.sse_keep_alive import (
|
||||
def test_disabled_returns_same_generator():
|
||||
"""For 0, negative, and NaN intervals the exact same generator object must
|
||||
be returned (identity), proving no wrapper or per-chunk overhead when
|
||||
disabled.
|
||||
"""
|
||||
disabled."""
|
||||
|
||||
async def gen():
|
||||
yield "x"
|
||||
@@ -34,8 +33,7 @@ def test_disabled_returns_same_generator():
|
||||
@pytest.mark.asyncio
|
||||
async def test_keep_alive_during_silent_start():
|
||||
"""A silent start (queue wait / prefill) emits keep-alive comments before
|
||||
the first real chunk.
|
||||
"""
|
||||
the first real chunk."""
|
||||
release = asyncio.Event()
|
||||
|
||||
async def gen():
|
||||
@@ -57,8 +55,7 @@ async def test_keep_alive_during_silent_start():
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_keep_alive_when_streaming_is_fast():
|
||||
"""If chunks arrive faster than the interval, no comment is emitted and all
|
||||
data is forwarded.
|
||||
"""
|
||||
data is forwarded."""
|
||||
|
||||
async def gen():
|
||||
for chunk in ["a", "b", "c"]:
|
||||
@@ -72,8 +69,7 @@ async def test_no_keep_alive_when_streaming_is_fast():
|
||||
@pytest.mark.asyncio
|
||||
async def test_keep_alive_between_chunks_does_not_drop_data():
|
||||
"""Keep-alives interleaved with real chunks must never drop or reorder
|
||||
data.
|
||||
"""
|
||||
data."""
|
||||
release_b = asyncio.Event()
|
||||
release_c = asyncio.Event()
|
||||
|
||||
@@ -139,8 +135,7 @@ async def test_upstream_cancelled_error_propagates():
|
||||
@pytest.mark.asyncio
|
||||
async def test_aclose_closes_upstream():
|
||||
"""Closing a wrapper that is mid-stream must complete promptly and close
|
||||
the upstream generator (regression for the earlier leaked-task hang).
|
||||
"""
|
||||
the upstream generator (regression for the earlier leaked-task hang)."""
|
||||
finalized = asyncio.Event()
|
||||
|
||||
async def gen():
|
||||
@@ -159,8 +154,7 @@ async def test_aclose_closes_upstream():
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_disconnect_cancels_upstream():
|
||||
"""Cancelling the consumer mid-stream must cancel the pending anext and
|
||||
close the upstream generator so its cleanup runs.
|
||||
"""
|
||||
close the upstream generator so its cleanup runs."""
|
||||
finalized = asyncio.Event()
|
||||
|
||||
async def gen():
|
||||
@@ -187,8 +181,7 @@ async def test_client_disconnect_cancels_upstream():
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_response_emits_comment():
|
||||
"""A keep-alive comment must reach the client through the real
|
||||
StreamingResponse used by the routers.
|
||||
"""
|
||||
StreamingResponse used by the routers."""
|
||||
release = asyncio.Event()
|
||||
|
||||
async def gen():
|
||||
@@ -227,8 +220,7 @@ async def test_streaming_response_emits_comment():
|
||||
async def test_streaming_response_disconnect_closes_upstream():
|
||||
"""On the ASGI 2.3 disconnect path used by the pinned Uvicorn, Starlette
|
||||
cancels streaming and the wrapper's finally must close the upstream
|
||||
generator.
|
||||
"""
|
||||
generator."""
|
||||
finalized = asyncio.Event()
|
||||
|
||||
async def gen():
|
||||
|
||||
@@ -295,8 +295,7 @@ async def test_session_update_invalid_model_returns_error(model_name):
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_commit_without_session_update_returns_error(model_name):
|
||||
"""Test that committing before validating the model returns an error
|
||||
and does not fall through to processing.
|
||||
"""
|
||||
and does not fall through to processing."""
|
||||
server_args = ["--enforce-eager", "--max-model-len", "2048"]
|
||||
|
||||
if model_name.startswith("mistralai"):
|
||||
|
||||
@@ -51,8 +51,7 @@ async def generate_mock_outputs(
|
||||
num_turns, prompt_token_counts, output_token_counts, cached_token_counts=None
|
||||
):
|
||||
"""Generate a sequence of mock RequestOutput objects to simulate multiple
|
||||
turns.
|
||||
"""
|
||||
turns."""
|
||||
if cached_token_counts is None:
|
||||
cached_token_counts = [0] * num_turns
|
||||
|
||||
@@ -271,8 +270,7 @@ def test_reasoning_tokens_counting():
|
||||
|
||||
def test_preamble_tokens_not_counted_as_reasoning():
|
||||
"""Preambles (commentary with no recipient) are visible user text,
|
||||
not hidden reasoning. They must NOT inflate num_reasoning_tokens.
|
||||
"""
|
||||
not hidden reasoning. They must NOT inflate num_reasoning_tokens."""
|
||||
context, parser = make_harmony_context()
|
||||
parser.enqueue_chunk_result(reasoning_token_count=0)
|
||||
|
||||
@@ -289,8 +287,7 @@ def test_preamble_tokens_not_counted_as_reasoning():
|
||||
|
||||
def test_commentary_with_recipient_counted_as_reasoning():
|
||||
"""Commentary directed at a tool (recipient != None) is hidden from
|
||||
the user, so it should still count as reasoning tokens.
|
||||
"""
|
||||
the user, so it should still count as reasoning tokens."""
|
||||
context, parser = make_harmony_context()
|
||||
parser.enqueue_chunk_result(reasoning_token_count=3)
|
||||
|
||||
@@ -349,8 +346,7 @@ async def test_single_turn_no_tool_output():
|
||||
@pytest.mark.asyncio
|
||||
async def test_negative_tool_tokens_edge_case():
|
||||
"""Test edge case where calculation could result in negative tool
|
||||
tokens. We should log an error and clamp the value to 0.
|
||||
"""
|
||||
tokens. We should log an error and clamp the value to 0."""
|
||||
# Use patch to check if logger.error was called
|
||||
with patch("vllm.entrypoints.openai.responses.context.logger.error") as mock_log:
|
||||
context, _ = make_harmony_context(available_tools=["browser"])
|
||||
@@ -690,8 +686,7 @@ def create_simple_context_output(
|
||||
finished=True,
|
||||
):
|
||||
"""Helper to create a RequestOutput with customizable text for
|
||||
SimpleContext tests.
|
||||
"""
|
||||
SimpleContext tests."""
|
||||
if token_ids is None:
|
||||
token_ids = []
|
||||
return RequestOutput(
|
||||
@@ -875,8 +870,7 @@ def test_simple_context_final_output():
|
||||
|
||||
def test_simple_context_output_messages_empty_text_with_tokens():
|
||||
"""output_messages should be returned when tokens exist even if text is
|
||||
empty (e.g. special tokens).
|
||||
"""
|
||||
empty (e.g. special tokens)."""
|
||||
context = SimpleContext()
|
||||
context.append_output(
|
||||
create_simple_context_output(
|
||||
@@ -894,8 +888,7 @@ def test_simple_context_output_messages_empty_text_with_tokens():
|
||||
|
||||
def test_simple_context_output_messages_no_mutation():
|
||||
"""Each call to output_messages returns a fresh list; callers can't
|
||||
corrupt internal state.
|
||||
"""
|
||||
corrupt internal state."""
|
||||
context = SimpleContext()
|
||||
context.append_output(
|
||||
create_simple_context_output(
|
||||
|
||||
@@ -121,8 +121,7 @@ def test_get_world_size_tp1(vllm_runner):
|
||||
@create_new_process_for_each_test()
|
||||
def test_init_weight_transfer_engine_calls_engine(vllm_runner):
|
||||
"""Test that init_weight_transfer_engine calls the engine's
|
||||
init_transfer_engine method.
|
||||
"""
|
||||
init_transfer_engine method."""
|
||||
if torch.accelerator.device_count() < 1:
|
||||
pytest.skip("Need at least 1 GPU for this test")
|
||||
|
||||
|
||||
@@ -86,8 +86,7 @@ def _all_kernel_classes() -> list[type]:
|
||||
|
||||
def _probe(cls: type):
|
||||
"""A bare kernel instance with a plausible config, so input_quant_key()
|
||||
can be queried without the hardware-gated constructor.
|
||||
"""
|
||||
can be queried without the hardware-gated constructor."""
|
||||
obj = cls.__new__(cls) # type: ignore[call-overload]
|
||||
if issubclass(cls, NvFp4LinearKernel):
|
||||
obj.config = NvFp4LinearLayerConfig()
|
||||
|
||||
@@ -247,8 +247,7 @@ def test_extend_attention_cpu_matches_reference(prefix_lens, extend_lens):
|
||||
|
||||
class _FakeLinear:
|
||||
"""Minimal duck-typed stand-in for kv_b_proj: only .weight and
|
||||
.quant_method are read by get_and_maybe_dequant_weights().
|
||||
"""
|
||||
.quant_method are read by get_and_maybe_dequant_weights()."""
|
||||
|
||||
def __init__(self, weight: torch.Tensor):
|
||||
self.weight = weight
|
||||
@@ -319,8 +318,7 @@ class _FakeAttnMetadata:
|
||||
def test_amx_mla_impl_forward_mqa_matches_reference(default_vllm_config):
|
||||
"""forward_mqa receives already-absorbed Q (as MLAAttention.forward_impl
|
||||
would produce via layer.W_UK_T) and must reproduce plain per-head MLA
|
||||
decode attention through W_UK/W_UV.
|
||||
"""
|
||||
decode attention through W_UK/W_UV."""
|
||||
set_random_seed(1)
|
||||
num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank = 4, 32, 32, 64
|
||||
impl, kv_b_weight = _make_amx_mla_impl(
|
||||
@@ -386,8 +384,7 @@ def test_amx_mla_impl_forward_mqa_matches_reference(default_vllm_config):
|
||||
def test_amx_mla_impl_forward_mha_matches_reference(default_vllm_config):
|
||||
"""forward_mha receives raw unabsorbed Q and must attend correctly for a
|
||||
prefill batch with a mix of fresh and continued (cached-prefix)
|
||||
sequences.
|
||||
"""
|
||||
sequences."""
|
||||
set_random_seed(2)
|
||||
num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank = 4, 32, 32, 64
|
||||
impl, kv_b_weight = _make_amx_mla_impl(
|
||||
|
||||
@@ -616,8 +616,7 @@ def test_non_causal_autoselect_backend():
|
||||
)
|
||||
def test_flash_attn_rejects_unhandled_kv_cache_dtypes(kv_cache_dtype: str):
|
||||
"""FlashAttentionBackend must not claim support for kv_cache dtypes
|
||||
that it cannot handle.
|
||||
"""
|
||||
that it cannot handle."""
|
||||
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
|
||||
|
||||
assert not FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype)
|
||||
@@ -628,8 +627,7 @@ def test_flash_attn_accepts_handled_fp8_variants(
|
||||
kv_cache_dtype: str, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
"""FlashAttentionBackend must accept the two fp8 dtypes it can actually
|
||||
handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'.
|
||||
"""
|
||||
handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'."""
|
||||
import vllm.v1.attention.backends.fa_utils as fa_utils_mod
|
||||
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
|
||||
|
||||
|
||||
@@ -225,8 +225,7 @@ def test_cutlass_mla_decode_cross_layer_view():
|
||||
"""The kernel must read the cache's page-dim stride instead of assuming
|
||||
pages are packed back-to-back. A per-layer view into a cross-layer
|
||||
(block-major) cache has stride(0) inflated by num_layers; outputs must
|
||||
match a contiguous cache holding the same data exactly.
|
||||
"""
|
||||
match a contiguous cache holding the same data exactly."""
|
||||
device = torch.device("cuda:0")
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
torch.set_default_device(device)
|
||||
|
||||
@@ -162,8 +162,7 @@ def test_fast_decode_plan_importable() -> None:
|
||||
@torch.inference_mode
|
||||
def test_fast_plan_decode_warmup_uses_full_plan(dtype: torch.dtype) -> None:
|
||||
"""On the first call fast_plan_decode must route through self.plan() and
|
||||
flip vllm_first_call to False on the wrapper object.
|
||||
"""
|
||||
flip vllm_first_call to False on the wrapper object."""
|
||||
from unittest.mock import patch
|
||||
|
||||
from vllm.v1.attention.backends.flashinfer import fast_plan_decode
|
||||
|
||||
@@ -134,8 +134,7 @@ def make_quantized_kv_cache(
|
||||
head_size: int,
|
||||
) -> tuple:
|
||||
"""Quantize kv_cache based on dtype. Returns (kv_cache, kv_cache_sf,
|
||||
kv_scale, ref_kv_cache, is_nvfp4_kv).
|
||||
"""
|
||||
kv_scale, ref_kv_cache, is_nvfp4_kv)."""
|
||||
is_nvfp4_kv = kv_quant_dtype == FP4_DTYPE
|
||||
if is_nvfp4_kv:
|
||||
data, scales, kv_scale, ref = make_nvfp4_kv_cache(
|
||||
|
||||
@@ -73,8 +73,7 @@ def _strided_context_inputs(
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""The real prefill-context layouts: ``k_nope``/``v`` as the two strided
|
||||
halves of one kv_b_proj output and ``k_pe`` as a column slice of the gather
|
||||
workspace (left in the fp8 cache layout for a plain fp8 cache).
|
||||
"""
|
||||
workspace (left in the fp8 cache layout for a plain fp8 cache)."""
|
||||
kv_nope = _randn(num_tokens, num_heads, 256, dtype=dtype)
|
||||
k_nope, v = kv_nope.split((128, 128), dim=-1)
|
||||
workspace = _randn(num_tokens, 576, dtype=dtype)
|
||||
|
||||
@@ -89,8 +89,7 @@ all_case_info: list[tuple] = []
|
||||
@pytest.mark.parametrize("output_dtype", [torch.float32, torch.half, torch.bfloat16])
|
||||
def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None:
|
||||
"""When a token is empty on both sides (both LSE -inf), the 0/0 softmax
|
||||
scales must not surface as NaN in the merged output.
|
||||
"""
|
||||
scales must not surface as NaN in the merged output."""
|
||||
num_tokens, num_heads, head_size = 6, 8, 128
|
||||
prefix_output = torch.zeros(
|
||||
num_tokens, num_heads, head_size, device=DEVICE, dtype=output_dtype
|
||||
|
||||
@@ -45,8 +45,7 @@ def kv_layout(request) -> KVCacheLayout:
|
||||
|
||||
def _layer_stride_order(layout: KVCacheLayout, ndim: int) -> tuple[int, ...]:
|
||||
"""Per-layer physical stride order for the given layout; the 3-dim
|
||||
indexer side cache (H=1) is contiguous, so identity.
|
||||
"""
|
||||
indexer side cache (H=1) is contiguous, so identity."""
|
||||
if ndim == 3:
|
||||
return (0, 1, 2)
|
||||
stride_order = layout.layer_view_order
|
||||
@@ -66,8 +65,7 @@ def _main_spec() -> FullAttentionSpec:
|
||||
|
||||
def _main_kv_logical_shape(num_pages: int) -> tuple[int, ...]:
|
||||
"""Standardized per-layer logical shape [B, H, N, C] for the main cache,
|
||||
derived the same way the production allocator does.
|
||||
"""
|
||||
derived the same way the production allocator does."""
|
||||
shape_bytes = compute_layer_kv_cache_shape_bytes(_main_spec(), num_pages)
|
||||
return (*shape_bytes[:-1], shape_bytes[-1] // DTYPE.itemsize)
|
||||
|
||||
@@ -77,8 +75,7 @@ def _allocate_main_kv_via_contract(
|
||||
) -> torch.Tensor:
|
||||
"""Build the main KV cache exactly as the production allocator does for the
|
||||
given layout: allocate the physical (permuted) tensor, then expose the
|
||||
inverse-permuted logical [B, H, N, C] view the backend sees.
|
||||
"""
|
||||
inverse-permuted logical [B, H, N, C] view the backend sees."""
|
||||
logical_shape = _main_kv_logical_shape(num_pages)
|
||||
stride_order = _layer_stride_order(layout, len(logical_shape))
|
||||
physical_shape = tuple(logical_shape[i] for i in stride_order)
|
||||
@@ -1379,8 +1376,7 @@ def test_amd_decode_index_topk_end_to_end(
|
||||
@pytest.mark.parametrize("num_idx_heads", [1, 4])
|
||||
def test_decode_index_topk_fp8(num_idx_heads: int):
|
||||
"""The standalone Triton path must score FP8 inputs in FP32 so its top-k
|
||||
matches a reference computed from the dequantized FP8 values.
|
||||
"""
|
||||
matches a reference computed from the dequantized FP8 values."""
|
||||
torch.manual_seed(0)
|
||||
topk, init_blocks, local_blocks, head_dim = 8, 0, 1, 128
|
||||
decode_query_len = 1
|
||||
@@ -1653,8 +1649,7 @@ def test_prefill_sparse_attention_correctness(
|
||||
def test_main_cache_layout_contract():
|
||||
"""The standardized per-layer logical shape is [B, H, N, C] with packed
|
||||
K/V content, and the legacy layout aliases resolve to the expected
|
||||
per-layer stride orders.
|
||||
"""
|
||||
per-layer stride orders."""
|
||||
nb, bs, h, d = 7, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
|
||||
logical = _main_kv_logical_shape(nb)
|
||||
assert logical == (nb, h, bs, 2 * d)
|
||||
@@ -1778,8 +1773,7 @@ def test_indexer_cache_squeezes_to_contiguous_3d():
|
||||
"""The indexer side cache is standardized 4D with H=1: under both layouts
|
||||
the allocator's logical view stays contiguous and squeezes (as
|
||||
`MiniMaxM3IndexerCache.bind_kv_cache` does) to the 3-dim
|
||||
[num_blocks, block_size, head_dim] cache the kernels consume.
|
||||
"""
|
||||
[num_blocks, block_size, head_dim] cache the kernels consume."""
|
||||
nb = 5
|
||||
ispec = MLAAttentionSpec(
|
||||
block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
|
||||
@@ -1799,8 +1793,7 @@ def test_indexer_cache_squeezes_to_contiguous_3d():
|
||||
|
||||
def test_hnd_allocation_is_packed_head_major():
|
||||
"""Under HND the backend-visible logical view is the packed head-major
|
||||
physical allocation.
|
||||
"""
|
||||
physical allocation."""
|
||||
nb, bs, h, d = 4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
|
||||
logical = _main_kv_logical_shape(nb)
|
||||
stride_order = KVCacheLayout.LBHNC.layer_view_order
|
||||
@@ -1828,8 +1821,7 @@ def test_hnd_allocation_is_packed_head_major():
|
||||
def test_main_cache_is_block_first_and_unpadded():
|
||||
"""The allocator's contiguous-view branch (not the padded-strided branch)
|
||||
is used for the main GQA cache: its spec is unpadded and the physical
|
||||
layout keeps num_blocks as the first dimension under both layouts.
|
||||
"""
|
||||
layout keeps num_blocks as the first dimension under both layouts."""
|
||||
spec = _main_spec()
|
||||
# Unpadded -> allocator uses kv_tensor.view(...) rather than as_strided().
|
||||
assert spec.page_size_padded is None
|
||||
@@ -1851,8 +1843,7 @@ def _build_decode_inputs(
|
||||
):
|
||||
"""Shared decode setup: uniform query tokens per request, a non-identity
|
||||
block table, and topk indices selecting the current block plus older causal
|
||||
blocks for each query token.
|
||||
"""
|
||||
blocks for each query token."""
|
||||
active_batch = len(seq_lens_list)
|
||||
batch = active_batch + num_padded_reqs
|
||||
pages_per_req = [(s + BLOCK_SIZE - 1) // BLOCK_SIZE for s in seq_lens_list]
|
||||
@@ -2188,8 +2179,7 @@ def test_decode_sparse_attention_correctness(
|
||||
):
|
||||
"""Decode (split-K) parity under both layouts: this is the only coverage of
|
||||
the decode-site cache feed, and the strided HND case fails if the kernel
|
||||
ignores the cache strides.
|
||||
"""
|
||||
ignores the cache strides."""
|
||||
torch.manual_seed(0)
|
||||
q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(
|
||||
seq_lens_list, decode_query_len, num_padded_reqs
|
||||
@@ -2239,8 +2229,7 @@ def test_decode_wrong_layout_breaks_parity():
|
||||
already contiguous-NHD (i.e. skipping the allocator's inverse permute)
|
||||
reorders the K/V content, so the decode output no longer matches the
|
||||
reference computed on the correct logical view. The mislabeled tensor keeps
|
||||
the same shape as the correct view, so the kernel stays in bounds.
|
||||
"""
|
||||
the same shape as the correct view, so the kernel stays in bounds."""
|
||||
torch.manual_seed(0)
|
||||
seq_lens_list = (130, 257)
|
||||
q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(seq_lens_list)
|
||||
@@ -2275,8 +2264,7 @@ def test_decode_wrong_layout_breaks_parity():
|
||||
def test_main_cache_byte_identical_through_production_allocator():
|
||||
"""AC-2: drive the real allocator (`create_kv_cache_views`) for the M3 main
|
||||
`FullAttentionSpec` under HND and assert the backend-visible view has the
|
||||
same shape, stride, and storage offset as the packed-HND allocation.
|
||||
"""
|
||||
same shape, stride, and storage offset as the packed-HND allocation."""
|
||||
nb = 4
|
||||
spec = _main_spec()
|
||||
raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8)
|
||||
@@ -2291,8 +2279,7 @@ def test_main_cache_byte_identical_through_production_allocator():
|
||||
def test_padded_main_cache_is_flagged():
|
||||
"""AC-2.1 negative: the M3 main cache relies on the allocator's
|
||||
contiguous-view branch (`page_size_padded is None`). A spec that sets
|
||||
`page_size_padded` is explicitly flagged rather than silently wrong-strided.
|
||||
"""
|
||||
`page_size_padded` is explicitly flagged rather than silently wrong-strided."""
|
||||
|
||||
def _require_unpadded_block_first(spec, stride_order):
|
||||
inv_order = [stride_order.index(i) for i in range(len(stride_order))]
|
||||
@@ -2329,8 +2316,7 @@ def test_padded_main_cache_is_flagged():
|
||||
def test_reshape_and_cache_flash_write_persists(kv_layout: KVCacheLayout):
|
||||
"""AC-5 write path: the `reshape_and_cache_flash` write site now consumes
|
||||
packed-content K/V split views. Writing through those views must persist
|
||||
into the bound storage under both layouts.
|
||||
"""
|
||||
into the bound storage under both layouts."""
|
||||
torch.manual_seed(0)
|
||||
num_pages = 4
|
||||
kv_cache = _allocate_main_kv_via_contract(num_pages, kv_layout)
|
||||
|
||||
@@ -23,8 +23,7 @@ pytestmark = pytest.mark.skipif(
|
||||
def test_concat_and_cache_mla_into_unified_slot_view():
|
||||
"""concat_and_cache_mla must write correctly into a per-layer view whose
|
||||
block stride is the full unified slot (block-major), with zero bleed into
|
||||
the other layers' segments of the same slot.
|
||||
"""
|
||||
the other layers' segments of the same slot."""
|
||||
from vllm import _custom_ops as ops
|
||||
|
||||
torch.manual_seed(0)
|
||||
@@ -86,8 +85,7 @@ def test_concat_and_cache_mla_into_unified_slot_view():
|
||||
def test_flashmla_dense_decode_unified_slot_view():
|
||||
"""FlashMLA dense decode (FLASHMLA backend, e.g. Kimi-K2-style dense MLA
|
||||
on Hopper) must read a unified-slot block-major view bit-identically to a
|
||||
contiguous per-layer cache.
|
||||
"""
|
||||
contiguous per-layer cache."""
|
||||
import vllm.v1.attention.ops.flashmla as fm
|
||||
|
||||
ok, reason = fm.is_flashmla_dense_supported()
|
||||
@@ -152,8 +150,7 @@ def test_flashmla_dense_decode_unified_slot_view():
|
||||
def test_flashinfer_mla_dense_decode_unified_slot_view():
|
||||
"""FlashInfer MLA dense decode must read a unified-slot block-major view
|
||||
(inflated stride(0), non-zero storage offset) bit-identically to a
|
||||
contiguous per-layer cache.
|
||||
"""
|
||||
contiguous per-layer cache."""
|
||||
try:
|
||||
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
|
||||
except ImportError:
|
||||
@@ -224,8 +221,7 @@ def test_flashinfer_mla_dense_decode_unified_slot_view():
|
||||
def test_flashmla_fp8_sparse_decode_unified_slot_view(next_n: int):
|
||||
"""FlashMLA fp8 sparse decode (DeepSeek V3.2/V4 DSA path) must read a
|
||||
unified-slot block-major view bit-identically to a contiguous fp8_ds_mla
|
||||
cache, with finite nonzero output.
|
||||
"""
|
||||
cache, with finite nonzero output."""
|
||||
import vllm.v1.attention.ops.flashmla as fm
|
||||
|
||||
ok, reason = fm.is_flashmla_sparse_supported()
|
||||
@@ -314,8 +310,7 @@ def test_flashmla_fp8_sparse_decode_unified_slot_view(next_n: int):
|
||||
def test_indexer_k_quant_and_cache_into_unified_slot_view():
|
||||
"""indexer_k_quant_and_cache (DeepSeek V3.2/V4 DSA indexer K write) must
|
||||
write correctly into a per-layer view whose block stride is the full
|
||||
unified slot, with zero bleed into the other layers' segments.
|
||||
"""
|
||||
unified slot, with zero bleed into the other layers' segments."""
|
||||
from vllm import _custom_ops as ops
|
||||
|
||||
torch.manual_seed(0)
|
||||
@@ -360,8 +355,7 @@ def test_indexer_k_quant_and_cache_into_unified_slot_view():
|
||||
|
||||
def test_flashattn_mla_dense_decode_unified_slot_view():
|
||||
"""FA3 decode (FLASH_ATTN_MLA backend) must read a unified-slot
|
||||
block-major view bit-identically to a contiguous per-layer cache.
|
||||
"""
|
||||
block-major view bit-identically to a contiguous per-layer cache."""
|
||||
try:
|
||||
from vllm.vllm_flash_attn import flash_attn_varlen_func
|
||||
except ImportError:
|
||||
@@ -434,8 +428,7 @@ def test_flashattn_mla_dense_decode_unified_slot_view():
|
||||
def test_flashmla_dense_fp8_decode_unified_slot_view():
|
||||
"""FlashMLA dense fp8 decode (FLASHMLA backend with quantized KV cache)
|
||||
must read a unified-slot block-major view bit-identically to a contiguous
|
||||
per-layer fp8 cache.
|
||||
"""
|
||||
per-layer fp8 cache."""
|
||||
import vllm.v1.attention.ops.flashmla as fm
|
||||
|
||||
ok, reason = fm.is_flashmla_dense_supported()
|
||||
@@ -506,8 +499,7 @@ def test_flashmla_dense_fp8_decode_unified_slot_view():
|
||||
def test_flashinfer_mla_dense_fp8_decode_unified_slot_view():
|
||||
"""FlashInfer MLA dense decode with an fp8 KV cache must read a
|
||||
unified-slot block-major view bit-identically to a contiguous per-layer
|
||||
cache.
|
||||
"""
|
||||
cache."""
|
||||
try:
|
||||
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
|
||||
except ImportError:
|
||||
|
||||
@@ -283,8 +283,7 @@ def _run_direct_flash_attn_case(
|
||||
# Backend contract tests --------------------------------------------------
|
||||
def test_aiter_mha_backend_contract():
|
||||
"""The ROCm backend advertises the dtypes, shapes, and attention types it
|
||||
is designed to handle.
|
||||
"""
|
||||
is designed to handle."""
|
||||
from vllm.v1.attention.backend import AttentionType
|
||||
from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
|
||||
|
||||
@@ -312,8 +311,7 @@ def test_aiter_mha_backend_contract():
|
||||
|
||||
def test_aiter_mha_backend_validates_kv_cache_block_size():
|
||||
"""The backend should reject KV cache specs that cannot be gathered
|
||||
correctly by the ROCm kernel.
|
||||
"""
|
||||
correctly by the ROCm kernel."""
|
||||
from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec
|
||||
|
||||
@@ -333,8 +331,7 @@ def test_aiter_mha_backend_validates_kv_cache_block_size():
|
||||
|
||||
def test_aiter_mha_backend_supports_compute_capability_matches_mi3xx_probe():
|
||||
"""The backend should trust the ROCm MI3xx probe instead of the raw torch
|
||||
capability tuple.
|
||||
"""
|
||||
capability tuple."""
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
|
||||
|
||||
@@ -359,8 +356,7 @@ def test_aiter_mha_env_flags_control_enablement(
|
||||
use_aiter, use_mha, expected, monkeypatch
|
||||
):
|
||||
"""Both the global AITER flag and the MHA-specific flag must be enabled
|
||||
before the ROCm MHA path is considered active.
|
||||
"""
|
||||
before the ROCm MHA path is considered active."""
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
|
||||
_assert_aiter_supported()
|
||||
@@ -378,8 +374,7 @@ def test_aiter_mha_env_flags_control_enablement(
|
||||
|
||||
def test_aiter_mha_platform_gate_matches_install_and_arch():
|
||||
"""The global AITER availability check should only open on ROCm MI3xx when
|
||||
the aiter package is installed.
|
||||
"""
|
||||
the aiter package is installed."""
|
||||
from vllm._aiter_ops import IS_AITER_FOUND, is_aiter_found_and_supported
|
||||
|
||||
assert is_aiter_found_and_supported() is (
|
||||
@@ -650,8 +645,7 @@ def test_aiter_mha_multi_batch(num_heads, head_size, dtype):
|
||||
@pytest.mark.parametrize("num_blocks", DIRECT_NUM_BLOCKS)
|
||||
def test_aiter_fa_large_block_table_matches_reference(num_blocks):
|
||||
"""The direct paged-KV path should stay stable for both normal and very
|
||||
large block tables.
|
||||
"""
|
||||
large block tables."""
|
||||
atol = 2e-2
|
||||
rtol = 2e-2
|
||||
_assert_aiter_supported()
|
||||
@@ -673,8 +667,7 @@ def test_aiter_fa_large_block_table_matches_reference(num_blocks):
|
||||
@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
|
||||
def test_aiter_fa_sliding_window_matches_reference():
|
||||
"""The direct kernel should respect the same sliding-window causal mask as
|
||||
the naive reference implementation.
|
||||
"""
|
||||
the naive reference implementation."""
|
||||
atol = 2e-2
|
||||
rtol = 2e-2
|
||||
_assert_aiter_supported()
|
||||
|
||||
@@ -111,8 +111,7 @@ def _make_mixed_common_metadata():
|
||||
|
||||
def _patch_build_deps(monkeypatch, events=None):
|
||||
"""Stub the aiter kernel, triton helper and CUDA sync so ``build()`` runs
|
||||
on CPU.
|
||||
"""
|
||||
on CPU."""
|
||||
|
||||
def fake_generate_sparse_seqlen_triton(
|
||||
query_lens, seq_lens, cu_query_lens, topk_token, num_tokens, max_query_len
|
||||
@@ -153,8 +152,7 @@ def test_build_populates_decode_only_split_fields(monkeypatch):
|
||||
|
||||
def test_build_populates_mixed_split_fields(monkeypatch):
|
||||
"""Mixed decode+prefill batch: split is reported, prefill fields stay
|
||||
default because this impl always runs the MQA path.
|
||||
"""
|
||||
default because this impl always runs the MQA path."""
|
||||
builder = _make_builder()
|
||||
_patch_build_deps(monkeypatch)
|
||||
|
||||
|
||||
@@ -125,8 +125,7 @@ def _quantize_to_fp8(tensor: torch.Tensor):
|
||||
@pytest.mark.parametrize("PAGE_SIZE", [1, 16])
|
||||
def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE):
|
||||
"""Test FP8 KV cache path: quantize K/V to FP8, run kernel with scales,
|
||||
and compare against BF16 reference output.
|
||||
"""
|
||||
and compare against BF16 reference output."""
|
||||
assert CACHE_SIZE % PAGE_SIZE == 0
|
||||
dtype = torch.bfloat16
|
||||
seq_len = L
|
||||
@@ -249,8 +248,7 @@ def test_decode_attention_cross_layer_view(H_Q, H_KV, D_QK, D_V, is_mla, PAGE_SI
|
||||
"""The kernel must honor the cache's page-dim stride, not assume pages are
|
||||
packed back-to-back. A per-layer view into a cross-layer (block-major)
|
||||
cache has stride(0) inflated by num_layers; outputs must match a
|
||||
contiguous cache holding the same data exactly.
|
||||
"""
|
||||
contiguous cache holding the same data exactly."""
|
||||
B = 3
|
||||
seq_len = 1027
|
||||
CACHE_SIZE = 16384
|
||||
|
||||
@@ -41,8 +41,7 @@ def _rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
|
||||
def test_fused_embed_norm_matches_reference():
|
||||
"""Main-model fusion: the residual is the exact gather and the second output
|
||||
is a correct RMSNorm. The norm matches a full-precision reference to ~2 bf16
|
||||
ulp (rtol 1e-2) -- that gap is bf16 rounding, not the kernel.
|
||||
"""
|
||||
ulp (rtol 1e-2) -- that gap is bf16 rounding, not the kernel."""
|
||||
set_random_seed(13)
|
||||
table = torch.randn(VOCAB, HIDDEN, dtype=DTYPE, device="cuda")
|
||||
ids = torch.randint(0, VOCAB, (NUM_TOKENS,), dtype=torch.int32, device="cuda")
|
||||
@@ -61,8 +60,7 @@ def test_fused_embed_norm_matches_reference():
|
||||
@torch.inference_mode()
|
||||
def test_fused_embed_eh_norm_matches_reference():
|
||||
"""MTP fusion (folded gather) is bit-exact vs gathering the embeds and
|
||||
feeding the untouched model-local ``fused_eh_norm``.
|
||||
"""
|
||||
feeding the untouched model-local ``fused_eh_norm``."""
|
||||
set_random_seed(13)
|
||||
table = torch.randn(VOCAB, HIDDEN, dtype=DTYPE, device="cuda")
|
||||
ids = torch.randint(0, VOCAB, (NUM_TOKENS,), dtype=torch.int32, device="cuda")
|
||||
|
||||
@@ -62,8 +62,7 @@ def test_fused_q_kv_rmsnorm_outputs_are_packed(num_tokens: int):
|
||||
a [1, q_size] slice of a [1, q_size + kv_size] buffer used to produce a
|
||||
qr_out with row stride q_size + kv_size. Downstream dispatchers that
|
||||
require packed row-major inputs then reject the tensor and silently fall
|
||||
back to a slower GEMM path on every decode step.
|
||||
"""
|
||||
back to a slower GEMM path on every decode step."""
|
||||
device = "cuda"
|
||||
dtype = torch.bfloat16
|
||||
q_size, kv_size = 192, 576
|
||||
@@ -89,8 +88,7 @@ def test_fused_q_kv_rmsnorm_outputs_are_packed(num_tokens: int):
|
||||
def test_fused_q_kv_rmsnorm_launches_past_grid_y_cap(num_tokens: int):
|
||||
"""Regression guard: grid used to be (2, num_tokens), hitting CUDA's
|
||||
65535 grid-y cap at num_tokens >= 65536. The new grid (num_tokens, 2)
|
||||
lifts that bound to 2**31-1.
|
||||
"""
|
||||
lifts that bound to 2**31-1."""
|
||||
device = "cuda"
|
||||
dtype = torch.bfloat16
|
||||
q_size, kv_size = 192, 576
|
||||
|
||||
@@ -2,8 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Tests that FusedRMSNormGated decomposes correctly under torch.compile,
|
||||
matching the eager triton kernel output.
|
||||
"""
|
||||
matching the eager triton kernel output."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
@@ -579,8 +579,7 @@ def test_rocm_aiter_mla_dual_rmsnorm_group_quant_vs_sequential(m):
|
||||
|
||||
def test_rocm_aiter_mla_dual_rmsnorm_group_quant_strided_split_view():
|
||||
"""Split-view inputs, as produced by the DeepSeek-V4 q/kv split, match
|
||||
contiguous inputs.
|
||||
"""
|
||||
contiguous inputs."""
|
||||
require_aiter()
|
||||
require_fp8()
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user