5225f1bacb · tduncan/Nvidia-Optimization-Stack-VLLM

tests/benchmarks/sweep/test_param_sweep.py

+8 -8

@@ -116,7 +116,7 @@ class TestParameterSweep:

116

117 def test_from_records_list(self):

118 """Test creating ParameterSweep from a list of records."""

119 - records = [

119 + records: list[dict[str, object]] = [

120 {"max_tokens": 100, "temperature": 0.7},

121 {"max_tokens": 200, "temperature": 0.9},

122 ]

@@ -127,7 +127,7 @@ class TestParameterSweep:

127

128 def test_read_from_dict(self):

129 """Test creating ParameterSweep from a dict format."""

130 - data = {

130 + data: dict[str, dict[str, object]] = {

131 "experiment1": {"max_tokens": 100, "temperature": 0.7},

132 "experiment2": {"max_tokens": 200, "temperature": 0.9},

133 }

@@ -149,7 +149,7 @@ class TestParameterSweep:

149

150 def test_read_json_list_format(self):

151 """Test reading JSON file with list format."""

152 - records = [

152 + records: list[dict[str, object]] = [

153 {"max_tokens": 100, "temperature": 0.7},

154 {"max_tokens": 200, "temperature": 0.9},

155 ]

@@ -168,7 +168,7 @@ class TestParameterSweep:

168

169 def test_read_json_dict_format(self):

170 """Test reading JSON file with dict format."""

171 - data = {

171 + data: dict[str, dict[str, object]] = {

172 "experiment1": {"max_tokens": 100, "temperature": 0.7},

173 "experiment2": {"max_tokens": 200, "temperature": 0.9},

174 }

@@ -190,7 +190,7 @@ class TestParameterSweep:

190 def test_unique_benchmark_names_validation(self):

191 """Test that duplicate _benchmark_name values raise an error."""

192 # Test with duplicate names in list format

193 - records = [

193 + records: list[dict[str, object]] = [

194 {"_benchmark_name": "exp1", "max_tokens": 100},

195 {"_benchmark_name": "exp1", "max_tokens": 200},

196 ]

@@ -200,7 +200,7 @@ class TestParameterSweep:

200

201 def test_unique_benchmark_names_multiple_duplicates(self):

202 """Test validation with multiple duplicate names."""

203 - records = [

203 + records: list[dict[str, object]] = [

204 {"_benchmark_name": "exp1", "max_tokens": 100},

205 {"_benchmark_name": "exp1", "max_tokens": 200},

206 {"_benchmark_name": "exp2", "max_tokens": 300},

@@ -212,7 +212,7 @@ class TestParameterSweep:

212

213 def test_no_benchmark_names_allowed(self):

214 """Test that records without _benchmark_name are allowed."""

215 - records = [

215 + records: list[dict[str, object]] = [

216 {"max_tokens": 100, "temperature": 0.7},

217 {"max_tokens": 200, "temperature": 0.9},

218 ]

@@ -221,7 +221,7 @@ class TestParameterSweep:

221

222 def test_mixed_benchmark_names_allowed(self):

223 """Test that mixing records with and without _benchmark_name is allowed."""

224 - records = [

224 + records: list[dict[str, object]] = [

225 {"_benchmark_name": "exp1", "max_tokens": 100},

226 {"max_tokens": 200, "temperature": 0.9},

227 ]

tests/benchmarks/test_audio_dataset.py

+10 -5

@@ -13,6 +13,7 @@ import soundfile as sf

13 import vllm.benchmarks.datasets.datasets as datasets_module

14 import vllm.benchmarks.lib.endpoint_request_func as request_func_module

15 from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput

16 +from vllm.tokenizers import TokenizerLike

17

18 pytestmark = pytest.mark.skip_global_cleanup

19

@@ -44,6 +45,10 @@ class _CohereNameOnlyTokenizer(_Tokenizer):

45 super().__init__("cohere/some-local-checkpoint")

46

47

48 +def _as_tokenizer(tokenizer: _Tokenizer) -> TokenizerLike:

49 + return cast(TokenizerLike, tokenizer)

50 +

51 +

52 def _write_wav(path: Path, duration_s: float = 0.1, sample_rate: int = 16_000) -> None:

53 num_samples = int(duration_s * sample_rate)

54 sf.write(path, np.zeros(num_samples, dtype=np.float32), sample_rate)

@@ -109,7 +114,7 @@ def test_asr_dataset_sample_handles_local_audio_paths(tmp_path: Path) -> None:

114 ]

115

116 samples = dataset.sample(

112 - tokenizer=_Tokenizer(),

117 + tokenizer=_as_tokenizer(_Tokenizer()),

118 num_requests=1,

119 output_len=32,

120 asr_min_audio_len_sec=0.0,

@@ -146,7 +151,7 @@ def test_asr_dataset_sample_handles_embedded_audio_bytes(

151 ]

152

153 samples = dataset.sample(

149 - tokenizer=_Tokenizer(),

154 + tokenizer=_as_tokenizer(_Tokenizer()),

155 num_requests=1,

156 output_len=32,

157 asr_min_audio_len_sec=0.0,

@@ -240,7 +245,7 @@ def _make_asr_dataset(tmp_path: Path) -> datasets_module.ASRDataset:

245 def test_asr_dataset_cohere_class_name_gets_decoder_prompt(tmp_path: Path) -> None:

246 dataset = _make_asr_dataset(tmp_path)

247 samples = dataset.sample(

243 - tokenizer=CohereAsrTokenizer(),

248 + tokenizer=_as_tokenizer(CohereAsrTokenizer()),

249 num_requests=1,

250 output_len=32,

251 asr_min_audio_len_sec=0.0,

@@ -255,7 +260,7 @@ def test_asr_dataset_cohere_name_or_path_fallback_gets_decoder_prompt(

260 ) -> None:

261 dataset = _make_asr_dataset(tmp_path)

262 samples = dataset.sample(

258 - tokenizer=_CohereNameOnlyTokenizer(),

263 + tokenizer=_as_tokenizer(_CohereNameOnlyTokenizer()),

264 num_requests=1,

265 output_len=32,

266 asr_min_audio_len_sec=0.0,

@@ -268,7 +273,7 @@ def test_asr_dataset_cohere_name_or_path_fallback_gets_decoder_prompt(

273 def test_asr_dataset_unknown_tokenizer_gets_empty_prompt(tmp_path: Path) -> None:

274 dataset = _make_asr_dataset(tmp_path)

275 samples = dataset.sample(

271 - tokenizer=_Tokenizer(name_or_path="some-other/asr-model"),

276 + tokenizer=_as_tokenizer(_Tokenizer(name_or_path="some-other/asr-model")),

277 num_requests=1,

278 output_len=32,

279 asr_min_audio_len_sec=0.0,

tests/benchmarks/test_bfcl_dataset.py

+2 -2

@@ -271,7 +271,7 @@ def test_bfcl_prompt_len_includes_tools(tmp_path: Path) -> None:

271 args.num_prompts = 1

272

273 with _patch_hf_api(fake_download):

274 - samples = get_samples(args, fake)

274 + samples = get_samples(args, fake) # type: ignore[arg-type]

275

276 assert len(samples) == 1

277 assert captured["tools"] is not None, (

@@ -311,7 +311,7 @@ def test_bfcl_prompt_len_falls_back_when_tokenizer_rejects_tools(

311 args.num_prompts = 1

312

313 with _patch_hf_api(fake_download):

314 - samples = get_samples(args, _LegacyTokenizer())

314 + samples = get_samples(args, _LegacyTokenizer()) # type: ignore[arg-type]

315

316 assert len(samples) == 1

317 assert samples[0].prompt_len > 0

tests/benchmarks/test_custom_dataset_chat_template_kwargs.py

+2 -2

@@ -58,7 +58,7 @@ def test_chat_template_kwargs_forwarded(tmp_path: Path) -> None:

58 _write_one(jsonl)

59

60 tok = _RecordingTokenizer()

61 - get_samples(_args(str(jsonl), {"thinking": True}), tok)

61 + get_samples(_args(str(jsonl), {"thinking": True}), tok) # type: ignore[arg-type]

62

63 assert tok.captured_kwargs == {"thinking": True}

64

@@ -70,6 +70,6 @@ def test_chat_template_kwargs_default_is_noop(tmp_path: Path) -> None:

70 _write_one(jsonl)

71

72 tok = _RecordingTokenizer()

73 - get_samples(_args(str(jsonl), None), tok)

73 + get_samples(_args(str(jsonl), None), tok) # type: ignore[arg-type]

74

75 assert tok.captured_kwargs == {}

tests/benchmarks/test_custom_image_dataset.py

+29 -15

@@ -4,7 +4,7 @@ import json

4 from argparse import Namespace

5 from io import BytesIO

6 from pathlib import Path

7 -from typing import Any

7 +from typing import Any, cast

8

9 import pybase64 as base64

10 import pytest

@@ -17,6 +17,7 @@ from vllm.benchmarks.lib.endpoint_request_func import (

17 _get_chat_content,

18 _get_chat_messages,

19 )

20 +from vllm.tokenizers import TokenizerLike

21

22 pytestmark = pytest.mark.skip_global_cleanup

23

@@ -31,6 +32,10 @@ class _Tokenizer:

32 return _TokenizedPrompt(prompt)

33

34

35 +def _tokenizer() -> TokenizerLike:

36 + return cast(TokenizerLike, _Tokenizer())

37 +

38 +

39 def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:

40 with path.open("w") as f:

41 for row in rows:

@@ -92,7 +97,7 @@ def test_get_samples_custom_image_cli_path_supports_multi_image_and_content(

97 ],

98 )

99

95 - samples = get_samples(_args_for_custom_image(jsonl), _Tokenizer())

100 + samples = get_samples(_args_for_custom_image(jsonl), _tokenizer())

101

102 assert len(samples) == 2

103 assert samples[0].request_id == "req-0"

@@ -106,6 +111,7 @@ def test_get_samples_custom_image_cli_path_supports_multi_image_and_content(

111 assert samples[1].multi_modal_data is None

112 assert isinstance(samples[1].prompt, list)

113 assert samples[1].prompt[0] == {"type": "text", "text": "Now compare "}

114 + assert isinstance(samples[1].prompt[1], dict)

115 assert samples[1].prompt[1]["image_url"]["url"] == f"file://{image_c}"

116

117

@@ -126,7 +132,7 @@ def test_custom_image_dataset_uses_all_image_files(tmp_path: Path) -> None:

132

133 dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)

134 samples = dataset.sample(

129 - tokenizer=_Tokenizer(),

135 + tokenizer=_tokenizer(),

136 num_requests=1,

137 output_len=32,

138 )

@@ -171,7 +177,7 @@ def test_custom_image_dataset_preserves_interleaved_content_order(

177

178 dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)

179 samples = dataset.sample(

174 - tokenizer=_Tokenizer(),

180 + tokenizer=_tokenizer(),

181 num_requests=1,

182 output_len=32,

183 )

@@ -181,14 +187,18 @@ def test_custom_image_dataset_preserves_interleaved_content_order(

187 assert sample.multi_modal_data is None

188 assert sample.prompt_len == 2

189 assert isinstance(sample.prompt, list)

184 - assert [part["type"] for part in sample.prompt] == [

190 + parts = []

191 + for part in sample.prompt:

192 + assert isinstance(part, dict)

193 + parts.append(part)

194 + assert [part["type"] for part in parts] == [

195 "text",

196 "image_url",

197 "text",

198 "image_url",

199 ]

190 - assert sample.prompt[1]["image_url"]["url"] == f"file://{image_a}"

191 - assert sample.prompt[3]["image_url"] == {

200 + assert parts[1]["image_url"]["url"] == f"file://{image_a}"

201 + assert parts[3]["image_url"] == {

202 "url": f"file://{image_b}",

203 "detail": "low",

204 }

@@ -223,7 +233,7 @@ def test_custom_image_dataset_wraps_interleaved_content_for_multimodal_chat(

233

234 dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)

235 samples = dataset.sample(

226 - tokenizer=_Tokenizer(),

236 + tokenizer=_tokenizer(),

237 num_requests=1,

238 output_len=32,

239 enable_multimodal_chat=True,

@@ -289,7 +299,7 @@ def test_custom_image_dataset_encodes_image_media_when_requested(

299

300 dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)

301 samples = dataset.sample(

292 - tokenizer=_Tokenizer(),

302 + tokenizer=_tokenizer(),

303 num_requests=1,

304 output_len=32,

305 ensure_client_side_data=True,

@@ -335,7 +345,7 @@ def test_custom_image_dataset_encodes_interleaved_image_media(

345

346 dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)

347 samples = dataset.sample(

338 - tokenizer=_Tokenizer(),

348 + tokenizer=_tokenizer(),

349 num_requests=1,

350 output_len=32,

351 ensure_client_side_data=True,

@@ -343,9 +353,13 @@ def test_custom_image_dataset_encodes_interleaved_image_media(

353

354 sample = samples[0]

355 assert isinstance(sample.prompt, list)

346 - _assert_png_data_url(sample.prompt[1]["image_url"]["url"])

347 - _assert_png_data_url(sample.prompt[2]["image_url"]["url"])

348 - assert sample.prompt[2]["image_url"]["detail"] == "low"

356 + parts = []

357 + for part in sample.prompt:

358 + assert isinstance(part, dict)

359 + parts.append(part)

360 + _assert_png_data_url(parts[1]["image_url"]["url"])

361 + _assert_png_data_url(parts[2]["image_url"]["url"])

362 + assert parts[2]["image_url"]["detail"] == "low"

363

364

365 @pytest.mark.benchmark

@@ -363,7 +377,7 @@ def test_custom_image_dataset_rejects_invalid_image_media(

377 dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)

378 with pytest.raises(ValueError, match="Invalid image URL"):

379 dataset.sample(

366 - tokenizer=_Tokenizer(),

380 + tokenizer=_tokenizer(),

381 num_requests=1,

382 output_len=32,

383 ensure_client_side_data=True,

@@ -380,7 +394,7 @@ def test_custom_image_dataset_rejects_invalid_content_part(

394 dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)

395 with pytest.raises(ValueError, match="type 'text', 'image', or 'image_url'"):

396 dataset.sample(

383 - tokenizer=_Tokenizer(),

397 + tokenizer=_tokenizer(),

398 num_requests=1,

399 output_len=32,

400 )

tests/benchmarks/test_random_dataset.py

+2

@@ -37,6 +37,7 @@ def random_dataset_params() -> Params:

37

38 def _fingerprint_sample(req: SampleRequest) -> tuple[str, int, int]:

39 """Project a SampleRequest into a comparable tuple."""

40 + assert isinstance(req.prompt, str)

41 return (req.prompt, req.prompt_len, req.expected_output_len)

42

43

@@ -153,6 +154,7 @@ def _mm_fingerprint_sample(

154 - count of multimodal items

155 - per-item type and URL prefix (e.g., 'data:image/jpeg;base64,')

156 """

157 + assert isinstance(req.prompt, str)

158 items = req.multi_modal_data or []

159 item_prefixes: list[str] = []

160 for it in items:

tests/benchmarks/test_sampling_params.py

+4 -4

@@ -19,7 +19,7 @@ class _FakeTokenizer(TokenizerLike):

19 # -- Properties required by TokenizerLike --

20

21 @classmethod

22 - def from_pretrained(cls, path_or_repo_id, *a, **kw): # type: ignore[override]

22 + def from_pretrained(cls, path_or_repo_id, *a, **kw):

23 return cls()

24

25 @property

@@ -80,16 +80,16 @@ class _FakeTokenizer(TokenizerLike):

80 def apply_chat_template(self, messages, **kw): # type: ignore[override]

81 raise NotImplementedError

82

83 - def convert_tokens_to_ids(self, tokens): # type: ignore[override]

83 + def convert_tokens_to_ids(self, tokens):

84 raise NotImplementedError

85

86 def convert_tokens_to_string(self, tokens: list[str]) -> str:

87 raise NotImplementedError

88

89 - def decode(self, ids, skip_special_tokens: bool = False) -> str: # type: ignore[override]

89 + def decode(self, ids, skip_special_tokens: bool = False) -> str:

90 raise NotImplementedError

91

92 - def convert_ids_to_tokens( # type: ignore[override]

92 + def convert_ids_to_tokens(

93 self, ids, skip_special_tokens: bool = False

94 ) -> list[str]:

95 raise NotImplementedError

tests/benchmarks/test_throughput_cli.py

+3 -2

@@ -9,6 +9,7 @@ serve-side dataset coverage.

9 """

10

11 import subprocess

12 +from argparse import Namespace

13 from types import SimpleNamespace

14

15 import pytest

@@ -289,7 +290,7 @@ def test_to_serve_args_supports_mm_processor_flags() -> None:

290 explicit 0 instead of coercing it to None (which later crashed

291 RandomMultiModalDataset.get_prefix on ``None <= 0``).

292 """

292 - args = SimpleNamespace(

293 + args = Namespace(

294 random_input_len=32,

295 random_output_len=1,

296 random_prefix_len=0,

@@ -306,7 +307,7 @@ def test_to_serve_args_supports_mm_processor_flags() -> None:

307

308 def test_to_serve_args_keeps_legacy_flag_fallback() -> None:

309 """Callers that only set the legacy flags still get them mapped."""

309 - args = SimpleNamespace(

310 + args = Namespace(

311 input_len=256,

312 output_len=8,

313 prefix_len=16,

tests/entrypoints/anthropic/test_anthropic_messages_conversion.py

+14 -6

@@ -355,9 +355,11 @@ class TestToolResultContent:

355 if m["role"] == "user" and isinstance(m.get("content"), list)

356 ]

357 assert len(follow_up) == 1

358 - assert follow_up[0]["content"][0]["image_url"]["url"] == (

359 - "data:image/jpeg;base64,QUFB"

360 - )

358 + content = follow_up[0]["content"]

359 + assert isinstance(content, list)

360 + image = content[0]

361 + assert image["type"] == "image_url"

362 + assert image["image_url"]["url"] == "data:image/jpeg;base64,QUFB"

363

364 def test_tool_result_with_multiple_images(self):

365 request = self._make_tool_result_request(

@@ -387,7 +389,12 @@ class TestToolResultContent:

389 if m["role"] == "user" and isinstance(m.get("content"), list)

390 ]

391 assert len(follow_up) == 1

390 - urls = [p["image_url"]["url"] for p in follow_up[0]["content"]]

392 + content = follow_up[0]["content"]

393 + assert isinstance(content, list)

394 + urls = []

395 + for part in content:

396 + assert part["type"] == "image_url"

397 + urls.append(part["image_url"]["url"])

398 assert urls == [

399 "data:image/png;base64,IMG1",

400 "https://example.com/img2.jpg",

@@ -584,7 +591,8 @@ class TestThinkingBlockConversion:

591 asst = asst_msgs[0]

592

593 assert asst.get("reasoning") == "I need to call the calculator."

587 - tool_calls = list(asst.get("tool_calls", []))

594 + tool_calls = asst.get("tool_calls")

595 + assert tool_calls is not None

596 assert len(tool_calls) == 1

597 assert tool_calls[0]["function"]["name"] == "calculator"

598 # No text content alongside reasoning + tool_use.

@@ -1439,7 +1447,7 @@ class TestDetectMergeInlineSystem:

1447 model_config=None,

1448 )

1449 assert (

1442 - AnthropicServingMessages._should_merge_inline_system(online_renderer)

1450 + AnthropicServingMessages._should_merge_inline_system(online_renderer) # type: ignore[arg-type]

1451 is expected

1452 )

1453 assert mock_logger.warning_once.called is expected

tests/entrypoints/multimodal/llm/test_mm_processor_kwargs.py

+69 -41

@@ -7,22 +7,26 @@ from unittest.mock import Mock

7 import pytest

8

9 from vllm import LLM, SamplingParams

10 +from vllm.entrypoints.chat_utils import ChatCompletionMessageParam

11 +from vllm.inputs import TextPrompt

12 +from vllm.outputs import RequestOutput

13

14

15 def _make_mock_llm() -> LLM:

16 llm = object.__new__(LLM)

14 - llm.model_config = SimpleNamespace(

17 + llm.model_config = SimpleNamespace( # type: ignore[assignment]

18 runner_type="generate", enable_prompt_embeds=False

19 )

20 return llm

21

22

20 -def test_generate_forwards_mm_processor_kwargs() -> None:

23 +def test_generate_forwards_mm_processor_kwargs(monkeypatch: pytest.MonkeyPatch) -> None:

24 llm = _make_mock_llm()

25 mm_processor_kwargs = {"num_crops": 4}

26 sampling_params = SamplingParams(max_tokens=1)

27

25 - llm._run_completion = Mock(return_value=["ok"])

28 + _run_completion = Mock(return_value=["ok"])

29 + monkeypatch.setattr(llm, "_run_completion", _run_completion)

30

31 outputs = llm.generate(

32 "prompt",

@@ -31,17 +35,18 @@ def test_generate_forwards_mm_processor_kwargs() -> None:

35 )

36

37 assert outputs == ["ok"]

34 - assert llm._run_completion.call_args.kwargs["mm_processor_kwargs"] == (

38 + assert _run_completion.call_args.kwargs["mm_processor_kwargs"] == (

39 mm_processor_kwargs

40 )

41

42

39 -def test_enqueue_forwards_mm_processor_kwargs() -> None:

43 +def test_enqueue_forwards_mm_processor_kwargs(monkeypatch: pytest.MonkeyPatch) -> None:

44 llm = _make_mock_llm()

45 mm_processor_kwargs = {"do_resize": False}

46 sampling_params = SamplingParams(max_tokens=1)

47

44 - llm._add_completion_requests = Mock(return_value=["req-0"])

48 + _add_completion_requests = Mock(return_value=["req-0"])

49 + monkeypatch.setattr(llm, "_add_completion_requests", _add_completion_requests)

50

51 request_ids = llm.enqueue(

52 "prompt",

@@ -51,18 +56,19 @@ def test_enqueue_forwards_mm_processor_kwargs() -> None:

56 )

57

58 assert request_ids == ["req-0"]

54 - assert llm._add_completion_requests.call_args.kwargs["mm_processor_kwargs"] == (

59 + assert _add_completion_requests.call_args.kwargs["mm_processor_kwargs"] == (

60 mm_processor_kwargs

61 )

62

63

59 -def test_chat_forwards_mm_processor_kwargs() -> None:

64 +def test_chat_forwards_mm_processor_kwargs(monkeypatch: pytest.MonkeyPatch) -> None:

65 llm = _make_mock_llm()

66 mm_processor_kwargs = {"do_pan_and_scan": True}

67 sampling_params = SamplingParams(max_tokens=1)

63 - messages = [{"role": "user", "content": "hello"}]

68 + messages: list[ChatCompletionMessageParam] = [{"role": "user", "content": "hello"}]

69

65 - llm._run_chat = Mock(return_value=["ok"])

70 + _run_chat = Mock(return_value=["ok"])

71 + monkeypatch.setattr(llm, "_run_chat", _run_chat)

72

73 outputs = llm.chat(

74 messages,

@@ -71,18 +77,19 @@ def test_chat_forwards_mm_processor_kwargs() -> None:

77 )

78

79 assert outputs == ["ok"]

74 - assert llm._run_chat.call_args.kwargs["mm_processor_kwargs"] == (

75 - mm_processor_kwargs

76 - )

80 + assert _run_chat.call_args.kwargs["mm_processor_kwargs"] == mm_processor_kwargs

81

82

79 -def test_enqueue_chat_forwards_mm_processor_kwargs() -> None:

83 +def test_enqueue_chat_forwards_mm_processor_kwargs(

84 + monkeypatch: pytest.MonkeyPatch,

85 +) -> None:

86 llm = _make_mock_llm()

87 mm_processor_kwargs = {"do_pan_and_scan": True}

88 sampling_params = SamplingParams(max_tokens=1)

83 - messages = [{"role": "user", "content": "hello"}]

89 + messages: list[ChatCompletionMessageParam] = [{"role": "user", "content": "hello"}]

90

85 - llm._add_chat_requests = Mock(return_value=["req-0"])

91 + _add_chat_requests = Mock(return_value=["req-0"])

92 + monkeypatch.setattr(llm, "_add_chat_requests", _add_chat_requests)

93

94 request_ids = llm.enqueue_chat(

95 messages,

@@ -92,67 +99,79 @@ def test_enqueue_chat_forwards_mm_processor_kwargs() -> None:

99 )

100

101 assert request_ids == ["req-0"]

95 - assert llm._add_chat_requests.call_args.kwargs["mm_processor_kwargs"] == (

102 + assert _add_chat_requests.call_args.kwargs["mm_processor_kwargs"] == (

103 mm_processor_kwargs

104 )

105

106

100 -def test_run_chat_forwards_mm_processor_kwargs() -> None:

107 +def test_run_chat_forwards_mm_processor_kwargs(monkeypatch: pytest.MonkeyPatch) -> None:

108 llm = _make_mock_llm()

109 mm_processor_kwargs = {"num_crops": 8}

110 sampling_params = SamplingParams(max_tokens=1)

104 - messages = [{"role": "user", "content": "hello"}]

111 + messages: list[ChatCompletionMessageParam] = [{"role": "user", "content": "hello"}]

112 sentinel_output = ["done"]

113

107 - llm._add_chat_requests = Mock()

108 - llm._run_engine = Mock(return_value=sentinel_output)

114 + _add_chat_requests = Mock()

115 + monkeypatch.setattr(llm, "_add_chat_requests", _add_chat_requests)

116 + _run_engine = Mock(return_value=sentinel_output)

117 + monkeypatch.setattr(llm, "_run_engine", _run_engine)

118

119 outputs = llm._run_chat(

120 messages=messages,

121 params=sampling_params,

113 - output_type=object,

122 + output_type=RequestOutput,

123 use_tqdm=False,

124 mm_processor_kwargs=mm_processor_kwargs,

125 )

126

127 assert outputs == sentinel_output

119 - assert llm._add_chat_requests.call_args.kwargs["mm_processor_kwargs"] == (

128 + assert _add_chat_requests.call_args.kwargs["mm_processor_kwargs"] == (

129 mm_processor_kwargs

130 )

131

132

124 -def test_run_completion_forwards_mm_processor_kwargs() -> None:

133 +def test_run_completion_forwards_mm_processor_kwargs(

134 + monkeypatch: pytest.MonkeyPatch,

135 +) -> None:

136 llm = _make_mock_llm()

137 mm_processor_kwargs = {"min_pixels": 4 * 28 * 28}

138 sampling_params = SamplingParams(max_tokens=1)

139 sentinel_output = ["done"]

140

130 - llm._add_completion_requests = Mock()

131 - llm._run_engine = Mock(return_value=sentinel_output)

141 + _add_completion_requests = Mock()

142 + monkeypatch.setattr(llm, "_add_completion_requests", _add_completion_requests)

143 + _run_engine = Mock(return_value=sentinel_output)

144 + monkeypatch.setattr(llm, "_run_engine", _run_engine)

145

146 outputs = llm._run_completion(

147 prompts=["prompt"],

148 params=sampling_params,

136 - output_type=object,

149 + output_type=RequestOutput,

150 use_tqdm=False,

151 mm_processor_kwargs=mm_processor_kwargs,

152 )

153

154 assert outputs == sentinel_output

142 - assert llm._add_completion_requests.call_args.kwargs["mm_processor_kwargs"] == (

155 + assert _add_completion_requests.call_args.kwargs["mm_processor_kwargs"] == (

156 mm_processor_kwargs

157 )

158

159

147 -def test_add_completion_requests_forwards_mm_processor_kwargs() -> None:

160 +def test_add_completion_requests_forwards_mm_processor_kwargs(

161 + monkeypatch: pytest.MonkeyPatch,

162 +) -> None:

163 llm = _make_mock_llm()

164 mm_processor_kwargs = {"max_dynamic_patch": 4}

165 sampling_params = SamplingParams(max_tokens=1)

166

152 - llm._params_to_seq = Mock(return_value=[sampling_params])

153 - llm._lora_request_to_seq = Mock(return_value=[None])

154 - llm._priority_to_seq = Mock(return_value=[0])

155 - llm._preprocess_cmpl_one = Mock(return_value={"prompt_token_ids": [1]})

167 + _params_to_seq = Mock(return_value=[sampling_params])

168 + monkeypatch.setattr(llm, "_params_to_seq", _params_to_seq)

169 + _lora_request_to_seq = Mock(return_value=[None])

170 + monkeypatch.setattr(llm, "_lora_request_to_seq", _lora_request_to_seq)

171 + _priority_to_seq = Mock(return_value=[0])

172 + monkeypatch.setattr(llm, "_priority_to_seq", _priority_to_seq)

173 + _preprocess_cmpl_one = Mock(return_value={"prompt_token_ids": [1]})

174 + monkeypatch.setattr(llm, "_preprocess_cmpl_one", _preprocess_cmpl_one)

175

176 captured_prompts = []

177

@@ -160,7 +179,8 @@ def test_add_completion_requests_forwards_mm_processor_kwargs() -> None:

179 captured_prompts.extend(prompts)

180 return ["req-0"]

181

163 - llm._render_and_add_requests = Mock(side_effect=fake_render_and_add_requests)

182 + _render_and_add_requests = Mock(side_effect=fake_render_and_add_requests)

183 + monkeypatch.setattr(llm, "_render_and_add_requests", _render_and_add_requests)

184

185 request_ids = llm._add_completion_requests(

186 prompts=["prompt"],

@@ -170,7 +190,7 @@ def test_add_completion_requests_forwards_mm_processor_kwargs() -> None:

190 )

191

192 assert request_ids == ["req-0"]

173 - llm._preprocess_cmpl_one.assert_called_once_with(

193 + _preprocess_cmpl_one.assert_called_once_with(

194 "prompt",

195 None,

196 mm_processor_kwargs=mm_processor_kwargs,

@@ -183,7 +203,7 @@ def test_preprocess_cmpl_applies_mm_processor_kwargs_to_renderer(

203 ) -> None:

204 llm = _make_mock_llm()

205 mm_processor_kwargs = {"num_crops": 8}

186 - prompt = {"prompt": "<image>", "multi_modal_data": {"image": object()}}

206 + prompt: TextPrompt = {"prompt": "<image>", "multi_modal_data": {"image": object()}}

207

208 renderer = Mock()

209 renderer.default_cmpl_tok_params = Mock()

@@ -213,7 +233,7 @@ def test_preprocess_cmpl_keeps_prompt_mm_processor_kwargs_when_no_override(

233 monkeypatch: pytest.MonkeyPatch,

234 ) -> None:

235 llm = _make_mock_llm()

216 - prompt = {

236 + prompt: TextPrompt = {

237 "prompt": "<image>",

238 "multi_modal_data": {"image": object()},

239 "mm_processor_kwargs": {"num_crops": 2},

@@ -243,7 +263,9 @@ def test_preprocess_cmpl_keeps_prompt_mm_processor_kwargs_when_no_override(

263 def test_preprocess_chat_applies_mm_processor_kwargs_to_renderer() -> None:

264 llm = _make_mock_llm()

265 mm_processor_kwargs = {"num_crops": 8}

246 - messages = [[{"role": "user", "content": "Describe this image."}]]

266 + messages: list[list[ChatCompletionMessageParam]] = [

267 + [{"role": "user", "content": "Describe this image."}]

268 + ]

269

270 renderer = Mock()

271 renderer.tokenizer = object()

@@ -269,7 +291,9 @@ def test_preprocess_chat_applies_mm_processor_kwargs_to_renderer() -> None:

291

292 def test_preprocess_chat_omits_mm_processor_kwargs_when_no_override() -> None:

293 llm = _make_mock_llm()

272 - messages = [[{"role": "user", "content": "Describe this image."}]]

294 + messages: list[list[ChatCompletionMessageParam]] = [

295 + [{"role": "user", "content": "Describe this image."}]

296 + ]

297

298 renderer = Mock()

299 renderer.tokenizer = object()

@@ -291,7 +315,9 @@ def test_preprocess_chat_omits_mm_processor_kwargs_when_no_override() -> None:

315 def test_preprocess_chat_defaults_add_special_tokens_to_false() -> None:

316 # Matches `ChatCompletionRequest.add_special_tokens` on the server.

317 llm = _make_mock_llm()

294 - messages = [[{"role": "user", "content": "hi"}]]

318 + messages: list[list[ChatCompletionMessageParam]] = [

319 + [{"role": "user", "content": "hi"}]

320 + ]

321

322 renderer = Mock()

323 renderer.tokenizer = object()

@@ -309,7 +335,9 @@ def test_preprocess_chat_defaults_add_special_tokens_to_false() -> None:

335

336 def test_preprocess_chat_tokenization_kwargs_override_add_special_tokens() -> None:

337 llm = _make_mock_llm()

312 - messages = [[{"role": "user", "content": "hi"}]]

338 + messages: list[list[ChatCompletionMessageParam]] = [

339 + [{"role": "user", "content": "hi"}]

340 + ]

341

342 renderer = Mock()

343 renderer.tokenizer = object()

tests/entrypoints/multimodal/openai/chat_completion/test_audio.py

+10 -8

@@ -47,18 +47,20 @@ async def client(server):

47

48 @pytest.fixture(scope="session")

49 def base64_encoded_audio() -> dict[str, str]:

50 - return {

51 - audio_url: encode_audio_base64(*fetch_audio(audio_url))

52 - for audio_url in TEST_AUDIO_URLS

53 - }

50 + encoded = {}

51 + for audio_url in TEST_AUDIO_URLS:

52 + audio, sample_rate = fetch_audio(audio_url)

53 + encoded[audio_url] = encode_audio_base64(audio, int(sample_rate))

54 + return encoded

55

56

57 @pytest.fixture(scope="session")

58 def url_encoded_audio() -> dict[str, str]:

58 - return {

59 - audio_url: encode_audio_url(*fetch_audio(audio_url))

60 - for audio_url in TEST_AUDIO_URLS

61 - }

59 + encoded = {}

60 + for audio_url in TEST_AUDIO_URLS:

61 + audio, sample_rate = fetch_audio(audio_url)

62 + encoded[audio_url] = encode_audio_url(audio, int(sample_rate))

63 + return encoded

64

65

66 def dummy_messages_from_audio_url(

tests/entrypoints/pooling/classify/test_offline.py

+3 -1

@@ -7,6 +7,7 @@ import torch

7

8 from tests.models.utils import softmax

9 from vllm import LLM, ClassificationRequestOutput, PoolingParams

10 +from vllm.inputs import PromptType

11 from vllm.tasks import PoolingTask

12

13 MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach"

@@ -53,7 +54,8 @@ def test_token_ids_prompts(llm: LLM):

54

55 @pytest.mark.skip_global_cleanup

56 def test_list_prompts(llm: LLM):

56 - outputs = llm.classify([prompt, prompt_token_ids], use_tqdm=False)

57 + prompts: list[PromptType] = [prompt, prompt_token_ids]

58 + outputs = llm.classify(prompts, use_tqdm=False)

59 assert len(outputs) == 2

60 for i in range(len(outputs)):

61 assert isinstance(outputs[i], ClassificationRequestOutput)

tests/entrypoints/pooling/embed/test_io_processor.py

+10 -9

@@ -21,6 +21,7 @@ from vllm.entrypoints.pooling.typing import (

21 PoolingEngineInput,

22 PoolingServeContext,

23 )

24 +from vllm.inputs import tokens_input

25 from vllm.outputs import PoolingOutput, PoolingRequestOutput

26

27

@@ -28,7 +29,7 @@ class TestEmbeddingRequestParsing:

29 """Unit tests for OpenAI embedding request parsing."""

30

31 def test_input_messages_parses_as_chat_request(self):

31 - request = TypeAdapter(EmbeddingRequest).validate_python(

32 + request: EmbeddingRequest = TypeAdapter(EmbeddingRequest).validate_python(

33 {

34 "model": "test",

35 "input": [{"role": "user", "content": "hello"}],

@@ -42,7 +43,7 @@ class TestEmbeddingRequestParsing:

43 assert request.chat_template_kwargs == {"instruction": "Represent the query: "}

44

45 def test_batched_input_messages_parses_as_batch_chat_input_request(self):

45 - request = TypeAdapter(EmbeddingRequest).validate_python(

46 + request: EmbeddingRequest = TypeAdapter(EmbeddingRequest).validate_python(

47 {

48 "model": "test",

49 "input": [

@@ -65,7 +66,7 @@ class TestEmbeddingRequestParsing:

66 assert request.chat_template_kwargs == {"instruction": "Represent the query: "}

67

68 def test_token_ids_still_parse_as_completion_request(self):

68 - request = TypeAdapter(EmbeddingRequest).validate_python(

69 + request: EmbeddingRequest = TypeAdapter(EmbeddingRequest).validate_python(

70 {

71 "model": "test",

72 "input": [[1, 2, 3], [4, 5]],

@@ -76,7 +77,7 @@ class TestEmbeddingRequestParsing:

77 assert request.input == [[1, 2, 3], [4, 5]]

78

79 def test_messages_still_parses_as_chat_request(self):

79 - request = TypeAdapter(EmbeddingRequest).validate_python(

80 + request: EmbeddingRequest = TypeAdapter(EmbeddingRequest).validate_python(

81 {

82 "model": "test",

83 "messages": [{"role": "user", "content": "hello"}],

@@ -89,7 +90,7 @@ class TestEmbeddingRequestParsing:

90 assert request.chat_template_kwargs == {"instruction": "Represent the query: "}

91

92 def test_batched_messages_parses_as_batch_chat_request(self):

92 - request = TypeAdapter(EmbeddingRequest).validate_python(

93 + request: EmbeddingRequest = TypeAdapter(EmbeddingRequest).validate_python(

94 {

95 "model": "test",

96 "messages": [

@@ -410,13 +411,13 @@ class TestChunkedEmbeddingProcessing:

411 @classmethod

412 def _make_handler(cls):

413 handler = object.__new__(EmbedIOProcessor)

413 - handler.model_config = cls._FakeModelConfig()

414 + handler.model_config = cls._FakeModelConfig() # type: ignore[assignment]

415 handler.enable_chunked_processing = True

416 return handler

417

418 @staticmethod

419 def _make_context() -> PoolingServeContext[EmbeddingCompletionRequest]:

419 - request = TypeAdapter(EmbeddingRequest).validate_python(

420 + request: EmbeddingRequest = TypeAdapter(EmbeddingRequest).validate_python(

421 {

422 "model": "test",

423 "input": [[0, 1, 2, 3, 4], [10, 11]],

@@ -431,13 +432,13 @@ class TestChunkedEmbeddingProcessing:

432 request_id="embd-client-prompt-999-chunk-888",

433 engine_inputs=[

434 PoolingEngineInput(

434 - prompts={"prompt_token_ids": [0, 1, 2, 3, 4]},

435 + prompts=tokens_input(prompt_token_ids=[0, 1, 2, 3, 4]),

436 params=pooling_params,

437 lora_requests=None,

438 priorities=0,

439 ),

440 PoolingEngineInput(

440 - prompts={"prompt_token_ids": [10, 11]},

441 + prompts=tokens_input(prompt_token_ids=[10, 11]),

442 params=pooling_params,

443 lora_requests=None,

444 priorities=0,

tests/entrypoints/pooling/embed/test_offline.py

+3 -1

@@ -7,6 +7,7 @@ import torch

7 import torch.nn.functional as F

8

9 from vllm import LLM, EmbeddingRequestOutput, PoolingParams

10 +from vllm.inputs import PromptType

11 from vllm.tasks import PoolingTask

12

13 MODEL_NAME = "intfloat/multilingual-e5-small"

@@ -54,7 +55,8 @@ def test_token_ids_prompts(llm: LLM):

55

56 @pytest.mark.skip_global_cleanup

57 def test_list_prompts(llm: LLM):

57 - outputs = llm.embed([prompt, prompt_token_ids], use_tqdm=False)

58 + prompts: list[PromptType] = [prompt, prompt_token_ids]

59 + outputs = llm.embed(prompts, use_tqdm=False)

60 assert len(outputs) == 2

61 for i in range(len(outputs)):

62 assert isinstance(outputs[i], EmbeddingRequestOutput)

tests/entrypoints/pooling/embed/test_online.py

+13 -9

@@ -12,9 +12,13 @@ import requests

12 import torch

13 import torch.nn.functional as F

14

15 -from tests.models.language.pooling.embed_utils import run_embedding_correctness_test

15 +from tests.models.language.pooling.embed_utils import (

16 + float_embeddings,

17 + run_embedding_correctness_test,

18 +)

19 from tests.models.utils import check_embeddings_close

20 from tests.utils import RemoteOpenAIServer

21 +from vllm.entrypoints.chat_utils import ChatCompletionMessageParam

22 from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse

23 from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse

24 from vllm.entrypoints.pooling.utils import (

@@ -145,7 +149,7 @@ async def test_completion_request(

149 assert embeddings.usage.prompt_tokens == len(input_tokens)

150 assert embeddings.usage.total_tokens == len(input_tokens)

151

148 - vllm_outputs = [d.embedding for d in embeddings.data]

152 + vllm_outputs = float_embeddings(embeddings)

153 run_embedding_correctness_test(hf_model, [input_text], vllm_outputs)

154

155 # test input: list[int]

@@ -165,7 +169,7 @@ async def test_completion_request(

169 assert embeddings.usage.prompt_tokens == len(input_tokens)

170 assert embeddings.usage.total_tokens == len(input_tokens)

171

168 - vllm_outputs = [d.embedding for d in embeddings.data]

172 + vllm_outputs = float_embeddings(embeddings)

173 run_embedding_correctness_test(hf_model, [input_text], vllm_outputs)

174

175

@@ -194,7 +198,7 @@ async def test_completion_request_batched(

198 assert embeddings.usage.prompt_tokens == len(input_tokens) * N

199 assert embeddings.usage.total_tokens == len(input_tokens) * N

200

197 - vllm_outputs = [d.embedding for d in embeddings.data]

201 + vllm_outputs = float_embeddings(embeddings)

202 run_embedding_correctness_test(hf_model, input_texts, vllm_outputs)

203

204 # test list[list[int]]

@@ -214,7 +218,7 @@ async def test_completion_request_batched(

218 assert embeddings.usage.prompt_tokens == len(input_tokens) * N

219 assert embeddings.usage.total_tokens == len(input_tokens) * N

220

217 - vllm_outputs = [d.embedding for d in embeddings.data]

221 + vllm_outputs = float_embeddings(embeddings)

222 run_embedding_correctness_test(hf_model, input_texts, vllm_outputs)

223

224

@@ -331,7 +335,7 @@ async def test_padding(client: openai.AsyncOpenAI, model_name: str):

335 async def test_chat_request(

336 server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str

337 ):

334 - messages = [

338 + messages: list[ChatCompletionMessageParam] = [

339 {

340 "role": "user",

341 "content": "The cat sat on the mat.",

@@ -382,8 +386,8 @@ async def test_chat_request(

386 assert chat_embeddings.created <= completion_embeddings.created

387 # Use tolerance-based comparison for embeddings

388 check_embeddings_close(

385 - embeddings_0_lst=[d.embedding for d in chat_embeddings.data],

386 - embeddings_1_lst=[d.embedding for d in completion_embeddings.data],

389 + embeddings_0_lst=float_embeddings(chat_embeddings),

390 + embeddings_1_lst=float_embeddings(completion_embeddings),

391 name_0="chat",

392 name_1="completion",

393 )

@@ -488,7 +492,7 @@ async def test_invocations_completion_request(

492

493 @pytest.mark.asyncio

494 async def test_invocations_chat_request(server: RemoteOpenAIServer):

491 - messages = [

495 + messages: list[ChatCompletionMessageParam] = [

496 {

497 "role": "user",

498 "content": "The cat sat on the mat.",

tests/entrypoints/pooling/embed/test_online_dimensions.py

+5 -2

@@ -6,7 +6,10 @@ import openai

6 import pytest

7

8 from tests.conftest import HfRunner

9 -from tests.models.language.pooling.embed_utils import run_embedding_correctness_test

9 +from tests.models.language.pooling.embed_utils import (

10 + float_embeddings,

11 + run_embedding_correctness_test,

12 +)

13 from tests.models.utils import EmbedModelInfo

14 from tests.utils import ROCM_EXTRA_ARGS, RemoteOpenAIServer

15 from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse

@@ -95,7 +98,7 @@ async def test_matryoshka(

98 if dimensions is not None:

99 assert len(embeddings.data[0].embedding) == dimensions

100

98 - vllm_outputs = [d.embedding for d in embeddings.data]

101 + vllm_outputs = float_embeddings(embeddings)

102 run_embedding_correctness_test(hf_model, prompts, vllm_outputs, dimensions)

103

104 if model_info.is_matryoshka:

tests/entrypoints/pooling/embed/test_online_long_text.py

+4

@@ -274,6 +274,7 @@ async def test_long_text_embedding_1500_chars(

274

275 # Verify embedding vector validity

276 embedding_vector = embeddings.data[0].embedding

277 + assert isinstance(embedding_vector, list)

278 assert all(isinstance(x, float) for x in embedding_vector), (

279 "Embedding vector should contain floats"

280 )

@@ -322,6 +323,7 @@ async def test_long_text_embedding_2500_chars(

323

324 # Verify embedding vector validity

325 embedding_vector = embeddings.data[0].embedding

326 + assert isinstance(embedding_vector, list)

327 assert all(isinstance(x, float) for x in embedding_vector), (

328 "Embedding vector should contain floats"

329 )

@@ -409,6 +411,7 @@ async def test_chunked_vs_normal_consistency(

411

412 # 验证embedding向量的有效性

413 embedding_vector = embeddings.data[0].embedding

414 + assert isinstance(embedding_vector, list)

415 assert all(isinstance(x, float) for x in embedding_vector)

416 assert not all(x == 0 for x in embedding_vector)

417

@@ -438,6 +441,7 @@ async def test_chunked_processing_response_format(

441

442 # Verify embedding vector properties

443 embedding_vector = embeddings.data[0].embedding

444 + assert isinstance(embedding_vector, list)

445 import math

446

447 vector_norm = math.sqrt(sum(x * x for x in embedding_vector))

tests/entrypoints/pooling/reward/test_token_reward_online.py

+3 -2

@@ -11,6 +11,7 @@ import torch

11

12 from tests.models.utils import check_embeddings_close

13 from tests.utils import RemoteOpenAIServer

14 +from vllm.entrypoints.chat_utils import ChatCompletionMessageParam

15 from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse

16 from vllm.entrypoints.pooling.utils import (

17 MetadataItem,

@@ -137,7 +138,7 @@ def test_completion_request_batched(server: RemoteOpenAIServer, model_name: str)

138 @pytest.mark.asyncio

139 @pytest.mark.parametrize("model_name", [MODEL_NAME])

140 async def test_chat_request(server: RemoteOpenAIServer, model_name: str):

140 - messages = [

141 + messages: list[ChatCompletionMessageParam] = [

142 {

143 "role": "user",

144 "content": "The cat sat on the mat.",

@@ -524,7 +525,7 @@ async def test_invocations_chat_request(server: RemoteOpenAIServer):

525

526 @pytest.mark.asyncio

527 async def test_invocations_conversation_chat_request(server: RemoteOpenAIServer):

527 - messages = [

528 + messages: list[ChatCompletionMessageParam] = [

529 {

530 "role": "user",

531 "content": "The cat sat on the mat.",

tests/entrypoints/pooling/scoring/test_io_processor_unit.py

+7 -3

@@ -47,7 +47,7 @@ def llm_reranker_processor() -> CrossEncoderIOProcessor:

47 )

48

49 processor = CrossEncoderIOProcessor.__new__(CrossEncoderIOProcessor)

50 - processor.model_config = SimpleNamespace(enable_prompt_embeds=False)

50 + processor.model_config = SimpleNamespace(enable_prompt_embeds=False) # type: ignore[assignment]

51 processor.tokenizer = TokenizersBackend(

52 tokenizer_object=backend,

53 unk_token="[UNK]",

@@ -95,12 +95,16 @@ def test_token_type_ids_stay_aligned_with_a_truncated_padded_prompt():

95 truncation_side="left",

96 )

97

98 - prompt = tok_params.apply_post_tokenization(tokenizer, prompt)

98 + processed_prompt = tok_params.apply_post_tokenization(

99 + tokenizer, # type: ignore[arg-type]

100 + prompt,

101 + )

102 + assert "prompt_token_ids" in processed_prompt

103 token_type_ids = _apply_post_tokenization_to_token_type_ids(

104 tokenizer, tok_params, token_type_ids

105 )

106

103 - prompt_token_ids = prompt["prompt_token_ids"]

107 + prompt_token_ids = processed_prompt["prompt_token_ids"]

108 assert len(token_type_ids) == len(prompt_token_ids)

109

110 # Keeping the last 40 tokens drops the first 10 query tokens, so 10 query

tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py

+15 -7

@@ -30,12 +30,18 @@ def test_online_forwards_truncate_prompt_tokens_to_proxy(monkeypatch):

30 silently disables truncate_prompt_tokens for Jina rerank/score.

31 """

32 proc = JinaRankingIOProcessor.__new__(JinaRankingIOProcessor)

33 - proc.valid_inputs_online = MagicMock(

34 - return_value=ScoringData(data_1=["query"], data_2=["doc"])

33 + monkeypatch.setattr(

34 + proc,

35 + "valid_inputs_online",

36 + MagicMock(return_value=ScoringData(data_1=["query"], data_2=["doc"])),

37 + )

38 + monkeypatch.setattr(proc, "_get_token_limits", MagicMock(return_value=(0, 0)))

39 + monkeypatch.setattr(

40 + proc, "ensure_str", MagicMock(side_effect=lambda data: list(data))

41 + )

42 + monkeypatch.setattr(

43 + proc, "format_docs_prompts_func", MagicMock(return_value="formatted prompt")

44 )

36 - proc._get_token_limits = MagicMock(return_value=(0, 0))

37 - proc.ensure_str = MagicMock(side_effect=lambda data: list(data))

38 - proc.format_docs_prompts_func = MagicMock(return_value="formatted prompt")

45

46 captured: dict[str, object] = {}

47

@@ -71,8 +77,10 @@ def offline_processor_and_context():

77 backend.pre_tokenizer = pre_tokenizers.WhitespaceSplit()

78 proc = JinaRankingIOProcessor.__new__(JinaRankingIOProcessor)

79 proc.tokenizer = TokenizersBackend(tokenizer_object=backend)

74 - proc.model_config = SimpleNamespace(max_model_len=1024, is_encoder_decoder=False)

75 - proc.renderer = SimpleNamespace(

80 + proc.model_config = SimpleNamespace( # type: ignore[assignment]

81 + max_model_len=1024, is_encoder_decoder=False

82 + )

83 + proc.renderer = SimpleNamespace( # type: ignore[assignment]

84 default_cmpl_tok_params=TokenizeParams(max_total_tokens=1024)

85 )

86 ctx = OfflineScoringInputsContext(

tests/entrypoints/pooling/scoring/test_late_interaction_serving.py

+20 -9

@@ -6,7 +6,7 @@ from unittest.mock import AsyncMock, Mock

6 import pytest

7

8 from vllm.entrypoints.pooling.scoring.serving import ServingScores

9 -from vllm.entrypoints.pooling.typing import PoolingServeContext

9 +from vllm.entrypoints.pooling.typing import PoolingEngineInput, PoolingServeContext

10 from vllm.inputs import tokens_input

11 from vllm.pooling_params import PoolingParams

12

@@ -21,8 +21,18 @@ def _make_context() -> PoolingServeContext:

21 priorities=None,

22 prompt_extras=None,

23 engine_inputs=[

24 - tokens_input(prompt_token_ids=[1]),

25 - tokens_input(prompt_token_ids=[2]),

24 + PoolingEngineInput(

25 + prompts=tokens_input(prompt_token_ids=[1]),

26 + params=PoolingParams(task="token_embed"),

27 + lora_requests=None,

28 + priorities=0,

29 + ),

30 + PoolingEngineInput(

31 + prompts=tokens_input(prompt_token_ids=[2]),

32 + params=PoolingParams(task="token_embed"),

33 + lora_requests=None,

34 + priorities=0,

35 + ),

36 ],

37 n_queries=1,

38 )

@@ -39,10 +49,13 @@ def _query_key(context: PoolingServeContext) -> str:

49

50

51 @pytest.mark.asyncio

42 -async def test_colliding_request_ids_use_distinct_query_cache_keys():

52 +async def test_colliding_request_ids_use_distinct_query_cache_keys(

53 + monkeypatch: pytest.MonkeyPatch,

54 +):

55 serving = object.__new__(ServingScores)

44 - serving._prepare_generators = AsyncMock()

45 - serving._collect_batch = AsyncMock()

56 + prepare_generators = AsyncMock()

57 + monkeypatch.setattr(serving, "_prepare_generators", prepare_generators)

58 + monkeypatch.setattr(serving, "_collect_batch", AsyncMock())

59

60 first = _make_context()

61 second = _make_context()

@@ -50,9 +63,7 @@ async def test_colliding_request_ids_use_distinct_query_cache_keys():

63 await serving._flash_late_interaction_encode_queries(context)

64 await serving._flash_late_interaction_encode_docs(context)

65

53 - prepared_contexts = [

54 - call.args[0] for call in serving._prepare_generators.await_args_list

55 - ]

66 + prepared_contexts = [call.args[0] for call in prepare_generators.await_args_list]

67 first_query_key = _query_key(prepared_contexts[0])

68 first_doc_key = _query_key(prepared_contexts[1])

69 second_query_key = _query_key(prepared_contexts[2])

tests/entrypoints/pooling/test_io_processor.py

+2 -2

@@ -53,8 +53,8 @@ def test_rejects_mismatched_lora_requests(processor: PoolingIOProcessor):

53

54

55 def test_rejects_conflicting_pooling_task(processor: PoolingIOProcessor):

56 - processor.model_config = SimpleNamespace(is_encoder_decoder=False)

57 - processor.renderer = SimpleNamespace(

56 + processor.model_config = SimpleNamespace(is_encoder_decoder=False) # type: ignore[assignment]

57 + processor.renderer = SimpleNamespace( # type: ignore[assignment] # Tokenization defaults only.

58 default_cmpl_tok_params=TokenizeParams(max_total_tokens=None)

59 )

60 ctx = OfflineEncodeInputsContext(

tests/entrypoints/pooling/token_classify/test_online.py

+1

@@ -45,6 +45,7 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: st

45 poolings = PoolingResponse.model_validate(response.json())

46 assert len(poolings.data) == 1

47 assert len(poolings.data[0].data) == 8

48 + assert isinstance(poolings.data[0].data[0], list)

49 assert len(poolings.data[0].data[0]) == 2

50

51

tests/entrypoints/pooling/token_embed/test_online.py

+1

@@ -71,6 +71,7 @@ async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):

71 assert poolings.id == f"pooling-{request_id}"

72 assert len(poolings.data) == 1

73 assert len(poolings.data[0].data) == len(input_tokens)

74 + assert isinstance(poolings.data[0].data[0], list)

75 assert len(poolings.data[0].data[0]) == 384

76

77

tests/kernels/attention/test_attention_selector.py

+12 -8

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from contextlib import contextmanager

5 +from typing import Literal

6 from unittest.mock import MagicMock, patch

7

8 import pytest

@@ -14,6 +15,7 @@ from vllm.config import (

15 VllmConfig,

16 set_current_vllm_config,

17 )

18 +from vllm.config.cache import CacheDType

19 from vllm.platforms import current_platform

20 from vllm.platforms.cpu import CpuPlatform

21 from vllm.platforms.interface import DeviceCapability

@@ -21,12 +23,12 @@ from vllm.platforms.interface import DeviceCapability

23 if current_platform.is_cuda():

24 from vllm.platforms.cuda import CudaPlatform

25 else:

24 - CudaPlatform = None

26 + CudaPlatform = None # type: ignore[assignment] # Unavailable platform import.

27

28 if current_platform.is_rocm():

29 from vllm.platforms.rocm import RocmPlatform

30 else:

29 - RocmPlatform = None

31 + RocmPlatform = None # type: ignore[misc, assignment] # Unavailable platform import.

32

33 from vllm.v1.attention.backend import AttentionType

34 from vllm.v1.attention.backends.registry import AttentionBackendEnum

@@ -354,14 +356,14 @@ def test_invalid_backend():

356 def test_auto_backend_string(auto_value: str):

357 """Test that 'auto' string value triggers automatic backend selection."""

358 # Using "auto" should result in backend=None (automatic selection)

357 - attention_config = AttentionConfig(backend=auto_value)

359 + attention_config = AttentionConfig(backend=auto_value) # type: ignore[arg-type]

360 assert attention_config.backend is None

361

362

363 def test_auto_backend_selection_behavior():

364 """Test that 'auto' backend behaves same as None (automatic selection)."""

365 # Create config with explicit "auto"

364 - auto_config = AttentionConfig(backend="auto")

366 + auto_config = AttentionConfig(backend="auto") # type: ignore[arg-type]

367

368 # Create config with None (default)

369 none_config = AttentionConfig(backend=None)

@@ -406,7 +408,7 @@ def test_auto_backend_selection_behavior():

408 reason="Attention backend FA3 is not supported on ROCm. This test can't succeed.",

409 )

410 def test_per_head_quant_scales_backend_selection(

409 - backend_name: str, flash_attn_version: int | None, should_succeed: bool

411 + backend_name: str, flash_attn_version: Literal[2, 3, 4] | None, should_succeed: bool

412 ):

413 """Test backend selection when use_per_head_quant_scales=True."""

414 # Clear cache to ensure fresh backend selection

@@ -559,7 +561,7 @@ def test_non_causal_autoselect_backend():

561 "int8_per_token_head",

562 ],

563 )

562 -def test_flash_attn_rejects_unhandled_kv_cache_dtypes(kv_cache_dtype: str):

564 +def test_flash_attn_rejects_unhandled_kv_cache_dtypes(kv_cache_dtype: CacheDType):

565 """FlashAttentionBackend must not claim support for kv_cache dtypes

566 that it cannot handle."""

567 from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend

@@ -569,7 +571,7 @@ def test_flash_attn_rejects_unhandled_kv_cache_dtypes(kv_cache_dtype: str):

571

572 @pytest.mark.parametrize("kv_cache_dtype", ["fp8", "fp8_e4m3"])

573 def test_flash_attn_accepts_handled_fp8_variants(

572 - kv_cache_dtype: str, monkeypatch: pytest.MonkeyPatch

574 + kv_cache_dtype: CacheDType, monkeypatch: pytest.MonkeyPatch

575 ):

576 """FlashAttentionBackend must accept the two fp8 dtypes it can actually

577 handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'."""

@@ -876,7 +878,9 @@ def test_rswa_selection_does_not_reuse_causal_result(blackwell_selection):

878 config = EngineArgs(

879 model="google/gemma-4-31B-it",

880 dtype="bfloat16",

879 - attention_config={"backend": "TRITON_FLASHINFER"},

881 + attention_config=AttentionConfig(

882 + backend=AttentionBackendEnum.TRITON_FLASHINFER

883 + ),

884 ).create_engine_config()

885 with set_current_vllm_config(config):

886 assert (

tests/kernels/attention/test_cascade_flash_attn.py

+1 -1

@@ -151,7 +151,7 @@ def test_cascade(

151 max_seqlen_k=max_kv_len,

152 softmax_scale=scale,

153 causal=True,

154 - window_size=window_size,

154 + window_size=list(window_size),

155 block_table=block_tables,

156 softcap=soft_cap if soft_cap is not None else 0,

157 )

tests/kernels/attention/test_cpu_attn.py

+1 -1

@@ -836,7 +836,7 @@ def varlen_encoder_zentorch_sdpa(

836 key,

837 value,

838 output,

839 - SimpleNamespace(query_start_loc=query_start_loc, causal=False),

839 + SimpleNamespace(query_start_loc=query_start_loc, causal=False), # type: ignore[arg-type]

840 scale,

841 sliding_window if sliding_window is not None else -1,

842 alibi_slopes,

tests/kernels/attention/test_flash_attn.py

+2 -2

@@ -187,7 +187,7 @@ def test_varlen_with_paged_kv(

187 max_seqlen_k=max_kv_len,

188 softmax_scale=scale,

189 causal=True,

190 - window_size=window_size,

190 + window_size=list(window_size),

191 block_table=block_tables,

192 softcap=soft_cap if soft_cap is not None else 0,

193 fa_version=fa_version,

@@ -274,7 +274,7 @@ def test_fa4_hd256_paged_call_shape(

274 max_seqlen_k=num_pages * block_size,

275 softmax_scale=scale,

276 causal=True,

277 - window_size=window_size,

277 + window_size=list(window_size),

278 block_table=block_tables[:, :num_pages],

279 num_splits=1,

280 fa_version=4,

tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py

+3 -6

@@ -197,11 +197,6 @@ def test_decode_concat_ignores_negative_slots_for_cache(cache_format: str) -> No

197 mixed_cache = initial_cache.clone()

198 reference_cache = _reference_cache(cache_format, inputs, initial_cache, slots)

199 scale_inv = torch.ones(1, device="cuda", dtype=torch.float32)

200 - kwargs = {

201 - "ds_mla": cache_format == "fp8_ds_mla",

202 - "q_scale_inv": scale_inv if cache_format == "fp8" else None,

203 - "cache_scale_inv": scale_inv if cache_format == "fp8" else None,

204 - }

200

201 output = fused_mla_decode_q_concat_kv_cache_insert(

202 inputs["ql_nope"],

@@ -210,7 +205,9 @@ def test_decode_concat_ignores_negative_slots_for_cache(cache_format: str) -> No

205 inputs["k_pe"],

206 mixed_cache,

207 slots,

213 - **kwargs,

208 + ds_mla=cache_format == "fp8_ds_mla",

209 + q_scale_inv=scale_inv if cache_format == "fp8" else None,

210 + cache_scale_inv=scale_inv if cache_format == "fp8" else None,

211 )

212

213 expected = _latent_query(inputs["ql_nope"], inputs["q_pe"])

tests/kernels/attention/test_mha_attn.py

+4 -4

@@ -277,7 +277,7 @@ def test_mha_attn_varlen_forward_flashinfer(

277 # Override vllm config so get_vit_attn_backend returns FLASHINFER (simulates

278 # --mm-encoder-attn-backend=FLASHINFER).

279 vllm_config = get_current_vllm_config()

280 - old_model_config = getattr(vllm_config, "model_config", None)

280 + old_model_config = vllm_config.model_config

281 minimal_model_config = type(

282 "MinimalModelConfig",

283 (),

@@ -376,7 +376,7 @@ def test_mha_attn_varlen_forward_aiter_fp8(

376 torch.set_default_dtype(dtype)

377

378 vllm_config = get_current_vllm_config()

379 - old_model_config = getattr(vllm_config, "model_config", None)

379 + old_model_config = vllm_config.model_config

380 minimal_model_config = type(

381 "MinimalModelConfig",

382 (),

@@ -442,7 +442,7 @@ def test_mha_attn_varlen_forward_aiter_fp8(

442 )

443 def test_mha_attn_aiter_fp8_rejects_unsupported_arch(default_vllm_config):

444 vllm_config = get_current_vllm_config()

445 - old_model_config = getattr(vllm_config, "model_config", None)

445 + old_model_config = vllm_config.model_config

446 vllm_config.model_config = type(

447 "MinimalModelConfig",

448 (),

@@ -468,7 +468,7 @@ def test_mha_attn_aiter_fp8_rejects_unsupported_arch(default_vllm_config):

468 )

469 def test_mha_attn_fp8_rejects_wrong_backend(default_vllm_config):

470 vllm_config = get_current_vllm_config()

471 - old_model_config = getattr(vllm_config, "model_config", None)

471 + old_model_config = vllm_config.model_config

472 vllm_config.model_config = type(

473 "MinimalModelConfig",

474 (),

tests/kernels/attention/test_minimax_m3.py

+17 -12

@@ -1689,8 +1689,10 @@ def test_aiter_sparse_pa_layout_contract(monkeypatch, layout):

1689 )

1690

1691 backend = sparse_attn_mod.MiniMaxM3SparseBackend

1692 - assert KVCacheLayout.LHBNC in backend.supported_kv_cache_layouts()

1693 - assert KVCacheLayout.LBHNC in backend.supported_kv_cache_layouts()

1692 + layouts = backend.supported_kv_cache_layouts()

1693 + assert layouts is not None

1694 + assert KVCacheLayout.LHBNC in layouts

1695 + assert KVCacheLayout.LBHNC in layouts

1696

1697 nb, h = 3, 1

1698 cache_dtype = current_platform.fp8_dtype()

@@ -1772,19 +1774,22 @@ def test_aiter_indexer_requires_the_aiter_attend(monkeypatch):

1774 import vllm.models.minimax_m3.amd.indexer_aiter as indexer_aiter_mod

1775

1776 monkeypatch.setattr(indexer_aiter_mod.current_platform, "is_rocm", lambda: True)

1775 - kwargs = dict(

1776 - topk_blocks=TOPK,

1777 - sparse_block_size=BLOCK_SIZE,

1778 - num_index_heads=1,

1779 - index_head_dim=HEAD_DIM,

1780 - indexer_kv_dtype="fp8_e4m3",

1781 - max_model_len=8192,

1782 - )

1777 +

1778 + def unsupported_reason() -> str | None:

1779 + return indexer_aiter_mod.aiter_indexer_unsupported_reason(

1780 + topk_blocks=TOPK,

1781 + sparse_block_size=BLOCK_SIZE,

1782 + num_index_heads=1,

1783 + index_head_dim=HEAD_DIM,

1784 + # The implementation also accepts this legacy fp8 alias.

1785 + indexer_kv_dtype="fp8_e4m3", # type: ignore[arg-type]

1786 + max_model_len=8192,

1787 + )

1788

1789 monkeypatch.setattr(

1790 indexer_aiter_mod, "_minimax_m3_aiter_sparse_pa_requested", lambda: False

1791 )

1787 - reason = indexer_aiter_mod.aiter_indexer_unsupported_reason(**kwargs)

1792 + reason = unsupported_reason()

1793 assert reason is not None and "AITER sparse PA attend" in reason

1794

1795 # With the attend asked for, the gate moves on to the kernel-contract

@@ -1792,7 +1797,7 @@ def test_aiter_indexer_requires_the_aiter_attend(monkeypatch):

1797 monkeypatch.setattr(

1798 indexer_aiter_mod, "_minimax_m3_aiter_sparse_pa_requested", lambda: True

1799 )

1795 - reason = indexer_aiter_mod.aiter_indexer_unsupported_reason(**kwargs)

1800 + reason = unsupported_reason()

1801 assert reason is None or "AITER sparse PA attend" not in reason

1802

1803

tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py

+4 -4

@@ -235,14 +235,14 @@ def test_msa_metadata_builder_prepares_cutlass_for_regular_decode(

235 builder.num_q_heads = 64

236 builder.num_kv_heads = 4

237 builder.topk_blocks = TOPK

238 - builder.kv_cache_spec = SimpleNamespace(num_kv_heads=4)

238 + builder.kv_cache_spec = SimpleNamespace(num_kv_heads=4) # type: ignore[assignment]

239 builder.kv_cache_dtype = "fp8_e4m3"

240 builder.decode_backend = "cutlass"

241 - builder.msa_cutlass_plan_cache = object()

241 + builder.msa_cutlass_plan_cache = object() # type: ignore[assignment] # Mocked plan-cache sentinel.

242

243 metadata = builder.build(

244 0,

245 - SimpleNamespace(

245 + SimpleNamespace( # type: ignore[arg-type]

246 seq_lens_cpu_upper_bound=torch.full((batch,), 257, dtype=torch.int32)

247 ),

248 )

@@ -377,7 +377,7 @@ def test_query_fp8_stays_valid_when_cutlass_plan_appears_on_replay(

377 with torch.cuda.stream(stream), override_forward_context(forward_context):

378 capture = BreakableCUDAGraphCapture()

379 with capture:

380 - query_fp8 = MiniMaxM3SparseAttention._allocate_query_fp8(attention, qkv)

380 + query_fp8 = MiniMaxM3SparseAttention._allocate_query_fp8(attention, qkv) # type: ignore[arg-type]

381 if query_fp8 is not None:

382 query_fp8.zero_()

383 run_attention(query_fp8)

tests/kernels/attention/test_rocm_aiter_fa.py

+3

@@ -783,6 +783,9 @@ def test_aiter_fa_shared_kv_matches_reference(

783 32 * BLOCK_SIZE, num_kv_heads, head_size, dtype=torch.bfloat16

784 )

785 value = torch.randn_like(key)

786 + from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionImpl

787 +

788 + assert isinstance(target.impl, AiterFlashAttentionImpl)

789 target.impl.do_kv_cache_update(

790 target,

791 key,

tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py

+4 -2

@@ -82,6 +82,7 @@ def _run_verify_block():

82 )

83 from vllm.config import SpeculativeConfig

84 from vllm.config.vllm import set_current_vllm_config

85 + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAImpl

86 from vllm.v1.attention.backends.registry import AttentionBackendEnum

87 from vllm.v1.kv_cache_interface import MLAAttentionSpec

88 from vllm.v1.worker.workspace import init_workspace_manager

@@ -106,7 +107,7 @@ def _run_verify_block():

107 vllm_config.speculative_config = SpeculativeConfig(

108 method="ngram", num_speculative_tokens=QLEN - 1

109 )

109 - vllm_config.model_config.get_num_attention_heads = types.MethodType(

110 + vllm_config.model_config.get_num_attention_heads = types.MethodType( # type: ignore[method-assign]

111 lambda self, parallel_config, arch_config=None: NUM_QUERY_HEADS,

112 vllm_config.model_config,

113 )

@@ -202,7 +203,8 @@ def _run_verify_block():

203 lambda: spy,

204 ),

205 ):

205 - impl.forward_mqa((q_nope, q_pe), kv_cache, metadata, layer=None)

206 + assert isinstance(impl, AiterMLAImpl)

207 + impl.forward_mqa((q_nope, q_pe), kv_cache, metadata, layer=None) # type: ignore[arg-type] # Spy decode path does not read layer.

208

209 return metadata, captured

210

tests/kernels/attention/test_rocm_aiter_mla_fp8_prefill.py

+2 -2

@@ -121,7 +121,7 @@ def _build_prefill_metadata(

121

122 builder = object.__new__(AiterMLAMetadataBuilder)

123 builder.num_heads = num_heads

124 - builder.mla_dims = SimpleNamespace(v_head_dim=V_HEAD_DIM)

124 + builder.mla_dims = SimpleNamespace(v_head_dim=V_HEAD_DIM) # type: ignore[assignment]

125 builder._init_fp8_prefill_ps_buffers(

126 max_num_reqs=len(seq_lens),

127 max_prefill_qlen=max_q,

@@ -136,7 +136,7 @@ def _build_prefill_metadata(

136 prefill = SimpleNamespace(query_start_loc=qo_indptr, max_query_len=max_q)

137 metadata = SimpleNamespace(prefill=prefill, num_decodes=0)

138 common = SimpleNamespace(query_start_loc_cpu=qo_indptr_cpu)

139 - builder._build_fp8_prefill_ps_metadata(metadata, common)

139 + builder._build_fp8_prefill_ps_metadata(metadata, common) # type: ignore[arg-type]

140 return metadata, total_q, builder

141

142

tests/kernels/attention/test_rocm_aiter_mla_sink.py

+9 -6

@@ -167,7 +167,10 @@ def test_sparse_mla_sink_matches_ragged_reference(

167 impl.sinks = sinks

168

169 output, lse = impl._forward_mla(

170 - SimpleNamespace(_q_scale=q_scale, _k_scale=kv_scale), q, kv, metadata

170 + SimpleNamespace(_q_scale=q_scale, _k_scale=kv_scale),

171 + q,

172 + kv,

173 + metadata, # type: ignore[arg-type]

174 )

175 kv_flat = kv_ref[:, 0]

176 references = []

@@ -241,7 +244,7 @@ def test_sparse_mla_sink_rejects_unsupported_aiter_dtypes(

244 SimpleNamespace(_q_scale=None, _k_scale=None),

245 torch.empty(1, 16, Q_HEAD_DIM, dtype=q_dtype),

246 torch.empty(1, 1, Q_HEAD_DIM, dtype=kv_dtype),

244 - metadata,

247 + metadata, # type: ignore[arg-type]

248 )

249

250

@@ -381,7 +384,7 @@ def test_sparse_mla_sink_matches_dense_attention_with_empty_rows_and_paged_cache

384 paged_kv_indices=torch.tensor(indices, dtype=torch.int32, device="cuda"),

385 paged_kv_indptr=lengths.cumsum(0).to(device="cuda", dtype=torch.int32),

386 )

384 - impl = ROCMAiterMLASparseImpl.__new__(ROCMAiterMLASparseImpl)

387 + impl = ROCMAiterMLASparseImpl.__new__(ROCMAiterMLASparseImpl) # type: ignore[arg-type]

388 impl.num_heads = num_heads

389 impl.head_size = head_dim

390 impl.kv_lora_rank = value_dim

@@ -438,9 +441,9 @@ def test_sparse_mla_backend_resolves_only_contiguous_layer_layouts(monkeypatch,

441 ]

442 if layout == "BLNHC":

443 with pytest.raises(ValueError, match="does not satisfy"):

441 - resolve_kv_cache_layout(config, supported)

444 + resolve_kv_cache_layout(config, supported) # type: ignore[arg-type]

445 else:

443 - assert resolve_kv_cache_layout(config, supported).name == layout

446 + assert resolve_kv_cache_layout(config, supported).name == layout # type: ignore[arg-type]

447

448

449 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])

@@ -491,7 +494,7 @@ def test_sparse_mla_sink_forward_mqa_preserves_split_query(dtype):

494 actual, _ = impl.forward_mqa(

495 (q[..., :V_HEAD_DIM], q[..., V_HEAD_DIM:]),

496 kv,

494 - metadata,

497 + metadata, # type: ignore[arg-type]

498 SimpleNamespace(_q_scale=None, _k_scale=None),

499 )

500 references = [

tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py

+2 -2

@@ -39,10 +39,10 @@ def _make_builder():

39 topk_tokens = 4

40

41 builder.device = torch.device("cpu")

42 - builder.kv_cache_spec = SimpleNamespace(block_size=1)

42 + builder.kv_cache_spec = SimpleNamespace(block_size=1) # type: ignore[assignment]

43 builder.model_dtype = torch.bfloat16

44 builder.kv_cache_dtype = "fp8"

45 - builder.mla_dims = SimpleNamespace(kv_lora_rank=512, qk_rope_head_dim=64)

45 + builder.mla_dims = SimpleNamespace(kv_lora_rank=512, qk_rope_head_dim=64) # type: ignore[assignment]

46 builder.topk_tokens = topk_tokens

47 builder.req_id_per_token_buffer = torch.zeros(

48 max_num_batched_tokens, dtype=torch.int32, device="cpu"

tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py

+4 -2

@@ -227,7 +227,9 @@ def _decode_metadata(case, rows, ratio, query_lens):

227 num_decode_tokens=len(rows),

228 num_prefills=0,

229 num_prefill_tokens=0,

230 - decode=types.SimpleNamespace(block_table=block_table, seq_lens=lens[:, None]),

230 + decode=types.SimpleNamespace( # type: ignore[arg-type]

231 + block_table=block_table, seq_lens=lens[:, None]

232 + ),

233 decode_row_lens=lens,

234 decode_block_ends=(lens + CAND_BLOCK - 1) // CAND_BLOCK,

235 decode_native=native,

@@ -278,7 +280,7 @@ def _prefill_metadata(

280 num_decode_tokens=0,

281 num_prefills=len(case.seq_lens),

282 num_prefill_tokens=len(rows),

281 - prefill=types.SimpleNamespace(chunks=chunks),

283 + prefill=types.SimpleNamespace(chunks=chunks), # type: ignore[arg-type]

284 prefill_plans=plans,

285 gather_launches=plan_gather_launches(chunks, plans, gather_rows),

286 )

tests/kernels/attention/test_rocm_triton_attn_dsv4.py

+2 -2

@@ -955,7 +955,7 @@ def test_rocm_capture_metadata_sets_adaptive_marker(monkeypatch) -> None:

955 )

956 builder = object.__new__(rocm_mod.DeepseekV4ROCMAiterMLASparseMetadataBuilder)

957

958 - actual = builder.build_for_cudagraph_capture(SimpleNamespace())

958 + actual = builder.build_for_cudagraph_capture(SimpleNamespace()) # type: ignore[arg-type] # Patched base method ignores metadata.

959

960 assert actual is metadata

961 assert actual.for_cudagraph_capture is _on_gfx950()

@@ -1563,7 +1563,7 @@ def test_dsv4_adaptive_mla_swa_metadata_graph_replay(monkeypatch) -> None:

1563 "sliding_window": window_size,

1564 },

1565 )

1566 - vllm_config.speculative_config = SimpleNamespace(

1566 + vllm_config.speculative_config = SimpleNamespace( # type: ignore[assignment]

1567 num_speculative_tokens=upper_query_len - 1,

1568 parallel_drafting=False,

1569 enable_adaptive_verification=True,

tests/kernels/attention/test_triton_unified_attention_diffkv.py

+1

@@ -153,6 +153,7 @@ def test_triton_unified_attn_diffkv_vs_reference(

153 else:

154 from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func

155

156 + assert fa_version is not None

157 # FA's TMA path needs aligned singleton strides (num_kv_heads == 1).

158 fa_k = canonicalize_singleton_dim_strides(key_cache)

159 fa_v = canonicalize_singleton_dim_strides(value_cache)

tests/kernels/helion/test_config_manager.py

+1 -1

@@ -79,7 +79,7 @@ class TestConfigSet:

79 config_set = ConfigSet("test_kernel")

80

81 with pytest.raises(KeyError, match="platform 'h100' not found"):

82 - config_set.get_config("h100", "nonexistent")

82 + config_set.get_config("h100", CaseKey({"batch": 32}))

83

84 config_data = {"num_warps": 8, "num_stages": 4}

85 data = {

tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py

+1 -1

@@ -81,7 +81,7 @@ class TestDynamicPerTokenScaledFp8QuantConfigPicker:

81 assert selected_key == CaseKey({"hidden_size": 2048, "num_tokens": 32})

82

83 def test_config_picker_no_configs(self):

84 - config_keys: list[dict] = []

84 + config_keys: list[CaseKey] = []

85

86 args = _generate_fake_input(16, 4096)

87 selected_key = pick_config(args, config_keys)

tests/kernels/helion/test_fused_qk_norm_rope.py

+1 -1

@@ -127,7 +127,7 @@ class TestFusedQkNormRopeConfigPicker:

127 )

128

129 def test_config_picker_no_configs(self):

130 - config_keys: list[dict] = []

130 + config_keys: list[CaseKey] = []

131

132 args = _generate_fake_input(16, 4096, 128)

133 selected_key = pick_config(args, config_keys)

tests/kernels/helion/test_pattern_matching.py

+4 -3

@@ -38,6 +38,7 @@ from torch._inductor.pattern_matcher import (

38 )

39 from torch.fx.experimental.proxy_tensor import make_fx

40

41 +from vllm.kernels.helion.case_key import CaseKey

42 from vllm.kernels.helion.config_manager import ConfigManager

43 from vllm.kernels.helion.register import HelionKernelWrapper

44

@@ -45,7 +46,7 @@ from vllm.kernels.helion.register import HelionKernelWrapper

46 @contextlib.contextmanager

47 def _helion_mock_context():

48 configs = {

48 - "default": helion.Config(block_sizes=[64], num_warps=2, num_stages=2),

49 + CaseKey.default(): helion.Config(block_sizes=[64], num_warps=2, num_stages=2),

50 }

51 mock_config_manager = Mock(spec=ConfigManager)

52 mock_config_manager.get_platform_configs = Mock(return_value=configs)

@@ -88,7 +89,7 @@ class TestMakeFxHop:

89 raw_kernel_func=raw_add_scale,

90 op_name="test_make_fx",

91 fake_impl=lambda *a, **kw: None,

91 - config_picker=lambda args, keys: "default",

92 + config_picker=lambda args, keys: CaseKey.default(),

93 )

94

95 def fn(x, y):

@@ -145,7 +146,7 @@ class TestMakeFxHop:

146 raw_kernel_func=raw_silu_mul,

147 op_name="test_pm_silu_mul",

148 fake_impl=lambda *a, **kw: None,

148 - config_picker=lambda args, keys: "default",

149 + config_picker=lambda args, keys: CaseKey.default(),

150 )

151

152 def pattern(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:

tests/kernels/helion/test_per_token_group_fp8_quant.py

+1 -1

@@ -107,7 +107,7 @@ class TestPerTokenGroupFp8QuantConfigPicker:

107 )

108

109 def test_config_picker_no_configs(self):

110 - config_keys: list[dict] = []

110 + config_keys: list[CaseKey] = []

111

112 args = _generate_fake_input(16, 4096, 128)

113 selected_key = pick_config(args, config_keys)

tests/kernels/helion/test_register.py

+1

@@ -397,6 +397,7 @@ class TestHelionKernelWrapper:

397 )

398

399 assert wrapper._disabled is True

400 + assert wrapper._disabled_reason is not None

401 assert "No configs available" in wrapper._disabled_reason

402

403 def test_disabled_wrapper_raises_on_call(self, sample_kernel):

tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py

+1 -1

@@ -89,7 +89,7 @@ class TestRmsNormDynamicPerTokenQuantConfigPicker:

89 assert selected_key == CaseKey({"hidden_size": 2048, "num_tokens": 32})

90

91 def test_config_picker_no_configs(self):

92 - config_keys: list[dict] = []

92 + config_keys: list[CaseKey] = []

93

94 args = _generate_fake_input(16, 4096)

95 selected_key = pick_config(args, config_keys)

tests/kernels/helion/test_rms_norm_per_block_quant.py

+1 -1

@@ -112,7 +112,7 @@ class TestRmsNormPerBlockQuantConfigPicker:

112 )

113

114 def test_config_picker_no_configs(self):

115 - config_keys: list[dict] = []

115 + config_keys: list[CaseKey] = []

116

117 args = _generate_fake_input(16, 4096, 128)

118 selected_key = pick_config(args, config_keys)

tests/kernels/helion/test_silu_and_mul_per_block_quant.py

+1 -1

@@ -96,7 +96,7 @@ class TestSiluAndMulPerBlockQuantConfigPicker:

96 )

97

98 def test_config_picker_no_configs(self):

99 - config_keys: list[dict] = []

99 + config_keys: list[CaseKey] = []

100

101 args = _generate_fake_input(16, 4096, 128)

102 selected_key = pick_config(args, config_keys)

tests/kernels/helion/test_silu_mul_fp8.py

+2 -1

@@ -83,7 +83,7 @@ class TestSiluMulFp8ConfigPicker:

83 assert selected_key == CaseKey({"intermediate": 4096, "numtokens": 256})

84

85 def test_config_picker_no_configs(self):

86 - config_keys: list[dict] = []

86 + config_keys: list[CaseKey] = []

87

88 input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")

89 scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")

@@ -355,6 +355,7 @@ class TestSiluMulFp8Integration:

355 registered_kernels = get_registered_kernels()

356 kernel_wrapper = registered_kernels["silu_mul_fp8"]

357 fake_impl = kernel_wrapper._fake_impl

358 + assert fake_impl is not None

359

360 fake_output = fake_impl(input_tensor, scale)

361

tests/models/language/pooling/embed_utils.py

+9

@@ -7,6 +7,15 @@ import pytest

7

8 from tests.conftest import HfRunner

9 from tests.models.utils import EmbedModelInfo, check_embeddings_close, matryoshka_fy

10 +from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse

11 +

12 +

13 +def float_embeddings(response: EmbeddingResponse) -> list[list[float]]:

14 + embeddings = []

15 + for d in response.data:

16 + assert isinstance(d.embedding, list)

17 + embeddings.append(d.embedding)

18 + return embeddings

19

20

21 def run_embedding_correctness_test(

tests/multimodal/media/test_video.py

+9 -9

@@ -53,12 +53,10 @@ def test_video_media_io_kwargs(monkeypatch: pytest.MonkeyPatch):

53 imageio = ImageMediaIO()

54

55 # Verify that different args pass/fail assertions as expected.

56 - videoio = VideoMediaIO(imageio, **{"num_frames": 10, "fps": 1.0})

56 + videoio = VideoMediaIO(imageio, num_frames=10, fps=1.0)

57 _ = videoio.load_bytes(b"test")

58

59 - videoio = VideoMediaIO(

60 - imageio, **{"num_frames": 10, "fps": 1.0, "not_used": "not_used"}

61 - )

59 + videoio = VideoMediaIO(imageio, num_frames=10, fps=1.0, not_used="not_used")

60 _ = videoio.load_bytes(b"test")

61

62 with pytest.raises(AssertionError, match="bad num_frames"):

@@ -66,11 +64,11 @@ def test_video_media_io_kwargs(monkeypatch: pytest.MonkeyPatch):

64 _ = videoio.load_bytes(b"test")

65

66 with pytest.raises(AssertionError, match="bad num_frames"):

69 - videoio = VideoMediaIO(imageio, **{"num_frames": 9, "fps": 1.0})

67 + videoio = VideoMediaIO(imageio, num_frames=9, fps=1.0)

68 _ = videoio.load_bytes(b"test")

69

70 with pytest.raises(AssertionError, match="bad fps"):

73 - videoio = VideoMediaIO(imageio, **{"num_frames": 10, "fps": 2.0})

71 + videoio = VideoMediaIO(imageio, num_frames=10, fps=2.0)

72 _ = videoio.load_bytes(b"test")

73

74

@@ -86,14 +84,14 @@ def test_opencv_video_io_colorspace(tmp_path, is_color: bool, fourcc: str, ext:

84 image = Image.open(image_path)

85

86 if not is_color:

89 - image_path = f"{tmp_path}/test_grayscale_image.png"

87 + image_path = Path(tmp_path) / "test_grayscale_image.png"

88 image = image.convert("L")

89 image.save(image_path)

90 # Convert to gray RGB for comparison

91 image = image.convert("RGB")

92 video_path = f"{tmp_path}/test_RGB_video.{ext}"

93 create_video_from_image(

96 - image_path,

94 + str(image_path),

95 video_path,

96 num_frames=2,

97 is_color=is_color,

@@ -385,7 +383,9 @@ def test_pynvvideocodec_unrelated_error_propagates(

383

384 with pytest.raises(RuntimeError) as exc_info:

385 PyNvVideoCodecVideoBackendMixin.decode_frames_pynvvideocodec(

388 - None, b"video", None

386 + None,

387 + b"video",

388 + None, # type: ignore[arg-type] # Fails before target metadata is used.

389 )

390

391 assert exc_info.value is original_error

tests/multimodal/test_audio.py

+13 -1

@@ -179,7 +179,7 @@ def test_audio_resampler_invalid_method(dummy_audio):

179 # Validated eagerly so a bad method fails at construction, not on the

180 # first audio request.

181 with pytest.raises(ValueError, match="Invalid resampling method"):

182 - AudioResampler(target_sr=22050, method="invalid")

182 + AudioResampler(target_sr=22050, method="invalid") # type: ignore[arg-type]

183

184

185 def test_audio_resampler_no_target_sr(dummy_audio):

@@ -378,6 +378,7 @@ class TestMultiModalDataParserChannelNormalization:

378 result = parser._parse_audio_data((stereo_audio, 16000))

379

380 # Check that result is mono (1D)

381 + assert result is not None

382 audio_item = result.get(0)

383 assert audio_item.ndim == 1, f"Expected 1D mono audio, got {audio_item.ndim}D"

384 assert audio_item.shape == (3,), f"Expected shape (3,), got {audio_item.shape}"

@@ -404,6 +405,7 @@ class TestMultiModalDataParserChannelNormalization:

405 result = parser._parse_audio_data((stereo_audio, 16000))

406

407 # Check that result preserves original shape (after resampling)

408 + assert result is not None

409 audio_item = result.get(0)

410 # When target_channels=None, stereo audio should be preserved

411 assert audio_item.ndim == 2, f"Expected 2D stereo audio, got {audio_item.ndim}D"

@@ -425,6 +427,7 @@ class TestMultiModalDataParserChannelNormalization:

427 result = parser._parse_audio_data((mono_audio, 16000))

428

429 # Check that result is still mono (1D)

430 + assert result is not None

431 audio_item = result.get(0)

432 assert audio_item.ndim == 1

433 assert audio_item.shape == (16000,)

@@ -446,6 +449,7 @@ class TestMultiModalDataParserChannelNormalization:

449 result = parser._parse_audio_data((surround_audio, 16000))

450

451 # Check that result is stereo (2 channels)

452 + assert result is not None

453 audio_item = result.get(0)

454 assert audio_item.ndim == 2

455 assert audio_item.shape[0] == 2 # 2 channels

@@ -483,6 +487,7 @@ class TestAudioPipelineE2E:

487

488 # Process audio through the parser

489 result = parser._parse_audio_data((stereo_torchaudio, 16000))

490 + assert result is not None

491 audio_output = result.get(0)

492

493 # Verify output is mono 1D

@@ -512,6 +517,7 @@ class TestAudioPipelineE2E:

517

518 # Process audio through the parser

519 result = parser._parse_audio_data((stereo_soundfile, 16000))

520 + assert result is not None

521 audio_output = result.get(0)

522

523 # Verify output is mono 1D

@@ -537,6 +543,7 @@ class TestAudioPipelineE2E:

543

544 # Process audio through the parser

545 result = parser._parse_audio_data((mono_pyav, 16000))

546 + assert result is not None

547 audio_output = result.get(0)

548

549 # Verify output is still mono 1D

@@ -572,6 +579,7 @@ class TestAudioPipelineE2E:

579

580 # Process audio through the parser

581 result = parser._parse_audio_data((surround_audio, 16000))

582 + assert result is not None

583 audio_output = result.get(0)

584

585 # Verify output is mono 1D

@@ -603,6 +611,7 @@ class TestAudioPipelineE2E:

611 # Process audio through the parser

612 # Note: Parser expects numpy, so we convert first (simulating real usage)

613 result = parser._parse_audio_data((stereo_torch.numpy(), 16000))

614 + assert result is not None

615 audio_output = result.get(0)

616

617 # Verify output is mono 1D numpy array

@@ -630,6 +639,7 @@ class TestAudioPipelineE2E:

639

640 # Process audio through the parser

641 result = parser._parse_audio_data((stereo_audio, 16000))

642 + assert result is not None

643 audio_output = result.get(0)

644

645 # Verify output preserves stereo (2D)

@@ -654,6 +664,7 @@ class TestAudioPipelineE2E:

664

665 # Process audio through the parser

666 result = parser._parse_audio_data((stereo_48k, 48000))

667 + assert result is not None

668 audio_output = result.get(0)

669

670 # Verify output is mono 1D at target sample rate

@@ -677,6 +688,7 @@ class TestAudioPipelineE2E:

688 )

689

690 result = parser._parse_audio_data((short_stereo, 16000))

691 + assert result is not None

692 audio_output = result.get(0)

693

694 # Should still produce mono output

tests/multimodal/test_cache.py

+13 -6

@@ -11,7 +11,6 @@ import torch

11 from vllm.config import ModelConfig, ParallelConfig, VllmConfig

12 from vllm.config.multimodal import MultiModalConfig

13 from vllm.multimodal.cache import (

14 - BaseMultiModalProcessorCache,

14 BaseMultiModalReceiverCache,

15 LruKeyReplicatedReceiverCache,

16 LruKeyReplicatedSenderCache,

@@ -145,6 +144,8 @@ def _compare_caches(

144 cache_1_p0 = processor_cache_from_config(config_1)

145 cache_1_p1 = engine_receiver_cache_from_config(config_1)

146

147 + assert config_0.model_config.multimodal_config is not None

148 + assert config_1.model_config.multimodal_config is not None

149 cache_size_gb = max(

150 config_0.model_config.multimodal_config.mm_processor_cache_gb,

151 config_1.model_config.multimodal_config.mm_processor_cache_gb,

@@ -280,6 +281,7 @@ def test_oversized_item_is_served_uncached():

281

282 p0 = LruKeyReplicatedSenderCache(model_config) # type: ignore[arg-type]

283 p1 = LruKeyReplicatedReceiverCache(model_config) # type: ignore[arg-type]

284 + assert item is not None

285 assert p0.get_and_update_item((item, []), "big")[0] is item

286 assert not p0.is_cached_item("big")

287 assert p1.get_and_update_item(item, "big") is item

@@ -318,6 +320,7 @@ def test_mm_cache_miss_raises_and_recovers():

320 assert p0.is_cached_item(mm_hash)

321 # On the next request P0 short-circuits to data=None -- the drift bug when

322 # P1 lacks the item.

323 + assert item is not None

324 hit = p0.get_and_update_item((item, []), mm_hash)

325 assert hit[0] is None

326

@@ -326,6 +329,7 @@ def test_mm_cache_miss_raises_and_recovers():

329 p0.invalidate(mm_hash)

330 assert not p0.is_cached_item(mm_hash)

331

332 + assert item is not None

333 resent = p0.get_and_update_item((item, []), mm_hash)

334 assert resent[0] is item # MISS again -> data is resent

335 assert p1.get_and_update_item(item, mm_hash) == item # P1 now caches it

@@ -459,6 +463,7 @@ def test_shm_receiver_handles_prefix_covered_items(monkeypatch):

463 address_item, _ = p0.get_and_update_item((item, []), mm_hash)

464 first = _feature(mm_hash, address_item)

465 p1.get_and_update_features([first])

466 + assert first.data is not None

467 assert torch.equal(first.data["dummy"].data, item["dummy"].data)

468

469 # Request 2 (identical, fully prefix-covered): the sender hit takes

@@ -476,6 +481,7 @@ def test_shm_receiver_handles_prefix_covered_items(monkeypatch):

481

482 assert covered_uncached.data is None

483 # The address item is resolved to the cached payload.

484 + assert covered_cached.data is not None

485 assert torch.equal(covered_cached.data["dummy"].data, item["dummy"].data)

486

487 # The hit's writer references were acknowledged by the worker, so the

@@ -488,7 +494,7 @@ def test_shm_receiver_handles_prefix_covered_items(monkeypatch):

494

495

496 def _run_test_cache_eviction_lru(

491 - p0_cache: BaseMultiModalProcessorCache,

497 + p0_cache: LruKeyReplicatedSenderCache,

498 p1_cache: BaseMultiModalReceiverCache,

499 base_item_size: int,

500 ):

@@ -603,7 +609,7 @@ def test_cache_eviction_lru_cache():

609 # image_B is protected from eviction then image_i cannot be added.

610 # This proving normal eviction and reuse behavior.

611 def _run_test_cache_eviction_shm(

606 - p0_cache: BaseMultiModalProcessorCache,

612 + p0_cache: ShmObjectStoreSenderCache,

613 p1_cache: BaseMultiModalReceiverCache,

614 base_item_size: int,

615 ):

@@ -860,9 +866,9 @@ async def test_release_kv_cache_resends_mm_payload(use_async, release_error):

866

867 async def call_release():

868 if use_async:

863 - await AsyncLLM.release_kv_cache_memory(engine)

869 + await AsyncLLM.release_kv_cache_memory(engine) # type: ignore[arg-type]

870 else:

865 - LLMEngine.release_kv_cache_memory(engine)

871 + LLMEngine.release_kv_cache_memory(engine) # type: ignore[arg-type]

872

873 if release_error:

874 with pytest.raises(RuntimeError, match=release_error):

@@ -894,9 +900,10 @@ def test_sleep_wake_preserves_mm_cache_consistency():

900 """Regression for vllm-project/vllm#42995."""

901 from vllm import LLM, SamplingParams

902 from vllm.assets.image import ImageAsset

903 + from vllm.inputs import TextPrompt

904

905 image = ImageAsset("stop_sign").pil_image

899 - prompt = {

906 + prompt: TextPrompt = {

907 "prompt": _SLEEP_VISION_PROMPT,

908 "multi_modal_data": {"image": image},

909 }

tests/multimodal/test_hasher.py

+2

@@ -131,6 +131,7 @@ def test_hash_collision_video_num_frames():

131 }

132 video = MediaWithBytes((frames, metadata), source)

133 items = MultiModalDataParser()._parse_video_data([video])

134 + assert items is not None

135 return items.get_all_items_for_hash()[0]

136

137 hasher = MultiModalHasher

@@ -155,6 +156,7 @@ def test_hash_video_tensor_frames():

156 }

157 video = MediaWithBytes((frames, metadata), source)

158 items = MultiModalDataParser()._parse_video_data([video])

159 + assert items is not None

160 return items.get_all_items_for_hash()[0]

161

162 np_frames = np.zeros((2, 8, 8, 3), dtype=np.uint8)

tests/multimodal/test_processing.py

+22 -13

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import time

5 +from collections.abc import Sequence

6 from contextlib import nullcontext

7 from copy import deepcopy

8 from types import SimpleNamespace

@@ -868,7 +869,7 @@ def test_limit_mm_per_prompt_apply(model_id, num_images, limit, is_valid):

869 ],

870 )

871 def test_budget_caps_prevent_dummy_input_validation_failure(

871 - model_id, user_limit, supported_limit

872 + model_id, user_limit, supported_limit, monkeypatch

873 ):

874 limit_mm_per_prompt = {"image": user_limit}

875

@@ -878,7 +879,9 @@ def test_budget_caps_prevent_dummy_input_validation_failure(

879 )

880

881 processor = MULTIMODAL_REGISTRY.create_processor(model_config)

881 - processor.info.get_supported_mm_limits = lambda: {"image": supported_limit}

882 + monkeypatch.setattr(

883 + processor.info, "get_supported_mm_limits", lambda: {"image": supported_limit}

884 + )

885

886 # This is what budget.py uses to derive mm_counts

887 allowed = processor.info.allowed_mm_limits

@@ -936,7 +939,7 @@ def test_hf_processor_init_kwargs(

939 )

940

941 processor = ctx.get_hf_processor(

939 - DummyProcessor, # type: ignore[arg-type]

942 + DummyProcessor,

943 **inference_kwargs,

944 )

945 assert processor.a == expected_kwargs["a"]

@@ -967,7 +970,7 @@ def test_hf_processor_call_kwargs(

970 tokenizer=None,

971 )

972

970 - processor = ctx.get_hf_processor(DummyProcessor) # type: ignore[arg-type]

973 + processor = ctx.get_hf_processor(DummyProcessor)

974

975 result = ctx.call_hf_processor(processor, {}, inference_kwargs)

976 assert result == expected_kwargs

@@ -1278,7 +1281,7 @@ def test_get_supported_mm_processor_kwargs_uses_supported_modalities(

1281 processor: SimpleNamespace,

1282 expected: dict[str, set[str]],

1283 ) -> None:

1281 - info = BaseProcessingInfo(SimpleNamespace())

1284 + info = BaseProcessingInfo(SimpleNamespace()) # type: ignore[arg-type]

1285 info.__dict__["supported_mm_limits"] = supported_mm_limits

1286 monkeypatch.setattr(info, "get_hf_processor", lambda **_: processor)

1287

@@ -1288,7 +1291,7 @@ def test_get_supported_mm_processor_kwargs_uses_supported_modalities(

1291 def test_supported_mm_processor_kwargs_is_cached(

1292 monkeypatch: pytest.MonkeyPatch,

1293 ) -> None:

1291 - info = BaseProcessingInfo(SimpleNamespace())

1294 + info = BaseProcessingInfo(SimpleNamespace()) # type: ignore[arg-type]

1295 expected = {"images_kwargs": {"size"}}

1296 calls = 0

1297

@@ -1541,7 +1544,10 @@ def test_processor_inputs_hashes_scope_kwargs_by_modality():

1544 "video": [np.zeros((2, 8, 8, 3), dtype=np.uint8)],

1545 }

1546 )

1544 - mm_uuid_items = {"image": ["image-uuid"], "video": ["video-uuid"]}

1547 + mm_uuid_items: dict[str, Sequence[str | None]] = {

1548 + "image": ["image-uuid"],

1549 + "video": ["video-uuid"],

1550 + }

1551

1552 def get_hashes(video_frames: int, image_size: int, video_size: int):

1553 return ProcessorInputs(

@@ -1621,7 +1627,7 @@ def test_processor_inputs_hashes_distinguish_kwargs_shapes(left, right):

1627 [(None, 491520, 491520), (True, 491520, 8192), (True, 128, 128), (False, 128, 128)],

1628 )

1629 def test_dummy_inputs_scheduler_budget(

1624 - chunked_prefill, max_model_len, expected_seq_len

1630 + chunked_prefill, max_model_len, expected_seq_len, monkeypatch

1631 ):

1632 ctx = build_model_context(

1633 "llava-hf/llava-v1.6-mistral-7b-hf",

@@ -1634,11 +1640,9 @@ def test_dummy_inputs_scheduler_budget(

1640 ctx.model_config,

1641 tokenizer=ctx.tokenizer,

1642 )

1637 - processor.apply = lambda *args, **kwargs: {"prompt_token_ids": [7]}

1638 -

1639 - kwargs = {}

1643 + scheduler_config = None

1644 if chunked_prefill is not None:

1641 - kwargs["scheduler_config"] = SchedulerConfig(

1645 + scheduler_config = SchedulerConfig(

1646 max_model_len=max_model_len,

1647 is_encoder_decoder=False,

1648 max_num_batched_tokens=8192,

@@ -1646,5 +1650,10 @@ def test_dummy_inputs_scheduler_budget(

1650 enable_chunked_prefill=chunked_prefill,

1651 )

1652

1649 - result = processor.get_dummy_mm_inputs({"image": 1}, **kwargs)

1653 + monkeypatch.setattr(

1654 + processor, "apply", lambda *args, **kwargs: {"prompt_token_ids": [7]}

1655 + )

1656 + result = processor.get_dummy_mm_inputs(

1657 + {"image": 1}, scheduler_config=scheduler_config

1658 + )

1659 assert len(result["prompt_token_ids"]) == expected_seq_len

tests/multimodal/test_registry.py

+1 -1

@@ -24,4 +24,4 @@ def test_create_processor_error_uses_served_model_name():

24 ValueError,

25 match="friendly-model-name is not a multimodal model",

26 ):

27 - MULTIMODAL_REGISTRY.create_processor(model_config)

27 + MULTIMODAL_REGISTRY.create_processor(model_config) # type: ignore[arg-type]

tests/multimodal/test_video.py

+26 -17

@@ -7,6 +7,7 @@ import sys

7 import threading

8 from contextlib import ExitStack, contextmanager

9 from pathlib import Path

10 +from typing import Any

11

12 import numpy as np

13 import numpy.typing as npt

@@ -83,24 +84,28 @@ def _fresh_decoder_pool():

84 @VIDEO_LOADER_REGISTRY.register("test_video_loader_1")

85 class TestVideoLoader1(VideoLoader):

86 @classmethod

86 - def load_bytes(cls, data: bytes, num_frames: int = -1) -> npt.NDArray:

87 - return FAKE_OUTPUT_1

87 + def load_bytes(

88 + cls, data: bytes, num_frames: int = -1, **kwargs

89 + ) -> tuple[npt.NDArray, dict[str, Any]]:

90 + return FAKE_OUTPUT_1, {}

91

92

93 @VIDEO_LOADER_REGISTRY.register("test_video_loader_2")

94 class TestVideoLoader2(VideoLoader):

95 @classmethod

93 - def load_bytes(cls, data: bytes, num_frames: int = -1) -> npt.NDArray:

94 - return FAKE_OUTPUT_2

96 + def load_bytes(

97 + cls, data: bytes, num_frames: int = -1, **kwargs

98 + ) -> tuple[npt.NDArray, dict[str, Any]]:

99 + return FAKE_OUTPUT_2, {}

100

101

102 def test_video_loader_registry():

103 custom_loader_1 = VIDEO_LOADER_REGISTRY.load("test_video_loader_1")

99 - output_1 = custom_loader_1.load_bytes(b"test")

104 + output_1, _ = custom_loader_1.load_bytes(b"test")

105 np.testing.assert_array_equal(output_1, FAKE_OUTPUT_1)

106

107 custom_loader_2 = VIDEO_LOADER_REGISTRY.load("test_video_loader_2")

103 - output_2 = custom_loader_2.load_bytes(b"test")

108 + output_2, _ = custom_loader_2.load_bytes(b"test")

109 np.testing.assert_array_equal(output_2, FAKE_OUTPUT_2)

110

111

@@ -488,7 +493,7 @@ def test_pynvvideocodec_failed_rebuild_invalidates_decoder_slot():

493

494 old_decoder = FakeDecoder()

495 slot = PyNvVideoCodecDecoderSlot(FakeStream())

491 - slot.decoder = old_decoder

496 + slot.decoder = old_decoder # type: ignore[assignment] # Fake retained GPU decoder.

497 slot.source_path = "valid.mp4"

498

499 class FakeNvc:

@@ -596,7 +601,9 @@ def test_pynvvideocodec_h200_recovers_after_unsupported_8k():

601 _pynv_decoder_pool.max_slots = old_max

602

603

599 -def test_pynvvideocodec_cross_subclass_shares_single_pool():

604 +def test_pynvvideocodec_cross_subclass_shares_single_pool(

605 + monkeypatch: pytest.MonkeyPatch,

606 +):

607 """Regression test for GHSA-j682-9xp5-rrf3.

608

609 Multiple subclasses of PyNvVideoCodecVideoBackendMixin must share the

@@ -623,11 +630,12 @@ def test_pynvvideocodec_cross_subclass_shares_single_pool():

630 with _fresh_decoder_pool() as pool:

631 pool.max_slots = 2

632

626 - orig_create = PyNvVideoCodecVideoBackendMixin._create_decoder_slot

627 - PyNvVideoCodecVideoBackendMixin._create_decoder_slot = classmethod(

628 - fake_create_slot

629 - )

630 - try:

633 + with monkeypatch.context() as m:

634 + m.setattr(

635 + PyNvVideoCodecVideoBackendMixin,

636 + "_create_decoder_slot",

637 + classmethod(fake_create_slot),

638 + )

639 with ExitStack() as stack:

640 stack.enter_context(MixinSubclassA._borrow_decoder_slot())

641 stack.enter_context(MixinSubclassB._borrow_decoder_slot())

@@ -652,8 +660,6 @@ def test_pynvvideocodec_cross_subclass_shares_single_pool():

660

661 assert create_count == 2

662 assert len(pool.slots) == 2

655 - finally:

656 - PyNvVideoCodecVideoBackendMixin._create_decoder_slot = orig_create

663

664

665 @pytest.mark.parametrize("hw_decoders", [0, -1, 1.5, True, "2"])

@@ -662,7 +668,7 @@ def test_pynvvideocodec_rejects_invalid_hw_decoders(hw_decoders: object):

668 VideoBackend.load_bytes(

669 b"fake video",

670 backend=PYNVVIDEOCODEC_VIDEO_BACKEND,

665 - hw_decoders=hw_decoders, # type: ignore[arg-type]

671 + hw_decoders=hw_decoders,

672 )

673

674

@@ -799,6 +805,7 @@ def test_video_processor_from_model_repo(

805 )

806

807 backend = get_video_loader_backend_for_processor(video_processor)

808 + assert backend is not None

809 loader = VIDEO_LOADER_REGISTRY.load(backend)

810 assert isinstance(loader, expected_loader_cls), (

811 f"{model_repo!r}: backend={backend!r} loaded "

@@ -1693,7 +1700,9 @@ def test_glm5next_backend_codec_parity(tmp_path, backend):

1700 path = _write_gray_video(tmp_path, total_frames, fps)

1701 # Dense default sampling (gap 5) and a sparse max_frames cap (gap 20).

1702 for max_frames in (None, 6):

1696 - kwargs = {} if max_frames is None else {"max_frames": max_frames}

1703 + kwargs: dict[str, Any] = (

1704 + {} if max_frames is None else {"max_frames": max_frames}

1705 + )

1706 expected = glm_sample_frame_indices(

1707 total_frames, float(fps), 12.0, max_frame_count=max_frames

1708 )

tests/quantization/test_auto_gptq.py

+4 -4

@@ -123,8 +123,8 @@ def test_auto_gptq_moe_creates_zero_initialized_expert_biases():

123 method = object.__new__(AutoGPTQMoEMethod)

124 method.quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {})

125 method.input_dtype = None

126 - method.experts_cls = None

127 - method.moe = SimpleNamespace(w13_num_shards=2)

126 + method.experts_cls = None # type: ignore[assignment] # Weight creation does not select experts.

127 + method.moe = SimpleNamespace(w13_num_shards=2) # type: ignore[assignment]

128 layer = torch.nn.Module()

129

130 method.create_weights(

@@ -170,7 +170,7 @@ def test_routed_experts_loads_per_expert_biases():

170 ("w1", torch.tensor([1.0, 2.0, 3.0, 4.0])),

171 ("w3", torch.tensor([5.0, 6.0, 7.0, 8.0])),

172 ):

173 - assert RoutedExperts.weight_loader(

173 + assert RoutedExperts.weight_loader( # type: ignore[call-overload]

174 loader,

175 w13_bias,

176 loaded,

@@ -180,7 +180,7 @@ def test_routed_experts_loads_per_expert_biases():

180 return_success=True,

181 )

182

183 - assert RoutedExperts.weight_loader(

183 + assert RoutedExperts.weight_loader( # type: ignore[call-overload]

184 loader,

185 w2_bias,

186 torch.tensor([9.0, 10.0, 11.0, 12.0]),

tests/quantization/test_auto_round.py

+9 -9

@@ -193,7 +193,7 @@ class DummyFusedMoE:

193

194

195 def make_config(**overrides) -> INCConfig:

196 - kwargs = {

196 + kwargs: dict[str, Any] = {

197 "weight_bits": 4,

198 "group_size": 128,

199 "sym": True,

@@ -208,7 +208,7 @@ def make_config(**overrides) -> INCConfig:

208

209

210 def make_layer_config(**overrides) -> INCLayerConfig:

211 - kwargs = {

211 + kwargs: dict[str, Any] = {

212 "bits": 4,

213 "group_size": 128,

214 "sym": True,

@@ -423,7 +423,7 @@ def test_inc_get_quant_method_unquantized_routed_experts_with_model_prefix(

423

424 config = make_config(extra_config={"model.moe_layer": {"bits": 16}})

425 layer = object.__new__(RoutedExperts)

426 - layer.moe_config = None

426 + layer.moe_config = None # type: ignore[assignment] # No MoE config used by this stub.

427

428 method = config.get_quant_method(layer, "moe_layer")

429

@@ -730,7 +730,7 @@ def test_wna16_cuda_low_bit_moe_routes_to_humming(monkeypatch, bits) -> None:

730 )

731

732 layer = object.__new__(RoutedExperts)

733 - layer.moe_config = DummyMoeConfig()

733 + layer.moe_config = DummyMoeConfig() # type: ignore[assignment]

734 layer_config = make_layer_config(bits=bits)

735 method = INCWna16Scheme().get_moe_method(

736 make_config(), layer, "model.layers.0.mlp", layer_config

@@ -861,7 +861,7 @@ def test_qwen3_30b_a3b_w4a16_autoround_routes_to_gptq_moe(

861

862 config = make_qwen3_autoround_config("qwen3_30b_a3b_w4a16")

863 layer = object.__new__(RoutedExperts)

864 - layer.moe_config = DummyMoeConfig()

864 + layer.moe_config = DummyMoeConfig() # type: ignore[assignment]

865

866 method = config.get_quant_method(layer, "model.layers.0.mlp")

867

@@ -890,7 +890,7 @@ def test_qwen3_30b_a3b_mxfp4_autoround_routes_to_mxfp4_moe(

890

891 config = make_qwen3_autoround_config("qwen3_30b_a3b_mxfp4")

892 layer = object.__new__(RoutedExperts)

893 - layer.moe_config = DummyMoeConfig()

893 + layer.moe_config = DummyMoeConfig() # type: ignore[assignment]

894

895 ignored_method = config.get_quant_method(

896 object.__new__(LinearBase), "model.layers.0.self_attn.q_proj"

@@ -1186,7 +1186,7 @@ def test_wna16_xpu_moe_routes_to_gptq_moe(monkeypatch) -> None:

1186 )

1187

1188 layer = object.__new__(RoutedExperts)

1189 - layer.moe_config = DummyMoeConfig()

1189 + layer.moe_config = DummyMoeConfig() # type: ignore[assignment]

1190 method = INCWna16Scheme().get_moe_method(

1191 make_config(),

1192 layer,

@@ -1669,7 +1669,7 @@ def test_inc_get_quant_method_unquantized_moe_returns_unquantized(

1669 when extra_config has bits >= 16."""

1670 config = make_config(extra_config={"layer": {"bits": 16}})

1671 layer = object.__new__(RoutedExperts)

1672 - layer.moe_config = None # UnquantizedFusedMoEMethod accepts moe_config

1672 + layer.moe_config = None # type: ignore[assignment] # No MoE config used by this stub.

1673

1674 class DummyUnquantizedFusedMoEMethod:

1675 def __init__(self, moe_config) -> None:

@@ -2094,7 +2094,7 @@ def _with_w4a8_kernel(monkeypatch) -> None:

2094

2095 def _dispatch(layer_config=None):

2096 return INCWna16Scheme().get_linear_method(

2097 - object(), object(), "layer", layer_config or make_layer_config()

2097 + make_config(), object(), "layer", layer_config or make_layer_config()

2098 )

2099

2100

tests/quantization/test_cutlass_w4a16.py

+1

@@ -99,6 +99,7 @@ def test_machete_rejects_invalid_config(

99 )

100 can_impl, reason = MacheteLinearKernel.can_implement(config)

101 assert not can_impl

102 + assert reason is not None

103 assert expected_reason in reason

104

105

tests/quantization/test_fp8.py

+27 -14

@@ -25,6 +25,7 @@ from vllm.model_executor.layers.attention.attention import (

25 set_default_quant_scales,

26 )

27 from vllm.model_executor.layers.fused_moe import FusedMoEFactory

28 +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig

29 from vllm.model_executor.layers.quantization.fp8 import (

30 Fp8Config,

31 Fp8LinearMethod,

@@ -79,17 +80,22 @@ def test_deepseek_v41_mxfp8_scale_loading(

80 for module in (model_module, linear_module):

81 monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)

82 monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: tp_rank)

82 - tp_args = dict(quant_config=quant_config, bias=False)

83 block = torch.nn.Module()

84 block.attn = torch.nn.Module()

85 - block.attn.wq_b = ColumnParallelLinear(128, 128, **tp_args)

86 - block.attn.fused_wqa_wkv = MergedColumnParallelLinear(128, [128, 64], **tp_args)

85 + block.attn.wq_b = ColumnParallelLinear(

86 + 128, 128, quant_config=quant_config, bias=False

87 + )

88 + block.attn.fused_wqa_wkv = MergedColumnParallelLinear(

89 + 128, [128, 64], quant_config=quant_config, bias=False

90 + )

91 block.ffn = torch.nn.Module()

92 block.ffn.shared_experts = torch.nn.Module()

93 block.ffn.shared_experts.gate_up_proj = MergedColumnParallelLinear(

90 - 128, [128, 128], **tp_args

94 + 128, [128, 128], quant_config=quant_config, bias=False

95 + )

96 + block.ffn.shared_experts.down_proj = RowParallelLinear(

97 + 128, 128, quant_config=quant_config, bias=False

98 )

92 - block.ffn.shared_experts.down_proj = RowParallelLinear(128, 128, **tp_args)

99 block.ffn.experts = torch.nn.Module()

100 block.ffn.experts.register_parameter(

101 "weight_scale", torch.nn.Parameter(torch.empty(3, 4, dtype=torch.uint8), False)

@@ -111,9 +117,11 @@ def test_deepseek_v41_mxfp8_scale_loading(

117 model.config = SimpleNamespace(num_attention_heads=2)

118 model.quant_config = quant_config

119 model.use_sequence_parallel = False

114 - model.get_expert_mapping = lambda: [

115 - ("experts.weight_scale", "experts.0.w1.weight_scale", 0, "w1")

116 - ]

120 + monkeypatch.setattr(

121 + model,

122 + "get_expert_mapping",

123 + lambda: [("experts.weight_scale", "experts.0.w1.weight_scale", 0, "w1")],

124 + )

125

126 checkpoints = [

127 ("attn.wq_b", "attn.wq_b", 128, 128, 0, None),

@@ -211,7 +219,10 @@ def test_deepseek_v41_vl_mapper_routes_linear_scales(

219 vllm_config = SimpleNamespace(

220 quant_config=SimpleNamespace(weight_block_size=weight_block_size)

221 )

214 - resolved = model_module._linear_scale_param_name(vllm_config, expert_dtype)

222 + resolved = model_module._linear_scale_param_name(

223 + vllm_config, # type: ignore[arg-type]

224 + expert_dtype,

225 + )

226 assert resolved == scale_name

227

228 mapper = vl_module._make_deepseek_v4_vl_weights_mapper(expert_dtype, resolved)

@@ -373,7 +384,7 @@ def test_deepgemm_mxfp8_preserves_weight_and_scale_values(

384 @pytest.mark.parametrize("config_source", ["deepseek", "mxfp8"])

385 @pytest.mark.parametrize("num_tokens", [1, 7, 128])

386 def test_mxfp8_bmm_loads_and_projects_grouped_weights(

376 - dist_init, default_vllm_config, prequantized, config_source, num_tokens

387 + dist_init, default_vllm_config, monkeypatch, prequantized, config_source, num_tokens

388 ):

389 """BMM metadata set after construction selects grouped weight processing."""

390 from vllm.model_executor.kernels.linear.mxfp8.deep_gemm import (

@@ -398,10 +409,11 @@ def test_mxfp8_bmm_loads_and_projects_grouped_weights(

409 pytest.skip("DeepGEMM MXFP8 BMM requires Blackwell")

410

411 default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16)

401 - quant_config = DeepseekV4FP8Config(

412 + deepseek_config = DeepseekV4FP8Config(

413 is_checkpoint_fp8_serialized=True, weight_block_size=[32, 32]

414 )

404 - quant_config._resolved_expert_dtype = "fp4"

415 + deepseek_config._resolved_expert_dtype = "fp4"

416 + quant_config: QuantizationConfig = deepseek_config

417 if config_source == "mxfp8":

418 quant_config = get_quantization_config("mxfp8").from_config(

419 {"quant_method": "mxfp8"}

@@ -432,7 +444,7 @@ def test_mxfp8_bmm_loads_and_projects_grouped_weights(

444 model.config = SimpleNamespace(num_attention_heads=2)

445 model.quant_config = quant_config

446 model.use_sequence_parallel = False

435 - model.get_expert_mapping = lambda: []

447 + monkeypatch.setattr(model, "get_expert_mapping", lambda: [])

448 model.load_weights(

449 [

450 ("layers.0.attn.wo_a.weight", weight),

@@ -443,6 +455,7 @@ def test_mxfp8_bmm_loads_and_projects_grouped_weights(

455 weight.float().reshape(8, 32, 16, 32)

456 * torch.exp2(scales.float() - 127)[:, None, :, None]

457 ).reshape(2, 128, 512)

458 + assert isinstance(linear.quant_method, ModelOptLinearMethod)

459 linear.quant_method.process_weights_after_loading(linear)

460 linear.quant_method.process_weights_after_loading(linear)

461 assert isinstance(linear.quant_method.kernel, DeepGemmMxfp8BmmLinearKernel)

@@ -785,7 +798,7 @@ def test_kv_cache_scale_sync_to_host_copies():

798 set_default_quant_scales(layer, register_buffer=True)

799 layer.kv_cache_dtype = "fp8"

800

788 - method = BaseKVCacheMethod(quant_config=None)

801 + method = BaseKVCacheMethod(quant_config=None) # type: ignore[arg-type] # Scale synchronization does not use quant_config.

802 method.create_weights(layer)

803 # 0.3 stays != 1.0 even after the fp8_fnuz x2 rescale.

804 checkpoint_scale = torch.tensor(0.3, dtype=torch.float32)

tests/quantization/test_gfx950_moe.py

+2 -1

@@ -8,6 +8,7 @@ These tests run on real hardware — no mocks. Skipped on non-GFX950 platforms.

8 import pytest

9 import torch

10

11 +from vllm.config.kernel import MoEBackend

12 from vllm.model_executor.layers.fused_moe.config import (

13 FusedMoEConfig,

14 FusedMoEParallelConfig,

@@ -53,7 +54,7 @@ def disable_rocm_aiter(monkeypatch: pytest.MonkeyPatch):

54 set_rocm_aiter(monkeypatch, False)

55

56

56 -def _make_w4a4_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:

57 +def _make_w4a4_moe_config(moe_backend: MoEBackend = "auto") -> FusedMoEConfig:

58 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

59

60 return FusedMoEConfig(

tests/quantization/test_int8_moe_oracle.py

+2 -1

@@ -10,6 +10,7 @@ CUDA (SM >= 7.5), ROCm, or XPU — not just gfx950.

10 import pytest

11 import torch

12

13 +from vllm.config.kernel import MoEBackend

14 from vllm.model_executor.layers.fused_moe.config import (

15 FusedMoEConfig,

16 FusedMoEParallelConfig,

@@ -44,7 +45,7 @@ requires_int8_moe = pytest.mark.skipif(

45 DEVICE = current_platform.device_type

46

47

47 -def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:

48 +def _make_int8_moe_config(moe_backend: MoEBackend = "auto") -> FusedMoEConfig:

49 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

50

51 return FusedMoEConfig(

tests/quantization/test_lm_head.py

+2

@@ -15,6 +15,7 @@ from tests.quantization.utils import load_model_without_vllm_runner

15 from vllm.model_executor.layers.linear import ColumnParallelLinear

16 from vllm.model_executor.layers.logits_processor import LogitsProcessor

17 from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQLinearMethod

18 +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig

19 from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501

20 CompressedTensorsConfig,

21 )

@@ -80,6 +81,7 @@ def test_quantized_lm_head_matches_linear(

81 dtype=torch.bfloat16, head_dtype=None

82 )

83 default_vllm_config.kernel_config.linear_backend = "humming"

84 + quant_config: QuantizationConfig

85 if quant_format is None:

86 quant_config = ModelOptNvFp4Config(

87 quant_method=preset,

tests/quantization/test_modelopt.py

+1

@@ -843,6 +843,7 @@ def test_modelopt_mixed_precision_builds_w4a16_sibling_config():

843 }

844 config = m.ModelOptMixedPrecisionConfig.from_config(hf_quant_config)

845

846 + assert isinstance(config, m.ModelOptMixedPrecisionConfig)

847 assert config.nvfp4_config.quant_method == "NVFP4"

848 assert config.w4a16_nvfp4_config.quant_method == "W4A16_NVFP4"

849

tests/quantization/test_moe_wna16.py

+7 -5

@@ -218,11 +218,13 @@ def test_gptq_nfirst_weights_pass_through_for_triton():

218

219 def test_moe_wna16_setup_forwards_selected_backend(monkeypatch):

220 method = object.__new__(MoeWNA16Method)

221 - method.experts_cls = object

221 + method.experts_cls = object # type: ignore[assignment] # Opaque mocked-expert sentinel.

222 method.wna16_backend = WNA16MoEBackend.HUMMING

223 - method.moe = object()

223 + method.moe = object() # type: ignore[assignment] # Opaque mocked-config sentinel.

224 quant_config = object()

225 - method.get_fused_moe_quant_config = lambda layer: quant_config

225 + monkeypatch.setattr(

226 + method, "get_fused_moe_quant_config", lambda layer: quant_config

227 + )

228 layer = SimpleNamespace(_expert_routing_tables=lambda: (None, None, None))

229 captured = {}

230 kernel = object()

@@ -233,7 +235,7 @@ def test_moe_wna16_setup_forwards_selected_backend(monkeypatch):

235

236 monkeypatch.setattr(moe_wna16, "make_wna16_moe_kernel", fake_make_wna16_moe_kernel)

237

236 - method._setup_kernel(layer)

238 + method._setup_kernel(layer) # type: ignore[arg-type] # Routing-table stub.

239

240 assert method.moe_kernel is kernel

241 assert captured["backend"] == WNA16MoEBackend.HUMMING

@@ -278,7 +280,7 @@ def test_moe_wna16_uses_humming_quant_config(monkeypatch):

280 ),

281 )

282

281 - assert method.get_fused_moe_quant_config(layer) is quant_config

283 + assert method.get_fused_moe_quant_config(layer) is quant_config # type: ignore[arg-type] # Identity sentinel.

284

285

286 @pytest.mark.skipif(

tests/quantization/test_online.py

+15 -14

@@ -16,6 +16,7 @@ from tests.quantization.utils import (

16 _test_online_quant_peak_mem_impl,

17 is_quant_method_supported,

18 load_model_without_vllm_runner,

19 + quant_config_args,

20 )

21 from vllm import _custom_ops as ops

22 from vllm._aiter_ops import rocm_aiter_ops

@@ -180,7 +181,7 @@ def test_qianfan_online_fp8_keeps_vision_layers_unquantized(

181 config = OnlineQuantizationConfig(args)

182 other_config = OnlineQuantizationConfig(args)

183 model = QianfanOCRForConditionalGeneration.__new__(

183 - QianfanOCRForConditionalGeneration

184 + QianfanOCRForConditionalGeneration # type: ignore[type-abstract]

185 )

186 model._patch_quant_config(

187 SimpleNamespace(vision_config=SimpleNamespace(num_hidden_layers=2)),

@@ -209,9 +210,9 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(

210 monkeypatch,

211 ) -> None:

212 method = object.__new__(Nvfp4OnlineMoEMethod)

212 - method.moe = SimpleNamespace(is_act_and_mul=True)

213 - method.nvfp4_backend = object()

214 - method.experts_cls = object

213 + method.moe = SimpleNamespace(is_act_and_mul=True) # type: ignore[assignment]

214 + method.nvfp4_backend = object() # type: ignore[assignment] # Opaque mocked-backend sentinel.

215 + method.experts_cls = object # type: ignore[assignment] # Opaque mocked-expert sentinel.

216 method.moe_quant_config = None

217 method.moe_kernel = None

218

@@ -239,7 +240,7 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(

240 )

241 make_kernel = Mock(return_value=kernel)

242 get_quant_config = Mock(return_value=quant_config)

242 - method.get_fused_moe_quant_config = get_quant_config

243 + monkeypatch.setattr(method, "get_fused_moe_quant_config", get_quant_config)

244

245 monkeypatch.setattr(

246 "vllm.model_executor.layers.quantization.online.nvfp4."

@@ -355,7 +356,7 @@ def test_online_prequantized_compatibility(

356 checkpoint_config = checkpoint_config_factory()

357

358 checkpoint_config.online_quantization_config = OnlineQuantizationConfig(

358 - QuantizationConfigArgs(linear="mxfp8")

359 + quant_config_args(linear="mxfp8")

360 )

361 config = checkpoint_config

362

@@ -408,7 +409,7 @@ def test_online_ignore_keeps_checkpoint_quantization_linear(

409 quant_config = _fully_quantized_quark_config()

410 prefix = "model.layers.0.self_attn.o_proj"

411 quant_config.online_quantization_config = OnlineQuantizationConfig(

411 - QuantizationConfigArgs(linear="mxfp8", ignore=[prefix])

412 + quant_config_args(linear="mxfp8", ignore=[prefix])

413 )

414 monkeypatch.setattr(

415 quant_config.online_quantization_config,

@@ -437,7 +438,7 @@ def test_online_quantization_rejects_prequantized_moe(

438 prefix = "model.layers.0.mlp.experts"

439 quant_config = _fully_quantized_quark_config()

440 quant_config.online_quantization_config = OnlineQuantizationConfig(

440 - QuantizationConfigArgs(linear="mxfp4", moe="mxfp4")

441 + quant_config_args(linear="mxfp4", moe="mxfp4")

442 )

443

444 with pytest.raises(ValueError, match="pre-quantized layer"):

@@ -460,7 +461,7 @@ def test_activation_only_override_applies_to_checkpoint_method(

461 model_config = ModelConfig(

462 model=str(tmp_path),

463 quantization="compressed-tensors",

463 - quantization_config=QuantizationConfigArgs(moe={"activation": "mxfp8"}),

464 + quantization_config=quant_config_args(moe={"activation": "mxfp8"}),

465 hf_overrides={

466 "quantization_config": {

467 "quant_method": "compressed-tensors",

@@ -514,7 +515,7 @@ def test_online_overlay_loads_checkpoint_config_file(tmp_path) -> None:

515 model_config = ModelConfig(

516 model=str(tmp_path),

517 quantization="awq",

517 - quantization_config=QuantizationConfigArgs(linear="mxfp8"),

518 + quantization_config=quant_config_args(linear="mxfp8"),

519 )

520

521 result = weight_utils.get_quant_config(model_config, LoadConfig())

@@ -574,7 +575,7 @@ def test_online_shorthand_selects_checkpoint_or_online_config(

575

576 def test_log_online_quantization_for_composable_config(monkeypatch) -> None:

577 """Composable configs log their nested online quantization results."""

577 - online_config = OnlineQuantizationConfig(QuantizationConfigArgs(linear="mxfp8"))

578 + online_config = OnlineQuantizationConfig(quant_config_args(linear="mxfp8"))

579 online_config.quantized_layers = {

580 "model.layers.0.self_attn.o_proj": ("linear", "mxfp8", None),

581 "model.layers.1.self_attn.o_proj": ("linear", "mxfp8", None),

@@ -1164,7 +1165,7 @@ def test_online_quantization_records_global_config(

1165 default_vllm_config, dist_init

1166 ) -> None:

1167 default_vllm_config.model_config = ModelConfig()

1167 - config = OnlineQuantizationConfig(QuantizationConfigArgs(linear="fp8_per_block"))

1168 + config = OnlineQuantizationConfig(quant_config_args(linear="fp8_per_block"))

1169 prefix = "model.layers.0.self_attn.o_proj"

1170 layer = ColumnParallelLinear(

1171 input_size=1,

@@ -1218,7 +1219,7 @@ def test_online_quantization_targets_reject_unsupported_layer() -> None:

1219

1220

1221 def test_log_online_quantization(default_vllm_config, monkeypatch) -> None:

1221 - config = OnlineQuantizationConfig(QuantizationConfigArgs(linear="fp8_per_tensor"))

1222 + config = OnlineQuantizationConfig(quant_config_args(linear="fp8_per_tensor"))

1223 config.quantized_layers = {

1224 "model.layers.0.mlp.down_proj": ("linear", "fp8_per_tensor", None),

1225 "model.layers.1.mlp.down_proj": ("linear", "fp8_per_tensor", None),

@@ -1431,7 +1432,7 @@ def test_online_int8_moe_w2_scale_matches_unsharded(monkeypatch) -> None:

1432 layer.w2_weight = torch.nn.Parameter(w2_in, requires_grad=False)

1433 layer.num_experts = layer.local_num_experts = w13.shape[0]

1434 method = SimpleNamespace(moe=SimpleNamespace(tp_size=moe_tp_size))

1434 - Int8OnlineMoEMethod._quantize_weights(method, layer)

1435 + Int8OnlineMoEMethod._quantize_weights(method, layer) # type: ignore[arg-type]

1436 return layer.w2_weight, layer.w2_scale

1437

1438 full_weight, full_scale = quantize(w2, 1)

tests/quantization/test_online_mxfp4.py

+2 -1

@@ -10,6 +10,7 @@ against the pure-torch reference in `reference_mxfp4.py`.

10 import pytest

11 import torch

12

13 +from vllm.config.kernel import MoEBackend

14 from vllm.config.model import ModelConfig

15 from vllm.model_executor.kernels.linear import _POSSIBLE_MXFP4_KERNELS

16 from vllm.model_executor.kernels.linear.mxfp4.aiter import (

@@ -282,7 +283,7 @@ def test_online_mxfp4_tp_weight_quant_matches_unsharded(tp_size: int):

283 ],

284 )

285 def test_online_mxfp4_moe_matches_quark(

285 - moe_backend: str,

286 + moe_backend: MoEBackend,

287 unpadded_hidden_size: int,

288 unpadded_intermediate_size: int,

289 default_vllm_config,

tests/quantization/test_quantization_config_args.py

+16 -11

@@ -6,6 +6,7 @@ from unittest.mock import Mock

6

7 import pytest

8

9 +from tests.quantization.utils import quant_config_args, quant_spec

10 from vllm.config.quantization import (

11 QUANT_KEY_NAMES,

12 QuantizationConfigArgs,

@@ -30,7 +31,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

31

32

33 def test_quant_spec_resolves_string_to_quant_key():

33 - spec = QuantSpec(weight="mxfp8", activation="fp8_per_token")

34 + spec = quant_spec(weight="mxfp8", activation="fp8_per_token")

35 assert spec.weight == kMxfp8Dynamic

36 assert spec.activation == kFp8DynamicTokenSym

37

@@ -42,12 +43,12 @@ def test_quant_spec_accepts_quant_key_directly():

43

44

45 def test_quant_spec_string_representation_uses_quantization_name():

45 - assert str(QuantSpec(weight="mxfp4")) == "mxfp4"

46 + assert str(quant_spec(weight="mxfp4")) == "mxfp4"

47

48

49 def test_quant_spec_rejects_unknown_name():

50 with pytest.raises(ValueError, match="unknown quantization name"):

50 - QuantSpec(weight="not_a_real_format")

51 + quant_spec(weight="not_a_real_format")

52

53

54 # ---- QuantizationConfigArgs string shorthand on linear/moe ----------------

@@ -55,7 +56,7 @@ def test_quant_spec_rejects_unknown_name():

56

57 def test_args_linear_string_resolves_via_quant_key_names():

58 # A bare QUANT_KEY_NAMES entry desugars to QuantSpec(weight=<key>).

58 - args = QuantizationConfigArgs(linear="fp8_per_block_static")

59 + args = quant_config_args(linear="fp8_per_block_static")

60 assert args.linear == QuantSpec(weight=kFp8Static128BlockSym)

61 assert args.moe is None

62

@@ -65,7 +66,7 @@ def test_args_moe_string_resolves_via_online_shorthand():

66 # (so `linear: "fp8_per_block"` and `moe: "fp8_per_block"` produce the

67 # same per-layer-kind spec the `--quantization fp8_per_block` shorthand

68 # would).

68 - args = QuantizationConfigArgs(moe="fp8_per_block")

69 + args = quant_config_args(moe="fp8_per_block")

70 assert args.moe == QuantSpec(weight=kFp8Static128BlockSym)

71

72

@@ -73,11 +74,11 @@ def test_args_string_shorthand_missing_slot_raises():

74 # int8_per_channel_weight_only sets only `moe`; using it on `linear`

75 # has no defined spec and should raise rather than silently no-op.

76 with pytest.raises(ValueError, match="does not define a linear spec"):

76 - QuantizationConfigArgs(linear="int8_per_channel_weight_only")

77 + quant_config_args(linear="int8_per_channel_weight_only")

78

79

80 def test_args_accepts_dict_form():

80 - args = QuantizationConfigArgs(moe={"activation": "mxfp8"})

81 + args = quant_config_args(moe={"activation": "mxfp8"})

82 assert args.moe == QuantSpec(weight=None, activation=kMxfp8Dynamic)

83

84

@@ -91,6 +92,7 @@ def test_targets_reject_non_string_keys():

92

93 def test_resolve_shorthand_only_populates_both_slots():

94 args = resolve_quantization_config("fp8_per_block", None)

95 + assert args is not None

96 assert args.linear == QuantSpec(weight=kFp8Static128BlockSym)

97 assert args.moe == QuantSpec(weight=kFp8Static128BlockSym)

98

@@ -105,6 +107,7 @@ def test_resolve_int8_shorthand_leaves_linear_unset():

107 # int8_per_channel_weight_only is MoE-only; linear stays None so that

108 # OnlineQuantizationConfig leaves Linear layers in full precision.

109 args = resolve_quantization_config("int8_per_channel_weight_only", None)

110 + assert args is not None

111 assert args.linear is None

112 assert args.moe == QuantSpec(weight=kInt8StaticChannelSym)

113

@@ -113,6 +116,7 @@ def test_resolve_quantization_config_only():

116 # When only `quantization_config` is given (e.g. for an already-quantized

117 # checkpoint that needs an activation override), it's returned as-is.

118 args = resolve_quantization_config(None, {"moe": {"activation": "mxfp8"}})

119 + assert args is not None

120 assert args.linear is None

121 assert args.moe == QuantSpec(weight=None, activation=kMxfp8Dynamic)

122

@@ -124,13 +128,14 @@ def test_resolve_merges_explicit_over_shorthand():

128 "fp8_per_tensor",

129 {"linear": "fp8_per_block"},

130 )

131 + assert args is not None

132 assert args.linear == QuantSpec(weight=kFp8Static128BlockSym)

133 assert args.moe == QuantSpec(weight=kFp8StaticTensorSym)

134

135

136 def test_resolve_quantization_config_with_checkpoint_quantization():

137 args = resolve_quantization_config("gptq", {"linear": "fp8_per_block"})

133 - assert args == QuantizationConfigArgs(linear="fp8_per_block")

138 + assert args == quant_config_args(linear="fp8_per_block")

139

140

141 # ---- QUANT_KEY_NAMES coverage --------------------------------------------

@@ -140,15 +145,15 @@ def test_quant_key_names_round_trip():

145 # Every advertised name should round-trip through QuantSpec without error

146 # and produce the same QuantKey it maps to.

147 for name, expected in QUANT_KEY_NAMES.items():

143 - assert QuantSpec(weight=name).weight == expected, name

144 - assert QuantSpec(activation=name).activation == expected, name

148 + assert quant_spec(weight=name).weight == expected, name

149 + assert quant_spec(activation=name).activation == expected, name

150

151

152 def test_static_block_weight_paired_with_dynamic_block_activation():

153 # The block-FP8 shorthand pair: 128x128 static weights + 1x128 dynamic

154 # activations. Pinning this so renames in QUANT_KEY_NAMES don't quietly

155 # rewire the kernel dispatch.

151 - spec = QuantSpec(weight="fp8_per_block_static", activation="fp8_per_block_dynamic")

156 + spec = quant_spec(weight="fp8_per_block_static", activation="fp8_per_block_dynamic")

157 assert spec.weight == kFp8Static128BlockSym

158 assert spec.activation == kFp8Dynamic128Sym

159

tests/quantization/test_quark.py

+8 -5

@@ -22,7 +22,7 @@ from packaging import version

22 from tests.quantization.utils import load_model_without_vllm_runner

23 from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops

24 from vllm.config import VllmConfig, set_current_vllm_config

25 -from vllm.config.cache import CacheConfig

25 +from vllm.config.cache import CacheConfig, CacheDType

26 from vllm.forward_context import set_forward_context

27 from vllm.model_executor import parameter

28 from vllm.model_executor.kernels.linear.scaled_mm.aiter import (

@@ -1081,6 +1081,7 @@ def test_quant_method_dispatch_mxfp8_moe_raises(default_vllm_config):

1081 )

1082 wk, ak, mcls = config.get_quant_method_target("experts", RoutedExperts)

1083 assert mcls is QuarkOCP_MX_MoEMethod

1084 + assert wk is not None

1085 # The OCP MX MoE constructor should fail loudly for MXFP8.

1086 fake_moe_config = MagicMock()

1087 with pytest.raises(ValueError, match="MXFP8 experts are not supported"):

@@ -1321,7 +1322,6 @@ def test_quark_fp8_ptpc_exposes_kernel_input_quant_key(monkeypatch):

1322 scheme = QuarkW8A8Fp8.__new__(QuarkW8A8Fp8)

1323 scheme.weight_qscheme = "per_channel"

1324 scheme.is_static_input_scheme = False

1324 - scheme.input_qscheme = "per_channel"

1325 scheme.activation_quant_key = kFp8DynamicTokenSym

1326 scheme.weight_quant_key = kFp8StaticChannelSym

1327 scheme.out_dtype = dtype

@@ -1343,7 +1343,7 @@ def test_quark_fp8_ptpc_exposes_kernel_input_quant_key(monkeypatch):

1343

1344 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])

1345 def test_quark_fp8_w_per_tensor_a_per_tensor(

1346 - kv_cache_dtype: str, monkeypatch, dist_init, workspace_init

1346 + kv_cache_dtype: CacheDType, monkeypatch, dist_init, workspace_init

1347 ):

1348 model_path = "amd/Llama-3.1-8B-Instruct-FP8-KV-Quark-test"

1349 checkpoint_scales = {}

@@ -2314,7 +2314,7 @@ class TestQuarkInt4Format:

2314 None,

2315 quant_config.quant_config["global_quant_config"]["weight"],

2316 quant_config.pack_method,

2317 - moe_config,

2317 + moe_config, # type: ignore[arg-type]

2318 )

2319

2320 layer = _FakeLayer(moe_config)

@@ -2325,7 +2325,10 @@ class TestQuarkInt4Format:

2325 torch.zeros(1, 16, 8 // tp_size, dtype=torch.uint8),

2326 requires_grad=False,

2327 )

2328 - loader = method.get_weight_loader(layer, weight_loader=None)

2328 + loader = method.get_weight_loader(

2329 + layer, # type: ignore[arg-type]

2330 + weight_loader=None,

2331 + )

2332 loader(

2333 param,

2334 raw_zp.clone(),

tests/quantization/test_register_quantization_config.py

+1 -1

@@ -75,7 +75,7 @@ class CustomQuantConfig(QuantizationConfig):

75

76 def get_name(self) -> QuantizationMethods:

77 """Name of the quantization method."""

78 - return "custom_quant"

78 + return "custom_quant" # type: ignore[return-value] # Runtime registration extends built-in names.

79

80 def get_supported_act_dtypes(self) -> list[torch.dtype]:

81 """List of supported activation dtypes."""

tests/quantization/test_torchao.py

+2

@@ -12,6 +12,7 @@ from vllm.config import set_current_vllm_config

12 from vllm.forward_context import set_forward_context

13 from vllm.model_executor.layers.attention import Attention

14 from vllm.model_executor.model_loader import get_model_loader

15 +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader

16 from vllm.platforms import current_platform

17

18 if current_platform.is_rocm():

@@ -126,6 +127,7 @@ def test_online_quant_config_dict_json(vllm_runner, enable_pickle):

127

128 def load_weights(model):

129 model_loader = get_model_loader(load_config)

130 + assert isinstance(model_loader, DefaultModelLoader)

131 weights_iterator = model_loader.get_all_weights(model_config, model)

132 model.load_weights(weights_iterator)

133

tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py

+1 -1

@@ -34,7 +34,7 @@ def test_shared_nvfp4_input_scales_have_writable_storage(monkeypatch):

34

35 outputs = prepare_nvfp4_moe_layer_for_fi_or_cutlass(

36 backend=NvFp4MoeBackend.FLASHINFER_CUTLASS,

37 - layer=layer,

37 + layer=layer, # type: ignore[arg-type]

38 w13=w13,

39 w13_scale=w13_scale,

40 w13_scale_2=weight_scale,

tests/quantization/test_turboquant.py

+1 -1

@@ -300,7 +300,7 @@ class TestTurboQuantKVCacheSpec:

300

301 # The layer builds an unpacked spec; the worker's spec-collection

302 # loop applies TQ slot packing via the backend's customize_spec hook.

303 - spec = Attention.get_kv_cache_spec(layer, vllm_config)

303 + spec = Attention.get_kv_cache_spec(layer, vllm_config) # type: ignore[arg-type]

304 assert isinstance(spec, FullAttentionSpec)

305 assert spec.kv_quant_mode.is_turboquant

306 assert spec.state_content_bytes is None

tests/quantization/utils.py

+16 -1

@@ -6,18 +6,31 @@ from typing import Any

6

7 import regex as re

8 import torch

9 +from pydantic import TypeAdapter

10

11 from vllm.config import (

12 CompilationConfig,

13 + CompilationMode,

14 ModelConfig,

15 VllmConfig,

16 set_current_vllm_config,

17 )

18 +from vllm.config.quantization import QuantizationConfigArgs, QuantSpec

19 from vllm.model_executor.layers.quantization import get_quantization_config

20 from vllm.model_executor.model_loader.default_loader import DefaultModelLoader

21 from vllm.platforms import current_platform

22

23

24 +def quant_spec(**kwargs: Any) -> QuantSpec:

25 + """Build a `QuantSpec` from user-facing names such as `"mxfp8"`."""

26 + return TypeAdapter(QuantSpec).validate_python(kwargs)

27 +

28 +

29 +def quant_config_args(**kwargs: Any) -> QuantizationConfigArgs:

30 + """Build `QuantizationConfigArgs` from user-facing shorthand values."""

31 + return TypeAdapter(QuantizationConfigArgs).validate_python(kwargs)

32 +

33 +

34 def _limit_num_hidden_layers(

35 model: torch.nn.Module, num_hidden_layers: int | None

36 ) -> None:

@@ -59,7 +72,9 @@ def load_model_without_vllm_runner(

72 **(model_config_kwargs or {}),

73 )

74 vllm_config_args = dict(vllm_config_kwargs or {})

62 - vllm_config_args.setdefault("compilation_config", CompilationConfig(mode=0))

75 + vllm_config_args.setdefault(

76 + "compilation_config", CompilationConfig(mode=CompilationMode.NONE)

77 + )

78 vllm_config = VllmConfig(model_config=model_config, **vllm_config_args)

79 hf_overrides = (model_config_kwargs or {}).get("hf_overrides") or {}

80 num_hidden_layers = hf_overrides.get("num_hidden_layers")

tests/reasoning/test_base_thinking_reasoning_parser.py

+6 -4

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import pytest

5 -from transformers import AutoTokenizer

5 +from transformers import AutoTokenizer, PreTrainedTokenizerBase

6

7 from tests.reasoning.utils import run_reasoning_extraction

8 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest

@@ -39,8 +39,10 @@ REASONING_MODEL_NAME = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"

39

40

41 @pytest.fixture(scope="module")

42 -def test_tokenizer():

43 - tokenizer = AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)

42 +def test_tokenizer() -> PreTrainedTokenizerBase:

43 + tokenizer: PreTrainedTokenizerBase = AutoTokenizer.from_pretrained(

44 + REASONING_MODEL_NAME

45 + )

46 # Add custom test tokens

47 test_tokens = ["<test:think>", "</test:think>", "<alt:start>", "<alt:end>"]

48 existing_tokens = set(tokenizer.get_vocab().keys())

@@ -66,7 +68,7 @@ class TestBaseThinkingReasoningParserInit:

68 def test_initialization_with_missing_tokenizer(self):

69 """Test that initialization fails without tokenizer."""

70 with pytest.raises(ValueError, match="model tokenizer must be passed"):

69 - TestThinkingReasoningParser(None)

71 + TestThinkingReasoningParser(None) # type: ignore[arg-type]

72

73 def test_initialization_with_missing_tokens(self, test_tokenizer):

74 """Test that initialization fails when tokens are not in vocabulary."""

tests/reasoning/test_deepseekr1_reasoning_parser.py

+6 -6

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import pytest

5 -from transformers import AutoTokenizer

5 +from transformers import AutoTokenizer, PreTrainedTokenizerBase

6

7 from tests.reasoning.utils import run_reasoning_extraction

8 from vllm.reasoning import ReasoningParser, ReasoningParserManager

@@ -15,7 +15,7 @@ REASONING_MODEL_NAME = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"

15

16

17 @pytest.fixture(scope="module")

18 -def deepseek_r1_qwen_tokenizer():

18 +def deepseek_r1_qwen_tokenizer() -> PreTrainedTokenizerBase:

19 return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)

20

21

@@ -279,10 +279,10 @@ def test_reasoning(

279

280 # Test extract_content

281 if param_dict["content"] is not None:

282 - content = parser.extract_content_ids(output_ids)

283 - assert content == deepseek_r1_qwen_tokenizer.convert_tokens_to_ids(

282 + content_ids = parser.extract_content_ids(output_ids)

283 + assert content_ids == deepseek_r1_qwen_tokenizer.convert_tokens_to_ids(

284 deepseek_r1_qwen_tokenizer.tokenize(param_dict["content"])

285 )

286 else:

287 - content = parser.extract_content_ids(output)

288 - assert content == []

287 + content_ids = parser.extract_content_ids(output)

288 + assert content_ids == []

tests/reasoning/test_deepseekv3_reasoning_parser.py

+2 -2

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import pytest

5 -from transformers import AutoTokenizer

5 +from transformers import AutoTokenizer, PreTrainedTokenizerBase

6

7 from vllm.entrypoints.generate.base.protocol import DeltaMessage

8 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest

@@ -17,7 +17,7 @@ REASONING_MODEL_NAME = "deepseek-ai/DeepSeek-V3.1"

17

18

19 @pytest.fixture(scope="module")

20 -def tokenizer():

20 +def tokenizer() -> PreTrainedTokenizerBase:

21 return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)

22

23

tests/reasoning/test_ernie45_reasoning_parser.py

+2 -2

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import pytest

5 -from transformers import AutoTokenizer

5 +from transformers import AutoTokenizer, PreTrainedTokenizerBase

6

7 from tests.reasoning.utils import run_reasoning_extraction

8 from vllm.reasoning import ReasoningParser, ReasoningParserManager

@@ -13,7 +13,7 @@ REASONING_MODEL_NAME = "baidu/ERNIE-4.5-21B-A3B-Thinking"

13

14

15 @pytest.fixture(scope="module")

16 -def ernie45_tokenizer():

16 +def ernie45_tokenizer() -> PreTrainedTokenizerBase:

17 return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)

18

19

tests/reasoning/test_glm4_moe_reasoning_parser.py

+13 -4

@@ -4,9 +4,11 @@

4 import copy

5

6 import pytest

7 -from transformers import AutoTokenizer

7 +from transformers import AutoTokenizer, PreTrainedTokenizerBase

8

9 from tests.reasoning.utils import run_reasoning_extraction

10 +from vllm.parser.engine.adapters import ParserEngineReasoningAdapter

11 +from vllm.parser.glm47_moe import Glm47MoeParser

12 from vllm.reasoning import ReasoningParser, ReasoningParserManager

13

14 parser_name = "glm45"

@@ -17,7 +19,7 @@ REASONING_MODEL_NAME = "zai-org/GLM-4.7"

19

20

21 @pytest.fixture(scope="module")

20 -def glm45_tokenizer():

22 +def glm45_tokenizer() -> PreTrainedTokenizerBase:

23 return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)

24

25

@@ -261,13 +263,20 @@ def glm53_style_tokenizer(glm45_tokenizer):

263 return tokenizer

264

265

266 +def _glm_engine(parser: ReasoningParser) -> Glm47MoeParser:

267 + assert isinstance(parser, ParserEngineReasoningAdapter)

268 + engine = parser._parser_engine

269 + assert isinstance(engine, Glm47MoeParser)

270 + return engine

271 +

272 +

273 @pytest.mark.parametrize(

274 "disable_kwargs", [{"enable_thinking": False}, {"thinking": False}]

275 )

276 def test_glm53_template_forces_reasoning(disable_kwargs: dict, glm53_style_tokenizer):

277 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)

278 parser = parser_cls(glm53_style_tokenizer, chat_template_kwargs=disable_kwargs)

270 - assert parser._parser_engine.thinking_enabled

279 + assert _glm_engine(parser).thinking_enabled

280

281 output = glm53_style_tokenizer.tokenize(GLM53_LEAK["output"])

282 output_tokens: list[str] = [

@@ -298,7 +307,7 @@ def test_glm47_template_honors_thinking_disable(glm45_tokenizer):

307 parser = ReasoningParserManager.get_reasoning_parser(parser_name)(

308 glm45_tokenizer, chat_template_kwargs={"enable_thinking": False}

309 )

301 - assert not parser._parser_engine.thinking_enabled

310 + assert not _glm_engine(parser).thinking_enabled

311

312 output = glm45_tokenizer.tokenize(GLM53_LEAK["output"])

313 output_tokens: list[str] = [

tests/reasoning/test_granite_thinking_reasoning_parser.py

+22 -21

@@ -1,10 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import TypedDict

4 +from typing import TypedDict, cast

5

6 import pytest

7 import regex as re

8 +from transformers import PreTrainedTokenizerBase

9

10 from tests.reasoning.utils import run_reasoning_extraction

11 from vllm.entrypoints.openai.chat_completion.protocol import (

@@ -52,8 +53,8 @@ class GraniteThinkingTokenizer:

53

54

55 @pytest.fixture

55 -def tokenizer():

56 - return GraniteThinkingTokenizer()

56 +def tokenizer() -> PreTrainedTokenizerBase:

57 + return cast(PreTrainedTokenizerBase, GraniteThinkingTokenizer())

58

59

60 # ── Basic reasoning extraction (non-streaming + streaming) ───────────

@@ -173,7 +174,7 @@ def tokenizer():

174 ],

175 )

176 def test_granite_thinking_reasoning(

176 - tokenizer: GraniteThinkingTokenizer,

177 + tokenizer: PreTrainedTokenizerBase,

178 streaming: bool,

179 param_dict: ReasoningCase,

180 ):

@@ -195,7 +196,7 @@ def test_granite_thinking_reasoning(

196

197

198 def test_granite_thinking_no_content_after_end_token(

198 - tokenizer: GraniteThinkingTokenizer,

199 + tokenizer: PreTrainedTokenizerBase,

200 ):

201 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)

202 parser = parser_cls(tokenizer)

@@ -215,7 +216,7 @@ def test_granite_thinking_no_content_after_end_token(

216

217 @pytest.mark.parametrize("streaming", [False, True])

218 def test_granite_thinking_whitespace_only_content(

218 - tokenizer: GraniteThinkingTokenizer,

219 + tokenizer: PreTrainedTokenizerBase,

220 streaming: bool,

221 ):

222 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)

@@ -235,7 +236,7 @@ def test_granite_thinking_whitespace_only_content(

236

237

238 def test_granite_thinking_unterminated_think_block(

238 - tokenizer: GraniteThinkingTokenizer,

239 + tokenizer: PreTrainedTokenizerBase,

240 ):

241 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)

242 parser = parser_cls(tokenizer)

@@ -260,7 +261,7 @@ def test_granite_thinking_unterminated_think_block(

261

262

263 def test_granite_thinking_disabled_moves_into_content(

263 - tokenizer: GraniteThinkingTokenizer,

264 + tokenizer: PreTrainedTokenizerBase,

265 ):

266 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)

267 parser = parser_cls(tokenizer)

@@ -285,7 +286,7 @@ def test_granite_thinking_disabled_moves_into_content(

286

287

288 def test_granite_thinking_disabled_with_leading_newline(

288 - tokenizer: GraniteThinkingTokenizer,

289 + tokenizer: PreTrainedTokenizerBase,

290 ):

291 # With enable_thinking=False, model output goes through the swap

292 # path: content is initially None (all text classified as

@@ -316,7 +317,7 @@ def test_granite_thinking_disabled_with_leading_newline(

317

318

319 def test_granite_thinking_force_nonempty_content_moves_into_content(

319 - tokenizer: GraniteThinkingTokenizer,

320 + tokenizer: PreTrainedTokenizerBase,

321 ):

322 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)

323 parser = parser_cls(tokenizer)

@@ -338,7 +339,7 @@ def test_granite_thinking_force_nonempty_content_moves_into_content(

339

340

341 def test_granite_thinking_force_nonempty_no_swap_when_newlines_only(

341 - tokenizer: GraniteThinkingTokenizer,

342 + tokenizer: PreTrainedTokenizerBase,

343 ):

344 # When </think> IS present and content is newlines-only, lstrip

345 # removes them but the swap should NOT fire — content was present,

@@ -363,7 +364,7 @@ def test_granite_thinking_force_nonempty_no_swap_when_newlines_only(

364

365

366 def test_granite_thinking_force_nonempty_swaps_when_content_absent(

366 - tokenizer: GraniteThinkingTokenizer,

367 + tokenizer: PreTrainedTokenizerBase,

368 ):

369 # When </think> IS present but content is truly absent (zero

370 # characters after </think>, e.g. max_tokens cut), swap fires.

@@ -387,7 +388,7 @@ def test_granite_thinking_force_nonempty_swaps_when_content_absent(

388

389

390 def test_granite_thinking_force_nonempty_keeps_real_content(

390 - tokenizer: GraniteThinkingTokenizer,

391 + tokenizer: PreTrainedTokenizerBase,

392 ):

393 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)

394 parser = parser_cls(tokenizer)

@@ -412,7 +413,7 @@ def test_granite_thinking_force_nonempty_keeps_real_content(

413

414

415 def test_granite_thinking_keeps_truncated_reasoning(

415 - tokenizer: GraniteThinkingTokenizer,

416 + tokenizer: PreTrainedTokenizerBase,

417 ):

418 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)

419 parser = parser_cls(tokenizer)

@@ -472,7 +473,7 @@ def _run_parse_delta(parser, tokenizer, text, request):

473

474

475 def test_granite_thinking_streaming_enable_thinking_false(

475 - tokenizer: GraniteThinkingTokenizer,

476 + tokenizer: PreTrainedTokenizerBase,

477 ):

478 # With enable_thinking=False, the parser (constructed without

479 # kwargs) starts in REASONING state. All text streams as reasoning

@@ -494,7 +495,7 @@ def test_granite_thinking_streaming_enable_thinking_false(

495

496

497 def test_granite_thinking_streaming_strips_leading_newline(

497 - tokenizer: GraniteThinkingTokenizer,

498 + tokenizer: PreTrainedTokenizerBase,

499 ):

500 request = ChatCompletionRequest(

501 model="test-model",

@@ -511,7 +512,7 @@ def test_granite_thinking_streaming_strips_leading_newline(

512

513

514 def test_granite_thinking_streaming_promotes_reasoning_to_content(

514 - tokenizer: GraniteThinkingTokenizer,

515 + tokenizer: PreTrainedTokenizerBase,

516 ):

517 request = ChatCompletionRequest(

518 model="test-model",

@@ -527,7 +528,7 @@ def test_granite_thinking_streaming_promotes_reasoning_to_content(

528

529

530 def test_granite_thinking_streaming_no_promotion_with_real_content(

530 - tokenizer: GraniteThinkingTokenizer,

531 + tokenizer: PreTrainedTokenizerBase,

532 ):

533 request = ChatCompletionRequest(

534 model="test-model",

@@ -548,7 +549,7 @@ def test_granite_thinking_streaming_no_promotion_with_real_content(

549

550

551 def test_granite_thinking_streaming_no_promotion_without_opt_in(

551 - tokenizer: GraniteThinkingTokenizer,

552 + tokenizer: PreTrainedTokenizerBase,

553 ):

554 request = ChatCompletionRequest(model="test-model", messages=[])

555 parser = _make_reasoning_parser(tokenizer)

@@ -564,7 +565,7 @@ def test_granite_thinking_streaming_no_promotion_without_opt_in(

565

566 @pytest.mark.parametrize("streaming", [False, True])

567 def test_granite_thinking_empty_think_block(

567 - tokenizer: GraniteThinkingTokenizer,

568 + tokenizer: PreTrainedTokenizerBase,

569 streaming: bool,

570 ):

571 # <think></think>\nHello — empty reasoning, content after newline.

@@ -592,7 +593,7 @@ def test_granite_thinking_empty_think_block(

593 ],

594 )

595 def test_granite_thinking_streaming_chunking_independent(

595 - tokenizer: GraniteThinkingTokenizer,

596 + tokenizer: PreTrainedTokenizerBase,

597 text: str,

598 expected_reasoning: str,

599 expected_content: str,

tests/reasoning/test_holo2_reasoning_parser.py

+2 -2

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import pytest

4 -from transformers import AutoTokenizer

4 +from transformers import AutoTokenizer, PreTrainedTokenizerBase

5

6 from tests.reasoning.utils import run_reasoning_extraction

7 from vllm.reasoning import ReasoningParser, ReasoningParserManager

@@ -15,7 +15,7 @@ REASONING_MODEL_NAME = "HCompany/Holo2-4B"

15

16

17 @pytest.fixture(scope="module")

18 -def tokenizer():

18 +def tokenizer() -> PreTrainedTokenizerBase:

19 return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)

20

21

tests/reasoning/test_k2_horizon_reasoning_parser.py

+13 -7

@@ -3,6 +3,7 @@

3

4 import string

5 from collections.abc import Sequence

6 +from typing import cast

7

8 import pytest

9

@@ -10,6 +11,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionReque

11 from vllm.parser.parser_manager import ParserManager

12 from vllm.reasoning import ReasoningParserManager

13 from vllm.reasoning.k2_horizon_reasoning_parser import K2HorizonReasoningParser

14 +from vllm.tokenizers import TokenizerLike

15

16 pytestmark = pytest.mark.skip_global_cleanup

17

@@ -59,6 +61,10 @@ class K2Tokenizer:

61 return "".join(inverse_vocab[token_id] for token_id in ids)

62

63

64 +def _k2_tokenizer() -> TokenizerLike:

65 + return cast(TokenizerLike, K2Tokenizer())

66 +

67 +

68 def _request(**kwargs) -> ChatCompletionRequest:

69 return ChatCompletionRequest(

70 model="k2-horizon",

@@ -70,7 +76,7 @@ def _request(**kwargs) -> ChatCompletionRequest:

76 @pytest.mark.parametrize("effort", ["high", "medium", "low"])

77 def test_reasoning_effort_selects_tokens(effort: str):

78 parser = K2HorizonReasoningParser(

73 - K2Tokenizer(),

79 + _k2_tokenizer(),

80 chat_template_kwargs={"reasoning_effort": effort},

81 )

82

@@ -78,7 +84,7 @@ def test_reasoning_effort_selects_tokens(effort: str):

84

85

86 def test_reasoning_effort_precedence_and_default():

81 - tokenizer = K2Tokenizer()

87 + tokenizer = _k2_tokenizer()

88 request = _request(

89 reasoning_effort="medium",

90 chat_template_kwargs={"reasoning_effort": "low"},

@@ -103,13 +109,13 @@ def test_reasoning_effort_precedence_and_default():

109 def test_invalid_reasoning_effort_raises(effort: str | None):

110 with pytest.raises(ValueError, match="Unsupported reasoning_effort"):

111 K2HorizonReasoningParser(

106 - K2Tokenizer(),

112 + _k2_tokenizer(),

113 chat_template_kwargs={"reasoning_effort": effort},

114 )

115

116

117 def test_non_streaming_explicit_and_implicit_boundaries():

112 - parser = K2HorizonReasoningParser(K2Tokenizer())

118 + parser = K2HorizonReasoningParser(_k2_tokenizer())

119 request = _request()

120

121 assert parser.extract_reasoning("<ifm|think>plan</ifm|think>answer", request) == (

@@ -126,7 +132,7 @@ def test_non_streaming_explicit_and_implicit_boundaries():

132

133 @pytest.mark.parametrize("reasoning", ["plan", " ", "\n", " \n"])

134 def test_character_split_reasoning_stream(reasoning: str):

129 - parser = K2HorizonReasoningParser(K2Tokenizer())

135 + parser = K2HorizonReasoningParser(_k2_tokenizer())

136 reasoning_parts: list[str] = []

137 content_parts: list[str] = []

138

@@ -150,7 +156,7 @@ def test_character_split_reasoning_stream(reasoning: str):

156

157

158 def test_latest_boundary_controls_reasoning_state():

153 - tokenizer = K2Tokenizer()

159 + tokenizer = _k2_tokenizer()

160 parser = K2HorizonReasoningParser(tokenizer)

161 vocab = tokenizer.get_vocab()

162

@@ -159,7 +165,7 @@ def test_latest_boundary_controls_reasoning_state():

165

166

167 def test_composed_streaming_reasoning_to_tool_handoff():

162 - tokenizer = K2Tokenizer()

168 + tokenizer = _k2_tokenizer()

169 tools = [

170 {

171 "type": "function",

tests/reasoning/test_kimi_k2_reasoning_parser.py

+13 -4

@@ -6,6 +6,7 @@ import pytest

6

7 from vllm.entrypoints.generate.base.protocol import DeltaMessage

8 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest

9 +from vllm.parser.kimi_k2 import KimiK2Parser

10 from vllm.reasoning.kimi_k2_reasoning_parser import KimiK2ReasoningParser

11 from vllm.tokenizers import get_tokenizer

12

@@ -65,9 +66,13 @@ def test_extract_reasoning_tool_section_ends_reasoning(kimi_k2_tokenizer):

66 def test_streaming_reasoning_then_content(kimi_k2_tokenizer):

67 """Token-by-token streaming: reasoning tokens then content after </think>."""

68 parser = KimiK2ReasoningParser(kimi_k2_tokenizer)

69 + parser_engine = parser._parser_engine

70 + assert isinstance(parser_engine, KimiK2Parser)

71

69 - think_id = parser._parser_engine._start_token_id

70 - end_think_id = parser._parser_engine._end_token_id

72 + think_id = parser_engine._start_token_id

73 + end_think_id = parser_engine._end_token_id

74 + assert think_id is not None

75 + assert end_think_id is not None

76 # Use a real token ID from the tokenizer for regular content

77 regular_id = kimi_k2_tokenizer.encode("hello", add_special_tokens=False)[0]

78

@@ -123,9 +128,13 @@ def test_streaming_reasoning_then_content(kimi_k2_tokenizer):

128 def test_streaming_tool_section_ends_reasoning(kimi_k2_tokenizer):

129 """<|tool_calls_section_begin|> in delta ends reasoning during streaming."""

130 parser = KimiK2ReasoningParser(kimi_k2_tokenizer)

131 + parser_engine = parser._parser_engine

132 + assert isinstance(parser_engine, KimiK2Parser)

133

127 - think_id = parser._parser_engine._start_token_id

128 - tool_begin_id = parser._parser_engine._tool_section_start_token_id

134 + think_id = parser_engine._start_token_id

135 + tool_begin_id = parser_engine._tool_section_start_token_id

136 + assert think_id is not None

137 + assert tool_begin_id is not None

138 regular_id = kimi_k2_tokenizer.encode("hello", add_special_tokens=False)[0]

139

140 # Tool section token arrives — should transition from reasoning to content

tests/reasoning/test_kimi_k3_reasoning_parser.py

+15 -4

@@ -1,6 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 +from typing import cast

5 +

6 import pytest

7

8 from vllm.entrypoints.generate.base.protocol import DeltaMessage

@@ -10,6 +12,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (

12 from vllm.parser.kimi_k3 import KimiK3Parser

13 from vllm.parser.parser_manager import ParserManager

14 from vllm.reasoning.kimi_k3_reasoning_parser import KimiK3ReasoningParser

15 +from vllm.tokenizers import TokenizerLike

16

17 pytestmark = pytest.mark.skip_global_cleanup

18

@@ -25,10 +28,13 @@ RESPONSE_OPEN_IDS = [ord(ch) for ch in RESPONSE_OPEN]

28

29

30 class DummyTokenizer:

31 + def __len__(self) -> int:

32 + return 1

33 +

34 def get_vocab(self) -> dict[str, int]:

35 return {}

36

31 - def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:

37 + def encode(self, text: str, *args, **kwargs) -> list[int]:

38 if text == THINK_OPEN:

39 return [1, 2, 3]

40 if text == THINK_CLOSE:

@@ -36,6 +42,10 @@ class DummyTokenizer:

42 return [ord(ch) for ch in text]

43

44

45 +def _dummy_tokenizer() -> TokenizerLike:

46 + return cast(TokenizerLike, DummyTokenizer())

47 +

48 +

49 class ReasoningOnlyParser(KimiK3Parser):

50 reasoning_parser_cls = KimiK3ReasoningParser

51

@@ -84,7 +94,7 @@ def test_extract_reasoning_with_generation_prefix_consumed():

94

95

96 def test_delegating_parser_strips_response_wrapper_without_tool_parser():

87 - parser = ReasoningOnlyParser(DummyTokenizer())

97 + parser = ReasoningOnlyParser(_dummy_tokenizer())

98 request = ChatCompletionRequest(model="test-model", messages=[])

99

100 reasoning, content, tool_calls = parser.parse(

@@ -165,7 +175,7 @@ def test_count_reasoning_tokens_is_zero_when_thinking_disabled():

175

176

177 def test_count_reasoning_tokens_through_delegating_parser():

168 - parser = ReasoningOnlyParser(DummyTokenizer())

178 + parser = ReasoningOnlyParser(_dummy_tokenizer())

179

180 assert parser.count_reasoning_tokens([*OPEN_IDS, 9, *CLOSE_IDS, 11]) == 1

181

@@ -253,7 +263,8 @@ def test_thinking_disabled_streams_content():

263

264 def test_delegating_parser_thinking_false_streams_response_content():

265 parser = ReasoningOnlyParser(

256 - DummyTokenizer(), chat_template_kwargs={"thinking": False}

266 + _dummy_tokenizer(),

267 + chat_template_kwargs={"thinking": False},

268 )

269 request = ChatCompletionRequest(

270 model="test-model",

tests/reasoning/test_minimax_m2_append_reasoning_parser.py

+2 -2

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import pytest

5 -from transformers import AutoTokenizer

5 +from transformers import AutoTokenizer, PreTrainedTokenizerBase

6

7 from tests.reasoning.utils import run_reasoning_extraction

8 from vllm.reasoning import ReasoningParser, ReasoningParserManager

@@ -15,7 +15,7 @@ REASONING_MODEL_NAME = "MiniMaxAI/MiniMax-M2"

15

16

17 @pytest.fixture(scope="module")

18 -def minimax_m2_tokenizer():

18 +def minimax_m2_tokenizer() -> PreTrainedTokenizerBase:

19 return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)

20

21

tests/reasoning/test_minimax_m2_reasoning_parser.py

+6 -6

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import pytest

5 -from transformers import AutoTokenizer

5 +from transformers import AutoTokenizer, PreTrainedTokenizerBase

6

7 from tests.reasoning.utils import run_reasoning_extraction

8 from vllm.reasoning import ReasoningParser, ReasoningParserManager

@@ -15,7 +15,7 @@ REASONING_MODEL_NAME = "MiniMaxAI/MiniMax-M2"

15

16

17 @pytest.fixture(scope="module")

18 -def minimax_m2_tokenizer():

18 +def minimax_m2_tokenizer() -> PreTrainedTokenizerBase:

19 return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)

20

21

@@ -195,10 +195,10 @@ def test_reasoning(

195

196 # Test extract_content

197 if param_dict["content"] is not None:

198 - content = parser.extract_content_ids(output_ids)

199 - assert content == minimax_m2_tokenizer.convert_tokens_to_ids(

198 + content_ids = parser.extract_content_ids(output_ids)

199 + assert content_ids == minimax_m2_tokenizer.convert_tokens_to_ids(

200 minimax_m2_tokenizer.tokenize(param_dict["content"])

201 )

202 else:

203 - content = parser.extract_content_ids(output)

204 - assert content == []

203 + content_ids = parser.extract_content_ids(output)

204 + assert content_ids == []

tests/reasoning/test_nemotron_v3_reasoning_parser.py

+12 -11

@@ -1,10 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -from typing import TypedDict

4 +from typing import TypedDict, cast

5

6 import pytest

7 import regex as re

8 +from transformers import PreTrainedTokenizerBase

9

10 from tests.reasoning.utils import run_reasoning_extraction

11 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest

@@ -48,8 +49,8 @@ class FakeNemotronTokenizer:

49

50

51 @pytest.fixture

51 -def tokenizer():

52 - return FakeNemotronTokenizer()

52 +def tokenizer() -> PreTrainedTokenizerBase:

53 + return cast(PreTrainedTokenizerBase, FakeNemotronTokenizer())

54

55

56 @pytest.mark.parametrize(

@@ -94,7 +95,7 @@ def tokenizer():

95 ],

96 )

97 def test_nemotron_v3_reasoning(

97 - tokenizer: FakeNemotronTokenizer,

98 + tokenizer: PreTrainedTokenizerBase,

99 streaming: bool,

100 param_dict: ReasoningCase,

101 ):

@@ -113,7 +114,7 @@ def test_nemotron_v3_reasoning(

114

115

116 def test_nemotron_v3_without_thinking_moves_into_content(

116 - tokenizer: FakeNemotronTokenizer,

117 + tokenizer: PreTrainedTokenizerBase,

118 ):

119 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)

120 parser = parser_cls(tokenizer)

@@ -137,7 +138,7 @@ def test_nemotron_v3_without_thinking_moves_into_content(

138

139

140 def test_nemotron_v3_force_nonempty_content_moves_into_content(

140 - tokenizer: FakeNemotronTokenizer,

141 + tokenizer: PreTrainedTokenizerBase,

142 ):

143 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)

144 parser = parser_cls(tokenizer)

@@ -159,7 +160,7 @@ def test_nemotron_v3_force_nonempty_content_moves_into_content(

160

161

162 def test_nemotron_v3_force_nonempty_keeps_real_content(

162 - tokenizer: FakeNemotronTokenizer,

163 + tokenizer: PreTrainedTokenizerBase,

164 ):

165 # When real content follows the closing tag nothing is promoted: the

166 # content after </think> is returned as-is and reasoning stays separate.

@@ -183,7 +184,7 @@ def test_nemotron_v3_force_nonempty_keeps_real_content(

184

185

186 def test_nemotron_v3_with_thinking_keeps_truncated_reasoning(

186 - tokenizer: FakeNemotronTokenizer,

187 + tokenizer: PreTrainedTokenizerBase,

188 ):

189 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)

190 parser = parser_cls(tokenizer)

@@ -242,7 +243,7 @@ def _run_parse_delta(parser, tokenizer, text, request):

243

244

245 def test_nemotron_v3_streaming_promotes_reasoning_to_content(

245 - tokenizer: FakeNemotronTokenizer,

246 + tokenizer: PreTrainedTokenizerBase,

247 ):

248 # Model never closes <think>: reasoning streams normally AND is duplicated

249 # into content on the terminal delta.

@@ -260,7 +261,7 @@ def test_nemotron_v3_streaming_promotes_reasoning_to_content(

261

262

263 def test_nemotron_v3_streaming_no_promotion_with_real_content(

263 - tokenizer: FakeNemotronTokenizer,

264 + tokenizer: PreTrainedTokenizerBase,

265 ):

266 request = ChatCompletionRequest(

267 model="test-model",

@@ -279,7 +280,7 @@ def test_nemotron_v3_streaming_no_promotion_with_real_content(

280

281

282 def test_nemotron_v3_streaming_no_promotion_without_opt_in(

282 - tokenizer: FakeNemotronTokenizer,

283 + tokenizer: PreTrainedTokenizerBase,

284 ):

285 # Without enable_thinking=False / force_nonempty_content the fallback must

286 # stay disabled: the response stays reasoning-only, content empty.

tests/reasoning/test_qwen3_reasoning_parser.py

+2 -2

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import pytest

5 -from transformers import AutoTokenizer

5 +from transformers import AutoTokenizer, PreTrainedTokenizerBase

6

7 from tests.reasoning.utils import (

8 StreamingReasoningReconstructor,

@@ -24,7 +24,7 @@ REASONING_MODEL_NAMES = [

24

25

26 @pytest.fixture(scope="module", params=REASONING_MODEL_NAMES)

27 -def qwen3_tokenizer(request):

27 +def qwen3_tokenizer(request) -> PreTrainedTokenizerBase:

28 return AutoTokenizer.from_pretrained(request.param)

29

30

tests/reasoning/test_step3p5_reasoning_parser.py

+8 -8

@@ -4,7 +4,8 @@

4 import pytest

5

6 from tests.reasoning.utils import run_reasoning_extraction

7 -from vllm.reasoning import ReasoningParser, ReasoningParserManager

7 +from vllm.reasoning import ReasoningParserManager

8 +from vllm.reasoning.step3p5_reasoning_parser import Step3p5ReasoningParser

9 from vllm.tokenizers import get_tokenizer

10

11 parser_name = "step3p5"

@@ -283,9 +284,8 @@ def test_reasoning(

284 output_tokens: list[str] = [

285 step3p5_tokenizer.convert_tokens_to_string([token]) for token in output

286 ]

286 - parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)(

287 - step3p5_tokenizer

288 - )

287 + parser = ReasoningParserManager.get_reasoning_parser(parser_name)(step3p5_tokenizer)

288 + assert isinstance(parser, Step3p5ReasoningParser)

289

290 reasoning, content = run_reasoning_extraction(

291 parser, output_tokens, streaming=streaming

@@ -317,7 +317,7 @@ def test_reasoning(

317

318 # Test extract_content

319 if param_dict["content"] is not None:

320 - content = parser.extract_content_ids(output_ids)

320 + content_ids = parser.extract_content_ids(output_ids)

321 # Fixed expected token ids for specific test cases

322 test_id = (

323 request.node.callspec.id if hasattr(request.node, "callspec") else None

@@ -332,7 +332,7 @@ def test_reasoning(

332 expected_content_ids = step3p5_tokenizer.convert_tokens_to_ids(

333 step3p5_tokenizer.tokenize(param_dict["content"])

334 )

335 - assert content == expected_content_ids

335 + assert content_ids == expected_content_ids

336 else:

337 - content = parser.extract_content_ids(output_ids)

338 - assert content == []

337 + content_ids = parser.extract_content_ids(output_ids)

338 + assert content_ids == []

tests/reasoning/utils.py

+7 -2

@@ -1,6 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 +from typing import cast

5 +

6 +from transformers import PreTrainedTokenizerBase

7 +

8 from vllm.entrypoints.generate.base.protocol import DeltaMessage

9 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest

10 from vllm.reasoning import ReasoningParser

@@ -102,10 +106,11 @@ def run_reasoning_extraction_streaming(

106 reconstructor = StreamingReasoningReconstructor()

107 previous_text = ""

108 previous_tokens: list[int] = []

109 + tokenizer = cast(PreTrainedTokenizerBase, reasoning_parser.model_tokenizer)

110 for delta in model_deltas:

111 token_delta = [

107 - reasoning_parser.vocab.get(token)

108 - for token in reasoning_parser.model_tokenizer.tokenize(delta)

112 + reasoning_parser.vocab[token]

113 + for token in tokenizer.tokenize(delta)

114 if token in reasoning_parser.vocab

115 ]

116 current_text = previous_text + delta

tests/v1/logits_processors/test_correctness.py

+54 -71

@@ -3,7 +3,7 @@

3

4 import random

5 from collections.abc import Callable

6 -from typing import NamedTuple, TypeAlias

6 +from typing import NamedTuple, NoReturn, TypeAlias

7

8 import numpy as np

9 import pytest

@@ -13,6 +13,7 @@ from tests.utils import create_new_process_for_each_test

13 from tests.v1.sample.utils import (

14 LogitsprocsTestFakes,

15 create_fake_logits,

16 + create_mock_reasoning_config,

17 create_penalty_tensor,

18 create_prompt_tokens_tensor,

19 fake_apply_logitsprocs,

@@ -103,14 +104,6 @@ class LogitsProcsRequestParams:

104 return f"MyClass({summ})"

105

106

106 -class MockReasoningConfig:

107 - """Minimal reasoning config for ``ThinkingBudgetStateHolder`` tests."""

108 -

109 - reasoning_start_token_ids = [THINK_START_TOKEN_ID]

110 - reasoning_end_token_ids = [THINK_END_TOKEN_ID]

111 - enabled = True

112 -

113 -

107 def _generate_fake_sampling_metadata(

108 num_output_tokens: int,

109 batch_size: int,

@@ -131,7 +124,9 @@ def _generate_fake_sampling_metadata(

124 )

125

126 vllm_config = VllmConfig()

134 - vllm_config.reasoning_config = MockReasoningConfig()

127 + vllm_config.reasoning_config = create_mock_reasoning_config(

128 + [THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID]

129 + )

130

131 logitsprocs = build_logitsprocs(

132 vllm_config=vllm_config,

@@ -242,7 +237,7 @@ def _raise_error_invalid(

237 request_params: LogitsProcsRequestParams,

238 step_idx: int,

239 err_cls: type[Exception] = ValueError,

245 -) -> None:

240 +) -> NoReturn:

241 raise err_cls(

242 f"Validation failed for step={step_idx}, "

243 f"batch_index={batch_index}, "

@@ -268,6 +263,7 @@ def _logit_bias_validate(

263 ) -> None:

264 """Validate logit bias logitproc applied correctly."""

265 logit_bias = request_params.params.logit_bias

266 + assert logit_bias is not None

267 logits_old = test_fakes.logits[persistent_batch[batch_index].workload_index].cpu()

268 logits_new = logits_new[batch_index].cpu()

269 for token_id in range(VOCAB_SIZE):

@@ -368,9 +364,7 @@ def _min_tokens_validate(

364 ref_num_out_tokens = len(request_params.out_tokens)

365 min_reached = ref_num_out_tokens >= MIN_TOKENS_LEN_THRESHOLD

366 ref_all_stop_token_ids = request_params.params.all_stop_token_ids

371 - mt_lp: MinTokensLogitsProcessor = next(

372 - test_fakes.get_logitsprocs_by_cls(MinTokensLogitsProcessor)

373 - )

367 + mt_lp = next(test_fakes.get_logitsprocs_by_cls(MinTokensLogitsProcessor))

368 assert isinstance(mt_lp, MinTokensLogitsProcessor)

369 min_tok = mt_lp.min_toks.get(batch_index, None)

370

@@ -689,7 +683,7 @@ class LogitsprocTestHelpers(NamedTuple):

683 gen_request_fxn: Callable | None = None

684

685

692 -logitsprocs_test_mapping = {

686 +logitsprocs_test_mapping: dict[LogitprocType, LogitsprocTestHelpers] = {

687 STR_NO_LOGITPROC: LogitsprocTestHelpers(eval_fxn=_none_validate),

688 LogitBiasLogitsProcessor: LogitsprocTestHelpers(

689 gen_request_fxn=_logit_bias_params, eval_fxn=_logit_bias_validate

@@ -706,22 +700,22 @@ logitsprocs_test_mapping = {

700 }

701

702

709 -def _get_test_cases() -> list[list[str]]:

703 +def _get_test_cases() -> list[list[LogitprocType]]:

704 """Each test case is a set of logitsprocs."""

711 - logitsprocs_types = list(logitsprocs_test_mapping.keys())

705 + logitsprocs_types: list[LogitprocType] = list(logitsprocs_test_mapping.keys())

706

707 # Isolate thinking-budget handling from other processors to avoid cross-talk.

708 thinking_id: LogitprocType = STR_THINKING_BUDGET

715 - other_processors = [

709 + other_processors: list[LogitprocType] = [

710 p for p in logitsprocs_types if p != STR_NO_LOGITPROC and p != thinking_id

711 ]

712

719 - return (

720 - [[STR_NO_LOGITPROC]]

721 - + [[logitproc_type, STR_NO_LOGITPROC] for logitproc_type in other_processors]

722 - + [other_processors]

723 - + [[thinking_id]]

713 + test_cases: list[list[LogitprocType]] = [[STR_NO_LOGITPROC]]

714 + test_cases.extend(

715 + [logitproc_type, STR_NO_LOGITPROC] for logitproc_type in other_processors

716 )

717 + test_cases.extend([other_processors, [thinking_id]])

718 + return test_cases

719

720

721 def _generate_fake_step_update(

@@ -774,6 +768,7 @@ def _generate_fake_step_update(

768 for add_req_params in workload_params[wdx : (wdx + num_step_add_replace)]:

769 # Replace as many removed requests as possible with added requests

770 add_remove_idx = batch_update_builder.pop_removed()

771 + assert add_remove_idx is not None

772 batch_update_builder.added.append(

773 (

774 add_remove_idx,

@@ -976,13 +971,6 @@ def test_logitsprocs(

971 step_idx += 1

972

973

979 -class MockReasoningNoEndTokens:

980 - """Reasoning config with no end token ids (disables enforcement in holder)."""

981 -

982 - reasoning_start_token_ids = [THINK_START_TOKEN_ID]

983 - reasoning_end_token_ids: list[int] = []

984 -

985 -

974 def test_maybe_create_thinking_budget_holder_without_reasoning():

975 cfg = VllmConfig()

976 assert cfg.reasoning_config is None

@@ -1000,7 +988,9 @@ def test_maybe_create_thinking_budget_holder_without_reasoning():

988

989 def test_thinking_budget_holder_has_tracked_after_sync_add():

990 vc = VllmConfig()

1003 - vc.reasoning_config = MockReasoningConfig()

991 + vc.reasoning_config = create_mock_reasoning_config(

992 + [THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID]

993 + )

994 h = ThinkingBudgetStateHolder(

995 vc.reasoning_config,

996 vc.scheduler_config.max_num_seqs,

@@ -1030,7 +1020,9 @@ def test_thinking_budget_holder_has_tracked_after_sync_add():

1020

1021 def test_thinking_budget_holder_sync_remove_clears_state():

1022 vc = VllmConfig()

1033 - vc.reasoning_config = MockReasoningConfig()

1023 + vc.reasoning_config = create_mock_reasoning_config(

1024 + [THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID]

1025 + )

1026 h = ThinkingBudgetStateHolder(

1027 vc.reasoning_config,

1028 vc.scheduler_config.max_num_seqs,

@@ -1060,7 +1052,9 @@ def test_thinking_budget_holder_sync_remove_clears_state():

1052

1053 def test_thinking_budget_holder_sync_add_without_budget_drops_row():

1054 vc = VllmConfig()

1063 - vc.reasoning_config = MockReasoningConfig()

1055 + vc.reasoning_config = create_mock_reasoning_config(

1056 + [THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID]

1057 + )

1058 h = ThinkingBudgetStateHolder(

1059 vc.reasoning_config,

1060 vc.scheduler_config.max_num_seqs,

@@ -1081,7 +1075,9 @@ def test_thinking_budget_holder_sync_add_without_budget_drops_row():

1075

1076 def test_thinking_budget_holder_swap_exchanges_state():

1077 vc = VllmConfig()

1084 - vc.reasoning_config = MockReasoningConfig()

1078 + vc.reasoning_config = create_mock_reasoning_config(

1079 + [THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID]

1080 + )

1081 h = ThinkingBudgetStateHolder(

1082 vc.reasoning_config,

1083 vc.scheduler_config.max_num_seqs,

@@ -1125,7 +1121,9 @@ def test_thinking_budget_holder_swap_exchanges_state():

1121

1122 def test_thinking_budget_holder_unidirectional_move():

1123 vc = VllmConfig()

1128 - vc.reasoning_config = MockReasoningConfig()

1124 + vc.reasoning_config = create_mock_reasoning_config(

1125 + [THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID]

1126 + )

1127 h = ThinkingBudgetStateHolder(

1128 vc.reasoning_config,

1129 vc.scheduler_config.max_num_seqs,

@@ -1163,7 +1161,9 @@ def test_thinking_budget_holder_unidirectional_move():

1161

1162 def test_thinking_budget_holder_update_state_repeat_indices_last_row_wins():

1163 vc = VllmConfig()

1166 - vc.reasoning_config = MockReasoningConfig()

1164 + vc.reasoning_config = create_mock_reasoning_config(

1165 + [THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID]

1166 + )

1167 h = ThinkingBudgetStateHolder(

1168 vc.reasoning_config,

1169 vc.scheduler_config.max_num_seqs,

@@ -1197,7 +1197,7 @@ def test_thinking_budget_holder_update_state_repeat_indices_last_row_wins():

1197

1198 def test_thinking_budget_holder_spec_mode_tensor_layout():

1199 h = ThinkingBudgetStateHolder(

1200 - MockReasoningConfig(),

1200 + create_mock_reasoning_config([THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID]),

1201 8,

1202 2,

1203 torch.device("cpu"),

@@ -1209,7 +1209,7 @@ def test_thinking_budget_holder_spec_mode_tensor_layout():

1209

1210 def test_thinking_budget_holder_empty_end_tokens_disables_row():

1211 vc = VllmConfig()

1212 - vc.reasoning_config = MockReasoningNoEndTokens()

1212 + vc.reasoning_config = create_mock_reasoning_config([THINK_START_TOKEN_ID], [])

1213 h = ThinkingBudgetStateHolder(

1214 vc.reasoning_config,

1215 vc.scheduler_config.max_num_seqs,

@@ -1253,7 +1253,9 @@ def test_thinking_budget_enforced_without_penalties():

1253 passing an empty list (the pre-fix behavior) prevents budget enforcement.

1254 """

1255 vc = VllmConfig()

1256 - vc.reasoning_config = MockReasoningConfig()

1256 + vc.reasoning_config = create_mock_reasoning_config(

1257 + [THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID]

1258 + )

1259 budget = 3 # allow 3 thinking tokens

1260

1261 h = ThinkingBudgetStateHolder(

@@ -1346,7 +1348,11 @@ def test_thinking_budget_long_thinking_section_end_marker_found_at_correct_index

1348 """Test thinking budget enforced for a long thinking run,

1349 then a natural end marker."""

1350 h = ThinkingBudgetStateHolder(

1349 - MockReasoningConfig(), 8, 0, torch.device("cpu"), False

1351 + create_mock_reasoning_config([THINK_START_TOKEN_ID], [THINK_END_TOKEN_ID]),

1352 + 8,

1353 + 0,

1354 + torch.device("cpu"),

1355 + False,

1356 )

1357 h.sync_batch(

1358 BatchUpdate(

@@ -1356,8 +1362,8 @@ def test_thinking_budget_long_thinking_section_end_marker_found_at_correct_index

1362 moved=(),

1363 )

1364 )

1359 - start = MockReasoningConfig.reasoning_start_token_ids

1360 - end = MockReasoningConfig.reasoning_end_token_ids

1365 + start = [THINK_START_TOKEN_ID]

1366 + end = [THINK_END_TOKEN_ID]

1367

1368 out: list[int] = list(start)

1369 h.update_state([out], None, None)

@@ -1396,13 +1402,9 @@ class TestThinkingBudgetReentry:

1402

1403 @staticmethod

1404 def _make_holder(end_token_ids: list[int]) -> ThinkingBudgetStateHolder:

1399 - class FakeReasoningConfig:

1400 - reasoning_start_token_ids = [TestThinkingBudgetReentry.THINK_START]

1401 - reasoning_end_token_ids: list[int] = []

1402 - enabled = True

1403 -

1404 - cfg = FakeReasoningConfig()

1405 - cfg.reasoning_end_token_ids = end_token_ids

1405 + cfg = create_mock_reasoning_config(

1406 + [TestThinkingBudgetReentry.THINK_START], end_token_ids

1407 + )

1408 return ThinkingBudgetStateHolder(

1409 reasoning_config=cfg,

1410 max_num_seqs=8,

@@ -1526,21 +1528,12 @@ class TestThinkingBudgetNaturalEndReentry:

1528

1529 @staticmethod

1530 def _make_holder(end_token_ids):

1529 - from dataclasses import dataclass

1530 -

1531 from vllm.v1.sample.thinking_budget_state import (

1532 ThinkingBudgetStateHolder,

1533 )

1534

1535 - @dataclass

1536 - class FakeReasoningConfig:

1537 - reasoning_start_token_ids: list[int]

1538 - reasoning_end_token_ids: list[int]

1539 - enabled: bool = True

1540 -

1541 - cfg = FakeReasoningConfig(

1542 - reasoning_start_token_ids=[TestThinkingBudgetNaturalEndReentry.THINK_START],

1543 - reasoning_end_token_ids=end_token_ids,

1535 + cfg = create_mock_reasoning_config(

1536 + [TestThinkingBudgetNaturalEndReentry.THINK_START], end_token_ids

1537 )

1538 return ThinkingBudgetStateHolder(

1539 reasoning_config=cfg,

@@ -1830,21 +1823,11 @@ class TestThinkingBudgetNaturalEndReentry:

1823 natural </think> exit, continue_thinking is cleared and Block 2

1824 enforcement works correctly.

1825 """

1833 - from dataclasses import dataclass

1826 from unittest.mock import MagicMock

1827

1828 from vllm.v1.sample.thinking_budget_state import ThinkingBudgetStateHolder

1829

1838 - @dataclass

1839 - class FakeReasoningConfig:

1840 - reasoning_start_token_ids: list[int]

1841 - reasoning_end_token_ids: list[int]

1842 - enabled: bool = True

1843 -

1844 - cfg = FakeReasoningConfig(

1845 - reasoning_start_token_ids=[self.THINK_START],

1846 - reasoning_end_token_ids=self.THINK_END_SINGLE,

1847 - )

1830 + cfg = create_mock_reasoning_config([self.THINK_START], self.THINK_END_SINGLE)

1831 holder = ThinkingBudgetStateHolder(

1832 reasoning_config=cfg,

1833 max_num_seqs=8,

tests/v1/logits_processors/utils.py

+10 -6

@@ -57,15 +57,20 @@ prompts = [

57 class DummyLogitsProcessor(LogitsProcessor):

58 """Fake logit processor to support unit testing and examples."""

59

60 - @classmethod

61 - def validate_params(cls, params: SamplingParams):

62 - target_token: int | None = params.extra_args and params.extra_args.get(

63 - "target_token"

60 + @staticmethod

61 + def _get_target_token(params: SamplingParams) -> int | None:

62 + target_token = (

63 + params.extra_args.get(DUMMY_LOGITPROC_ARG) if params.extra_args else None

64 )

65 if target_token is not None and not isinstance(target_token, int):

66 raise VLLMValidationError(

67 f"target_token value {target_token} {type(target_token)} is not int"

68 )

69 + return target_token

70 +

71 + @classmethod

72 + def validate_params(cls, params: SamplingParams):

73 + cls._get_target_token(params)

74

75 def __init__(

76 self, vllm_config: "VllmConfig", device: torch.device, is_pin_memory: bool

@@ -78,8 +83,7 @@ class DummyLogitsProcessor(LogitsProcessor):

83

84 def update_state(self, batch_update: BatchUpdate | None):

85 def extract_extra_arg(params: SamplingParams) -> int | None:

81 - self.validate_params(params)

82 - return params.extra_args and params.extra_args.get("target_token")

86 + return self._get_target_token(params)

87

88 process_dict_updates(

89 self.req_info,

tests/v1/sample/test_diffusion_gemma_reads.py

+1

@@ -341,6 +341,7 @@ def test_batch_allowed_needs_one_shared_set():

341 assert states.batch_allowed([]) is None

342

343 shared = states.batch_allowed([0])

344 + assert shared is not None

345 assert shared.tolist() == [1, 2, 3]

346 assert shared.dtype == torch.int64

347 assert states.batch_allowed([0, 0]) is shared

tests/v1/sample/test_diffusion_sampler_stats.py

+3

@@ -34,6 +34,8 @@ def test_stats_match_reference(vocab: int, dtype: torch.dtype):

34 )

35 assert torch.equal(argmax, ref_argmax)

36 torch.testing.assert_close(entropy, ref_entropy, atol=2e-3, rtol=1e-3)

37 + assert probs is not None

38 + assert ref_probs is not None

39 torch.testing.assert_close(probs.float(), ref_probs.float(), atol=5e-3, rtol=2e-2)

40 assert probs.dtype == torch.bfloat16

41 assert probs.shape == logits.shape

@@ -100,6 +102,7 @@ def test_masked_logits_keep_a_finite_entropy():

102 assert entropy[0].item() < 1e-6 and entropy[1].item() < 1e-6

103 assert abs(entropy[2].item() - math.log(2.0)) < 1e-4

104 assert torch.equal(argmax[[0, 1, 3]], torch.tensor([3, 3, 3], device="cuda"))

105 + assert probs is not None

106 assert probs[0, 3].item() == 1.0 and probs[0].sum().item() == 1.0

107 _, _, ref_entropy, _ = sample_row_stats_reference(logits, temps, 1, None)

108 assert torch.isfinite(ref_entropy).all()

tests/v1/sample/test_dry.py

+6 -3

@@ -340,7 +340,7 @@ def _make_v2_state(

340 device=device,

341 all_token_ids=SimpleNamespace(gpu=all_tokens),

342 )

343 - state = DryState(_make_vllm_config(), req_states)

343 + state = DryState(_make_vllm_config(), req_states) # type: ignore[arg-type]

344 if tokenizer is not None:

345 # Re-resolve the default set now that a tokenizer is installed.

346 state._tokenizer = tokenizer

@@ -446,7 +446,7 @@ def test_spec_decode_is_refused_at_construction():

446 )

447 config = _make_vllm_config(speculative_config=SimpleNamespace())

448 with pytest.raises(ValueError, match="speculative decoding"):

449 - DryState(config, req_states)

449 + DryState(config, req_states) # type: ignore[arg-type]

450

451

452 def test_v2_draft_expanded_logits_raise():

@@ -917,7 +917,10 @@ def test_loads_through_the_loader_by_fqcn():

917 total_len=None,

918 )

919 procs = loader.build_custom_logits_processors(

920 - _make_vllm_config(), req_states, False, [DRY_FQCN]

920 + _make_vllm_config(),

921 + req_states, # type: ignore[arg-type]

922 + False,

923 + [DRY_FQCN],

924 )

925 assert [type(p) for p in procs] == [DryState]

926

tests/v1/sample/test_head_dtype.py

+21 -7

@@ -16,6 +16,7 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor

16 from vllm.model_executor.layers.vocab_parallel_embedding import (

17 ParallelLMHead,

18 UnquantizedEmbeddingMethod,

19 + VocabParallelEmbedding,

20 )

21

22

@@ -32,10 +33,17 @@ class _FakeLmHead:

33 self.tp_size = 1

34

35

36 +def _fake_lm_head(

37 + weight: torch.Tensor, quantized: bool = False, shard_indices: object | None = None

38 +) -> VocabParallelEmbedding:

39 + # The dtype tests only need the projection fields, not distributed setup.

40 + return _FakeLmHead(weight, quantized, shard_indices) # type: ignore[return-value]

41 +

42 +

43 def _build_processor(vocab_size: int) -> LogitsProcessor:

44 lp = LogitsProcessor(vocab_size)

45 # The TP gather is orthogonal to the dtype behavior under test.

38 - lp._gather_logits = lambda logits: logits

46 + lp._gather_logits = lambda logits: logits # type: ignore[method-assign]

47 return lp

48

49

@@ -47,8 +55,9 @@ def test_fp32_head_runs_projection_in_fp32(default_vllm_config):

55 hidden_states = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16)

56 weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16)

57

50 - logits = lp._get_logits(hidden_states, _FakeLmHead(weight), None)

58 + logits = lp._get_logits(hidden_states, _fake_lm_head(weight), None)

59

60 + assert logits is not None

61 assert logits.dtype == torch.float32

62 assert torch.isfinite(logits).all()

63 expected = torch.nn.functional.linear(hidden_states.float(), weight.float())

@@ -66,8 +75,9 @@ def test_non_fp32_head_dtype_uses_cast_path(default_vllm_config):

75 hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16)

76 weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16)

77

69 - logits = lp._get_logits(hidden_states, _FakeLmHead(weight), None)

78 + logits = lp._get_logits(hidden_states, _fake_lm_head(weight), None)

79

80 + assert logits is not None

81 assert logits.dtype == torch.float16

82 expected = torch.nn.functional.linear(hidden_states.half(), weight.half())

83 torch.testing.assert_close(logits, expected)

@@ -82,7 +92,7 @@ def test_head_dtype_equal_to_model_dtype_uses_quant_method(default_vllm_config):

92

93 hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16)

94 weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16)

85 - lm_head = _FakeLmHead(weight)

95 + lm_head = _fake_lm_head(weight)

96

97 with mock.patch.object(

98 lm_head.quant_method,

@@ -94,6 +104,7 @@ def test_head_dtype_equal_to_model_dtype_uses_quant_method(default_vllm_config):

104 logits = lp._get_logits(hidden_states, lm_head, None)

105

106 apply_mock.assert_called_once()

107 + assert logits is not None

108 assert logits.dtype == torch.bfloat16

109

110

@@ -120,9 +131,10 @@ def test_fp32_head_uses_mm_fast_path_on_device(default_vllm_config):

131 with mock.patch(

132 "vllm.model_executor.layers.logits_processor.F.linear"

133 ) as linear_mock:

123 - logits = lp._get_logits(hidden_states, _FakeLmHead(weight), None)

134 + logits = lp._get_logits(hidden_states, _fake_lm_head(weight), None)

135

136 linear_mock.assert_not_called()

137 + assert logits is not None

138 assert logits.dtype == torch.float32

139 expected = torch.nn.functional.linear(hidden_states.float(), weight.float())

140 torch.testing.assert_close(logits, expected)

@@ -131,7 +143,7 @@ def test_fp32_head_uses_mm_fast_path_on_device(default_vllm_config):

143 def test_fp32_head_rejects_quantized_lm_head(default_vllm_config):

144 lp = _build_processor(64)

145 lp.head_dtype = torch.float32

134 - lm_head = _FakeLmHead(torch.randn(64, 16, dtype=torch.bfloat16), quantized=True)

146 + lm_head = _fake_lm_head(torch.randn(64, 16, dtype=torch.bfloat16), quantized=True)

147

148 with pytest.raises(ValueError, match="unquantized"):

149 lp._get_logits(torch.randn(4, 16, dtype=torch.bfloat16), lm_head, None)

@@ -197,7 +209,7 @@ def test_get_top_tokens_honors_head_dtype(default_vllm_config):

209

210 hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16)

211 weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16)

200 - lm_head = _FakeLmHead(

212 + lm_head = _fake_lm_head(

213 weight,

214 shard_indices=types.SimpleNamespace(

215 num_org_vocab_padding=0, org_vocab_start_index=0

@@ -236,7 +248,9 @@ def test_fp32_head_e2e_no_nan():

248

249 for output in outputs:

250 for completion in output.outputs:

251 + assert completion.logprobs is not None

252 for token_id, position in zip(completion.token_ids, completion.logprobs):

253 + assert position is not None

254 # The sampled token survived filtering, so its logprob is finite.

255 assert math.isfinite(position[token_id].logprob)

256 # No returned logprob is NaN.

tests/v1/sample/test_logprobs.py

+33 -17

@@ -24,6 +24,7 @@ from vllm import SamplingParams

24 from vllm.config.model import LogprobsMode

25 from vllm.distributed import cleanup_dist_env_and_memory

26 from vllm.exceptions import VLLMValidationError

27 +from vllm.logprobs import Logprob

28 from vllm.platforms import current_platform

29 from vllm.v1.engine.input_processor import InputProcessor

30

@@ -145,7 +146,7 @@ def _repeat_logprob_config(

146 def _run_and_validate(

147 vllm_model: VllmRunner,

148 test_prompts: list[str],

148 - vllm_sampling_params: SamplingParams,

149 + vllm_sampling_params: SamplingParams | list[SamplingParams],

150 hf_logprobs: list[list[torch.Tensor]],

151 hf_outputs: list[tuple[list[int], str]],

152 logprob_prompt_logprob_list: BatchLogprobsSpecType,

@@ -163,11 +164,12 @@ def _run_and_validate(

164 # Extract request-level (prompt)logprobs config

165 num_top_logprobs, num_top_prompt_logprobs = logprob_prompt_logprob

166

167 + assert vllm_result.prompt_token_ids is not None

168 # Test whether sampled token output is consistent between vLLM and HF

169 # vLLM prompt+completion should match HF output

170 if temperature == 0.0:

171 assert (

170 - vllm_result.prompt_token_ids + vllm_result.outputs[0].token_ids

172 + vllm_result.prompt_token_ids + list(vllm_result.outputs[0].token_ids)

173 == hf_output[0]

174 )

175 else:

@@ -191,7 +193,9 @@ def _run_and_validate(

193

194 # Confirm that the output token appears among the logprobs

195 assert token_id in logprobs

194 - token_in_topk = logprobs[token_id].rank <= num_top_logprobs

196 + rank = logprobs[token_id].rank

197 + assert rank is not None

198 + token_in_topk = rank <= num_top_logprobs

199

200 # If the output token is not included in the top K

201 # logprob, it can return 1 more data

@@ -207,8 +211,10 @@ def _run_and_validate(

211

212 output_text = vllm_result.outputs[0].text

213 output_string_from_most_likely_tokens_lst: list[str] = []

210 - for top_logprobs in vllm_result.outputs[0].logprobs:

211 - top_logprob = next(iter(top_logprobs.values()))

214 + for output_logprobs in vllm_result.outputs[0].logprobs:

215 + assert output_logprobs is not None

216 + top_logprob = next(iter(output_logprobs.values()))

217 + assert top_logprob.decoded_token is not None

218 output_string_from_most_likely_tokens_lst.append(

219 top_logprob.decoded_token

220 )

@@ -226,8 +232,9 @@ def _run_and_validate(

232

233 # Compare vLLM sample logprobs to HF

234 vllm_sample_logprobs = vllm_result.outputs[0].logprobs

229 - for i, top_logprobs in enumerate(vllm_sample_logprobs):

230 - for token_id, sample_logprob in top_logprobs.items():

235 + for i, sample_logprobs in enumerate(vllm_sample_logprobs):

236 + assert sample_logprobs is not None

237 + for token_id, sample_logprob in sample_logprobs.items():

238 if temperature == 0.0 or i == 0:

239 logprob = sample_logprob.logprob

240 torch.testing.assert_close(

@@ -271,9 +278,9 @@ def _run_and_validate(

278

279 # Confirm that the prompt token appears among the logprobs

280 assert prompt_token_id in prompt_logprobs

274 - token_in_topk = (

275 - prompt_logprobs[prompt_token_id].rank <= num_top_prompt_logprobs

276 - )

281 + prompt_rank = prompt_logprobs[prompt_token_id].rank

282 + assert prompt_rank is not None

283 + token_in_topk = prompt_rank <= num_top_prompt_logprobs

284

285 # If the prompt token is not included in the top K

286 # logprob, it can return 1 more data

@@ -294,9 +301,10 @@ def _run_and_validate(

301 # 1:.

302 vllm_prompt_logprobs = vllm_result.prompt_logprobs[1:]

303 for i, vllm_prompt_logprob_dict in enumerate(vllm_prompt_logprobs):

297 - for token_id, logprob in vllm_prompt_logprob_dict.items():

304 + assert vllm_prompt_logprob_dict is not None

305 + for token_id, prompt_lp in vllm_prompt_logprob_dict.items():

306 torch.testing.assert_close(

299 - logprob.logprob,

307 + prompt_lp.logprob,

308 hf_logprob[0][i][token_id].item(),

309 atol=2e-2,

310 rtol=2e-2,

@@ -478,12 +486,12 @@ def test_prompt_logprob_token_ids_require_v2_model_runner():

486 params = SamplingParams(prompt_logprob_token_ids=[1, 2])

487 with patch.object(SamplingParams, "verify"):

488 with pytest.raises(VLLMValidationError, match="V2 model runner"):

481 - InputProcessor._validate_params(processor, params, ("generate",))

489 + InputProcessor._validate_params(processor, params, ("generate",)) # type: ignore[arg-type]

490 processor.vllm_config.use_v2_model_runner = True

483 - InputProcessor._validate_params(processor, params, ("generate",))

491 + InputProcessor._validate_params(processor, params, ("generate",)) # type: ignore[arg-type]

492 processor.vllm_config.cache_config.kv_sharing_fast_prefill = True

493 with pytest.raises(VLLMValidationError, match="fast-prefill"):

486 - InputProcessor._validate_params(processor, params, ("generate",))

494 + InputProcessor._validate_params(processor, params, ("generate",)) # type: ignore[arg-type]

495

496

497 def test_none_logprobs(vllm_model, example_prompts):

@@ -607,7 +615,9 @@ def test_logprobs_mode(logprobs_mode: LogprobsMode):

615 total_token_with_logprobs = 0

616 positive_values = 0

617 for output in results[0].outputs:

618 + assert output.logprobs is not None

619 for logprobs in output.logprobs:

620 + assert logprobs is not None

621 for token_id in logprobs:

622 logprob = logprobs[token_id]

623 if logprobs_mode in ("raw_logprobs", "processed_logprobs"):

@@ -648,6 +658,7 @@ def test_prompt_logprobs_mode():

658 )

659 assert results[0].prompt_logprobs is not None

660 assert results[0].prompt_logprobs[1] is not None

661 + assert results[0].prompt_token_ids is not None

662 tok_id = results[0].prompt_token_ids[1]

663 values[mode] = results[0].prompt_logprobs[1][tok_id].logprob

664 finally:

@@ -1228,10 +1239,12 @@ def test_spec_decode_logprobs(

1239 [prompt, prompt], [sampling_params, penalty_sampling_params]

1240 )

1241 # Collect logprobs outputs from reference LLM.

1231 - ref_logprobs = []

1242 + ref_logprobs: list[Logprob] = []

1243 for results in ref_results:

1244 for output in results.outputs:

1245 + assert output.logprobs is not None

1246 for logprobs in output.logprobs:

1247 + assert logprobs is not None

1248 ref_logprobs.extend(logprobs.values())

1249 del ref_llm

1250 torch.accelerator.empty_cache()

@@ -1249,10 +1262,12 @@ def test_spec_decode_logprobs(

1262 [prompt, prompt], [sampling_params, penalty_sampling_params]

1263 )

1264 # Collect logprobs outputs from spec decode LLM.

1252 - spec_logprobs = []

1265 + spec_logprobs: list[Logprob] = []

1266 for results in spec_results:

1267 for output in results.outputs:

1268 + assert output.logprobs is not None

1269 for logprobs in output.logprobs:

1270 + assert logprobs is not None

1271 spec_logprobs.extend(logprobs.values())

1272 del spec_llm

1273 torch.accelerator.empty_cache()

@@ -1326,6 +1341,7 @@ def test_prompt_logprobs_with_chunking_and_preemption():

1341 "Unexpected number of prompt logprob positions"

1342 )

1343

1344 + assert sampling_params.prompt_logprobs is not None

1345 # Each position should have the requested number of logprobs

1346 for pos, logprobs_dict in enumerate(prompt_logprobs):

1347 if logprobs_dict is not None: # First token may be None

tests/v1/sample/test_sampler.py

+1

@@ -333,6 +333,7 @@ def test_sampler_repetition_penalty(

333 for batch_idx in range(batch_size):

334 non_penalized_token_id = logits[batch_idx].argmax().item()

335 penalized_token_id = logits[batch_idx].argmin().item()

336 + assert sampling_metadata.prompt_token_ids is not None

337 prompt_tokens = sampling_metadata.prompt_token_ids[batch_idx][:].tolist()

338 output_tokens = sampling_metadata.output_token_ids[batch_idx]

339 if repetition_penalty > 1.0:

tests/v1/sample/test_thinking_budget_state.py

+2 -6

@@ -4,6 +4,7 @@

4

5 import torch

6

7 +from tests.v1.sample.utils import create_mock_reasoning_config

8 from vllm.sampling_params import SamplingParams

9 from vllm.v1.sample.logits_processor.interface import (

10 BatchUpdate,

@@ -12,14 +13,9 @@ from vllm.v1.sample.logits_processor.interface import (

13 from vllm.v1.sample.thinking_budget_state import ThinkingBudgetStateHolder

14

15

15 -class _MockReasoningConfig:

16 - reasoning_start_token_ids = [151667]

17 - reasoning_end_token_ids = [151668]

18 -

19 -

16 def _make_holder() -> ThinkingBudgetStateHolder:

17 return ThinkingBudgetStateHolder(

22 - _MockReasoningConfig(),

18 + create_mock_reasoning_config([151667], [151668]),

19 8,

20 0,

21 torch.device("cpu"),

tests/v1/sample/test_trace_replay_params.py

+2 -1

@@ -104,7 +104,8 @@ def test_trace_decode_token_ids_rejects_speculative_decoding():

104 params = SamplingParams(trace_decode_token_ids=[1])

105 with pytest.raises(ValueError, match="not supported with speculative decoding"):

106 params._validate_trace_replay(

107 - _make_model_config(vocab_size=100), speculative_config=object()

107 + _make_model_config(vocab_size=100),

108 + speculative_config=object(), # type: ignore[arg-type] # Only non-None presence is tested.

109 )

110

111

tests/v1/sample/utils.py

+17 -1

@@ -4,16 +4,28 @@

4 from collections.abc import Iterator

5 from enum import Enum

6 from typing import NamedTuple

7 +from unittest.mock import MagicMock

8

9 import regex as re

10 import torch

11

12 from vllm import CompletionOutput

13 +from vllm.config import ReasoningConfig

14 from vllm.utils.torch_utils import make_tensor_with_pad

15 from vllm.v1.sample.logits_processor import BatchUpdate, LogitsProcessor

16 from vllm.v1.sample.metadata import SamplingMetadata

17

18

19 +def create_mock_reasoning_config(

20 + start_token_ids: list[int], end_token_ids: list[int]

21 +) -> ReasoningConfig:

22 + return MagicMock(

23 + spec=ReasoningConfig,

24 + reasoning_start_token_ids=start_token_ids,

25 + reasoning_end_token_ids=end_token_ids,

26 + )

27 +

28 +

29 class BatchLogprobsComposition(Enum):

30 """Types of logprobs configs to include in test batch."""

31

@@ -141,7 +153,11 @@ def compute_correct_cumulative_logprob(completion_output: CompletionOutput) -> f

153 token_ids = completion_output.token_ids

154 logprobs = completion_output.logprobs

155 assert logprobs is not None

144 - return sum([lp[tok_id].logprob for tok_id, lp in zip(token_ids, logprobs)])

156 + values: list[float] = []

157 + for tok_id, lp in zip(token_ids, logprobs):

158 + assert lp is not None

159 + values.append(lp[tok_id].logprob)

160 + return sum(values)

161

162

163 def create_fake_logits(batch_size: int, vocab_size: int) -> torch.Tensor:

tests/v1/spec_decode/test_adaptive_verification.py

+16 -8

@@ -25,9 +25,9 @@ def make_manager(

25 num_reqs, num_steps = confidences.shape

26 manager = AdaptiveVerificationManager.__new__(AdaptiveVerificationManager)

27 manager.num_speculative_steps = num_steps

28 - manager._stale_confidences = [SimpleNamespace(np=confidences)]

28 + manager._stale_confidences = [SimpleNamespace(np=confidences)] # type: ignore[list-item] # CPU-only confidence buffer.

29 manager._stale_idx = 0

30 - manager.req_states = SimpleNamespace(

30 + manager.req_states = SimpleNamespace( # type: ignore[assignment] # CPU-only request fixture.

31 req_id_to_index={"low": 0, "high": 1},

32 num_computed_tokens_np=np.ones(num_reqs, dtype=np.int32),

33 prefill_len=SimpleNamespace(np=np.ones(num_reqs, dtype=np.int32)),

@@ -153,11 +153,11 @@ def test_manager_checks_target_varlen_cudagraph_bound(

153 manager = maybe_create_adaptive_verification_manager(

154 enable_adaptive_verification=True,

155 attn_groups=groups,

156 - req_states=SimpleNamespace(num_speculative_steps=7),

156 + req_states=SimpleNamespace(num_speculative_steps=7), # type: ignore[arg-type]

157 query_start_loc=object(),

158 num_bonus_tokens=1,

159 max_total_logits=1,

160 - vllm_config=None,

160 + vllm_config=None, # type: ignore[arg-type]

161 target_layer_names={"target"},

162 additional_attn_cg_support=additional,

163 )

@@ -176,6 +176,7 @@ def test_budget_stops_where_marginal_drafts_stop_paying_for_themselves():

176 {"low": 3, "high": 3},

177 {"low": [1, 2], "high": [3, 4]},

178 )

179 + assert manager._batch_budget is not None

180 valid_drafts, num_non_draft_tokens, draft_budget = manager._batch_budget

181

182 assert draft_budget == 1

@@ -183,14 +184,18 @@ def test_budget_stops_where_marginal_drafts_stop_paying_for_themselves():

184 assert num_non_draft_tokens == {"low": 1, "high": 1}

185

186

186 -def test_profiled_batches_seed_cost_curves_via_consumer():

187 +def test_profiled_batches_seed_cost_curves_via_consumer(monkeypatch):

188 manager = AdaptiveVerificationManager.__new__(AdaptiveVerificationManager)

188 - manager.req_states = SimpleNamespace(max_num_batched_tokens=4096, max_num_reqs=64)

189 + manager.req_states = SimpleNamespace( # type: ignore[assignment] # CPU-only request fixture.

190 + max_num_batched_tokens=4096, max_num_reqs=64

191 + )

192 manager.num_speculative_steps = 7

193 manager.num_bonus_tokens = 1

194 curves: dict[str, list[tuple[int, float]]] = {}

192 - manager.set_cost_curves = lambda draft, verify: curves.update(

193 - draft=draft, verify=verify

195 + monkeypatch.setattr(

196 + manager,

197 + "set_cost_curves",

198 + lambda draft, verify: curves.update(draft=draft, verify=verify),

199 )

200

201 timings = [

@@ -258,6 +263,7 @@ def test_budget_caps_at_one_rejection_sampler_chunk():

263 {"low": 3, "high": 3},

264 {"low": [1, 2], "high": [3, 4]},

265 )

266 + assert manager._batch_budget is not None

267 _, _, draft_budget = manager._batch_budget

268 assert draft_budget <= 1

269

@@ -285,6 +291,7 @@ def test_zero_budget_rebuilds_cpu_cu_num_logits():

291 {"low": 3, "high": 3, "prefill": 40},

292 {"low": [1, 2], "high": [3, 4]},

293 )

294 + assert manager._batch_budget is not None

295 _, _, draft_budget = manager._batch_budget

296 assert draft_budget == 0

297

@@ -323,6 +330,7 @@ def test_zero_budget_keeps_one_grammar_row_per_scheduled_draft():

330 manager.get_num_tokens(

331 {"low": 3, "high": 3, "prefill": 40}, scheduled_spec_decode_tokens

332 )

333 + assert manager._batch_budget is not None

334 assert manager._batch_budget[2] == 0

335

336 req_ids = ["low", "high", "prefill"]

tests/v1/spec_decode/test_dflash2.py

+10 -8

@@ -98,7 +98,7 @@ def test_draft_quant_exclusions_include_global_layer_indices():

98 ]

99 )

100

101 - _add_global_draft_layer_exclusions(quant_config, 88, 5)

101 + _add_global_draft_layer_exclusions(quant_config, 88, 5) # type: ignore[arg-type]

102

103 assert "layers.88.mlp_conv*" in quant_config.exclude_modules

104 assert "*layers.92.self_attn.q_proj" in quant_config.exclude_modules

@@ -171,7 +171,7 @@ def test_selector_leaves_greedy_drafting_without_proposal_logits(monkeypatch):

171 allocating one here would claim a proposal the walk never sampled from.

172 """

173 _stub_base(monkeypatch, None)

174 - speculator = DFlash2Speculator(None, torch.device("cpu"))

174 + speculator = DFlash2Speculator(None, torch.device("cpu")) # type: ignore[arg-type] # Base initialization is stubbed.

175

176 assert speculator.draft_logits is None

177

@@ -183,7 +183,7 @@ def test_selector_asks_for_fp32_proposal_logits():

183 candidate row often enough that the walk and the rejection sampler checking it

184 would no longer read the same distribution.

185 """

186 - dtype, fill = DFlash2Speculator.draft_logits_spec(None, None)

186 + dtype, fill = DFlash2Speculator.draft_logits_spec(None, None) # type: ignore[arg-type] # Constant spec ignores self/config.

187

188 assert dtype is torch.float32

189 assert fill == float("-inf")

@@ -309,9 +309,9 @@ def test_candidate_model_decoder_layer_cls(monkeypatch, variant):

309 vllm_config.compilation_config = mock_current_vllm_config.compilation_config

310

311 # 3. Instantiate the model under meta device to avoid parameter allocation issues

312 - with set_current_vllm_config(mock_current_vllm_config), torch.device("meta"):

312 + with set_current_vllm_config(mock_current_vllm_config), torch.device("meta"): # type: ignore[arg-type]

313 model_cls = DFlash2Qwen3Model if variant == "dflash2" else LiLiCorr

314 - model = model_cls(vllm_config=vllm_config)

314 + model = model_cls(vllm_config=vllm_config) # type: ignore[arg-type]

315

316 # 4. Assert that the layers are DFlash2Qwen3DecoderLayer (the subclass)

317 assert len(model.layers) == 2

@@ -350,7 +350,7 @@ def test_conv_projections_use_draft_quant_config(monkeypatch):

350 quant_method="W4A16_NVFP4", is_checkpoint_nvfp4_serialized=True

351 )

352 layer = DFlash2Qwen3DecoderLayer(

353 - SimpleNamespace(

353 + SimpleNamespace( # type: ignore[arg-type]

354 speculative_config=SimpleNamespace(num_speculative_tokens=7),

355 model_config=SimpleNamespace(dtype=torch.bfloat16),

356 ),

@@ -434,13 +434,15 @@ def test_context_kv_uses_quantized_projection_fallback(monkeypatch):

434 for projection in projections

435 ]

436 qwen3_dflash.DFlashQwen3Model._build_context_kv_buffers(

437 - model, layers_attn, has_bias=False

437 + model, # type: ignore[arg-type]

438 + layers_attn,

439 + has_bias=False,

440 )

441 assert model._fused_kv_weight is None

442 assert all(not hasattr(projection, "weight") for projection in projections)

443

444 actual_k, actual_v = qwen3_dflash.DFlashQwen3Model._project_context_kv(

443 - model,

445 + model, # type: ignore[arg-type]

446 context_states,

447 num_ctx=2,

448 num_layers=2,

tests/v1/spec_decode/test_dflash_lookahead.py

+4 -4

@@ -141,9 +141,9 @@ def test_dflash_drafter_window_reserves_bonus_token():

141 speculative_config=_dflash_speculative_config(NUM_SPECULATIVE_TOKENS),

142 )

143 # window = 4, so 96 fits (96 + 4 == 100) but 97 does not (97 + 4 == 101)

144 - assert input_fits_in_drafter(dflash_runner, SimpleNamespace(max_seq_len=96))

145 - assert not input_fits_in_drafter(dflash_runner, SimpleNamespace(max_seq_len=97))

146 - assert not input_fits_in_drafter(dflash_runner, None) # no metadata

144 + assert input_fits_in_drafter(dflash_runner, SimpleNamespace(max_seq_len=96)) # type: ignore[arg-type]

145 + assert not input_fits_in_drafter(dflash_runner, SimpleNamespace(max_seq_len=97)) # type: ignore[arg-type]

146 + assert not input_fits_in_drafter(dflash_runner, None) # type: ignore[arg-type] # No metadata.

147

148 # Other drafters don't reserve the bonus token, so 97 fits (97 + 3 == 100).

149 plain_runner = SimpleNamespace(

@@ -151,4 +151,4 @@ def test_dflash_drafter_window_reserves_bonus_token():

151 effective_drafter_max_model_len=100,

152 speculative_config=SimpleNamespace(use_dflash=lambda: False),

153 )

154 - assert input_fits_in_drafter(plain_runner, SimpleNamespace(max_seq_len=97))

154 + assert input_fits_in_drafter(plain_runner, SimpleNamespace(max_seq_len=97)) # type: ignore[arg-type]

tests/v1/spec_decode/test_dsv4_mtp_local_argmax.py

+4 -1

@@ -56,6 +56,7 @@ def _make_predictor(vocab_size: int) -> DeepSeekV4MultiTokenPredictor:

56 ):

57 shared_head = SharedHead(config=config, prefix="shared_head")

58

59 + assert shared_head.head is not None

60 generator = torch.Generator().manual_seed(7)

61 with torch.no_grad():

62 shared_head.head.weight.normal_(0.0, 0.02, generator=generator)

@@ -157,7 +158,9 @@ def test_wrapper_forwards_spec_step_idx():

158

159 hidden_states = torch.empty(0)

160 wrapper = object.__new__(DeepSeekV4MTP)

160 - wrapper.model = SimpleNamespace(get_top_tokens=lambda *args: args)

161 + wrapper.model = SimpleNamespace( # type: ignore[assignment]

162 + get_top_tokens=lambda *args: args

163 + )

164

165 assert DeepSeekV4MTP.get_top_tokens(wrapper, hidden_states, 2) == (hidden_states, 2)

166 assert DeepSeekV4MTP.get_top_tokens(wrapper, hidden_states) == (hidden_states, 0)

tests/v1/spec_decode/test_dynamic_sd_cug.py

+4 -4

@@ -25,7 +25,7 @@ def _create_vllm_config_for_dsd(

25 max_num_seqs: int,

26 max_spec_tokens: int,

27 *,

28 - cudagraph_mode: str = "FULL_AND_PIECEWISE",

28 + cudagraph_mode: CUDAGraphMode = CUDAGraphMode.FULL_AND_PIECEWISE,

29 use_dynamic_sd: bool = True,

30 num_spec_per_batch_size: list[tuple[int, int, int]] | None = None,

31 ) -> MagicMock:

@@ -199,7 +199,7 @@ def test_dynamic_sd_non_uniform_batch_falls_back_to_piecewise(monkeypatch):

199 vllm_config = _create_vllm_config_for_dsd(

200 max_num_seqs=512,

201 max_spec_tokens=max_spec_tokens,

202 - cudagraph_mode="FULL_AND_PIECEWISE",

202 + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,

203 use_dynamic_sd=True,

204 )

205 manager = gpu_cudagraph_utils.CudaGraphManager(

@@ -316,7 +316,7 @@ def test_basic_sd_does_not_capture_shorter_full_decode_shapes(monkeypatch):

316 vllm_config = _create_vllm_config_for_dsd(

317 max_num_seqs=max_num_seqs,

318 max_spec_tokens=max_spec_tokens,

319 - cudagraph_mode="FULL_AND_PIECEWISE",

319 + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,

320 use_dynamic_sd=False,

321 )

322 manager = gpu_cudagraph_utils.CudaGraphManager(

@@ -382,7 +382,7 @@ def test_dynamic_sd_only_captures_scheduled_query_lengths(monkeypatch):

382 vllm_config = _create_vllm_config_for_dsd(

383 max_num_seqs=max_num_seqs,

384 max_spec_tokens=max_spec_tokens,

385 - cudagraph_mode="FULL_AND_PIECEWISE",

385 + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,

386 use_dynamic_sd=True,

387 num_spec_per_batch_size=num_spec_per_batch_size,

388 )

tests/v1/spec_decode/test_eagle.py

+17 -11

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 +from typing import Literal

5 from unittest import mock

6

7 import numpy as np

@@ -31,6 +32,7 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum

32 from vllm.v1.spec_decode.dflash import DFlashProposer

33 from vllm.v1.spec_decode.draft_model import DraftModelProposer

34 from vllm.v1.spec_decode.eagle import EagleProposer

35 +from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer

36 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

37 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch

38

@@ -46,13 +48,13 @@ DEVICE_TYPE = current_platform.device_type

48

49

50 def _create_proposer(

49 - method: str,

51 + method: Literal["eagle", "eagle3", "draft_model", "dflash"],

52 num_speculative_tokens: int,

53 attention_backend: str | None = None,

54 parallel_drafting: bool = False,

53 - rejection_sample_method: str = "standard",

54 - draft_sample_method: str = "greedy",

55 -) -> EagleProposer:

55 + rejection_sample_method: Literal["standard", "synthetic", "block"] = "standard",

56 + draft_sample_method: Literal["greedy", "probabilistic"] = "greedy",

57 +) -> SpecDecodeBaseProposer:

58 # Method-dependent setup

59 if method == "eagle":

60 target_model_dir = model_dir

@@ -102,9 +104,14 @@ def _create_proposer(

104 max_model_len=model_config.max_model_len,

105 is_encoder_decoder=model_config.is_encoder_decoder,

106 ),

105 - attention_config=AttentionConfig(backend=attention_backend),

107 + attention_config=AttentionConfig(

108 + backend=AttentionBackendEnum[attention_backend]

109 + if attention_backend

110 + else None

111 + ),

112 )

113

114 + proposer: SpecDecodeBaseProposer

115 if method == "dflash":

116 proposer = DFlashProposer(vllm_config=vllm_config, device=device)

117 elif "eagle" in method:

@@ -135,7 +142,7 @@ def test_prepare_next_token_ids():

142 )

143

144 mock_num_scheduled_tokens = {req_id: 0 for req_id in req_ids}

138 - mock_requests = {}

145 + mock_requests: dict[str, CachedRequestState] = {}

146 for req_id in req_ids:

147 mock_request = mock.MagicMock(spec=CachedRequestState)

148 # Each request will have a backup next token id of 10, 20, 30, 40

@@ -948,13 +955,12 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):

955

956 attn_metadata_builder = attn_metadata_builder_cls(

957 kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),

951 - layer_names=proposer._draft_attn_layer_names,

958 + layer_names=sorted(proposer._draft_attn_layer_names),

959 vllm_config=proposer.vllm_config,

960 device=device,

961 )

962

956 - # Mock runner and draft_attn_groups for attention metadata building

957 - proposer.runner = mock.MagicMock()

963 + # Mock draft_attn_groups for attention metadata building.

964 mock_attn_group = mock.MagicMock()

965 mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder

966 mock_attn_group.layer_names = list(proposer._draft_attn_layer_names)

@@ -1054,11 +1060,10 @@ def test_propose_stores_probabilistic_draft_probs(attn_backend, monkeypatch):

1060 )

1061 attn_metadata_builder = attn_metadata_builder_cls(

1062 kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),

1057 - layer_names=proposer._draft_attn_layer_names,

1063 + layer_names=sorted(proposer._draft_attn_layer_names),

1064 vllm_config=proposer.vllm_config,

1065 device=device,

1066 )

1061 - proposer.runner = mock.MagicMock()

1067 mock_attn_group = mock.MagicMock()

1068 mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder

1069 mock_attn_group.layer_names = list(proposer._draft_attn_layer_names)

@@ -1133,6 +1138,7 @@ def test_set_inputs_first_pass_dflash():

1138

1139 num_speculative_tokens = 3

1140 proposer = _create_proposer("dflash", num_speculative_tokens)

1141 + assert isinstance(proposer, DFlashProposer)

1142 mask_token_id = proposer.parallel_drafting_token_id

1143

1144 # Setup batch with 3 requests

tests/v1/spec_decode/test_eagle_draft_attn_metadata.py

+1 -1

@@ -93,7 +93,7 @@ def _run_build(

93

94 with patch.object(base_speculator, "build_attn_metadata", fake_build_attn_metadata):

95 EagleSpeculator._build_uniform_attn_metadata(

96 - fake, # type: ignore[arg-type]

96 + fake,

97 batch_desc=batch_desc,

98 num_reqs=num_reqs,

99 num_query_per_req=num_query_per_req,

tests/v1/spec_decode/test_extract_hidden_states.py

+5 -5

@@ -51,7 +51,7 @@ def _create_proposer(

51 target_parallel_config=ParallelConfig(),

52 method="extract_hidden_states",

53 num_speculative_tokens=num_speculative_tokens,

54 - draft_model_config={

54 + draft_model_config={ # type: ignore[arg-type]

55 "hf_config": {

56 "eagle_aux_hidden_state_layer_ids": layer_ids,

57 }

@@ -103,7 +103,7 @@ def test_proposer_initialization_missing_layer_ids():

103 target_parallel_config=ParallelConfig(),

104 method="extract_hidden_states",

105 num_speculative_tokens=1,

106 - draft_model_config={

106 + draft_model_config={ # type: ignore[arg-type]

107 "hf_config": {} # Missing eagle_aux_hidden_state_layer_ids

108 },

109 )

@@ -146,7 +146,7 @@ def test_prepare_next_token_ids_padded():

146 mock_input_batch.vocab_size = 100

147 mock_input_batch.num_tokens_no_spec = np.array([5] * num_requests)

148

149 - mock_requests = {}

149 + mock_requests: dict[str, CachedRequestState] = {}

150 for req_id in req_ids:

151 mock_request = mock.MagicMock(spec=CachedRequestState)

152 # Each request will have a backup next token id of 10, 20, 30, 40

@@ -363,7 +363,7 @@ def test_extract_hidden_states_text_only_config_regression():

363 target_parallel_config=ParallelConfig(),

364 method="extract_hidden_states",

365 num_speculative_tokens=1,

366 - draft_model_config={

366 + draft_model_config={ # type: ignore[arg-type]

367 "hf_config": {

368 "eagle_aux_hidden_state_layer_ids": [1, 2, 3, 4],

369 }

@@ -450,7 +450,7 @@ def test_extract_hidden_states_speculative_config_vlm():

450 target_parallel_config=ParallelConfig(),

451 method="extract_hidden_states",

452 num_speculative_tokens=1,

453 - draft_model_config={

453 + draft_model_config={ # type: ignore[arg-type]

454 "hf_config": {

455 "eagle_aux_hidden_state_layer_ids": [1, 2],

456 }

tests/v1/spec_decode/test_lilicorr.py

+38 -20

@@ -11,6 +11,7 @@ from torch import nn

11

12 from vllm.config import (

13 CompilationConfig,

14 + CompilationMode,

15 DeviceConfig,

16 LoadConfig,

17 VllmConfig,

@@ -122,7 +123,7 @@ def test_lilicorr_matches_exported_head(head_width, slots, dtype):

123 set_current_vllm_config(

124 VllmConfig(

125 device_config=DeviceConfig("cuda"),

125 - compilation_config=CompilationConfig(mode=0),

126 + compilation_config=CompilationConfig(mode=CompilationMode.NONE),

127 )

128 ),

129 ):

@@ -183,23 +184,23 @@ def test_runtime_length_keeps_trained_checkpoint_geometry(monkeypatch):

184 with set_current_vllm_config(

185 VllmConfig(

186 device_config=DeviceConfig("cpu"),

186 - compilation_config=CompilationConfig(mode=0),

187 + compilation_config=CompilationConfig(mode=CompilationMode.NONE),

188 )

189 ):

190 for slots in (1, 3):

191 spec.num_speculative_tokens = slots

192 head = LiLiCorr(

192 - vllm_config=SimpleNamespace(speculative_config=spec)

193 + vllm_config=SimpleNamespace(speculative_config=spec) # type: ignore[arg-type]

194 ).lilicorr

195 assert head.slot_embedding.shape == (1, 1, 3, 1, 8)

196 assert head.relative_slot_bias.shape == (2, 7)

197 spec.draft_model_config.hf_config.dflash_config["lilicorr_enabled"] = False

198 with pytest.raises(ValueError, match="lilicorr_enabled"):

198 - LiLiCorr(vllm_config=SimpleNamespace(speculative_config=spec))

199 + LiLiCorr(vllm_config=SimpleNamespace(speculative_config=spec)) # type: ignore[arg-type]

200 spec.draft_model_config.hf_config.dflash_config["lilicorr_enabled"] = True

201 spec.num_speculative_tokens = 4

202 with pytest.raises(ValueError, match="trained block_size"):

202 - LiLiCorr(vllm_config=SimpleNamespace(speculative_config=spec))

203 + LiLiCorr(vllm_config=SimpleNamespace(speculative_config=spec)) # type: ignore[arg-type]

204

205

206 @pytest.mark.parametrize("tp_size", [1, 2])

@@ -268,7 +269,11 @@ def test_candidates_use_global_partition_and_exclude_padding(monkeypatch, tp_siz

269 )

270 monkeypatch.setattr(logits_processor, "_topk", lambda x, k: x.topk(k, dim=-1))

271 ids, values = logits_processor.LogitsProcessor.get_top_k_tokens(

271 - fake, lm_head, torch.empty(2, 1), 2, return_log_probs=True

272 + fake, # type: ignore[arg-type]

273 + lm_head, # type: ignore[arg-type]

274 + torch.empty(2, 1),

275 + 2,

276 + return_log_probs=True,

277 )

278 expected = torch.cat(valid_shards, -1).log_softmax(-1).topk(2)

279 torch.testing.assert_close(ids, expected.indices)

@@ -286,13 +291,13 @@ def test_context_anchor_is_last_committed_normalized_feature():

291 anchor_valid=torch.ones(4, dtype=torch.bool),

292 )

293 batch = SimpleNamespace(num_reqs=2, query_start_loc=torch.tensor([0, 2, 6]))

289 - LiLiCorrSpeculator.prepare_context_anchor(spec, batch, torch.tensor([0, 2]))

294 + LiLiCorrSpeculator.prepare_context_anchor(spec, batch, torch.tensor([0, 2])) # type: ignore[arg-type]

295 torch.testing.assert_close(spec.anchor_hidden[:2], norm(hidden[[1, 3]]))

296 assert spec.anchor_valid.tolist() == [True, True, False, False]

297 assert not spec.anchor_hidden[2:].any()

298 # A smaller reordered batch must not retain an old request's anchor.

299 batch = SimpleNamespace(num_reqs=1, query_start_loc=torch.tensor([0, 3]))

295 - LiLiCorrSpeculator.prepare_context_anchor(spec, batch, torch.tensor([1]))

300 + LiLiCorrSpeculator.prepare_context_anchor(spec, batch, torch.tensor([1])) # type: ignore[arg-type]

301 torch.testing.assert_close(spec.anchor_hidden[0], norm(hidden[1]))

302 assert not spec.anchor_hidden[1:].any()

303

@@ -316,7 +321,7 @@ def test_candidate_generation_routes_scores_and_adaptive_inputs(monkeypatch, lil

321

322 monkeypatch.setattr(DFlashSpeculator, "__init__", init_base)

323 cls = LiLiCorrSpeculator if lilicorr else DFlash2Speculator

319 - spec = cls(None, torch.device("cpu"))

324 + spec = cls(None, torch.device("cpu")) # type: ignore[arg-type]

325 hidden = torch.arange(36).float().view(9, 4)

326 spec._run_model = lambda *args: hidden

327 spec.sample_indices = torch.tensor([1, 2, 4, 5, 7, 8])

@@ -329,7 +334,7 @@ def test_candidate_generation_routes_scores_and_adaptive_inputs(monkeypatch, lil

334 spec.draft_logits = None

335 spec.use_fp64_gumbel = False

336 spec.enable_adaptive_verification = True

332 - spec.input_buffers = SimpleNamespace(input_ids=torch.arange(9))

337 + spec.input_buffers = SimpleNamespace(input_ids=torch.arange(9)) # type: ignore[assignment]

338 candidates = torch.arange(24).view(6, 4)

339 log_probs = candidates.float().log_softmax(-1)

340 scores = torch.randn(3, 2, 4, 4)

@@ -349,22 +354,35 @@ def test_candidate_generation_routes_scores_and_adaptive_inputs(monkeypatch, lil

354 compute_candidates=lambda h: (candidates, log_probs),

355 model=SimpleNamespace(lilicorr=score, candidate_selector=score),

356 )

352 - spec.target_embeddings = lambda ids: ids.float().unsqueeze(-1).expand(*ids.shape, 4)

357 if lilicorr:

358 + assert isinstance(spec, LiLiCorrSpeculator)

359 + monkeypatch.setattr(

360 + spec,

361 + "target_embeddings",

362 + lambda ids: ids.float().unsqueeze(-1).expand(*ids.shape, 4),

363 + raising=False,

364 + )

365 spec.anchor_hidden.copy_(torch.arange(12).view(3, 4))

366 spec.anchor_valid[:2] = True

356 - spec.candidate_sampler.sample = sample

357 - spec._maybe_predict_acceptance = lambda *args: adaptive_inputs.extend(args)

367 + monkeypatch.setattr(spec.candidate_sampler, "sample", sample)

368 + monkeypatch.setattr(

369 + spec, "_maybe_predict_acceptance", lambda *args: adaptive_inputs.extend(args)

370 + )

371 spec._generate_draft(3, 9, None, None, None)

372

373 expected_ids = candidates.view(3, 2, 4)

361 - expected_first = spec.target_embeddings(expected_ids) if lilicorr else expected_ids

374 + expected_first = (

375 + spec.target_embeddings(expected_ids)

376 + if isinstance(spec, LiLiCorrSpeculator)

377 + else expected_ids

378 + )

379 torch.testing.assert_close(scoring_inputs[0], expected_first)

380 torch.testing.assert_close(scoring_inputs[1], log_probs.view(3, 2, 4))

381 torch.testing.assert_close(

382 scoring_inputs[2], hidden[spec.sample_indices].view(3, 2, 4)

383 )

384 if lilicorr:

385 + assert isinstance(spec, LiLiCorrSpeculator)

386 torch.testing.assert_close(scoring_inputs[3], spec.anchor_hidden)

387 torch.testing.assert_close(scoring_inputs[4], spec.anchor_valid)

388 else:

@@ -472,7 +490,7 @@ def test_checkpoint_coverage_rejects_incomplete_or_wrong_heads(

490 ):

491 from vllm.model_executor.models.lilicorr import LiLiCorr, LiLiCorrForCausalLM

492

475 - wrapper = LiLiCorrForCausalLM.__new__(LiLiCorrForCausalLM)

493 + wrapper = LiLiCorrForCausalLM.__new__(LiLiCorrForCausalLM) # type: ignore[type-abstract]

494 nn.Module.__init__(wrapper)

495 wrapper.model = LiLiCorr.__new__(LiLiCorr)

496 nn.Module.__init__(wrapper.model)

@@ -480,7 +498,7 @@ def test_checkpoint_coverage_rejects_incomplete_or_wrong_heads(

498 with set_current_vllm_config(

499 VllmConfig(

500 device_config=DeviceConfig("cpu"),

483 - compilation_config=CompilationConfig(mode=0),

501 + compilation_config=CompilationConfig(mode=CompilationMode.NONE),

502 )

503 ):

504 wrapper.model.lilicorr = LiLiCorrHead(

@@ -615,10 +633,10 @@ def owned_head_model(monkeypatch):

633 with set_current_vllm_config(

634 VllmConfig(

635 device_config=DeviceConfig("cpu"),

618 - compilation_config=CompilationConfig(mode=0),

636 + compilation_config=CompilationConfig(mode=CompilationMode.NONE),

637 )

638 ):

621 - model = LiLiCorrForCausalLM(vllm_config=vllm_config)

639 + model = LiLiCorrForCausalLM(vllm_config=vllm_config) # type: ignore[arg-type]

640 return model, vllm_config

641

642 return build

@@ -720,7 +738,7 @@ def test_quantized_head_calls_methods_without_reading_packed_weights(monkeypatch

738 with set_current_vllm_config(

739 VllmConfig(

740 device_config=DeviceConfig("cpu"),

723 - compilation_config=CompilationConfig(mode=0),

741 + compilation_config=CompilationConfig(mode=CompilationMode.NONE),

742 )

743 ):

744 head = LiLiCorrHead(

@@ -728,7 +746,7 @@ def test_quantized_head_calls_methods_without_reading_packed_weights(monkeypatch

746 block_size=4,

747 rms_norm_eps=1e-6,

748 config=_config(hidden_size=8),

731 - quant_config=quant_config,

749 + quant_config=quant_config, # type: ignore[arg-type]

750 prefix="model.lilicorr",

751 )

752 with torch.no_grad():

tests/v1/spec_decode/test_llm_base_proposer.py

+7 -4

@@ -16,6 +16,7 @@ from types import SimpleNamespace

16 import pytest

17

18 import vllm.v1.spec_decode.llm_base_proposer as llm_base_proposer

19 +from vllm.v1.kv_cache_interface import KVCacheConfig

20 from vllm.v1.spec_decode.eagle import EagleProposer

21

22 SCHEDULER_BLOCK_SIZE = 256

@@ -52,7 +53,7 @@ def _make_proposer(

53 monkeypatch.setattr(llm_base_proposer, "AttentionGroup", _FakeAttentionGroup)

54

55 proposer = EagleProposer.__new__(EagleProposer)

55 - proposer.vllm_config = None

56 + proposer.vllm_config = None # type: ignore[assignment] # Patched backend discovery does not read config.

57 proposer.device = None

58 proposer._draft_attn_layer_names = set(layer_names)

59 proposer.kv_cache_gid = -1

@@ -61,10 +62,10 @@ def _make_proposer(

62 return proposer

63

64

64 -def _make_kv_cache_config(layer_names: set[str]) -> SimpleNamespace:

65 +def _make_kv_cache_config(layer_names: set[str]) -> KVCacheConfig:

66 spec = SimpleNamespace(block_size=SCHEDULER_BLOCK_SIZE)

67 group = SimpleNamespace(layer_names=list(layer_names), kv_cache_spec=spec)

67 - return SimpleNamespace(kv_cache_groups=[group])

68 + return SimpleNamespace(kv_cache_groups=[group]) # type: ignore[return-value]

69

70

71 def test_block_size_uses_kernel_block_size(monkeypatch: pytest.MonkeyPatch):

@@ -81,7 +82,9 @@ def test_block_size_uses_kernel_block_size(monkeypatch: pytest.MonkeyPatch):

82 assert proposer.block_size == KERNEL_BLOCK_SIZE

83 assert proposer.block_size != SCHEDULER_BLOCK_SIZE

84 # The metadata builder keeps receiving the kernel block size as well.

84 - assert proposer.draft_attn_groups[0].kernel_block_size == KERNEL_BLOCK_SIZE

85 + group = proposer.draft_attn_groups[0]

86 + assert isinstance(group, _FakeAttentionGroup)

87 + assert group.kernel_block_size == KERNEL_BLOCK_SIZE

88

89

90 def test_block_size_falls_back_to_kv_cache_spec(monkeypatch: pytest.MonkeyPatch):

tests/v1/spec_decode/test_llm_base_proposer_sampling.py

+1 -1

@@ -85,7 +85,7 @@ def test_compute_probs_and_sample_next_token_uses_fp64_exponential_race():

85 def test_mtp_model_returns_tuple(architecture: str, expected: bool):

86 proposer = object.__new__(SpecDecodeBaseProposer)

87 proposer.method = "mtp"

88 - proposer.draft_model_config = SimpleNamespace(

88 + proposer.draft_model_config = SimpleNamespace( # type: ignore[assignment] # Architecture-only config.

89 hf_config=SimpleNamespace(architectures=[architecture])

90 )

91

tests/v1/spec_decode/test_mtp.py

-1

@@ -220,7 +220,6 @@ def test_mtp_propose(num_speculative_tokens, monkeypatch):

220 device=device,

221 )

222

223 - proposer.runner = mock.MagicMock()

223 mock_attn_group = mock.MagicMock()

224 mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder

225 mock_attn_group.layer_names = list(proposer._draft_attn_layer_names)

tests/v1/spec_decode/test_mtp_structured_output.py

+7 -1

@@ -2,6 +2,8 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Backend and engine-reasoner checks that are not manager-flow tests."""

4

5 +from typing import Literal

6 +

7 import pytest

8 from transformers import AutoTokenizer

9

@@ -17,7 +19,9 @@ TOKENIZER = "gpt2"

19 NUM_SPEC_TOKENS = 4

20

21

20 -def _make_manager_and_request(backend: str, prompt_str: str = '{"a": "b"}'):

22 +def _make_manager_and_request(

23 + backend: Literal["xgrammar", "guidance"], prompt_str: str = '{"a": "b"}'

24 +):

25 tokenizer = AutoTokenizer.from_pretrained(TOKENIZER)

26 prompt = tokenizer.encode(prompt_str)

27

@@ -33,6 +37,7 @@ def _make_manager_and_request(backend: str, prompt_str: str = '{"a": "b"}'):

37 sampling_params = SamplingParams(

38 structured_outputs=StructuredOutputsParams(json='{"type": "object"}'),

39 )

40 + assert sampling_params.structured_outputs is not None

41 sampling_params.structured_outputs._backend = backend

42 sampling_params.update_from_generation_config({}, tokenizer.eos_token_id)

43

@@ -43,6 +48,7 @@ def _make_manager_and_request(backend: str, prompt_str: str = '{"a": "b"}'):

48 pooling_params=None,

49 )

50 manager.grammar_init(request)

51 + assert request.structured_output_request is not None

52 while not request.structured_output_request._check_grammar_completion():

53 continue

54

tests/v1/spec_decode/test_rejection_sampler_utils.py

+4 -6

@@ -168,7 +168,7 @@ def test_watermarked_bonus_uses_target_key():

168 watermarker=GumbelWatermarker(key=42, context_width=2),

169 )

170 expected = GumbelWatermarker(key=42, context_width=2).sample(

171 - target_logits.unsqueeze(0), contexts[1:], lambda _: None

171 + target_logits.unsqueeze(0), contexts[1:]

172 )

173

174 assert num_sampled.item() == 2

@@ -196,7 +196,7 @@ def test_watermarked_recovery_uses_target_key():

196 watermarker=GumbelWatermarker(key=42, context_width=2),

197 )

198 expected = GumbelWatermarker(key=42, context_width=2).sample(

199 - target_logits.unsqueeze(0), contexts[:1], lambda _: None

199 + target_logits.unsqueeze(0), contexts[:1]

200 )

201

202 assert num_sampled.item() == 1

@@ -225,9 +225,7 @@ def test_watermarked_recovery_supports_smaller_draft_vocabulary():

225 watermarking=torch.tensor([True], device="cuda"),

226 watermarker=watermarker,

227 )

228 - expected = watermarker.sample(

229 - target_logits[:4].unsqueeze(0), contexts[:1], lambda _: None

230 - )

228 + expected = watermarker.sample(target_logits[:4].unsqueeze(0), contexts[:1])

229

230 assert num_sampled.item() == 1

231 assert sampled[0, 0] == expected.token_ids[0]

@@ -702,7 +700,7 @@ def test_verify_rejects_unproposed_drafts():

700 first_pos = inputs["pos"].view(num_trials, K + 1)[:, 0]

701 prefill_len = first_pos + 1

702 rejection_sampler = object.__new__(RejectionSampler)

705 - rejection_sampler.sampler = SimpleNamespace(

703 + rejection_sampler.sampler = SimpleNamespace( # type: ignore[assignment]

704 apply_sampling_params=lambda logits, *args: logits,

705 sampling_states=SimpleNamespace(

706 temperature=SimpleNamespace(gpu=inputs["temperature"]),

tools/pre_commit/mypy.py

+12 -12

@@ -28,27 +28,36 @@ import regex as re

28 # Fixing a directory means moving it from SEPARATE_GROUPS to here. Without that

29 # move the fixes are not enforced because "tests" claims every file below it.

30 SILENT_GROUPS = [

31 + "tests/benchmarks",

32 "tests/compile/correctness_e2e",

33 "tests/compile/fullgraph",

34 "tests/compile/fusions_e2e",

35 "tests/compile/passes",

36 "tests/config",

37 + "tests/entrypoints/anthropic",

38 "tests/entrypoints/generate",

39 + "tests/entrypoints/multimodal",

40 + "tests/entrypoints/pooling",

41 "tests/entrypoints/serve",

42 "tests/entrypoints/tool_parsers",

43 "tests/entrypoints/unit_tests",

44 "tests/entrypoints/weight_transfer",

45 + "tests/kernels/attention",

46 "tests/kernels/core",

47 + "tests/kernels/helion",

48 "tests/kernels/mamba",

49 "tests/kernels/quantization",

50 "tests/model_executor/layers",

51 "tests/models/language",

52 "tests/models/quantization",

53 + "tests/multimodal",

54 "tests/plugins/bge_m3_sparse_plugin",

55 "tests/plugins/prithvi_io_processor_plugin",

56 "tests/plugins/vllm_add_dummy_platform",

57 "tests/plugins/vllm_add_dummy_stat_logger",

58 "tests/plugins_tests",

59 + "tests/quantization",

60 + "tests/reasoning",

61 "tests/samplers",

62 "tests/spec_decode",

63 "tests/tokenizers_",

@@ -59,8 +68,11 @@ SILENT_GROUPS = [

68 "tests/v1/determinism",

69 "tests/v1/distributed",

70 "tests/v1/executor",

71 + "tests/v1/logits_processors",

72 + "tests/v1/sample",

73 "tests/v1/shutdown",

74 "tests/v1/simple_kv_offload",

75 + "tests/v1/spec_decode",

76 "tests/v1/structured_output",

77 ]

78

@@ -68,18 +80,12 @@ SILENT_GROUPS = [

80 # from "skip" to "silent", move its directory to SILENT_GROUPS.

81 SEPARATE_GROUPS = [

82 "tests",

71 - "tests/benchmarks",

83 "tests/compile",

84 "tests/distributed",

74 - "tests/entrypoints/anthropic",

85 "tests/entrypoints/llm",

76 - "tests/entrypoints/multimodal",

86 "tests/entrypoints/openai",

78 - "tests/entrypoints/pooling",

87 "tests/entrypoints/speech_to_text",

88 "tests/kernels",

81 - "tests/kernels/attention",

82 - "tests/kernels/helion",

89 "tests/kernels/moe",

90 "tests/lora",

91 "tests/model_executor",

@@ -87,10 +93,7 @@ SEPARATE_GROUPS = [

93 "tests/models",

94 "tests/models/test_initialization.py",

95 "tests/models/multimodal",

90 - "tests/multimodal",

96 "tests/parser",

92 - "tests/quantization",

93 - "tests/reasoning",

97 "tests/renderers",

98 "tests/tool_parsers",

99 "tests/v1",

@@ -101,10 +104,7 @@ SEPARATE_GROUPS = [

104 "tests/v1/engine",

105 "tests/v1/kv_connector",

106 "tests/v1/kv_offload",

104 - "tests/v1/logits_processors",

107 "tests/v1/metrics",

106 - "tests/v1/sample",

107 - "tests/v1/spec_decode",

108 "tests/v1/streaming_input",

109 "tests/v1/worker",

110 ]

[MyPy][3/N] Fix MyPy errors in test groups (part 3) (#55939)