[Refactor] Remove dead or duplicate tests (#58446)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Sep 24, 2026 at 10:29 UTC
3c42385d7dcce60cc231bcd9127b56379f43f823
11 files changed
-646
tests/entrypoints/pooling/scoring/test_bi_encoder_online.py
-29
@@ -126,35 +126,6 @@ async def test_score_api_queries_str_n_documents_str_n(
126
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
127
128
129
-@pytest.mark.asyncio
130
-async def test_score_api_queries_vs_documents(hf_model, server: RemoteOpenAIServer):
131
- text_pairs = [
132
- [TEXTS_1[0], TEXTS_2[0]],
133
- [TEXTS_1[1], TEXTS_2[1]],
134
- ]
135
-
136
- score_response = requests.post(
137
- server.url_for("score"),
138
- json={
139
- "model": MODEL_NAME,
140
- "queries": TEXTS_1,
141
- "documents": TEXTS_2,
142
- },
143
- )
144
- score_response.raise_for_status()
145
- score = ScoreResponse.model_validate(score_response.json())
146
-
147
- assert score.id is not None
148
- assert score.data is not None
149
- assert len(score.data) == 2
150
-
151
- vllm_outputs = [d.score for d in score.data]
152
- hf_outputs = hf_model.predict(text_pairs).tolist()
153
-
154
- for i in range(len(vllm_outputs)):
155
- assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
156
-
157
-
129
@pytest.mark.asyncio
130
async def test_score_api_queries_vs_items(hf_model, server: RemoteOpenAIServer):
131
text_pairs = [
tests/entrypoints/pooling/scoring/test_cross_encoder_online.py
-29
@@ -143,35 +143,6 @@ async def test_score_api_queries_str_n_documents_str_n(
143
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
144
145
146
-@pytest.mark.asyncio
147
-async def test_score_api_queries_vs_documents(hf_model, server: RemoteOpenAIServer):
148
- text_pairs = [
149
- [TEXTS_1[0], TEXTS_2[0]],
150
- [TEXTS_1[1], TEXTS_2[1]],
151
- ]
152
-
153
- score_response = requests.post(
154
- server.url_for("score"),
155
- json={
156
- "model": MODEL_NAME,
157
- "queries": TEXTS_1,
158
- "documents": TEXTS_2,
159
- },
160
- )
161
- score_response.raise_for_status()
162
- score = ScoreResponse.model_validate(score_response.json())
163
-
164
- assert score.id is not None
165
- assert score.data is not None
166
- assert len(score.data) == 2
167
-
168
- vllm_outputs = [d.score for d in score.data]
169
- hf_outputs = hf_model.predict(text_pairs).tolist()
170
-
171
- for i in range(len(vllm_outputs)):
172
- assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
173
-
174
-
146
@pytest.mark.asyncio
147
async def test_score_api_queries_vs_items(hf_model, server: RemoteOpenAIServer):
148
text_pairs = [
tests/kernels/attention/test_attention_selector.py
-55
@@ -341,61 +341,6 @@ def test_fp32_fallback(device: str):
341
get_attn_backend(16, torch.float32, None)
342
343
344
-def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
345
- """Test FlashAttn validation."""
346
- pytest.skip(
347
- "Skipping as current backend selector does not "
348
- "handle fallbacks when a backend is explicitly set."
349
- )
350
-
351
- attention_config = AttentionConfig(backend=AttentionBackendEnum.FLASH_ATTN)
352
- cache_config = CacheConfig(block_size=16)
353
- vllm_config = VllmConfig(
354
- attention_config=attention_config, cache_config=cache_config
355
- )
356
-
357
- with set_current_vllm_config(vllm_config):
358
- # Unsupported CUDA arch
359
- monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _=None: (7, 5))
360
- backend = get_attn_backend(16, torch.float16, None)
361
- assert backend.get_name() != "FLASH_ATTN"
362
-
363
- # Reset the monkeypatch for subsequent tests
364
- monkeypatch.undo()
365
-
366
- # Unsupported data type
367
- backend = get_attn_backend(16, torch.float8_e4m3fn, None)
368
- assert backend.get_name() != "FLASH_ATTN"
369
-
370
- # Unsupported kv cache data type
371
- backend = get_attn_backend(16, torch.float16, "fp8")
372
- assert backend.get_name() != "FLASH_ATTN"
373
-
374
- # Unsupported block size
375
- vllm_config.cache_config.block_size = 8
376
- backend = get_attn_backend(16, torch.float16, None)
377
- assert backend.get_name() != "FLASH_ATTN"
378
-
379
- # flash-attn is not installed
380
- import sys
381
-
382
- vllm_config.cache_config.block_size = 16
383
- original_module = sys.modules.get("vllm_flash_attn")
384
- monkeypatch.setitem(sys.modules, "vllm_flash_attn", None)
385
- backend = get_attn_backend(16, torch.float16, None)
386
- assert backend.get_name() != "FLASH_ATTN"
387
-
388
- # Restore the original module if it existed
389
- if original_module is not None:
390
- monkeypatch.setitem(sys.modules, "vllm_flash_attn", original_module)
391
- else:
392
- monkeypatch.delitem(sys.modules, "vllm_flash_attn", raising=False)
393
-
394
- # Unsupported head size
395
- backend = get_attn_backend(17, torch.float16, None)
396
- assert backend.get_name() != "FLASH_ATTN"
397
-
398
-
344
def test_invalid_backend():
345
"""Test that invalid attention backend names raise ValueError."""
346
with (
tests/kernels/attention/test_rocm_attention_selector.py
deleted
-73
@@ -1,73 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import pytest
5
-import torch
6
-
7
-from vllm.config import AttentionConfig, VllmConfig, set_current_vllm_config
8
-from vllm.platforms.rocm import RocmPlatform
9
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
10
-from vllm.v1.attention.selector import _cached_get_attn_backend, get_attn_backend
11
-
12
-
13
-@pytest.fixture(autouse=True)
14
-def clear_cache():
15
- """Clear lru cache to ensure each test case runs without caching."""
16
- _cached_get_attn_backend.cache_clear()
17
-
18
-
19
-@pytest.mark.skip(reason="Skipped for now. Should be revisited.")
20
-def test_selector(monkeypatch: pytest.MonkeyPatch):
21
- # Set the current platform to ROCm using monkeypatch
22
- monkeypatch.setattr("vllm.v1.attention.selector.current_platform", RocmPlatform())
23
-
24
- # Test standard ROCm attention
25
- attention_config = AttentionConfig(backend=AttentionBackendEnum.ROCM_ATTN)
26
- vllm_config = VllmConfig(attention_config=attention_config)
27
-
28
- with set_current_vllm_config(vllm_config):
29
- backend = get_attn_backend(16, torch.float16, torch.float16, 16, False)
30
- assert backend.get_name() == "ROCM_FLASH" or backend.get_name() == "TRITON_ATTN"
31
-
32
- # MLA test for deepseek related
33
- # Change the attention backend to triton MLA
34
- attention_config = AttentionConfig(backend=AttentionBackendEnum.TRITON_MLA)
35
- vllm_config = VllmConfig(attention_config=attention_config)
36
-
37
- with set_current_vllm_config(vllm_config):
38
- backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False, use_mla=True)
39
- assert backend.get_name() == "TRITON_MLA"
40
-
41
- # If attention backend is None
42
- # If use_mla is true
43
- # The selected backend is triton MLA
44
- attention_config = AttentionConfig(backend=None)
45
- vllm_config = VllmConfig(attention_config=attention_config)
46
-
47
- with set_current_vllm_config(vllm_config):
48
- backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False, use_mla=True)
49
- assert backend.get_name() == "TRITON_MLA"
50
-
51
- # Change the attention backend to AITER MLA
52
- attention_config = AttentionConfig(backend=AttentionBackendEnum.ROCM_AITER_MLA)
53
- vllm_config = VllmConfig(attention_config=attention_config)
54
-
55
- with set_current_vllm_config(vllm_config):
56
- backend = get_attn_backend(576, torch.bfloat16, "auto", 1, False, use_mla=True)
57
- assert backend.get_name() == "ROCM_AITER_MLA"
58
-
59
- # If attention backend is None
60
- # If use_mla is true
61
- # If VLLM_ROCM_USE_AITER is enabled
62
- # The selected backend is ROCM_AITER_MLA
63
- with monkeypatch.context() as m:
64
- m.setenv("VLLM_ROCM_USE_AITER", "1")
65
-
66
- attention_config = AttentionConfig(backend=None)
67
- vllm_config = VllmConfig(attention_config=attention_config)
68
-
69
- with set_current_vllm_config(vllm_config):
70
- backend = get_attn_backend(
71
- 576, torch.bfloat16, "auto", 1, False, use_mla=True
72
- )
73
- assert backend.get_name() == "ROCM_AITER_MLA"
tests/quantization/test_auto_round.py
-13
@@ -1653,19 +1653,6 @@ def test_wna16_linear_gptq_unsupported_config_raises() -> None:
1653
INCWNA16LinearScheme(make_layer_config(sym=False))
1654
1655
1656
-def test_wna16_xpu_unsupported_config_still_raises(monkeypatch) -> None:
1657
- monkeypatch.setattr(current_platform, "is_xpu", lambda: True)
1658
- monkeypatch.setattr(current_platform, "is_cpu", lambda: False)
1659
-
1660
- with pytest.raises(NotImplementedError, match="unsupported config"):
1661
- INCWna16Scheme().get_linear_method(
1662
- make_config(weight_bits=2, sym=False),
1663
- object(),
1664
- "layer",
1665
- make_layer_config(bits=2, sym=False),
1666
- )
1667
-
1668
-
1656
def test_inc_get_quant_method_unquantized_linear_returns_unquantized() -> None:
1657
config = make_config(extra_config={"layer": {"bits": 16}})
1658
layer = object.__new__(LinearBase)
tests/quantization/test_gptq_v2.py
deleted
-106
@@ -1,106 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Tests whether vllm correctly load and run gptq_v2 format checkpoints.
4
-
5
-Run `pytest tests/quantization/test_gptq_v2.py --forked`.
6
-
7
-Note: 2/3-bit GPTQ models are no longer supported after the consolidation
8
-to Marlin kernels. Only 4/8-bit symmetric GPTQ models are supported.
9
-"""
10
-
11
-import pytest
12
-import torch
13
-from transformers import AutoTokenizer
14
-
15
-from vllm import SamplingParams
16
-from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQLinearMethod
17
-
18
-# A dummy small model quantized by GPTQModel, stored in GPTQ v2 format
19
-# Note: This is a 2-bit model which is no longer supported with Marlin kernels
20
-MODELS = ["XXXXyu/Qwen3-1.7B-w2g64-gptq_v2"]
21
-
22
-# Generate multiple sequences for testing, because an 1.7B 2-bit model
23
-# cannot always generate normal texts.
24
-N_SEQ = 5
25
-
26
-
27
-@pytest.mark.skip(reason="2-bit GPTQ is no longer supported after Marlin consolidation")
28
-@pytest.mark.parametrize("model_id", MODELS)
29
-def test_model_load(vllm_runner, model_id, monkeypatch):
30
- # `LLM.apply_model` requires pickling a function.
31
- monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
32
-
33
- with vllm_runner(model_id, dtype=torch.float16, max_model_len=512) as llm:
34
-
35
- def check_model(model_id):
36
- for name, submodule in model_id.named_modules():
37
- # Could check more modules if necessary
38
- if name == "model_id.layers.0.self_attn.qkv_proj":
39
- assert isinstance(submodule.quant_method, AutoGPTQLinearMethod)
40
- # Just break since currently we only check 1 module
41
- break
42
-
43
- # Check if gptq_v2 format is correctly loaded
44
- llm.apply_model(check_model)
45
-
46
-
47
-@pytest.mark.skip(reason="2-bit GPTQ is no longer supported after Marlin consolidation")
48
-@pytest.mark.parametrize("model_id", MODELS)
49
-def test_model_inference(vllm_runner, model_id):
50
- # Prepare prompt to test the model's generation result.
51
- prompt = "What is the meaning of life?"
52
- messages = [
53
- {"role": "system", "content": "You are a helpful assistant."},
54
- {"role": "user", "content": prompt},
55
- ]
56
- tokenizer = AutoTokenizer.from_pretrained(model_id)
57
- text = tokenizer.apply_chat_template(
58
- messages,
59
- tokenize=False,
60
- add_generation_prompt=True,
61
- enable_thinking=False, # If thinking model, set it to false
62
- )
63
- sampling_params = SamplingParams(
64
- n=N_SEQ,
65
- max_tokens=128,
66
- temperature=0.7,
67
- top_p=0.8,
68
- top_k=20,
69
- min_p=0,
70
- presence_penalty=2,
71
- )
72
-
73
- with vllm_runner(model_id, dtype=torch.float16, max_model_len=512) as llm:
74
- # Generate a response to verify inference correctness
75
- output = llm.generate(text, sampling_params)
76
-
77
- # Make sure the output exists
78
- assert output
79
- assert output[0][1]
80
- assert len(output[0][1]) == N_SEQ
81
-
82
- def has_normal_char_distribution(texts, min_len):
83
- for text in texts:
84
- # Response too short
85
- if len(text) < min_len:
86
- return False
87
-
88
- # Basic ratio checks
89
- letters = sum(c.isalpha() for c in text)
90
- spaces = sum(c.isspace() for c in text)
91
- total = len(text)
92
-
93
- letter_ratio = letters / total
94
- space_ratio = spaces / total
95
-
96
- # At least 1 normal text should exist within output sequences
97
- # Normal text should be mostly letters with reasonable spacing
98
- # Some magic numbers, could be adjusted
99
- if 0.5 <= letter_ratio <= 0.9 and 0.01 <= space_ratio <= 0.3:
100
- return True
101
- # No sequence contains normal text, output might be broken
102
- return False
103
-
104
- # Apply some simple checks for giberish output
105
- # Print the output sequences if failed
106
- assert has_normal_char_distribution(output[0][1], 5), output[0][1]
tests/v1/core/test_async_scheduler.py
-36
@@ -171,42 +171,6 @@ def test_connector_metadata_precedes_async_placeholder_advance(monkeypatch):
171
assert observed_placeholders == [0, 1]
172
173
174
-def test_preempt():
175
- scheduler = create_scheduler(async_scheduling=True)
176
- requests = create_requests(num_requests=10, max_tokens=20)
177
-
178
- for req in requests:
179
- scheduler.add_request(req)
180
-
181
- sched_outputs: deque[SchedulerOutput] = deque()
182
- sched_outputs.append(scheduler.schedule())
183
- sched_outputs.append(scheduler.schedule())
184
-
185
- abort_order = [0, 8, 3, 1, 6, 4, 2, 5, 7, 9]
186
- abort_order_copy = abort_order.copy()
187
-
188
- def abort_request():
189
- if not abort_order:
190
- return
191
- req = requests[abort_order.pop(0)]
192
- scheduler.finish_requests(req.request_id, RequestStatus.FINISHED_ABORTED)
193
-
194
- while sched_outputs:
195
- # Abort a scheduled request.
196
- abort_request()
197
- sched_output = sched_outputs.popleft()
198
- model_runner_output = _make_model_runner_output(sched_output)
199
- scheduler.update_from_output(sched_output, model_runner_output)
200
-
201
- sched_output = scheduler.schedule()
202
- if sched_output.num_scheduled_tokens:
203
- sched_outputs.append(sched_output)
204
-
205
- for i, req in enumerate(requests):
206
- assert req.status == RequestStatus.FINISHED_ABORTED
207
- assert req.num_output_tokens == abort_order_copy.index(i)
208
-
209
-
174
def test_prefix_caching_for_prefill_dedup():
175
CHUNK_SIZE = 1000
176
BLOCK_SIZE = 16
tests/v1/core/test_kv_sharing.py
-32
@@ -16,38 +16,6 @@ def new_kv_cache_spec():
16
)
17
18
19
-def test_initialize_kv_cache_for_kv_sharing_different_attn_groups():
20
- """Test initializing KV cache sharing with different attention groups.
21
- Layers in the same KV cache group might be placed in different attn groups
22
- if they have different attention backends.
23
- """
24
- shared_kv_cache_layers = {
25
- "model.layers.2": "model.layers.0",
26
- "model.layers.3": "model.layers.1",
27
- }
28
-
29
- # Layers 0 and 1 both belong in KV cache group 0
30
- # However, if they have different attention backends, they will be
31
- # placed in different attention groups for KV cache group 0
32
- kv_cache_groups = [
33
- KVCacheGroupSpec(["model.layers.0", "model.layers.1"], new_kv_cache_spec()),
34
- ]
35
-
36
- add_kv_sharing_layers_to_kv_cache_groups(
37
- shared_kv_cache_layers=shared_kv_cache_layers,
38
- kv_cache_groups=kv_cache_groups,
39
- )
40
-
41
- # Check that the layers were added to the correct KV cache group
42
- assert len(kv_cache_groups) == 1
43
- assert kv_cache_groups[0].layer_names == [
44
- "model.layers.0",
45
- "model.layers.1",
46
- "model.layers.2",
47
- "model.layers.3",
48
- ]
49
-
50
-
19
def test_initialize_kv_cache_for_kv_sharing_same_attn_groups():
20
"""Test case assuming that all layers in the same KV cache group have the same
21
attention backends. This is true for most models.
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
-53
@@ -3440,59 +3440,6 @@ class TestEagle:
3440
# whole request.
3441
assert sched._lookup(req_status) == 4
3442
3443
- def test_eagle_verified_survives_eagle_tighten(self, request_runner):
3444
- """Eagle group tightening does NOT clear eagle_verified.
3445
-
3446
- Groups: 0=non-eagle full-attn, 1=eagle full-attn.
3447
- Group 0 finds 3 hits (max_hit=12). Group 1 finds 3 hits, pops to 2
3448
- (max_hit=8). Since group 1 IS eagle, eagle_verified is NOT cleared.
3449
- Result: 8 tokens (eagle only pops once).
3450
- """
3451
- block_size = 4
3452
- groups = [
3453
- KVCacheGroupSpec(
3454
- ["layer0"],
3455
- FullAttentionSpec(
3456
- block_size=block_size,
3457
- num_kv_heads=1,
3458
- head_size=1,
3459
- dtype=torch.float32,
3460
- ),
3461
- is_eagle_group=False,
3462
- ),
3463
- KVCacheGroupSpec(
3464
- ["layer1"],
3465
- FullAttentionSpec(
3466
- block_size=block_size,
3467
- num_kv_heads=2,
3468
- head_size=1,
3469
- dtype=torch.float32,
3470
- ),
3471
- is_eagle_group=True,
3472
- ),
3473
- ]
3474
- runner = request_runner(
3475
- block_size=block_size,
3476
- num_gpu_blocks=100,
3477
- async_scheduling=False,
3478
- kv_cache_groups=groups,
3479
- )
3480
- runner.manager.lookup.side_effect = lambda key, req_context: (
3481
- LookupResult.HIT
3482
- if int(get_offload_block_hash(key).decode()) in {1, 2, 3}
3483
- else LookupResult.MISS
3484
- )
3485
- sched = runner.connector_scheduler
3486
- req_status = self._make_req_status(
3487
- sched,
3488
- num_tokens=12,
3489
- offload_keys_per_group=[[1, 2, 3], [1, 2, 3]],
3490
- )
3491
- # Group 0: 3 hits → max_hit=12, num_hit=12
3492
- # Group 1 (eagle): 3 hits, pop to 2 → max_hit=8, num_hit=8
3493
- # Tightened but IS eagle → no clear. No re-iteration.
3494
- assert sched._lookup(req_status) == 8
3495
-
3443
# -------------------------------------------------------------------
3444
# Integration tests: store and load via request_runner
3445
# -------------------------------------------------------------------
tests/v1/sample/test_topk_topp_sampler.py
-73
@@ -167,79 +167,6 @@ def test_topk_impl_equivalence():
167
assert torch.allclose(result1, result2)
168
169
170
-@pytest.mark.skip(
171
- reason="FlashInfer top-k/top-p renorm comparison fails; "
172
- "needs investigation of tolerance threshold or "
173
- "interface differences between Python and FlashInfer implementations"
174
-)
175
-def test_flashinfer_sampler():
176
- """This test verifies that the FlashInfer top-k and top-p sampling
177
- implementation produces the same results as the Python implementation.
178
-
179
- NOTE: FlashInfer did not directly expose an interface for fused top-k and
180
- top-p prob renorm (it did provide fused sampling but we cannot compare
181
- sampling results due to randomness), so we will compare the probability
182
- renormed consequently by top-k and then top-p of FlashInfer implementation.
183
- """
184
- try:
185
- from flashinfer.sampling import top_k_renorm_probs, top_p_renorm_probs
186
-
187
- is_flashinfer_available = True
188
- except ImportError:
189
- is_flashinfer_available = False
190
-
191
- FLASHINFER_ENABLED = current_platform.is_cuda() and is_flashinfer_available
192
-
193
- if not FLASHINFER_ENABLED:
194
- pytest.skip("FlashInfer not installed or not available on this platform.")
195
-
196
- torch.set_default_device(DEVICE_TYPE)
197
- generator = Generator(device=DEVICE_TYPE).manual_seed(42)
198
-
199
- # Generate random logits
200
- logits = torch.rand((BATCH_SIZE, VOCAB_SIZE), generator=generator)
201
-
202
- # Generate various top-k and top-p values
203
- k_values = torch.randint(1, 1000, (BATCH_SIZE,), generator=generator)
204
- p_values = (
205
- torch.rand((BATCH_SIZE,), generator=generator) * 0.5 + 0.5
206
- ) # range in [0.5, 1.0]
207
-
208
- # Sometimes disable top-k (k=vocab_size)
209
- k_values.masked_fill_(
210
- torch.randint(0, 2, (BATCH_SIZE,), generator=generator, dtype=torch.bool),
211
- VOCAB_SIZE,
212
- )
213
-
214
- # Sometimes disable top-p (p=1.0)
215
- p_values.masked_fill_(
216
- torch.randint(0, 2, (BATCH_SIZE,), generator=generator, dtype=torch.bool), 1.0
217
- )
218
-
219
- python_logits = apply_top_k_top_p_pytorch(
220
- logits=logits.clone(),
221
- k=k_values,
222
- p=p_values,
223
- )
224
- python_probs = torch.softmax(python_logits, dim=-1)
225
-
226
- # FlashInfer only exposed renorm interfaces for probs so convert first
227
- flashinfer_probs = torch.softmax(logits.clone(), dim=-1)
228
- flashinfer_probs = top_k_renorm_probs(
229
- probs=flashinfer_probs,
230
- top_k=k_values,
231
- )
232
- flashinfer_probs = top_p_renorm_probs(
233
- probs=flashinfer_probs,
234
- top_p=p_values,
235
- )
236
-
237
- # Compare the results
238
- assert torch.allclose(python_probs, flashinfer_probs, atol=2e-2), (
239
- "FlashInfer and Python sampling implementations do not match!"
240
- )
241
-
242
-
170
@pytest.mark.skipif(
171
not FLASHINFER_TOPK_TOPP_SUPPORTED,
172
reason="Requires CUDA and flashinfer",
tests/v1/spec_decode/test_backup_token_async_spec.py
deleted
-147
@@ -1,147 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Regression tests for the backup token fix in prepare_next_token_ids_padded.
4
-
5
-Fixes #38098: with async scheduling, seq_lens_cpu is inflated by unaccepted
6
-draft token placeholders, causing get_token_id() to return -1.
7
-"""
8
-
9
-from __future__ import annotations
10
-
11
-import numpy as np
12
-import pytest
13
-import torch
14
-
15
-
16
-class _FakeRequest:
17
- def __init__(self, prompt_tokens: list[int], output_tokens: list[int]):
18
- self.num_prompt_tokens = len(prompt_tokens)
19
- self._prompt = prompt_tokens
20
- self._output = output_tokens
21
-
22
- @property
23
- def num_tokens(self) -> int:
24
- return self.num_prompt_tokens + len(self._output)
25
-
26
- def get_token_id(self, idx: int) -> int:
27
- if idx < self.num_prompt_tokens:
28
- return self._prompt[idx]
29
- out_idx = idx - self.num_prompt_tokens
30
- if out_idx < len(self._output):
31
- return self._output[out_idx]
32
- return -1 # out of range
33
-
34
-
35
-class _FakeInputBatch:
36
- def __init__(
37
- self,
38
- req_ids: list[str],
39
- num_tokens_no_spec: list[int],
40
- vocab_size: int = 32000,
41
- ):
42
- self.req_ids = req_ids
43
- self.num_reqs = len(req_ids)
44
- self.vocab_size = vocab_size
45
- self.num_tokens_no_spec = np.array(num_tokens_no_spec, dtype=np.int64)
46
-
47
-
48
-def _make_requests(
49
- req_ids: list[str],
50
- prompt_lens: list[int],
51
- output_lens: list[int],
52
-) -> dict[str, _FakeRequest]:
53
- requests = {}
54
- for rid, plen, olen in zip(req_ids, prompt_lens, output_lens):
55
- requests[rid] = _FakeRequest(list(range(plen)), list(range(1000, 1000 + olen)))
56
- return requests
57
-
58
-
59
-def _backup_buggy(
60
- seq_lens_cpu: torch.Tensor,
61
- requests: dict[str, _FakeRequest],
62
- batch: _FakeInputBatch,
63
-) -> list[int]:
64
- """Old logic: uses seq_lens_cpu directly (may be inflated)."""
65
- n = batch.num_reqs
66
- return [
67
- requests[batch.req_ids[i]].get_token_id(int(seq_lens_cpu[i])) for i in range(n)
68
- ]
69
-
70
-
71
-def _backup_fixed(
72
- requests: dict[str, _FakeRequest],
73
- batch: _FakeInputBatch,
74
-) -> list[int]:
75
- """New logic: uses num_tokens_no_spec - 1 (last committed token)."""
76
- n = batch.num_reqs
77
- idx = (batch.num_tokens_no_spec[:n] - 1).tolist()
78
- return [requests[batch.req_ids[i]].get_token_id(int(idx[i])) for i in range(n)]
79
-
80
-
81
-class TestBackupTokenAsyncSpec:
82
- def test_no_inflation_fixed_returns_last_token(self):
83
- req_ids = ["r0", "r1"]
84
- requests = _make_requests(req_ids, [3, 3], [2, 2])
85
- batch = _FakeInputBatch(req_ids, [5, 5])
86
- # idx = 5-1 = 4 → output[1] = 1001
87
- assert _backup_fixed(requests, batch) == [1001, 1001]
88
-
89
- def test_inflation_buggy_returns_placeholder(self):
90
- req_ids = ["r0", "r1"]
91
- requests = _make_requests(req_ids, [3, 3], [2, 2])
92
- batch = _FakeInputBatch(req_ids, [5, 5])
93
- # inflated by 3 spec tokens → idx 8 is out of range
94
- seq_lens = torch.tensor([8, 8], dtype=torch.int64)
95
- assert _backup_buggy(seq_lens, requests, batch) == [-1, -1]
96
-
97
- def test_inflation_fixed_returns_correct_token(self):
98
- req_ids = ["r0", "r1"]
99
- requests = _make_requests(req_ids, [3, 3], [2, 2])
100
- batch = _FakeInputBatch(req_ids, [5, 5])
101
- assert _backup_fixed(requests, batch) == [1001, 1001]
102
-
103
- def test_mixed_inflation_per_request(self):
104
- req_ids = ["r0", "r1", "r2"]
105
- requests = {
106
- "r0": _FakeRequest([0, 1], [1000, 1001, 1002]),
107
- "r1": _FakeRequest([0, 1, 2, 3], [2000]),
108
- "r2": _FakeRequest([0], [3000, 3001, 3002, 3003]),
109
- }
110
- batch = _FakeInputBatch(req_ids, [5, 5, 5])
111
- seq_lens = torch.tensor([7, 9, 5], dtype=torch.int64)
112
-
113
- assert _backup_buggy(seq_lens, requests, batch) == [-1, -1, -1]
114
- assert _backup_fixed(requests, batch) == [1002, 2000, 3003]
115
-
116
- def test_prefill_only_request(self):
117
- """No output tokens yet — backup should be the last prompt token."""
118
- req_ids = ["r0"]
119
- requests = {"r0": _FakeRequest([10, 20, 30], [])}
120
- batch = _FakeInputBatch(req_ids, [3])
121
- # idx = 3-1 = 2 → prompt[2] = 30
122
- assert _backup_fixed(requests, batch) == [30]
123
-
124
- @pytest.mark.parametrize("num_spec_tokens", [1, 2, 3, 4, 5])
125
- def test_various_spec_token_counts(self, num_spec_tokens: int):
126
- req_ids = ["r0"]
127
- requests = {"r0": _FakeRequest([0, 1, 2], list(range(1000, 1005)))}
128
- batch = _FakeInputBatch(req_ids, [8])
129
- # idx = 8-1 = 7 → output[4] = 1004
130
- assert _backup_fixed(requests, batch) == [1004]
131
-
132
- def test_buggy_code_was_always_off_by_one(self):
133
- """The original code used seq_len as index, which is always one past
134
- the end of output_token_ids even without async inflation."""
135
- req_ids = ["r0"]
136
- requests = {"r0": _FakeRequest([0, 1, 2], [1000, 1001])}
137
- batch = _FakeInputBatch(req_ids, [5])
138
-
139
- # no inflation: seq_len == num_tokens == 5 → idx 5 is out of range
140
- seq_lens = torch.tensor([5], dtype=torch.int64)
141
- assert _backup_buggy(seq_lens, requests, batch) == [-1]
142
- assert _backup_fixed(requests, batch) == [1001]
143
-
144
- # with inflation: still -1, fixed still correct
145
- seq_lens_inf = torch.tensor([8], dtype=torch.int64)
146
- assert _backup_buggy(seq_lens_inf, requests, batch) == [-1]
147
- assert _backup_fixed(requests, batch) == [1001]