@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a289cc1dde
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 13, 02:00 UTC

[Test] Batch Invariant: Rename and organize tests (#27421)

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

Wentao Ye committed Nov 17, 2025 at 18:09 UTC a289cc1dde4a1aeee05492bbe4cc39a18f070135
5 files changed +248 -80
tests/v1/determinism/conftest.py new
+11
@@ -0,0 +1,11 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +
7 +@pytest.fixture(autouse=True)
8 +def enable_batch_invariant_mode(monkeypatch: pytest.MonkeyPatch):
9 + """Automatically enable batch invariant kernel overrides for all tests."""
10 + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
11 + yield
tests/v1/determinism/test_batch_invariance.py renamed
+1 -74
@@ -6,66 +6,9 @@ import random
6
7 import pytest
8 import torch
9 +from utils import _extract_step_logprobs, _random_prompt, skip_unsupported
10
11 from vllm import LLM, SamplingParams
11 -from vllm.platforms import current_platform
12 -
13 -skip_unsupported = pytest.mark.skipif(
14 - not (current_platform.is_cuda() and current_platform.has_device_capability(90)),
15 - reason="Requires CUDA and >= Hopper (SM90)",
16 -)
17 -
18 -
19 -@pytest.fixture(autouse=True)
20 -def enable_batch_invariant_mode(monkeypatch: pytest.MonkeyPatch):
21 - """Automatically enable batch invariant kernel overrides for all tests."""
22 - monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
23 - yield
24 -
25 -
26 -def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str:
27 - # Generate more realistic prompts that will actually produce varied tokens
28 - # Use a mix of common English text patterns
29 -
30 - prompt_templates = [
31 - # Question-answer style
32 - "Question: What is the capital of France?\nAnswer: The capital of France is",
33 - "Q: How does photosynthesis work?\nA: Photosynthesis is the process by which",
34 - "User: Can you explain quantum mechanics?\nAssistant: Quantum mechanics is",
35 - # Story/narrative style
36 - "Once upon a time in a distant galaxy, there lived",
37 - "The old man walked slowly down the street, remembering",
38 - "In the year 2157, humanity finally discovered",
39 - # Technical/code style
40 - "To implement a binary search tree in Python, first we need to",
41 - "The algorithm works by iterating through the array and",
42 - "Here's how to optimize database queries using indexing:",
43 - # Factual/informative style
44 - "The Renaissance was a period in European history that",
45 - "Climate change is caused by several factors including",
46 - "The human brain contains approximately 86 billion neurons which",
47 - # Conversational style
48 - "I've been thinking about getting a new laptop because",
49 - "Yesterday I went to the store and bought",
50 - "My favorite thing about summer is definitely",
51 - ]
52 -
53 - # Pick a random template
54 - base_prompt = random.choice(prompt_templates)
55 -
56 - if max_words < min_words:
57 - max_words = min_words
58 - target_words = random.randint(min_words, max_words)
59 -
60 - if target_words > 50:
61 - # For longer prompts, repeat context
62 - padding_text = (
63 - " This is an interesting topic that deserves more explanation. "
64 - * (target_words // 50)
65 - )
66 - base_prompt = base_prompt + padding_text
67 -
68 - return base_prompt
12
13
14 @skip_unsupported
@@ -204,22 +147,6 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
147 llm_bsN.shutdown()
148
149
207 -def _extract_step_logprobs(request_output):
208 - if getattr(request_output, "outputs", None):
209 - inner = request_output.outputs[0]
210 - if hasattr(inner, "logprobs") and inner.logprobs is not None:
211 - t = torch.tensor(
212 - [
213 - inner.logprobs[i][tid].logprob
214 - for i, tid in enumerate(inner.token_ids)
215 - ],
216 - dtype=torch.float32,
217 - )
218 - return t, inner.token_ids
219 -
220 - return None, None
221 -
222 -
150 @skip_unsupported
151 @pytest.mark.parametrize(
152 "backend",
tests/v1/determinism/test_online_batch_invariance.py new
+161
@@ -0,0 +1,161 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +HTTP-based batch invariance test: send requests to a running
5 +vLLM server and compare BS=1 vs BS=N results (tokens and per-step logprobs).
6 +
7 +Environment variables:
8 + - VLLM_TEST_MODEL: served model name (e.g., Qwen/Qwen3-1.7B / DeepSeek-R1)
9 + - VLLM_TP_SIZE: tensor parallelism size (e.g., 4)
10 +
11 +"""
12 +
13 +import os
14 +import random
15 +import sys
16 +from typing import Any
17 +
18 +import openai
19 +from utils import _random_prompt, skip_unsupported
20 +
21 +from tests.utils import RemoteOpenAIServer
22 +
23 +
24 +def _request_completion(
25 + client: openai.OpenAI,
26 + model: str,
27 + prompt: Any,
28 + sp: dict[str, Any],
29 + max_retries: int = 3,
30 + retry_backoff: float = 0.5,
31 +) -> dict[str, Any] | None:
32 + payload: dict[str, Any] = {"model": model, "prompt": prompt}
33 + payload.update(sp)
34 +
35 + for attempt in range(max_retries + 1):
36 + try:
37 + completion = client.completions.create(**payload)
38 + # Convert to plain dict so downstream logic can keep using
39 + # dict-style access just like with raw HTTP JSON.
40 + return completion.model_dump()
41 + except Exception as e: # pragma: no cover
42 + if attempt < max_retries:
43 + import time as _t
44 +
45 + _t.sleep(retry_backoff * (2**attempt))
46 + continue
47 + sys.stderr.write(f"Error: {e}\n")
48 + return None
49 + return None
50 +
51 +
52 +def _extract_tokens_and_logprobs(
53 + choice: dict[str, Any],
54 +) -> tuple[list[Any], list[float] | None]:
55 + tokens: list[Any] = []
56 + token_logprobs: list[float] | None = None
57 + lp = choice.get("logprobs")
58 + if lp and isinstance(lp, dict):
59 + tokens = lp.get("token_ids") or lp.get("tokens") or []
60 + token_logprobs = lp.get("token_logprobs", None)
61 + return tokens, token_logprobs
62 +
63 +
64 +def _compare_bs1_vs_bsn_single_process(
65 + prompts: list[str],
66 + sp_kwargs: dict[str, Any],
67 + client: openai.OpenAI,
68 + model_name: str,
69 +) -> None:
70 + # BS=1
71 + bs1_tokens_per_prompt: list[list[Any]] = []
72 + bs1_logprobs_per_prompt: list[list[float] | None] = []
73 + for p in prompts:
74 + resp = _request_completion(client, model_name, p, sp_kwargs)
75 + if resp is None or not resp.get("choices"):
76 + raise AssertionError("BS=1 empty/failed response")
77 + choice = resp["choices"][0]
78 + toks, lps = _extract_tokens_and_logprobs(choice)
79 + if lps is None:
80 + raise AssertionError(
81 + "logprobs not returned; ensure server supports 'logprobs'"
82 + )
83 + bs1_tokens_per_prompt.append(list(toks))
84 + bs1_logprobs_per_prompt.append(list(lps))
85 +
86 + # BS=N
87 + bsN_tokens_per_prompt: list[list[Any]] = [None] * len(prompts) # type: ignore[list-item]
88 + bsN_logprobs_per_prompt: list[list[float] | None] = [None] * len(prompts)
89 + resp = _request_completion(client, model_name, prompts, sp_kwargs)
90 + if resp is None or not resp.get("choices"):
91 + raise AssertionError("BS=N empty/failed batched response")
92 + choices = resp.get("choices", [])
93 + if len(choices) != len(prompts):
94 + raise AssertionError(
95 + f"BS=N choices length {len(choices)} != num prompts {len(prompts)}"
96 + )
97 + for idx, choice in enumerate(choices):
98 + toks, lps = _extract_tokens_and_logprobs(choice)
99 + if lps is None:
100 + raise AssertionError(f"BS=N missing logprobs for prompt {idx}")
101 + bsN_tokens_per_prompt[idx] = list(toks)
102 + bsN_logprobs_per_prompt[idx] = list(lps)
103 +
104 + # compare
105 + for i, (tokens_bs1, tokens_bsN, logprobs_bs1, logprobs_bsN) in enumerate(
106 + zip(
107 + bs1_tokens_per_prompt,
108 + bsN_tokens_per_prompt,
109 + bs1_logprobs_per_prompt,
110 + bsN_logprobs_per_prompt,
111 + )
112 + ):
113 + if tokens_bs1 != tokens_bsN:
114 + raise AssertionError(
115 + f"Prompt {i} (sampling): Different tokens sampled. "
116 + f"BS=1 tokens: {tokens_bs1} BS=N tokens: {tokens_bsN}"
117 + )
118 + if logprobs_bs1 is None or logprobs_bsN is None:
119 + raise AssertionError(f"Prompt {i}: Missing logprobs in one of the runs")
120 + if len(logprobs_bs1) != len(logprobs_bsN):
121 + raise AssertionError(
122 + f"Prompt {i}: Different number of steps: "
123 + f"{len(logprobs_bs1)} (BS=1) vs {len(logprobs_bsN)} (BS=N)."
124 + )
125 + for t, (a, b) in enumerate(zip(logprobs_bs1, logprobs_bsN)):
126 + if a != b:
127 + diff = abs(a - b)
128 + raise AssertionError(
129 + f"Prompt {i} Step {t}: Bitwise mismatch "
130 + f"(abs diff={diff:.6e}). "
131 + f"BS=1 tokens: {tokens_bs1} BS=N tokens: {tokens_bsN}"
132 + )
133 +
134 +
135 +@skip_unsupported
136 +def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN():
137 + random.seed(int(os.getenv("VLLM_TEST_SEED", "12345")))
138 + model_name = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
139 + prompts_all = [_random_prompt(10, 50) for _ in range(32)]
140 +
141 + sp_kwargs: dict[str, Any] = {
142 + "temperature": 0.6,
143 + "top_p": 1.0,
144 + "max_tokens": 8,
145 + "seed": 42,
146 + "logprobs": 5,
147 + }
148 +
149 + tp_size = os.getenv("VLLM_TP_SIZE", "1")
150 + server_args: list[str] = []
151 + if tp_size:
152 + server_args += ["-tp", tp_size]
153 +
154 + with RemoteOpenAIServer(model_name, server_args) as server:
155 + client = server.get_client()
156 + _compare_bs1_vs_bsn_single_process(
157 + prompts=prompts_all,
158 + sp_kwargs=sp_kwargs,
159 + client=client,
160 + model_name=model_name,
161 + )
tests/v1/determinism/test_rms_norm_batch_invariant.py renamed
+1 -6
@@ -9,15 +9,10 @@ with the standard CUDA-based implementation to ensure numerical accuracy.
9
10 import pytest
11 import torch
12 +from utils import skip_unsupported
13
14 from vllm.model_executor.layers.batch_invariant import rms_norm as triton_rms_norm
15 from vllm.model_executor.layers.layernorm import RMSNorm
15 -from vllm.platforms import current_platform
16 -
17 -skip_unsupported = pytest.mark.skipif(
18 - not (current_platform.is_cuda() and current_platform.has_device_capability(90)),
19 - reason="Requires CUDA and >= Hopper (SM90)",
20 -)
16
17
18 @skip_unsupported
tests/v1/determinism/utils.py new
+74
@@ -0,0 +1,74 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import random
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm.platforms import current_platform
9 +
10 +skip_unsupported = pytest.mark.skipif(
11 + not (current_platform.is_cuda() and current_platform.has_device_capability(90)),
12 + reason="Requires CUDA and >= Hopper (SM90)",
13 +)
14 +
15 +
16 +def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str:
17 + # Generate more realistic prompts that will actually produce varied tokens
18 + # Use a mix of common English text patterns
19 +
20 + prompt_templates = [
21 + # Question-answer style
22 + "Question: What is the capital of France?\nAnswer: The capital of France is",
23 + "Q: How does photosynthesis work?\nA: Photosynthesis is the process by which",
24 + "User: Can you explain quantum mechanics?\nAssistant: Quantum mechanics is",
25 + # Story/narrative style
26 + "Once upon a time in a distant galaxy, there lived",
27 + "The old man walked slowly down the street, remembering",
28 + "In the year 2157, humanity finally discovered",
29 + # Technical/code style
30 + "To implement a binary search tree in Python, first we need to",
31 + "The algorithm works by iterating through the array and",
32 + "Here's how to optimize database queries using indexing:",
33 + # Factual/informative style
34 + "The Renaissance was a period in European history that",
35 + "Climate change is caused by several factors including",
36 + "The human brain contains approximately 86 billion neurons which",
37 + # Conversational style
38 + "I've been thinking about getting a new laptop because",
39 + "Yesterday I went to the store and bought",
40 + "My favorite thing about summer is definitely",
41 + ]
42 +
43 + # Pick a random template
44 + base_prompt = random.choice(prompt_templates)
45 +
46 + if max_words < min_words:
47 + max_words = min_words
48 + target_words = random.randint(min_words, max_words)
49 +
50 + if target_words > 50:
51 + # For longer prompts, repeat context
52 + padding_text = (
53 + " This is an interesting topic that deserves more explanation. "
54 + * (target_words // 50)
55 + )
56 + base_prompt = base_prompt + padding_text
57 +
58 + return base_prompt
59 +
60 +
61 +def _extract_step_logprobs(request_output):
62 + if getattr(request_output, "outputs", None):
63 + inner = request_output.outputs[0]
64 + if hasattr(inner, "logprobs") and inner.logprobs is not None:
65 + t = torch.tensor(
66 + [
67 + inner.logprobs[i][tid].logprob
68 + for i, tid in enumerate(inner.token_ids)
69 + ],
70 + dtype=torch.float32,
71 + )
72 + return t, inner.token_ids
73 +
74 + return None, None