5dd693122c · tduncan/Nvidia-Optimization-Stack-VLLM

tests/compile/passes/distributed/test_async_tp.py

+1 -1

@@ -285,7 +285,7 @@ def test_async_tp_pass_requires_full_graph_compilation():

285 ]

286

287 async_tp_pass = object.__new__(AsyncTPPass)

288 - async_tp_pass.compilation_config = vllm_config.compilation_config

288 + async_tp_pass.compilation_config = vllm_config.compilation_config # type: ignore[assignment]

289

290 with pytest.raises(

291 AssertionError, match="AsyncTPPass requires full-graph compilation"

tests/compile/passes/distributed/test_fusion_all_reduce.py

+2 -2

@@ -111,7 +111,7 @@ def test_fused_ar_workspace_hidden_dim(

111 target_hidden: int, draft_hidden: int | None, expected: int

112 ):

113 config = _make_workspace_hidden_config(target_hidden, draft_hidden)

114 - assert _fused_ar_workspace_hidden_dim(config) == expected

114 + assert _fused_ar_workspace_hidden_dim(config) == expected # type: ignore[arg-type]

115

116

117 def test_fused_ar_workspace_hidden_dim_spec_without_draft():

@@ -120,7 +120,7 @@ def test_fused_ar_workspace_hidden_dim_spec_without_draft():

120 model_config=SimpleNamespace(get_hidden_size=lambda: 2048),

121 speculative_config=SimpleNamespace(draft_model_config=None),

122 )

123 - assert _fused_ar_workspace_hidden_dim(config) == 2048

123 + assert _fused_ar_workspace_hidden_dim(config) == 2048 # type: ignore[arg-type]

124

125

126 class TestAllReduceRMSNormModel(torch.nn.Module):

tests/compile/passes/distributed/test_sequence_parallelism.py

+1 -1

@@ -222,7 +222,7 @@ def test_sequence_parallelism_pass_requires_full_graph_compilation():

222 ]

223

224 sequence_parallelism_pass = object.__new__(SequenceParallelismPass)

225 - sequence_parallelism_pass.compilation_config = vllm_config.compilation_config

225 + sequence_parallelism_pass.compilation_config = vllm_config.compilation_config # type: ignore[assignment]

226 sequence_parallelism_pass.min_token_num = 1

227

228 with pytest.raises(

tests/compile/passes/test_fusion_attn.py

+3 -1

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

33 VllmConfig,

34 set_current_vllm_config,

35 )

36 +from vllm.config.cache import CacheDType

37 from vllm.forward_context import get_forward_context, set_forward_context

38 from vllm.model_executor.layers.attention import Attention

39 from vllm.model_executor.layers.quantization.utils.quant_utils import (

@@ -198,6 +199,7 @@ class TestAttentionAiterFp8StaticQuantPatternModel(

199 def __init__(self, *args, **kwargs):

200 super().__init__(*args, **kwargs)

201 self.fp8_linear.kernel.quant_fp8.use_aiter = True

202 + assert self.fp8_linear.input_scale is not None

203 self.fp8_linear.input_scale.fill_(0.125)

204

205

@@ -307,7 +309,7 @@ def test_attention_quant_pattern(

309 head_size: int,

310 batch_size: int,

311 dtype: torch.dtype,

310 - kv_cache_dtype: str,

312 + kv_cache_dtype: CacheDType,

313 custom_ops: str,

314 model_name: str,

315 model_class: type[AttentionQuantPatternModel],

tests/compile/passes/test_mla_rope_kvcache_cat_fusion.py

+2 -1

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

24 PassConfig,

25 VllmConfig,

26 )

27 +from vllm.config.cache import CacheDType

28 from vllm.forward_context import get_forward_context, set_forward_context

29 from vllm.model_executor.layers.attention import MLAAttention

30 from vllm.model_executor.layers.linear import ColumnParallelLinear

@@ -254,7 +255,7 @@ def test_mla_rope_kvcache_cat_fusion(

255 block_size: int,

256 is_neox: bool,

257 dtype: torch.dtype,

257 - kv_cache_dtype: str,

258 + kv_cache_dtype: CacheDType,

259 monkeypatch: pytest.MonkeyPatch,

260 ):

261 torch.set_default_device("cuda")

tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py

+4 -3

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

39 PassConfig,

40 VllmConfig,

41 )

42 +from vllm.config.cache import CacheDType

43 from vllm.forward_context import get_forward_context, set_forward_context

44 from vllm.model_executor.layers.attention import Attention

45 from vllm.model_executor.layers.layernorm import RMSNorm

@@ -342,7 +343,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(

343 use_shuffle_kv_layout: str,

344 kv_layout: KVCacheLayout,

345 dtype: torch.dtype,

345 - kv_cache_dtype: str,

346 + kv_cache_dtype: CacheDType,

347 rms_norm_eps: float,

348 custom_op: str,

349 monkeypatch: pytest.MonkeyPatch,

@@ -611,7 +612,7 @@ def test_qk_norm_rope_kvcache_fusion(

612 kv_layout: KVCacheLayout,

613 block_size: int,

614 dtype: torch.dtype,

614 - kv_cache_dtype: str,

615 + kv_cache_dtype: CacheDType,

616 rms_norm_eps: float,

617 custom_op: str,

618 monkeypatch: pytest.MonkeyPatch,

@@ -663,7 +664,7 @@ def test_qk_norm_mrope_kvcache_fusion(

664 mrope_interleaved: bool,

665 is_neox: bool,

666 kv_layout: KVCacheLayout,

666 - kv_cache_dtype: str,

667 + kv_cache_dtype: CacheDType,

668 monkeypatch: pytest.MonkeyPatch,

669 ):

670 _run_qk_norm_rope_kvcache_fusion_test(

tests/compile/passes/test_rope_kvcache_fusion.py

+3 -2

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

30 PassConfig,

31 VllmConfig,

32 )

33 +from vllm.config.cache import CacheDType

34 from vllm.config.utils import Range

35 from vllm.forward_context import get_forward_context, set_forward_context

36 from vllm.model_executor.layers.attention import Attention

@@ -324,7 +325,7 @@ def test_rope_kvcache_fusion(

325 block_size: int,

326 is_neox: bool,

327 dtype: torch.dtype,

327 - kv_cache_dtype: str,

328 + kv_cache_dtype: CacheDType,

329 monkeypatch: pytest.MonkeyPatch,

330 ):

331 torch.set_default_device("cuda")

@@ -468,7 +469,7 @@ def test_rope_static_qquant_kvcache_fusion(

469 block_size: int,

470 is_neox: bool,

471 dtype: torch.dtype,

471 - kv_cache_dtype: str,

472 + kv_cache_dtype: CacheDType,

473 monkeypatch: pytest.MonkeyPatch,

474 ):

475 torch.set_default_device("cuda")

tests/compile/passes/test_silu_mul_quant_fusion.py

+6 -3

@@ -60,7 +60,7 @@ class TestSiluMulFp8QuantModel(torch.nn.Module):

60 def __init__(

61 self,

62 hidden_size: int,

63 - force_kernel: FP8ScaledMMLinearKernel,

63 + force_kernel: type[FP8ScaledMMLinearKernel],

64 dtype: torch.dtype,

65 **kwargs,

66 ):

@@ -282,7 +282,7 @@ def test_fusion_silu_and_mul_quant(

282 ],

283 enable_silu_mul_custom_op: bool,

284 enable_quant_fp8_custom_op: bool,

285 - force_kernel: FP8ScaledMMLinearKernel | None,

285 + force_kernel: type[FP8ScaledMMLinearKernel] | None,

286 monkeypatch: pytest.MonkeyPatch,

287 ):

288 if model_class is TestSiluMulNvfp4QuantModel and not is_nvfp4_supported():

@@ -330,7 +330,10 @@ def test_fusion_silu_and_mul_quant(

330 passes = [NoOpEliminationPass(config), *fusion_passes, PostCleanupPass(config)]

331 backend = TestBackend(*passes)

332 model = model_class(

333 - hidden_size=hidden_size, force_kernel=force_kernel, x=x, dtype=dtype

333 + hidden_size=hidden_size,

334 + force_kernel=force_kernel, # type: ignore[arg-type]

335 + x=x,

336 + dtype=dtype,

337 )

338

339 # First dimension dynamic

tests/compile/passes/test_silu_mul_quant_manual_fusion.py

+1 -1

@@ -84,7 +84,7 @@ def test_manual_fusion_fp8_static_with_linear(

84 num_tokens: int,

85 hidden_size: int,

86 dtype: torch.dtype,

87 - force_kernel: FP8ScaledMMLinearKernel,

87 + force_kernel: type[FP8ScaledMMLinearKernel],

88 ):

89 """Test manual fusion with real FP8 linear layer (kFp8StaticTensorSym).

90

tests/entrypoints/launchers/api_server/test_api_server_process_manager.py

+4 -3

@@ -307,7 +307,8 @@ def test_external_process_monitoring(api_server_args):

307 def run_with_exception_capture():

308 try:

309 wait_for_completion_or_failure(

310 - api_server_manager=manager, coordinator=mock_coordinator

310 + api_server_manager=manager,

311 + coordinator=mock_coordinator, # type: ignore[arg-type]

312 )

313 except Exception as e:

314 result["exception"] = e

@@ -376,7 +377,7 @@ def test_gather_actual_addresses_end_to_end():

377 manager = APIServerProcessManager(

378 listen_address=f"tcp://{host}:0",

379 sock=sock,

379 - args="test_args",

380 + args="test_args", # type: ignore[arg-type]

381 num_servers=num_servers,

382 input_addresses=placeholder_inputs,

383 output_addresses=placeholder_outputs,

@@ -423,7 +424,7 @@ def test_gather_actual_addresses_child_crash_before_report():

424 manager = APIServerProcessManager(

425 listen_address=f"tcp://{host}:0",

426 sock=sock,

426 - args="test_args",

427 + args="test_args", # type: ignore[arg-type]

428 num_servers=num_servers,

429 input_addresses=placeholder_inputs,

430 output_addresses=placeholder_outputs,

tests/entrypoints/serve/lora/test_serving_models.py

+2 -2

@@ -108,8 +108,8 @@ async def test_unload_lora_adapter_success():

108 response = await serving_models.load_lora_adapter(request)

109 assert len(serving_models.lora_requests) == 1

110

111 - request = UnloadLoRAAdapterRequest(lora_name="adapter1")

112 - response = await serving_models.unload_lora_adapter(request)

111 + unload_request = UnloadLoRAAdapterRequest(lora_name="adapter1")

112 + response = await serving_models.unload_lora_adapter(unload_request)

113 assert response == LORA_UNLOADING_SUCCESS_MESSAGE.format(lora_name="adapter1")

114 assert len(serving_models.lora_requests) == 0

115

tests/entrypoints/serve/middleware/test_authentication_middleware.py

+4 -1

@@ -84,7 +84,10 @@ def task_routes(request, monkeypatch) -> tuple[str, list[tuple[str, list[str]]]]

84

85 # Register routers for this specific task (development mode already enabled).

86 register_api_routers(

87 - args, app, supported_tasks=(task,), model_config=MockModelConfig()

87 + args,

88 + app,

89 + supported_tasks=(task,),

90 + model_config=MockModelConfig(), # type: ignore[arg-type]

91 )

92

93 routes = get_all_http_routes(app)

tests/entrypoints/serve/tokenize/test_serving_tokenization.py

+5 -4

@@ -19,6 +19,7 @@ from vllm.entrypoints.serve.tokenize.api_router import attach_router

19 from vllm.entrypoints.serve.tokenize.protocol import (

20 TokenizeChatRequest,

21 TokenizeCompletionRequest,

22 + TokenizeResponse,

23 )

24 from vllm.entrypoints.serve.tokenize.serving import ServingTokenization

25 from vllm.exceptions import VLLMNotFoundError, VLLMValidationError

@@ -95,7 +96,7 @@ async def test_tokenize_chat_skips_mm_cache_for_renderer_only_path():

96 mock_engine.renderer = MagicMock()

97

98 serving = _build_serving_tokenization(mock_engine)

98 - serving.online_renderer.preprocess_chat = AsyncMock(

99 + serving.online_renderer.preprocess_chat = AsyncMock( # type: ignore[method-assign]

100 return_value=(

101 [{"role": "user", "content": "Test"}],

102 [{"prompt_token_ids": [1, 2, 3]}],

@@ -109,7 +110,7 @@ async def test_tokenize_chat_skips_mm_cache_for_renderer_only_path():

110

111 response = await serving.create_tokenize(request, MagicMock(headers={}))

112

112 - assert response.tokens == [1, 2, 3]

113 + assert isinstance(response, TokenizeResponse) and response.tokens == [1, 2, 3]

114 assert (

115 serving.online_renderer.preprocess_chat.call_args.kwargs["skip_mm_cache"]

116 is True

@@ -125,7 +126,7 @@ async def test_tokenize_completion_skips_mm_cache_for_renderer_only_path():

126 mock_engine.renderer = MagicMock()

127

128 serving = _build_serving_tokenization(mock_engine)

128 - serving.online_renderer.preprocess_completion = AsyncMock(

129 + serving.online_renderer.preprocess_completion = AsyncMock( # type: ignore[method-assign]

130 return_value=[{"prompt_token_ids": [1, 2, 3]}]

131 )

132

@@ -136,7 +137,7 @@ async def test_tokenize_completion_skips_mm_cache_for_renderer_only_path():

137

138 response = await serving.create_tokenize(request, MagicMock(headers={}))

139

139 - assert response.tokens == [1, 2, 3]

140 + assert isinstance(response, TokenizeResponse) and response.tokens == [1, 2, 3]

141 assert (

142 serving.online_renderer.preprocess_completion.call_args.kwargs["skip_mm_cache"]

143 is True

tests/entrypoints/serve/tokenize/test_tokenization.py

+4 -4

@@ -101,13 +101,13 @@ async def test_tokenize_chat(

101 if continue_final:

102 conversation.append({"role": "assistant", "content": "Sure,"})

103

104 - prompt = tokenizer.apply_chat_template(

104 + prompt = tokenizer.apply_chat_template( # type: ignore[call-arg]

105 add_generation_prompt=add_generation,

106 continue_final_message=continue_final,

107 conversation=conversation,

108 tokenize=False,

109 )

110 - tokens = tokenizer.encode(prompt, add_special_tokens=add_special)

110 + tokens = tokenizer.encode(prompt, add_special_tokens=add_special) # type: ignore[arg-type]

111

112 response = requests.post(

113 server.url_for("tokenize"),

@@ -169,14 +169,14 @@ async def test_tokenize_chat_with_tools(

169 if continue_final:

170 conversation.append({"role": "assistant", "content": "Sure,"})

171

172 - prompt = tokenizer.apply_chat_template(

172 + prompt = tokenizer.apply_chat_template( # type: ignore[call-arg]

173 add_generation_prompt=add_generation,

174 continue_final_message=continue_final,

175 conversation=conversation,

176 tools=tools,

177 tokenize=False,

178 )

179 - tokens = tokenizer.encode(prompt, add_special_tokens=add_special)

179 + tokens = tokenizer.encode(prompt, add_special_tokens=add_special) # type: ignore[arg-type]

180

181 response = requests.post(

182 server.url_for("tokenize"),

tests/entrypoints/serve/utils/test_fingerprint.py

+6 -3

@@ -18,7 +18,7 @@ def _cfg(tp=1, pp=1, dp=1, ep=False, digest="a3b21f94deadbeef"):

18 enable_expert_parallel=ep,

19 )

20 )

21 - c.compute_hash = lambda: digest # type: ignore[attr-defined]

21 + c.compute_hash = lambda: digest

22 return c

23

24

@@ -59,6 +59,7 @@ def test_get_respects_set_default():

59

60 fp.set_default_fingerprint_mode("hash")

61 hashed = fp.get_system_fingerprint(cfg)

62 + assert hashed is not None

63 assert hashed != full

64 assert "tp8" not in hashed

65

@@ -72,5 +73,7 @@ def test_get_respects_set_default():

73 def test_compute_hash_failure_does_not_raise():

74 cfg = _cfg()

75 cfg.compute_hash = lambda: (_ for _ in ()).throw(RuntimeError("boom"))

75 - assert fp.build_system_fingerprint(cfg, "full").endswith("-nohash")

76 - assert fp.build_system_fingerprint(cfg, "hash").endswith("-nohash")

76 + full_fp = fp.build_system_fingerprint(cfg, "full")

77 + assert full_fp is not None and full_fp.endswith("-nohash")

78 + hash_fp = fp.build_system_fingerprint(cfg, "hash")

79 + assert hash_fp is not None and hash_fp.endswith("-nohash")

tests/entrypoints/unit_tests/test_chat_utils.py

+18 -16

@@ -17,6 +17,7 @@ from vllm.config import ModelConfig

17 from vllm.entrypoints.chat_utils import (

18 MEDIA_CONNECTOR_REGISTRY,

19 AsyncMultiModalItemTracker,

20 + ChatCompletionMessageParam,

21 ConversationMessage,

22 _load_embeds_dict,

23 _parse_metadata_array,

@@ -215,7 +216,7 @@ def video_url():

216 @pytest.fixture(scope="module")

217 def audio_url():

218 audio = AudioAsset("mary_had_lamb")

218 - return encode_audio_url(*audio.audio_and_sample_rate)

219 + return encode_audio_url(*audio.audio_and_sample_rate) # type: ignore[arg-type]

220

221

222 def _assert_mm_data_is_image_input(

@@ -769,7 +770,9 @@ async def test_text_only_chat_does_not_initialize_media_connector(

770 ):

771 load_connector = MagicMock()

772 monkeypatch.setattr(MEDIA_CONNECTOR_REGISTRY, "load", load_connector)

772 - messages = [{"role": "user", "content": "Who are you?"}]

773 + messages: list[ChatCompletionMessageParam] = [

774 + {"role": "user", "content": "Who are you?"}

775 + ]

776

777 parse_chat_messages(

778 messages,

@@ -2457,7 +2460,7 @@ def test_parse_chat_messages_image_vision_chunk(

2460 kimi_k2_5_model_config,

2461 image_url,

2462 ):

2460 - messages = [

2463 + messages: list[ChatCompletionMessageParam] = [

2464 {

2465 "role": "user",

2466 "content": [

@@ -2493,7 +2496,7 @@ def test_parse_chat_messages_video_vision_chunk(

2496 kimi_k2_5_model_config,

2497 video_url,

2498 ):

2496 - messages = [

2499 + messages: list[ChatCompletionMessageParam] = [

2500 {

2501 "role": "user",

2502 "content": [

@@ -2530,7 +2533,7 @@ def test_parse_chat_messages_image_vision_chunk_with_uuid(

2533 image_url,

2534 ):

2535 image_uuid = "image_123"

2533 - messages = [

2536 + messages: list[ChatCompletionMessageParam] = [

2537 {

2538 "role": "user",

2539 "content": [

@@ -2568,7 +2571,7 @@ def test_parse_chat_messages_video_vision_chunk_with_uuid(

2571 video_url,

2572 ):

2573 video_uuid = "video_456"

2571 - messages = [

2574 + messages: list[ChatCompletionMessageParam] = [

2575 {

2576 "role": "user",

2577 "content": [

@@ -2606,7 +2609,7 @@ def test_parse_chat_messages_mixed_vision_chunk(

2609 image_url,

2610 video_url,

2611 ):

2609 - messages = [

2612 + messages: list[ChatCompletionMessageParam] = [

2613 {

2614 "role": "user",

2615 "content": [

@@ -2655,7 +2658,7 @@ def test_parse_chat_messages_mixed_vision_chunk_with_uuid(

2658 ):

2659 image_uuid = "image_123"

2660 video_uuid = "video_456"

2658 - messages = [

2661 + messages: list[ChatCompletionMessageParam] = [

2662 {

2663 "role": "user",

2664 "content": [

@@ -2707,7 +2710,7 @@ async def test_parse_chat_messages_mixed_vision_chunk_async(

2710 image_url,

2711 video_url,

2712 ):

2710 - messages = [

2713 + messages: list[ChatCompletionMessageParam] = [

2714 {

2715 "role": "user",

2716 "content": [

@@ -2757,7 +2760,7 @@ async def test_parse_chat_messages_mixed_vision_chunk_with_uuid_async(

2760 ):

2761 image_uuid = "image_123"

2762 video_uuid = "video_456"

2760 - messages = [

2763 + messages: list[ChatCompletionMessageParam] = [

2764 {

2765 "role": "user",

2766 "content": [

@@ -2808,7 +2811,7 @@ async def test_parse_chat_messages_image_vision_chunk_async(

2811 kimi_k2_5_model_config,

2812 image_url,

2813 ):

2811 - messages = [

2814 + messages: list[ChatCompletionMessageParam] = [

2815 {

2816 "role": "user",

2817 "content": [

@@ -2845,7 +2848,7 @@ async def test_parse_chat_messages_video_vision_chunk_async(

2848 kimi_k2_5_model_config,

2849 video_url,

2850 ):

2848 - messages = [

2851 + messages: list[ChatCompletionMessageParam] = [

2852 {

2853 "role": "user",

2854 "content": [

@@ -2883,7 +2886,7 @@ async def test_parse_chat_messages_image_vision_chunk_with_uuid_async(

2886 image_url,

2887 ):

2888 image_uuid = "image_123"

2886 - messages = [

2889 + messages: list[ChatCompletionMessageParam] = [

2890 {

2891 "role": "user",

2892 "content": [

@@ -2922,7 +2925,7 @@ async def test_parse_chat_messages_video_vision_chunk_with_uuid_async(

2925 video_url,

2926 ):

2927 video_uuid = "video_456"

2925 - messages = [

2928 + messages: list[ChatCompletionMessageParam] = [

2929 {

2930 "role": "user",

2931 "content": [

@@ -2997,8 +3000,7 @@ async def test_resolve_items_runs_modalities_concurrently_and_preserves_order():

3000 finally:

3001 active_fetches -= 1

3002

3000 - tracker = AsyncMultiModalItemTracker(MagicMock())

3001 - tracker._model_config.is_multimodal_model = True

3003 + tracker = AsyncMultiModalItemTracker(MagicMock(is_multimodal_model=True))

3004 tracker.__dict__["mm_processor"] = MagicMock()

3005 tracker._items_by_modality["video"] = [

3006 lambda: _fetch("video-0", 0.02),

tests/kernels/quantization/test_block_fp8.py

+3 -2

@@ -37,7 +37,8 @@ from vllm.utils.flashinfer import (

37 )

38 from vllm.utils.import_utils import has_deep_gemm

39

40 -if current_platform.get_device_capability() < (9, 0):

40 +capability = current_platform.get_device_capability()

41 +if capability is None or capability < (9, 0):

42 pytest.skip("FP8 Triton requires CUDA 9.0 or higher", allow_module_level=True)

43

44 vllm_config = VllmConfig()

@@ -258,7 +259,7 @@ def test_w8a8_block_fp8_torch_scaled_mm_matmul():

259 stub = BlockWiseTorchFP8ScaledMMLinearKernel.__new__(

260 BlockWiseTorchFP8ScaledMMLinearKernel

261 )

261 - stub.config = types.SimpleNamespace(out_dtype=out_dtype)

262 + stub.config = types.SimpleNamespace(out_dtype=out_dtype) # type: ignore[assignment]

263 out = stub.apply_block_scaled_mm(

264 A_fp8_cuda.cuda(), B_fp8.cuda(), As_cuda.cuda(), Bs.cuda()

265 )

tests/kernels/quantization/test_cutlass_scaled_mm.py

+3 -2

@@ -62,7 +62,8 @@ PER_TOKEN_GROUP_SHAPE = (1, -1)

62 PER_OUT_CH_GROUP_SHAPE = (-1, 1)

63

64 capability = current_platform.get_device_capability()

65 -capability = capability[0] * 10 + capability[1]

65 +# Fall back to 0 so import doesn't fail on unknown capability

66 +capability_int = (capability[0] * 10 + capability[1]) if capability is not None else 0

67

68

69 def rand_int8(shape: tuple, device: str = "cuda"):

@@ -639,7 +640,7 @@ def test_cutlass_cuda_graph(per_act_token: bool, per_out_ch: bool):

640

641

642 def test_cutlass_support_opcheck():

642 - opcheck(torch.ops._C.cutlass_scaled_mm_supports_fp8, (capability,))

643 + opcheck(torch.ops._C.cutlass_scaled_mm_supports_fp8, (capability_int,))

644

645

646 @pytest.mark.parametrize("num_experts", [8, 64])

tests/kernels/quantization/test_cutlass_w4a8.py

+3 -2

@@ -29,7 +29,8 @@ if not current_platform.is_cuda():

29 # `is_quant_method_supported` conflates kernels with quantization methods

30 # an assumption which is breaking down as quantizations methods can have

31 # have kernels and some kernels support multiple quantization methods.

32 -IS_SUPPORTED_BY_GPU = current_platform.get_device_capability()[0] >= 9

32 +capability = current_platform.get_device_capability()

33 +IS_SUPPORTED_BY_GPU = capability is not None and capability[0] >= 9

34

35 MNK_SHAPES = [

36 (1, 128, 128),

@@ -207,7 +208,7 @@ def mm_test_helper(

208 a=tensors.a,

209 b_q=tensors.w_q,

210 b_group_scales=tensors.w_g_s,

210 - b_group_size=group_size,

211 + b_group_size=group_size, # type: ignore[arg-type]

212 b_channel_scales=tensors.w_ch_s,

213 a_token_scales=tensors.w_tok_s,

214 )

tests/kernels/quantization/test_cutlass_w4a8_moe.py

+2 -1

@@ -17,8 +17,9 @@ from vllm.platforms import current_platform

17 from vllm.scalar_type import ScalarType, scalar_types

18 from vllm.utils.torch_utils import set_random_seed

19

20 +capability = current_platform.get_device_capability()

21 IS_SUPPORTED_BY_GPU = (

21 - current_platform.is_cuda() and current_platform.get_device_capability()[0] >= 9

22 + current_platform.is_cuda() and capability is not None and capability[0] >= 9

23 )

24

25

tests/kernels/quantization/test_fp8_quant.py

+6 -2

@@ -12,6 +12,7 @@ from tests.kernels.quant_utils import (

12 )

13 from tests.kernels.utils import opcheck

14 from vllm.model_executor.layers.quantization.utils.quant_utils import (

15 + GroupShape,

16 scaled_quantize,

17 )

18 from vllm.platforms import current_platform

@@ -208,7 +209,10 @@ def test_static_fp8_quant_group_2d(

209

210 x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda")

211 ref_out, scale = scaled_quantize(

211 - x, group_shape, current_platform.fp8_dtype(), compute_dtype=torch.float32

212 + x,

213 + GroupShape(*group_shape),

214 + current_platform.fp8_dtype(),

215 + compute_dtype=torch.float32,

216 )

217 ops_out, ops_scale = ops.scaled_fp8_quant(x, scale=scale, group_shape=group_shape)

218

@@ -236,7 +240,7 @@ def test_static_fp8_quant_1d_scale(

240

241 x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda")

242 ref_out, scale_2d = scaled_quantize(

239 - x, group_shape, FP8_DTYPE, compute_dtype=torch.float32

243 + x, GroupShape(*group_shape), FP8_DTYPE, compute_dtype=torch.float32

244 )

245

246 # Flatten scale to 1D for testing 1D scale path

tests/kernels/quantization/test_machete_mm.py

+4 -3

@@ -38,7 +38,8 @@ CUDA_DEVICES = [

38 # `is_quant_method_supported` conflates kernels with quantization methods

39 # an assumption which is breaking down as quantizations methods can have

40 # have kernels and some kernels support multiple quantization methods.

41 -IS_SUPPORTED_BY_GPU = current_platform.get_device_capability()[0] >= 9

41 +capability = current_platform.get_device_capability()

42 +IS_SUPPORTED_BY_GPU = capability is not None and capability[0] >= 9

43

44 MNK_SHAPES = [

45 (1, 128, 128),

@@ -305,7 +306,7 @@ def test_machete_all_schedules(shape, types: TypeConfig):

306 if types.group_scale_type is None:

307 group_sizes = [None]

308 else:

308 - group_sizes = query_machete_supported_group_sizes(types.act_type)

309 + group_sizes = list(query_machete_supported_group_sizes(types.act_type))

310

311 for group_size in group_sizes:

312 if not group_size_valid(shape, group_size):

@@ -334,7 +335,7 @@ def test_machete_heuristic(shape, types: TypeConfig):

335 if types.group_scale_type is None:

336 group_sizes = [None]

337 else:

337 - group_sizes = query_machete_supported_group_sizes(types.act_type)

338 + group_sizes = list(query_machete_supported_group_sizes(types.act_type))

339

340 for group_size in group_sizes:

341 if not group_size_valid(shape, group_size):

tests/kernels/quantization/test_marlin_gemm.py

+1 -1

@@ -342,7 +342,7 @@ def marlin_generate_valid_test_cases():

342 [quant_test_config["b_type"]],

343 quant_test_config.get("c_type", f16_types),

344 quant_test_config["group_blocks"],

345 - )

345 + ) # type: ignore[call-overload]

346

347 for sub_case in inner_combinations:

348 if (

tests/kernels/quantization/test_marlin_tile_padding.py

+5 -2

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

11 import torch

12

13 from vllm import _custom_ops as ops

14 +from vllm.model_executor.layers.linear import LinearBase

15 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

16 GPTQ_MARLIN_TILE,

17 apply_gptq_marlin_linear,

@@ -524,13 +525,15 @@ def test_check_marlin_supports_layer_allow_tile_padding():

525 )

526

527 # Tile-misaligned but group-aligned: rejected strictly, allowed w/ padding

527 - layer = _FakeLinear(4640, 512, input_size=2048)

528 + layer: LinearBase = _FakeLinear( # type: ignore[assignment]

529 + 4640, 512, input_size=2048

530 + )

531 assert not check_marlin_supports_layer(layer, 128)

532 assert check_marlin_supports_layer(layer, 128, allow_tile_padding=True)

533 assert check_marlin_supports_layer(layer, -1, allow_tile_padding=True)

534

535 # A group straddling the TP shard cannot be fixed by padding

533 - layer = _FakeLinear(4608, 4672, input_size=18688)

536 + layer = _FakeLinear(4608, 4672, input_size=18688) # type: ignore[assignment]

537 assert not check_marlin_supports_layer(layer, 128, allow_tile_padding=True)

538

539

tests/kernels/quantization/test_nvfp4_emulation.py

+1 -1

@@ -206,7 +206,7 @@ def test_nvfp4_emulation_support_check_rejects_bias_and_lora(

206 in_dtype=torch.bfloat16,

207 device="cuda",

208 routing_method=RoutingMethodType.TopK,

209 - **config_kwargs,

209 + **config_kwargs, # type: ignore[arg-type]

210 )

211

212 supported, reason = Nvfp4QuantizationEmulationTritonExperts.is_supported_config(

tests/kernels/quantization/test_scaled_mm_kernel_selection.py

+1 -1

@@ -35,7 +35,7 @@ pytestmark = pytest.mark.cpu_test

35

36

37 def test_auto_linear_backend_is_unchanged_without_overrides():

38 - kernels = [object]

38 + kernels: list[type] = [object]

39 quantizations = (

40 "fp8_block_w8a8",

41 "fp8_w8a8",

tests/model_executor/layers/test_fused_shared_expert.py

+14 -8

@@ -17,6 +17,7 @@ import vllm.config as vllm_config_module

17 from vllm.config import VllmConfig, set_current_vllm_config

18 from vllm.model_executor.layers.fused_moe import utils as fused_moe_utils

19 from vllm.model_executor.layers.fused_moe.layer import determine_expert_counts

20 +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig

21 from vllm.model_executor.layers.quantization.fp8 import Fp8Config

22 from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig

23 from vllm.model_executor.layers.quantization.utils.config_utils import (

@@ -140,10 +141,15 @@ def get_deepseek_v4_quark_config(exclude: list[str]) -> dict[str, Any]:

141 return quantization_config

142

143

144 +def _stub_quant_config() -> QuantizationConfig:

145 + return cast(QuantizationConfig, object())

146 +

147 +

148 def get_fse_test_model_config(

149 model_type: str,

150 quantization_config: dict[str, Any],

151 ) -> tuple[object, type[nn.Module]]:

152 + config: object

153 if model_type == "minimax_m3":

154 config = MiniMaxM3TextConfig(

155 hidden_size=128,

@@ -301,7 +307,7 @@ def test_resolve_layer_fused_shared_expert_skips_compatibility_when_disabled(

307 )

308

309 assert not fused_moe_utils.resolve_layer_fused_shared_expert(

304 - object(), "model.layers.0.mlp"

310 + _stub_quant_config(), "model.layers.0.mlp"

311 )

312

313

@@ -316,7 +322,7 @@ def test_resolve_layer_fused_shared_expert_normalizes_unavailable_aiter(

322

323 assert (

324 fused_moe_utils.resolve_layer_fused_shared_expert(

319 - object(), "model.layers.0.mlp"

325 + _stub_quant_config(), "model.layers.0.mlp"

326 )

327 is False

328 )

@@ -325,7 +331,7 @@ def test_resolve_layer_fused_shared_expert_normalizes_unavailable_aiter(

331 def test_resolve_layer_fused_shared_expert_passes_module_prefixes(

332 monkeypatch: pytest.MonkeyPatch,

333 ) -> None:

328 - quant_config = object()

334 + quant_config = _stub_quant_config()

335 monkeypatch.setattr(

336 fused_moe_utils.rocm_aiter_ops,

337 "is_fusion_moe_shared_experts_enabled",

@@ -374,7 +380,7 @@ def test_resolve_layer_fused_shared_expert_rejects_incompatible_quantization(

380 )

381

382 assert not fused_moe_utils.resolve_layer_fused_shared_expert(

377 - object(), "model.layers.0.mlp"

383 + _stub_quant_config(), "model.layers.0.mlp"

384 )

385 assert "shared experts are excluded" in caplog.text

386

@@ -407,7 +413,7 @@ def test_deepseek_v4_shared_expert_fse_uses_mtp_quantization_config_prefix(

413 )

414

415 compatible, reason = is_shared_expert_quant_fse_compatible(

410 - DeepseekV4Config(),

416 + DeepseekV4Config(), # type: ignore[arg-type]

417 "model.layers.2.ffn.experts",

418 "model.layers.2.ffn.shared_experts",

419 )

@@ -818,7 +824,7 @@ def test_models_fse_init(

824 model_type, quantization_config

825 )

826 vllm_config = VllmConfig()

821 - vllm_config.model_config = SimpleNamespace(

827 + vllm_config.model_config = SimpleNamespace( # type: ignore[assignment]

828 hf_config=config,

829 hf_text_config=config,

830 dtype=torch.bfloat16,

@@ -887,7 +893,7 @@ def test_models_fse_init(

893 DeepSeekV4MTP,

894 )

895

890 - vllm_config.speculative_config = SimpleNamespace(

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

897 draft_model_config=SimpleNamespace(hf_config=config),

898 method="mtp",

899 parallel_drafting=False,

@@ -1214,7 +1220,7 @@ def test_quark_packed_layer_config_must_match_global_config() -> None:

1220

1221 def test_non_quark_shared_expert_fse_is_incompatible() -> None:

1222 compatible, reason = is_shared_expert_quant_fse_compatible(

1217 - object(),

1223 + _stub_quant_config(),

1224 "model.layers.0.mlp.experts",

1225 "model.layers.0.mlp.shared_experts",

1226 )

tests/model_executor/layers/test_mla_short_prefill_indexer.py

+4 -4

@@ -46,7 +46,7 @@ def test_sparse_attention_refreshes_batch_state_inside_eager_segment(

46

47 with pytest.raises(BatchStateRefreshed):

48 DeepseekV32Attention._sparse_indexer_and_attn(

49 - layer,

49 + layer, # type: ignore[arg-type]

50 torch.empty(1, dtype=torch.long),

51 torch.empty(1, 1),

52 torch.empty(1, 1, 1),

@@ -80,7 +80,7 @@ def make_indexer_metadata(

80 num_decode_tokens=num_decode_tokens,

81 num_prefills=num_prefills,

82 num_prefill_tokens=num_prefill_tokens,

83 - prefill=SimpleNamespace(chunks=[]) if num_prefills else None,

83 + prefill=SimpleNamespace(chunks=[]) if num_prefills else None, # type: ignore[arg-type]

84 )

85

86

@@ -129,7 +129,7 @@ def test_short_prefill_updates_k_cache_before_scoring_decision(

129 slot_mapping=slot_mapping,

130 )

131 if indexer_metadata.num_decodes:

132 - indexer_metadata.decode = object()

132 + indexer_metadata.decode = object() # type: ignore[assignment]

133 mla_metadata = make_mla_metadata(

134 use_dense_mha=batch_kind != "force_mqa",

135 num_decode_tokens=mla_num_decode_tokens,

@@ -307,7 +307,7 @@ def test_deepseek_v32_dispatches_selected_mha(

307 output = torch.empty(2, 2)

308

309 DeepseekV32Attention._sparse_indexer_and_attn(

310 - layer,

310 + layer, # type: ignore[arg-type]

311 torch.arange(2),

312 torch.empty(2, 2),

313 q_nope,

tests/model_executor/layers/test_pooler_activations.py

+2 -2

@@ -224,11 +224,11 @@ class TestResolveClassifierActFn:

224 model_config = SimpleNamespace(

225 hf_config=SimpleNamespace(num_labels=3, problem_type="")

226 )

227 - result = resolve_classifier_act_fn(model_config, act_fn=None)

227 + result = resolve_classifier_act_fn(model_config, act_fn=None) # type: ignore[arg-type]

228 assert isinstance(result, PoolerClassify)

229 assert result.num_labels == 3

230

231 def test_passes_through_provided_act_fn(self):

232 custom = PoolerIdentity()

233 - result = resolve_classifier_act_fn(None, act_fn=custom)

233 + result = resolve_classifier_act_fn(None, act_fn=custom) # type: ignore[arg-type]

234 assert result is custom

tests/model_executor/layers/test_pooler_heads.py

+9 -1

@@ -15,6 +15,7 @@ from vllm.model_executor.layers.pooler.tokwise.heads import (

15 TokenEmbeddingPoolerHead,

16 )

17 from vllm.pooling_params import PoolingParams

18 +from vllm.tasks import PoolingTask

19 from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates

20

21 _HIDDEN = 16

@@ -24,7 +25,7 @@ _BATCH = 3

25 def _make_params(

26 n: int,

27 *,

27 - task: str = "embed",

28 + task: PoolingTask = "embed",

29 dimensions: int | None = None,

30 use_activation: bool | None = None,

31 ) -> list[PoolingParams]:

@@ -320,6 +321,7 @@ class TestTokenEmbeddingPoolerHead:

321 x = torch.randn(5, _HIDDEN)

322 param = PoolingParams(task="token_embed")

323 out = head.forward_chunk(x, param)

324 + assert out is not None

325 assert out.dtype == torch.float16

326

327 def test_projector(self):

@@ -328,6 +330,7 @@ class TestTokenEmbeddingPoolerHead:

330 x = torch.randn(5, _HIDDEN)

331 param = PoolingParams(task="token_embed")

332 out = head.forward_chunk(x, param)

333 + assert out is not None

334 assert out.shape == (5, 8)

335 assert torch.allclose(out, proj(x))

336

@@ -336,6 +339,7 @@ class TestTokenEmbeddingPoolerHead:

339 x = torch.randn(5, _HIDDEN)

340 param = PoolingParams(task="token_embed", dimensions=4)

341 out = head.forward_chunk(x, param)

342 + assert out is not None

343 assert out.shape == (5, 4)

344 assert torch.equal(out, x[..., :4])

345

@@ -360,6 +364,7 @@ class TestTokenEmbeddingPoolerHead:

364 x = torch.randn(5, _HIDDEN)

365 param = PoolingParams(task="token_embed", dimensions=4)

366 out = head.forward_chunk(x, param)

367 + assert out is not None

368 assert out.shape == (5, 4)

369 assert torch.equal(out, proj(x)[..., :4])

370

@@ -368,6 +373,7 @@ class TestTokenEmbeddingPoolerHead:

373 x = torch.randn(5, _HIDDEN)

374 param = PoolingParams(task="token_embed", dimensions=4, use_activation=True)

375 out = head.forward_chunk(x, param)

376 + assert out is not None

377 assert out.shape == (5, 4)

378 norms = torch.linalg.norm(out, dim=-1)

379 assert torch.allclose(norms, torch.ones(5), atol=1e-5)

@@ -416,6 +422,7 @@ class TestTokenClassifierPoolerHead:

422 x = torch.randn(5, _HIDDEN)

423 param = PoolingParams(task="token_classify")

424 out = head.forward_chunk(x, param)

425 + assert out is not None

426 assert out.dtype == torch.float16

427

428 def test_classifier(self):

@@ -424,6 +431,7 @@ class TestTokenClassifierPoolerHead:

431 x = torch.randn(5, _HIDDEN)

432 param = PoolingParams(task="token_classify")

433 out = head.forward_chunk(x, param)

434 + assert out is not None

435 assert out.shape == (5, 3)

436 assert torch.allclose(out, clf(x))

437

tests/model_executor/layers/test_pooler_methods.py

+3 -2

@@ -26,6 +26,7 @@ from vllm.model_executor.layers.pooler.tokwise.methods import (

26 get_tok_pooling_method,

27 )

28 from vllm.pooling_params import PoolingParams

29 +from vllm.tasks import PoolingTask

30 from vllm.v1.pool.metadata import PoolingCursor, PoolingMetadata, PoolingStates

31

32 _CPU = torch.device("cpu")

@@ -66,7 +67,7 @@ def _make_pooling_cursor(

67 def _make_metadata(

68 prompt_lens: list[int],

69 *,

69 - tasks: list[str] | None = None,

70 + tasks: list[PoolingTask] | None = None,

71 token_ids: list[list[int]] | None = None,

72 pooling_params: list[PoolingParams] | None = None,

73 num_scheduled_tokens: list[int] | None = None,

@@ -319,7 +320,7 @@ def test_dispatch_seq_cls_honors_token_pooling_type(tok_pooling_type):

320 ),

321 )

322 classifier = torch.nn.Linear(4, 3)

322 - with set_current_vllm_config(config):

323 + with set_current_vllm_config(config): # type: ignore[arg-type]

324 pooler = DispatchPooler.for_seq_cls(pooler_config, classifier=classifier)

325

326 token_ids = [[2, 99, 3, 99], [4, 99], [5, 6]]

tests/plugins_tests/test_bge_m3_sparse_io_processor_plugins.py

+2 -1

@@ -114,7 +114,8 @@ async def test_bge_m3_sparse_plugin_online(

114 response = ret.json()

115

116 # Verify the request response is in the correct format

117 - assert (parsed_response := IOProcessorResponse(**response).data)

117 + parsed_response: object = IOProcessorResponse(**response).data

118 + assert parsed_response

119

120 # Verify the output is formatted as expected for this plugin

121 assert _get_attr_or_val(parsed_response, "data")

tests/plugins_tests/test_colbert_query_io_processor_plugins.py

+1 -1

@@ -79,7 +79,7 @@ def _post_pooling(server: RemoteOpenAIServer, data: dict):

79 ret = requests.post(server.url_for("pooling"), json=request_payload)

80 ret.raise_for_status()

81 response = ret.json()

82 - parsed_response = IOProcessorResponse(**response).data

82 + parsed_response: object = IOProcessorResponse(**response).data

83 assert parsed_response

84 return parsed_response

85

tests/plugins_tests/test_endpoint_plugins.py

+2 -2

@@ -172,7 +172,7 @@ async def test_init_state_is_noop_without_phase_a(monkeypatch: pytest.MonkeyPatc

172 monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin")

173

174 state = State()

175 - await init_endpoint_plugins_state(_FakeEngineClient(), state, _build_args())

175 + await init_endpoint_plugins_state(_FakeEngineClient(), state, _build_args()) # type: ignore[arg-type]

176

177 assert not hasattr(state, "dummy_engine_client")

178

@@ -223,7 +223,7 @@ async def test_endpoint_plugin_end_to_end(monkeypatch: pytest.MonkeyPatch):

223 )

224

225 fake_engine_client = _FakeEngineClient(rpc_result=["cfg-a", "cfg-b"])

226 - await init_endpoint_plugins_state(fake_engine_client, app.state, args)

226 + await init_endpoint_plugins_state(fake_engine_client, app.state, args) # type: ignore[arg-type]

227

228 assert app.state.dummy_engine_client is fake_engine_client

229

tests/plugins_tests/test_oot_registration_offline.py

+1 -1

@@ -92,7 +92,7 @@ def test_oot_registration_multimodal(

92 )

93

94 first_token = llm.get_tokenizer().decode(0)

95 - outputs = llm.generate(prompts, sampling_params)

95 + outputs = llm.generate(prompts, sampling_params) # type: ignore[arg-type]

96

97 for output in outputs:

98 generated_text = output.outputs[0].text

tests/plugins_tests/test_stats_logger_plugins.py

+5 -3

@@ -23,7 +23,8 @@ def test_stat_logger_plugin_is_discovered(monkeypatch: pytest.MonkeyPatch):

23 # instantiate and confirm the right type

24 vllm_config = VllmConfig()

25 instance = factories[0](vllm_config)

26 - assert isinstance(instance, DummyStatLogger)

26 + # The `is` check above narrows DummyStatLogger to a Callable

27 + assert isinstance(instance, DummyStatLogger) # type: ignore[arg-type]

28

29

30 def test_no_plugins_loaded_if_env_empty(monkeypatch: pytest.MonkeyPatch):

@@ -66,10 +67,11 @@ async def test_stat_logger_plugin_integration_with_engine(

67

68 engine = AsyncLLM.from_engine_args(engine_args=engine_args)

69

70 + assert engine.logger_manager is not None

71 assert len(engine.logger_manager.stat_loggers) == 2

70 - assert len(engine.logger_manager.stat_loggers[0].per_engine_stat_loggers) == 1

72 + assert len(engine.logger_manager.stat_loggers[0].per_engine_stat_loggers) == 1 # type: ignore[attr-defined]

73 assert isinstance(

72 - engine.logger_manager.stat_loggers[0].per_engine_stat_loggers[0],

74 + engine.logger_manager.stat_loggers[0].per_engine_stat_loggers[0], # type: ignore[attr-defined]

75 DummyStatLogger,

76 )

77

tests/plugins_tests/test_terratorch_io_processor_plugins.py

+2 -1

@@ -95,7 +95,8 @@ async def test_prithvi_mae_plugin_online(

95 response = ret.json()

96

97 # verify the request response is in the correct format

98 - assert (parsed_response := IOProcessorResponse(**response))

98 + parsed_response: IOProcessorResponse = IOProcessorResponse(**response)

99 + assert parsed_response

100

101 # verify the output is formatted as expected for this plugin

102 plugin_data = parsed_response.data

tests/samplers/test_beam_search.py

+6 -4

@@ -16,6 +16,7 @@ from transformers import AutoModelForSeq2SeqLM

16 from vllm import CompletionOutput, RequestOutput

17 from vllm.assets.audio import AudioAsset

18 from vllm.entrypoints.llm import LLM

19 +from vllm.inputs import TokensInput

20 from vllm.logprobs import Logprob, SampleLogprobs

21 from vllm.platforms import current_platform

22 from vllm.sampling_params import BeamSearchParams, StructuredOutputsParams

@@ -107,11 +108,12 @@ def test_beam_search_abort_returns_partial_outputs_and_continues_other_prompts(

108 llm.renderer = Mock(get_tokenizer=Mock(return_value=tokenizer))

109 monkeypatch.setattr(llm, "_preprocess_cmpl", lambda prompts: prompts)

110 monkeypatch.setattr(llm, "_render_and_run_requests", run_requests)

110 - prompts = [

111 + # Engine inputs, since _preprocess_cmpl is patched to the identity

112 + prompts: list[TokensInput] = [

113 {"type": "token", "prompt_token_ids": [token]} for token in [prompt_token, 2]

114 ]

115 params = BeamSearchParams(beam_width=2, max_tokens=3)

114 - outputs = llm.beam_search(prompts, params)

116 + outputs = llm.beam_search(prompts, params) # type: ignore[arg-type]

117

118 aborted, normal = outputs

119 expected_tokens = (

@@ -324,7 +326,7 @@ def test_beam_search_structured_output(

326

327 llm = LLM(

328 model=model,

327 - dtype=dtype,

329 + dtype=dtype, # type: ignore[arg-type]

330 max_model_len=512,

331 structured_outputs_config=dict(

332 backend="xgrammar",

@@ -343,7 +345,7 @@ def test_beam_search_structured_output(

345 "Generate a JSON object for a person with name and age:",

346 ]

347

346 - outputs = llm.beam_search(prompts, params)

348 + outputs = llm.beam_search(prompts, params) # type: ignore[arg-type]

349

350 assert len(outputs) == len(prompts)

351 for output in outputs:

tests/samplers/test_beam_search_online.py

+10 -5

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

5

6 from vllm import CompletionOutput, RequestOutput

7 from vllm.entrypoints.generate.beam_search.online import BeamSearchOnlineMixin

8 +from vllm.inputs import TokensInput

9 from vllm.logprobs import Logprob, SampleLogprobs

10 from vllm.sampling_params import BeamSearchParams

11

@@ -56,13 +57,13 @@ class _EngineClient:

57

58

59 class _Serving(BeamSearchOnlineMixin):

59 - renderer = _Renderer()

60 - engine_client = _EngineClient()

60 + renderer = _Renderer() # type: ignore[assignment]

61 + engine_client = _EngineClient() # type: ignore[assignment]

62

63

64 @pytest.mark.asyncio

65 async def test_beam_search_handles_extra_logprob_candidates() -> None:

65 - prompt = {

66 + prompt: TokensInput = {

67 "type": "token",

68 "prompt": "prompt",

69 "prompt_token_ids": [1],

@@ -90,7 +91,7 @@ async def test_beam_search_handles_extra_logprob_candidates() -> None:

91 async def test_beam_search_respects_skip_special_tokens(

92 skip_special_tokens: bool, expected_text: str

93 ) -> None:

93 - prompt = {

94 + prompt: TokensInput = {

95 "type": "token",

96 "prompt": "prompt",

97 "prompt_token_ids": [1],

@@ -140,7 +141,11 @@ async def test_beam_search_abort_returns_partial_outputs(

141

142 serving = _Serving()

143 monkeypatch.setattr(serving.engine_client, "generate", generate)

143 - prompt = {"type": "token", "prompt": "prompt", "prompt_token_ids": [prompt_token]}

144 + prompt: TokensInput = {

145 + "type": "token",

146 + "prompt": "prompt",

147 + "prompt_token_ids": [prompt_token],

148 + }

149 outputs = [

150 output

151 async for output in serving.beam_search(

tests/samplers/test_no_bad_words.py

+9 -4

@@ -8,11 +8,12 @@ Run `pytest tests/samplers/test_no_bad_words.py`.

8

9 from transformers import AutoTokenizer

10

11 -from vllm import LLM, SamplingParams

11 +from tests.conftest import VllmRunner

12 +from vllm import SamplingParams

13

14

15 def _generate(

15 - llm: LLM,

16 + llm: VllmRunner,

17 prompt: str,

18 num_prompt_tokens: int,

19 temperature: float = 0,

@@ -56,7 +57,9 @@ class TestOneTokenBadWord:

57 output_token_ids = self._generate(llm, bad_words=[self.TARGET_TOKEN])

58 assert self.target_token_id not in output_token_ids

59

59 - def _generate(self, llm: LLM, bad_words: list[str] | None = None) -> list[int]:

60 + def _generate(

61 + self, llm: VllmRunner, bad_words: list[str] | None = None

62 + ) -> list[int]:

63 return _generate(

64 llm=llm,

65 prompt=self.PROMPT,

@@ -151,7 +154,9 @@ class TestTwoTokenBadWord:

154 self.neighbour_token_id2 in output_token_ids

155 )

156

154 - def _generate(self, llm: LLM, bad_words: list[str] | None = None) -> list[int]:

157 + def _generate(

158 + self, llm: VllmRunner, bad_words: list[str] | None = None

159 + ) -> list[int]:

160 return _generate(

161 llm=llm,

162 prompt=self.PROMPT,

tests/tokenizers_/test_detokenize.py

+23 -5

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

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

3

4 from collections.abc import Generator

5 -from typing import Any

5 +from typing import Any, cast

6

7 import pytest

8 from transformers import AutoTokenizer, PythonBackend, TokenizersBackend

@@ -10,6 +10,7 @@ from transformers import AutoTokenizer, PythonBackend, TokenizersBackend

10 from vllm.sampling_params import SamplingParams

11 from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens

12 from vllm.tokenizers.mistral import MistralTokenizer

13 +from vllm.tokenizers.protocol import TokenizerLike

14 from vllm.v1.engine import EngineCoreRequest

15 from vllm.v1.engine.detokenizer import (

16 FastIncrementalDetokenizer,

@@ -281,9 +282,26 @@ class _MockTokenizer:

282 return "".join(self._decoded[tid] for tid in ids)

283

284

285 +def _mock_tokenizer(

286 + raw_tokens: dict[int, str],

287 + decoded_tokens: dict[int, str],

288 + pre_tokenizer_type: str = "Metaspace",

289 + replacement: str | None = "▁",

290 +) -> TokenizerLike:

291 + return cast(

292 + TokenizerLike,

293 + _MockTokenizer(

294 + raw_tokens=raw_tokens,

295 + decoded_tokens=decoded_tokens,

296 + pre_tokenizer_type=pre_tokenizer_type,

297 + replacement=replacement,

298 + ),

299 + )

300 +

301 +

302 def test_sentencepiece_leading_space_preserved():

303 """▁true and true must produce distinct strings."""

286 - tok = _MockTokenizer(

304 + tok = _mock_tokenizer(

305 raw_tokens={0: "▁true", 1: "true", 2: "▁false", 3: "false"},

306 decoded_tokens={0: "true", 1: "true", 2: "false", 3: "false"},

307 )

@@ -297,7 +315,7 @@ def test_sentencepiece_leading_space_preserved():

315

316 def test_whitespace_run_tokens_stay_distinct():

317 """▁, ▁▁, ▁▁▁ must produce different-length space strings."""

300 - tok = _MockTokenizer(

318 + tok = _mock_tokenizer(

319 raw_tokens={0: "▁", 1: "▁▁", 2: "▁▁▁"},

320 decoded_tokens={0: "", 1: " ", 2: " "},

321 )

@@ -307,7 +325,7 @@ def test_whitespace_run_tokens_stay_distinct():

325

326 def test_bpe_leading_space_already_preserved():

327 """GPT-2 BPE: Ġtrue already decodes to ' true', no fix needed."""

310 - tok = _MockTokenizer(

328 + tok = _mock_tokenizer(

329 raw_tokens={0: "Ġtrue", 1: "true"},

330 decoded_tokens={0: " true", 1: "true"},

331 pre_tokenizer_type="ByteLevel",

@@ -319,7 +337,7 @@ def test_bpe_leading_space_already_preserved():

337

338 def test_logprobs_count_stable_across_k():

339 """logprobs=4 and logprobs=10 must return 4 and 10 entries."""

322 - tok = _MockTokenizer(

340 + tok = _mock_tokenizer(

341 raw_tokens={

342 0: "▁true",

343 1: "a",

tests/tokenizers_/test_hf.py

+2 -2

@@ -61,7 +61,7 @@ def test_thread_pool_tokenizer_pickle(model_id: str):

61 unpickled_tokenizer = pickle.loads(pickle.dumps(pooled_tokenizer))

62 assert unpickled_tokenizer is not None

63 assert isinstance(unpickled_tokenizer, ThreadSafeHFTokenizerMixin)

64 - assert unpickled_tokenizer.encode("prompt") == reference_tokenizer.encode("prompt")

64 + assert unpickled_tokenizer.encode("prompt") == reference_tokenizer.encode("prompt") # type: ignore[attr-defined]

65

66 # Idempotence: wrapping an already-pooled tokenizer returns it unchanged.

67 - assert maybe_make_thread_pool(pooled_tokenizer) is pooled_tokenizer

67 + assert maybe_make_thread_pool(pooled_tokenizer) is pooled_tokenizer # type: ignore[type-var]

tests/tokenizers_/test_mistral.py

+26 -12

@@ -10,6 +10,7 @@ from mistral_common.exceptions import InvalidMessageStructureException

10 from mistral_common.guidance.grammar_factory import GrammarFactory

11 from mistral_common.tokens.tokenizers.base import SpecialTokenPolicy, SpecialTokens

12

13 +from vllm.entrypoints.chat_utils import ChatCompletionMessageParam

14 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest

15 from vllm.exceptions import VLLMValidationError

16 from vllm.tokenizers.mistral import (

@@ -21,7 +22,9 @@ from vllm.tokenizers.mistral import (

22

23 def test_validate_apply_chat_template_args():

24 # add_generation_prompt with assistant last message → error

24 - messages = [{"role": "assistant", "content": "Hello"}]

25 + messages: list[ChatCompletionMessageParam] = [

26 + {"role": "assistant", "content": "Hello"}

27 + ]

28 with pytest.raises(ValueError):

29 _validate_apply_chat_template_args(messages, add_generation_prompt=True)

30

@@ -32,7 +35,9 @@ def test_validate_apply_chat_template_args():

35 # both add_generation_prompt and continue_final_message → error

36 with pytest.raises(ValueError):

37 _validate_apply_chat_template_args(

35 - messages, add_generation_prompt=True, continue_final_message=True

38 + messages,

39 + add_generation_prompt=True,

40 + continue_final_message=True,

41 )

42

43 # continue_final_message with assistant last message → ok

@@ -835,7 +840,9 @@ class TestMistralTokenizer:

840 )

841

842 def test_apply_chat_template_error(self, mistral_tokenizer: MistralTokenizer):

838 - messages = [{"role": "user", "content": "Hello world !"}]

843 + messages: list[ChatCompletionMessageParam] = [

844 + {"role": "user", "content": "Hello world !"}

845 + ]

846

847 with pytest.raises(ValueError):

848 mistral_tokenizer.apply_chat_template(

@@ -968,7 +975,7 @@ class TestMistralTokenizer:

975 mistral_tokenizer: MistralTokenizer,

976 ):

977 assert mistral_tokenizer.batch_decode(

971 - [[]],

978 + [[]], # type: ignore[arg-type]

979 ) == [""]

980

981 def test_convert_tokens_to_string(self, mistral_tokenizer: MistralTokenizer):

@@ -2167,13 +2174,15 @@ class TestMistralTokenizer:

2174 def test_apply_chat_template_tool_optional_fields(

2175 self,

2176 mistral_tokenizer: MistralTokenizer,

2170 - messages: list[dict[str, Any]],

2177 + messages: list[ChatCompletionMessageParam],

2178 tools: list[dict[str, Any]],

2179 tekken_expected_substrings: list[str],

2180 spm_expected_substrings: list[str],

2181 ) -> None:

2182 output = mistral_tokenizer.apply_chat_template(

2176 - messages, tools=tools, add_generation_prompt=True

2183 + messages,

2184 + tools=tools,

2185 + add_generation_prompt=True,

2186 )

2187 decoded = mistral_tokenizer.tokenizer.decode(output, SpecialTokenPolicy.KEEP)

2188

@@ -2188,7 +2197,7 @@ class TestMistralTokenizer:

2197 def test_apply_chat_template_tools_not_mutated(

2198 self, mistral_tokenizer: MistralTokenizer

2199 ) -> None:

2191 - messages: list[dict[str, Any]] = [

2200 + messages: list[ChatCompletionMessageParam] = [

2201 {"role": "user", "content": "Hello"},

2202 ]

2203 tools: list[dict[str, Any]] = [

@@ -2209,7 +2218,9 @@ class TestMistralTokenizer:

2218 original_tools = copy.deepcopy(tools)

2219

2220 mistral_tokenizer.apply_chat_template(

2212 - messages, tools=tools, add_generation_prompt=True

2221 + messages,

2222 + tools=tools,

2223 + add_generation_prompt=True,

2224 )

2225

2226 assert tools == original_tools

@@ -2235,7 +2246,8 @@ class TestMistralTokenizer:

2246 ]

2247

2248 output = mistral_tokenizer.apply_chat_template(

2238 - messages, add_generation_prompt=True

2249 + messages, # type: ignore[arg-type]

2250 + add_generation_prompt=True,

2251 )

2252 decoded = mistral_tokenizer.tokenizer.decode(output, SpecialTokenPolicy.KEEP)

2253

@@ -2310,6 +2322,7 @@ def v15_mistral_tokenizer() -> MistralTokenizer:

2322 return MistralTokenizer.from_pretrained("mistralai/Mistral-Small-4-119B-2603")

2323 except Exception:

2324 pytest.skip("v15 tokenizer unavailable")

2325 + raise AssertionError("unreachable")

2326

2327

2328 @pytest.fixture(scope="module")

@@ -2319,6 +2332,7 @@ def v13_mistral_tokenizer() -> MistralTokenizer:

2332 return MistralTokenizer.from_pretrained("mistralai/Magistral-Small-2509")

2333 except Exception:

2334 pytest.skip("v13 tokenizer unavailable")

2335 + raise AssertionError("unreachable")

2336

2337

2338 def test_v15_apply_chat_template_passes_reasoning_effort_high(

@@ -2338,7 +2352,7 @@ def test_v15_apply_chat_template_passes_reasoning_effort_high(

2352 )

2353

2354 v15_mistral_tokenizer.apply_chat_template(

2341 - messages=_PASSTHROUGH_MESSAGES,

2355 + messages=_PASSTHROUGH_MESSAGES, # type: ignore[arg-type]

2356 reasoning_effort="high",

2357 )

2358

@@ -2361,7 +2375,7 @@ def test_v15_apply_chat_template_passes_reasoning_effort_none_by_default(

2375 fake_apply_chat_template,

2376 )

2377

2364 - v15_mistral_tokenizer.apply_chat_template(messages=_PASSTHROUGH_MESSAGES)

2378 + v15_mistral_tokenizer.apply_chat_template(messages=_PASSTHROUGH_MESSAGES) # type: ignore[arg-type]

2379

2380 assert "reasoning_effort" in captured_kwargs[-1]

2381 assert captured_kwargs[-1]["reasoning_effort"] is None

@@ -2384,7 +2398,7 @@ def test_pre_v15_apply_chat_template_omits_reasoning_effort(

2398 )

2399

2400 v13_mistral_tokenizer.apply_chat_template(

2387 - messages=_PASSTHROUGH_MESSAGES,

2401 + messages=_PASSTHROUGH_MESSAGES, # type: ignore[arg-type]

2402 reasoning_effort="high",

2403 )

2404

tests/tokenizers_/test_registry.py

+1 -1

@@ -205,7 +205,7 @@ def test_cached_tokenizer_from_config_registers_local_config(tmp_path: Path):

205 side_effect=lambda tokenizer: tokenizer,

206 ),

207 ):

208 - tokenizer = cached_tokenizer_from_config(model_config)

208 + tokenizer = cached_tokenizer_from_config(model_config) # type: ignore[arg-type]

209

210 assert tokenizer.is_fast is True

211 finally:

tests/tool_use/test_gemma4_responses_adjust_request.py

+1 -1

@@ -148,7 +148,7 @@ def test_tool_parser_adjust_request_builds_valid_response_text_config() -> None:

148 description string that the previous two-step construction injected.

149 """

150 parser = ToolParser.__new__(ToolParser)

151 - parser.model_tokenizer = None

151 + parser.model_tokenizer = None # type: ignore[assignment]

152

153 request = _build_responses_request(tool_choice="required")

154 ToolParser.adjust_request(parser, request)

tests/tool_use/test_kimi_k3_tool_parser.py

+28 -21

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

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

3

4 import json

5 +from typing import cast

6

7 import pytest

8

@@ -15,6 +16,7 @@ from vllm.exceptions import VLLMValidationError

16 from vllm.parser.kimi_k3 import KimiK3Parser

17 from vllm.parser.parser_manager import ParserManager

18 from vllm.reasoning.kimi_k3_reasoning_parser import KimiK3ReasoningParser

19 +from vllm.tokenizers.protocol import TokenizerLike

20 from vllm.tool_parsers.kimi_k3_tool_parser import KimiK3ToolParser

21

22 OPEN = "<|open|>"

@@ -37,6 +39,10 @@ class DummyTokenizer:

39 return [ord(ch) for ch in text]

40

41

42 +def _dummy_tokenizer() -> TokenizerLike:

43 + return cast(TokenizerLike, DummyTokenizer())

44 +

45 +

46 class KimiK3DelegatingParser(KimiK3Parser):

47 reasoning_parser_cls = KimiK3ReasoningParser

48 tool_parser_cls = KimiK3ToolParser

@@ -124,7 +130,7 @@ def _tools(*calls: str) -> str:

130

131

132 def test_extract_tool_calls_with_response_and_typed_arguments():

127 - parser = KimiK3ToolParser(DummyTokenizer())

133 + parser = KimiK3ToolParser(_dummy_tokenizer())

134

135 output = _response("answer") + _tools(

136 _call(

@@ -150,7 +156,7 @@ def test_extract_tool_calls_with_response_and_typed_arguments():

156

157

158 def test_delegating_parser_preserves_tool_calls_after_reasoning():

153 - parser = KimiK3DelegatingParser(DummyTokenizer())

159 + parser = KimiK3DelegatingParser(_dummy_tokenizer())

160 output = (

161 f"{THINK_OPEN}step{THINK_CLOSE}"

162 + _response("answer")

@@ -172,7 +178,7 @@ def test_delegating_parser_preserves_tool_calls_after_reasoning():

178

179

180 def test_delegating_parser_required_tool_choice_uses_xtml_parser():

175 - parser = KimiK3DelegatingParser(DummyTokenizer())

181 + parser = KimiK3DelegatingParser(_dummy_tokenizer())

182 request = _request().model_copy(update={"tool_choice": "required"})

183 output = (

184 f"{THINK_OPEN}step{THINK_CLOSE}"

@@ -195,7 +201,7 @@ def test_delegating_parser_required_tool_choice_uses_xtml_parser():

201

202

203 def test_delegating_parser_named_tool_choice_uses_xtml_parser():

198 - parser = KimiK3DelegatingParser(DummyTokenizer())

204 + parser = KimiK3DelegatingParser(_dummy_tokenizer())

205 output = (

206 f"{THINK_OPEN}step{THINK_CLOSE}"

207 + _response("")

@@ -218,7 +224,7 @@ def test_delegating_parser_named_tool_choice_uses_xtml_parser():

224

225 def test_delegating_parser_auto_no_call_strips_consumed_response_prefix():

226 parser = KimiK3DelegatingParser(

221 - DummyTokenizer(), chat_template_kwargs={"thinking": False}

227 + _dummy_tokenizer(), chat_template_kwargs={"thinking": False}

228 )

229 request = _request().model_copy(

230 update={"chat_template_kwargs": {"thinking": False}}

@@ -237,7 +243,7 @@ def test_delegating_parser_auto_no_call_strips_consumed_response_prefix():

243

244 def test_delegating_parser_required_call_strips_consumed_response_prefix():

245 parser = KimiK3DelegatingParser(

240 - DummyTokenizer(), chat_template_kwargs={"thinking": False}

246 + _dummy_tokenizer(), chat_template_kwargs={"thinking": False}

247 )

248 request = _request().model_copy(

249 update={

@@ -263,7 +269,7 @@ def test_delegating_parser_required_call_strips_consumed_response_prefix():

269

270 def test_delegating_parser_truncated_tools_do_not_leak_xtml():

271 parser = KimiK3DelegatingParser(

266 - DummyTokenizer(), chat_template_kwargs={"thinking": False}

272 + _dummy_tokenizer(), chat_template_kwargs={"thinking": False}

273 )

274 request = _request().model_copy(

275 update={

@@ -284,7 +290,7 @@ def test_delegating_parser_truncated_tools_do_not_leak_xtml():

290

291

292 def test_extract_tool_calls_unescapes_attributes():

287 - parser = KimiK3ToolParser(DummyTokenizer())

293 + parser = KimiK3ToolParser(_dummy_tokenizer())

294

295 output = _tools(_call("a&b"c", 1, _arg("k&q", "string", "v")))

296 extracted = parser.extract_tool_calls(output, _request())

@@ -295,7 +301,7 @@ def test_extract_tool_calls_unescapes_attributes():

301

302

303 def test_extract_tool_calls_allows_less_than_in_attributes():

298 - parser = KimiK3ToolParser(DummyTokenizer())

304 + parser = KimiK3ToolParser(_dummy_tokenizer())

305

306 output = _tools(_call("calc<beta", 1, _arg("foo<bar", "string", "raw")))

307 extracted = parser.extract_tool_calls(output, _request())

@@ -306,7 +312,7 @@ def test_extract_tool_calls_allows_less_than_in_attributes():

312

313

314 def test_extract_content_from_whitespace_degraded_markers():

309 - parser = KimiK3ToolParser(DummyTokenizer())

315 + parser = KimiK3ToolParser(_dummy_tokenizer())

316

317 extracted = parser.extract_tool_calls(

318 f"{OPEN} response {SEP}answer{CLOSE} response {SEP}",

@@ -318,7 +324,7 @@ def test_extract_content_from_whitespace_degraded_markers():

324

325

326 def test_streaming_split_markers_do_not_leak():

321 - parser = KimiK3ToolParser(DummyTokenizer())

327 + parser = KimiK3ToolParser(_dummy_tokenizer())

328 request = _request()

329 previous_text = ""

330 previous_ids: list[int] = []

@@ -362,21 +368,22 @@ def test_streaming_split_markers_do_not_leak():

368 assert OPEN not in content

369 assert SEP not in content

370 assert len(tool_deltas) == 1

371 + assert tool_deltas[0].function is not None

372 assert tool_deltas[0].function.name == "calc"

366 - assert json.loads(tool_deltas[0].function.arguments) == {"x": 1}

373 + assert json.loads(tool_deltas[0].function.arguments) == {"x": 1} # type: ignore[arg-type]

374

375

376 def test_tool_call_ids_are_unique_across_messages():

377 output = _tools(_call("calc", 1))

378

372 - first = KimiK3ToolParser(DummyTokenizer()).extract_tool_calls(output, _request())

373 - second = KimiK3ToolParser(DummyTokenizer()).extract_tool_calls(output, _request())

379 + first = KimiK3ToolParser(_dummy_tokenizer()).extract_tool_calls(output, _request())

380 + second = KimiK3ToolParser(_dummy_tokenizer()).extract_tool_calls(output, _request())

381

382 assert first.tool_calls[0].id != second.tool_calls[0].id

383

384

385 def test_streaming_consumed_response_prefix_no_call_keeps_content():

379 - parser = KimiK3ToolParser(DummyTokenizer())

386 + parser = KimiK3ToolParser(_dummy_tokenizer())

387 request = _request()

388 previous_text = ""

389 previous_ids: list[int] = []

@@ -406,7 +413,7 @@ def test_streaming_consumed_response_prefix_no_call_keeps_content():

413

414 def test_delegating_parser_tool_choice_none_strips_xtml_and_suppresses_calls():

415 parser = KimiK3DelegatingParser(

409 - DummyTokenizer(), chat_template_kwargs={"thinking": False}

416 + _dummy_tokenizer(), chat_template_kwargs={"thinking": False}

417 )

418 request = _request().model_copy(

419 update={

@@ -443,7 +450,7 @@ def test_delegating_parser_tool_choice_none_strips_xtml_and_suppresses_calls():

450

451

452 def test_adjust_request_keeps_xtml_markers_contiguous():

446 - parser = KimiK3ToolParser(DummyTokenizer())

453 + parser = KimiK3ToolParser(_dummy_tokenizer())

454 request = _request()

455

456 adjusted = parser.adjust_request(request)

@@ -455,7 +462,7 @@ def test_adjust_request_keeps_xtml_markers_contiguous():

462

463

464 def test_adjust_request_required_uses_xtml_parser_not_json_guidance():

458 - parser = KimiK3ToolParser(DummyTokenizer())

465 + parser = KimiK3ToolParser(_dummy_tokenizer())

466 request = _request().model_copy(update={"tool_choice": "required"})

467

468 adjusted = parser.adjust_request(request)

@@ -476,7 +483,7 @@ def test_adjust_request_required_uses_xtml_parser_not_json_guidance():

483 ],

484 )

485 def test_adjust_request_rejects_named_tool_choice(tool_request):

479 - parser = KimiK3ToolParser(DummyTokenizer())

486 + parser = KimiK3ToolParser(_dummy_tokenizer())

487

488 with pytest.raises(VLLMValidationError) as exc_info:

489 parser.adjust_request(tool_request)

@@ -512,7 +519,7 @@ def test_responses_chat_params_keeps_template_tool_choice_when_api_auto():

519

520 def test_responses_required_tool_choice_uses_xtml_parser():

521 parser = KimiK3DelegatingParser(

515 - DummyTokenizer(), chat_template_kwargs={"thinking": False}

522 + _dummy_tokenizer(), chat_template_kwargs={"thinking": False}

523 )

524 request = _responses_request(tool_choice="required").model_copy(

525 update={"chat_template_kwargs": {"thinking": False}}

@@ -538,7 +545,7 @@ def test_responses_required_tool_choice_uses_xtml_parser():

545

546 def test_responses_named_tool_choice_uses_xtml_parser():

547 parser = KimiK3DelegatingParser(

541 - DummyTokenizer(), chat_template_kwargs={"thinking": False}

548 + _dummy_tokenizer(), chat_template_kwargs={"thinking": False}

549 )

550 request = _responses_request(

551 tool_choice={"type": "function", "name": "calc"}

tests/tool_use/test_muse_glimmer.py

+36 -20

@@ -24,9 +24,11 @@ a real tokenizer live in ``test_muse_glimmer_parse_delta.py``.

24

25 import json

26 from types import SimpleNamespace

27 +from typing import cast

28

29 import pytest

30

31 +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest

32 from vllm.reasoning.muse_glimmer_reasoning_parser import MuseGlimmerReasoningParser

33 from vllm.tool_parsers.muse_glimmer_tool_parser import MuseGlimmerToolParser

34

@@ -61,10 +63,21 @@ class _FakeReq:

63 tools = None

64

65

64 -def _req(*names):

66 +def _fake_req() -> ChatCompletionRequest:

67 + return cast(ChatCompletionRequest, _FakeReq())

68 +

69 +

70 +# The parsers accept a missing request at runtime.

71 +_NO_REQ = cast(ChatCompletionRequest, None)

72 +

73 +

74 +def _req(*names: str) -> ChatCompletionRequest:

75 """A request with ``names`` registered as tools."""

66 - return SimpleNamespace(

67 - tools=[SimpleNamespace(function=SimpleNamespace(name=n)) for n in names]

76 + return cast(

77 + ChatCompletionRequest,

78 + SimpleNamespace(

79 + tools=[SimpleNamespace(function=SimpleNamespace(name=n)) for n in names]

80 + ),

81 )

82

83

@@ -90,7 +103,7 @@ def test_single_tool_call_after_reasoning():

103 '<atem:parameter name="units">celsius</atem:parameter>\n'

104 "</atem:invoke>\n</atem:function_calls><|eot|>"

105 )

93 - out = MuseGlimmerToolParser.extract_tool_calls(T, raw, None)

106 + out = MuseGlimmerToolParser.extract_tool_calls(T, raw, _NO_REQ)

107 assert out.tools_called and len(out.tool_calls) == 1

108 assert out.tool_calls[0].function.name == "weather.get"

109 assert json.loads(out.tool_calls[0].function.arguments) == {

@@ -112,7 +125,7 @@ def test_parallel_calls_across_eom_boundaries():

125 '<atem:parameter name="b">4</atem:parameter>\n'

126 "</atem:invoke>\n</atem:function_calls><|eot|>"

127 )

115 - out = MuseGlimmerToolParser.extract_tool_calls(T, raw, None)

128 + out = MuseGlimmerToolParser.extract_tool_calls(T, raw, _NO_REQ)

129 assert out.tools_called and len(out.tool_calls) == 2, len(out.tool_calls)

130 assert [t.function.name for t in out.tool_calls] == ["math.add", "math.mul"]

131 # JSON-typed values decode to ints

@@ -127,14 +140,14 @@ def test_echoed_invoke_in_reasoning_is_not_parsed():

140 "but I will not.<|eom|>"

141 "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>"

142 )

130 - out = MuseGlimmerToolParser.extract_tool_calls(T, raw, None)

143 + out = MuseGlimmerToolParser.extract_tool_calls(T, raw, _NO_REQ)

144 assert not out.tools_called, "channel scoping failed -- echoed invoke parsed!"

145 assert out.content == "The answer is 42.", repr(out.content)

146

147

148 def test_plain_answer_yields_no_tool_calls():

149 out = MuseGlimmerToolParser.extract_tool_calls(

137 - T, "to=user<|message|>Just a plain answer.<|eot|>", None

150 + T, "to=user<|message|>Just a plain answer.<|eot|>", _NO_REQ

151 )

152 assert not out.tools_called

153

@@ -147,7 +160,7 @@ def test_json_object_array_and_bool_params_decode():

160 '<atem:parameter name="flag">true</atem:parameter>\n'

161 "</atem:invoke>\n</atem:function_calls><|eot|>"

162 )

150 - out = MuseGlimmerToolParser.extract_tool_calls(T, raw, None)

163 + out = MuseGlimmerToolParser.extract_tool_calls(T, raw, _NO_REQ)

164 assert json.loads(out.tool_calls[0].function.arguments) == {

165 "payload": {"nested": [1, 2, 3]},

166 "flag": True,

@@ -166,10 +179,10 @@ def test_reasoning_to_toolcall_handoff():

179 '<atem:parameter name="city">Paris</atem:parameter>\n'

180 "</atem:invoke>\n</atem:function_calls>"

181 )

169 - reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(R, raw, None)

182 + reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(R, raw, _NO_REQ)

183 assert reasoning == "Let me call the tool.", repr(reasoning)

184 assert content is not None and "<atem:invoke" in content, repr(content)

172 - out = MuseGlimmerToolParser.extract_tool_calls(T, content, None)

185 + out = MuseGlimmerToolParser.extract_tool_calls(T, content, _NO_REQ)

186 assert out.tools_called and len(out.tool_calls) == 1

187 assert out.tool_calls[0].function.name == "weather.get"

188

@@ -179,15 +192,17 @@ def test_reasoning_then_user_answer():

192 " to=self<|message|>thinking<|eom|>"

193 "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>"

194 )

182 - reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(R, raw, None)

195 + reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(R, raw, _NO_REQ)

196 assert reasoning == "thinking", repr(reasoning)

197 assert content == "The answer is 42.", repr(content)

185 - assert not MuseGlimmerToolParser.extract_tool_calls(T, content, None).tools_called

198 + assert not MuseGlimmerToolParser.extract_tool_calls(

199 + T, content, _NO_REQ

200 + ).tools_called

201

202

203 def test_plain_content_without_framing_passes_through():

204 reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(

190 - R, "Just a direct answer.", None

205 + R, "Just a direct answer.", _NO_REQ

206 )

207 assert reasoning is None and content == "Just a direct answer.", (

208 reasoning,

@@ -207,9 +222,10 @@ def test_reasoning_then_parallel_calls():

222 '<atem:parameter name="a">3</atem:parameter>\n</atem:invoke>\n'

223 "</atem:function_calls><|eot|>"

224 )

210 - reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(R, raw, None)

225 + reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(R, raw, _NO_REQ)

226 assert reasoning == "need two calls", repr(reasoning)

212 - out = MuseGlimmerToolParser.extract_tool_calls(T, content, None)

227 + assert content is not None

228 + out = MuseGlimmerToolParser.extract_tool_calls(T, content, _NO_REQ)

229 assert [t.function.name for t in out.tool_calls] == ["math.add", "math.mul"], (

230 out.tool_calls

231 )

@@ -231,7 +247,7 @@ def _stream(raw: str, chunk: int):

247 R, prev, cur, delta, [], [], []

248 )

249 if dm is not None:

234 - if getattr(dm, "reasoning", None):

250 + if dm.reasoning:

251 reasoning.append(dm.reasoning)

252 content_delta = getattr(dm, "content", None)

253 # Tool-channel content is an internal handoff to the tool parser,

@@ -239,7 +255,7 @@ def _stream(raw: str, chunk: int):

255 if content_delta and "<atem:function_calls>" not in content_delta:

256 content.append(content_delta)

257 dt = MuseGlimmerToolParser.extract_tool_calls_streaming(

242 - T, prev, cur, delta, [], [], [], _FakeReq()

258 + T, prev, cur, delta, [], [], [], _fake_req()

259 )

260 if dt is not None and dt.tool_calls:

261 toolcalls.extend(dt.tool_calls)

@@ -319,11 +335,11 @@ def test_streaming_reasoning_then_content():

335

336

337 def test_truncated_cot_no_toolcall_nonstreaming():

322 - out = MuseGlimmerToolParser.extract_tool_calls(T, RAW_TRUNCATED, _FakeReq())

338 + out = MuseGlimmerToolParser.extract_tool_calls(T, RAW_TRUNCATED, _fake_req())

339 assert not out.tools_called and out.tool_calls == []

340 # partial reasoning must still be recovered by the reasoning parser

341 reasoning, _ = MuseGlimmerReasoningParser.extract_reasoning(

326 - R, RAW_TRUNCATED, _FakeReq()

342 + R, RAW_TRUNCATED, _fake_req()

343 )

344 assert reasoning and "Maybe I should call" in reasoning, repr(reasoning)

345

@@ -374,7 +390,7 @@ def test_trailing_segment_ambiguous_left_alone():

390

391 def test_no_registered_tools_passthrough():

392 out = MuseGlimmerToolParser.extract_tool_calls(

377 - T, _call("get_weather.get_weather"), None

393 + T, _call("get_weather.get_weather"), _NO_REQ

394 )

395 assert out.tool_calls[0].function.name == "get_weather.get_weather"

396

tests/tool_use/test_responses_request_validations.py

+1

@@ -203,6 +203,7 @@ def test_input_image_detail_defaults_or_preserves(detail: str | None) -> None:

203

204 request = ResponsesRequest.model_validate({"input": [item]})

205

206 + assert isinstance(request.input, list)

207 assert request.input[0]["content"][0]["detail"] == (

208 "auto" if detail is None else detail

209 )

tests/tool_use/test_tool_choice_required.py

+5 -6

@@ -306,18 +306,17 @@ def _collect_required_tool_streaming_json(output_json: str, delta_len: int) -> s

306

307 combined_messages = "["

308 for message in messages:

309 - if message.tool_calls[0].function.name:

309 + fn = message.tool_calls[0].function

310 + assert fn is not None

311 + if fn.name:

312 if len(combined_messages) > 1:

313 combined_messages += "},"

314

315 combined_messages += (

314 - '{"name": "'

315 - + message.tool_calls[0].function.name

316 - + '", "parameters": '

317 - + message.tool_calls[0].function.arguments

316 + '{"name": "' + fn.name + '", "parameters": ' + (fn.arguments or "")

317 )

318 else:

320 - combined_messages += message.tool_calls[0].function.arguments

319 + combined_messages += fn.arguments or ""

320 combined_messages += "}]"

321 return combined_messages

322

tests/utils_/test_argparse_utils.py

+9 -9

@@ -476,32 +476,32 @@ def test_compilation_config_mode_validator():

476 """Test that CompilationConfig.mode field validator converts strings to integers."""

477 from vllm.config.compilation import CompilationConfig, CompilationMode

478

479 - config = CompilationConfig(mode=0)

479 + config = CompilationConfig(mode=0) # type: ignore[arg-type]

480 assert config.mode == CompilationMode.NONE

481

482 - config = CompilationConfig(mode=3)

482 + config = CompilationConfig(mode=3) # type: ignore[arg-type]

483 assert config.mode == CompilationMode.VLLM_COMPILE

484

485 - config = CompilationConfig(mode="NONE")

485 + config = CompilationConfig(mode="NONE") # type: ignore[arg-type]

486 assert config.mode == CompilationMode.NONE

487

488 - config = CompilationConfig(mode="STOCK_TORCH_COMPILE")

488 + config = CompilationConfig(mode="STOCK_TORCH_COMPILE") # type: ignore[arg-type]

489 assert config.mode == CompilationMode.STOCK_TORCH_COMPILE

490

491 - config = CompilationConfig(mode="DYNAMO_TRACE_ONCE")

491 + config = CompilationConfig(mode="DYNAMO_TRACE_ONCE") # type: ignore[arg-type]

492 assert config.mode == CompilationMode.DYNAMO_TRACE_ONCE

493

494 - config = CompilationConfig(mode="VLLM_COMPILE")

494 + config = CompilationConfig(mode="VLLM_COMPILE") # type: ignore[arg-type]

495 assert config.mode == CompilationMode.VLLM_COMPILE

496

497 - config = CompilationConfig(mode="none")

497 + config = CompilationConfig(mode="none") # type: ignore[arg-type]

498 assert config.mode == CompilationMode.NONE

499

500 - config = CompilationConfig(mode="vllm_compile")

500 + config = CompilationConfig(mode="vllm_compile") # type: ignore[arg-type]

501 assert config.mode == CompilationMode.VLLM_COMPILE

502

503 with pytest.raises(ValidationError, match="Invalid compilation mode"):

504 - CompilationConfig(mode="INVALID_MODE")

504 + CompilationConfig(mode="INVALID_MODE") # type: ignore[arg-type]

505

506

507 def test_flat_product():

tests/utils_/test_cache.py

+1 -1

@@ -128,7 +128,7 @@ def test_lru_cache():

128

129

130 def test_lru_cache_put_if_fits():

131 - cache = LRUCache(10, getsizeof=lambda x: x)

131 + cache = LRUCache[str, int](10, getsizeof=lambda x: x)

132

133 assert cache.put_if_fits("ok", 4) is True

134 assert cache["ok"] == 4

tests/utils_/test_network_utils.py

+2 -2

@@ -206,8 +206,8 @@ def test_make_zmq_socket_ipv6():

206

207

208 def test_make_zmq_path():

209 - assert make_zmq_path("tcp", "127.0.0.1", "5555") == "tcp://127.0.0.1:5555"

210 - assert make_zmq_path("tcp", "::1", "5555") == "tcp://[::1]:5555"

209 + assert make_zmq_path("tcp", "127.0.0.1", 5555) == "tcp://127.0.0.1:5555"

210 + assert make_zmq_path("tcp", "::1", 5555) == "tcp://[::1]:5555"

211

212

213 def test_get_tcp_uri():

tests/v1/cudagraph/test_cudagraph_dispatch.py

+3 -3

@@ -257,7 +257,7 @@ class TestCudagraphDispatcher:

257 def test_get_capture_descs_empty_when_not_initialized(self):

258 """Test that get_capture_descs returns empty list when keys not initialized."""

259 comp_config = CompilationConfig(

260 - cudagraph_mode="FULL",

260 + cudagraph_mode=CUDAGraphMode.FULL,

261 mode=CompilationMode.NONE,

262 cudagraph_capture_sizes=[1, 8],

263 )

@@ -419,7 +419,7 @@ def _run_and_monitor_call(

419 def test_capture_replay_bypass_logic():

420 comp_config = CompilationConfig(

421 mode=CompilationMode.VLLM_COMPILE,

422 - cudagraph_mode="FULL",

422 + cudagraph_mode=CUDAGraphMode.FULL,

423 cudagraph_capture_sizes=[1, 2],

424 )

425 vllm_config = _create_vllm_config(comp_config)

@@ -489,7 +489,7 @@ def test_nested_wrappers():

489 """Tests a scenario with a PIECEWISE wrapper inside a FULL one."""

490 comp_config = CompilationConfig(

491 mode=CompilationMode.VLLM_COMPILE,

492 - cudagraph_mode="FULL",

492 + cudagraph_mode=CUDAGraphMode.FULL,

493 cudagraph_capture_sizes=[1],

494 )

495 vllm_config = _create_vllm_config(comp_config)

tests/v1/cudagraph/test_cudagraph_manager.py

+5 -4

@@ -38,7 +38,7 @@ def _reset_graph_pool_id():

38

39 def _create_vllm_config() -> MagicMock:

40 compilation_config = CompilationConfig(

41 - cudagraph_mode="FULL",

41 + cudagraph_mode=CUDAGraphMode.FULL,

42 cudagraph_capture_sizes=[4],

43 )

44 compilation_config.max_cudagraph_capture_size = 4

@@ -134,7 +134,7 @@ def test_piecewise_capture_uses_pcp_dummy_slot_mappings():

134 device=torch.device("cpu"),

135 max_num_reqs=num_reqs,

136 max_num_tokens=num_tokens,

137 - block_tables=pcp_block_tables,

137 + block_tables=pcp_block_tables, # type: ignore[arg-type]

138 )

139

140 block_tables = MagicMock()

@@ -224,7 +224,7 @@ def test_speculator_capture_preserves_decode_query_bounds(

224 )

225 manager = SpeculatorCudaGraphManager.__new__(SpeculatorCudaGraphManager)

226 manager.max_num_reqs = num_reqs

227 - manager.vllm_config = SimpleNamespace(

227 + manager.vllm_config = SimpleNamespace( # type: ignore[assignment]

228 parallel_config=SimpleNamespace(data_parallel_size=1)

229 )

230 buffers = InputBuffers(num_reqs, num_tokens, torch.device("cpu"))

@@ -264,7 +264,7 @@ def _create_decode_vllm_config(

264 use_kda_recoverssm: bool = False,

265 ) -> MagicMock:

266 compilation_config = CompilationConfig(

267 - cudagraph_mode="FULL_AND_PIECEWISE",

267 + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,

268 cudagraph_capture_sizes=capture_sizes,

269 )

270 compilation_config.max_cudagraph_capture_size = capture_sizes[-1]

@@ -473,6 +473,7 @@ def test_dynamic_spec_decode_shared_token_count_stays_reachable(monkeypatch):

473 assert any(len(descs) > 1 for descs in by_num_tokens.values())

474

475 for desc in full_descs:

476 + assert desc.num_reqs is not None

477 assert desc in manager._candidates[(desc.num_tokens, 0)], desc

478 assert (

479 manager.dispatch(

tests/v1/cudagraph/test_encoder_cudagraph.py

+57 -7

@@ -14,7 +14,7 @@ Test organization:

14

15 from collections.abc import Hashable

16 from types import SimpleNamespace

17 -from typing import Any

17 +from typing import Any, cast

18 from unittest.mock import MagicMock, patch

19

20 import pytest

@@ -100,6 +100,50 @@ class _MockModel(SupportsEncoderCudaGraph):

100 def get_encoder_cudagraph_budget_range(self, vllm_config):

101 return (self._min_budget, self._max_budget)

102

103 + def get_max_frames_per_video(self) -> int:

104 + return 0

105 +

106 + def get_encoder_cudagraph_item_specs(

107 + self, mm_kwargs: dict[str, Any]

108 + ) -> list[EncoderItemSpec]:

109 + return []

110 +

111 + def prepare_encoder_cudagraph_capture_inputs(

112 + self,

113 + token_budget: int,

114 + max_batch_size: int,

115 + max_frames_per_batch: int,

116 + device: torch.device,

117 + dtype: torch.dtype,

118 + path: str = "default",

119 + axis_keys: tuple[Hashable, ...] | None = None,

120 + ) -> EncoderCudaGraphCaptureInputs:

121 + return EncoderCudaGraphCaptureInputs(values={})

122 +

123 + def prepare_encoder_cudagraph_replay_buffers(

124 + self,

125 + mm_kwargs: dict[str, Any],

126 + max_batch_size: int,

127 + max_frames_per_batch: int,

128 + path: str = "default",

129 + ) -> EncoderCudaGraphReplayBuffers:

130 + return EncoderCudaGraphReplayBuffers(values={})

131 +

132 + def encoder_cudagraph_forward(

133 + self, inputs: dict[str, torch.Tensor], path: str = "default"

134 + ) -> torch.Tensor:

135 + return torch.zeros(0)

136 +

137 + def encoder_eager_forward(

138 + self, mm_kwargs: dict[str, Any], path: str = "default"

139 + ) -> torch.Tensor:

140 + return torch.zeros(0)

141 +

142 + def select_encoder_cudagraph_items(

143 + self, mm_kwargs: dict[str, Any], indices: list[int]

144 + ) -> dict[str, Any]:

145 + return {}

146 +

147

148 def _make_manager_with_budgets(budgets: list[int]) -> EncoderCudaGraphManager:

149 """Create a minimal EncoderCudaGraphManager with only token_budgets set.

@@ -594,6 +638,7 @@ def test_eonly_capture_preserves_outputs_across_replay_and_fallback(profile_only

638 ECConnectorBase,

639 ECConnectorMetadata,

640 )

641 + from vllm.v1.core.sched.output import SchedulerOutput

642 from vllm.v1.worker.gpu.ec_connector import ActiveECConnector

643 from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner

644 from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner

@@ -607,7 +652,7 @@ def test_eonly_capture_preserves_outputs_across_replay_and_fallback(profile_only

652 encoder.device = device

653 encoder.cudagraph_manager = manager

654 runner = object.__new__(MMEncoderModelRunner)

610 - runner.model_state = SimpleNamespace(encoder_runner=encoder)

655 + runner.model_state = SimpleNamespace(encoder_runner=encoder) # type: ignore[assignment]

656 # No decoder manager is installed: capture must be encoder-only.

657 with patch(

658 "vllm.v1.worker.mm_encoder_model_runner.lock_workspace", wraps=lock_workspace

@@ -628,9 +673,12 @@ def test_eonly_capture_preserves_outputs_across_replay_and_fallback(profile_only

673 with patch(

674 "vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=connector

675 ):

631 - ec = ActiveECConnector(SimpleNamespace(), cache)

632 - scheduled = SimpleNamespace(

633 - ec_connector_metadata=ECConnectorMetadata(), finished_req_ids=frozenset()

676 + ec = ActiveECConnector(SimpleNamespace(), cache) # type: ignore[arg-type]

677 + scheduled = cast(

678 + SchedulerOutput,

679 + SimpleNamespace(

680 + ec_connector_metadata=ECConnectorMetadata(), finished_req_ids=frozenset()

681 + ),

682 )

683 saved_outputs: list[tuple[torch.Tensor, torch.Tensor]] = []

684 # Mixed sizes exercise packing order; 64 is the boundary, 81 falls back.

@@ -643,6 +691,7 @@ def test_eonly_capture_preserves_outputs_across_replay_and_fallback(profile_only

691 outputs = manager.execute(inputs)

692 assert outputs is not None

693 cache[str(len(cache))] = outputs[0]

694 + assert ec_output is not None

695 assert ec_output.finished_sending is None

696 assert outputs is not None

697 for actual, eager in zip(outputs, expected):

@@ -656,6 +705,7 @@ def test_eonly_capture_preserves_outputs_across_replay_and_fallback(profile_only

705 connector.get_finished.return_value = (set(cache), None)

706 with ec.maybe_get_output(scheduled) as ec_output:

707 pass

708 + assert ec_output is not None

709 assert ec_output.finished_sending == set(cache)

710

711

@@ -666,7 +716,7 @@ def test_eonly_without_encoder_graph_skips_capture():

716 encoder = object.__new__(EncoderRunner)

717 encoder.cudagraph_manager = None

718 runner = object.__new__(MMEncoderModelRunner)

669 - runner.model_state = SimpleNamespace(encoder_runner=encoder)

719 + runner.model_state = SimpleNamespace(encoder_runner=encoder) # type: ignore[assignment]

720 assert runner.capture_model() == 0

721

722

@@ -973,7 +1023,7 @@ class TestInitInvariantValidation:

1023 vllm_config = _MockVllmConfig(token_budgets, max_mm_items)

1024 model = _MockModel(min_budget, max_budget)

1025 return EncoderCudaGraphManager(

976 - vllm_config=vllm_config,

1026 + vllm_config=vllm_config, # type: ignore[arg-type]

1027 device=torch.device("cpu"),

1028 dtype=torch.float32,

1029 model=model,

tests/v1/determinism/test_batch_invariance.py

+5 -12

@@ -155,7 +155,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(

155 # Ensure engines are shutdown to free GPU/VRAM across test sessions

156 if llm is not None:

157 with contextlib.suppress(Exception):

158 - llm.shutdown()

158 + llm.llm_engine.engine_core.shutdown()

159

160

161 @skip_unsupported

@@ -456,7 +456,7 @@ def test_logprobs_bitwise_batch_invariance_ragged_chunked_prefill(backend):

456 )

457 finally:

458 with contextlib.suppress(Exception):

459 - llm.shutdown()

459 + llm.llm_engine.engine_core.shutdown()

460

461

462 @skip_unsupported

@@ -505,7 +505,7 @@ def test_simple_generation(backend):

505

506 finally:

507 with contextlib.suppress(Exception):

508 - llm.shutdown()

508 + llm.llm_engine.engine_core.shutdown()

509

510

511 @skip_unsupported

@@ -815,15 +815,8 @@ def test_decode_logprobs_match_prefill_logprobs(

815 if token_idx == 0:

816 prefix_prompt = prompt

817 else:

818 - # Use the partial output text up to this token

819 - # We'll need to construct this from the full output

820 - prefix_output = decode_output.outputs[0]

821 - # Get the text for tokens 0 to token_idx-1

822 - # Unfortunately, we don't have per-token text, so we'll use

823 - # a different approach: run prefill with prompt + tokens[0:token_idx]

824 -

825 - # Actually, we need to get the actual text. Let's use a workaround:

826 - # Run a generation with max_tokens = token_idx to get that prefix

818 + # No per-token text, so regenerate with max_tokens = token_idx

819 + # to get the prefix text

820 prefix_sp = SamplingParams(

821 temperature=0.0,

822 max_tokens=token_idx,

tests/v1/determinism/test_batch_invariance_vlm.py

+5 -4

@@ -12,6 +12,7 @@ from PIL import Image

12 from utils import _extract_step_logprobs, skip_if_not_cuda

13

14 from vllm import LLM, SamplingParams

15 +from vllm.inputs import TextPrompt

16

17 VLM_TEST_MODEL = os.getenv("VLLM_VLM_TEST_MODEL", "Qwen/Qwen3-VL-2B-Instruct")

18

@@ -30,8 +31,8 @@ def _make_image(seed: int, size: int) -> Image.Image:

31 return Image.fromarray(rng.integers(0, 255, (size, size, 3), dtype=np.uint8))

32

33

33 -def _make_inputs(input_type: str, num_reqs: int) -> list[dict]:

34 - inputs = []

34 +def _make_inputs(input_type: str, num_reqs: int) -> list[TextPrompt]:

35 + inputs: list[TextPrompt] = []

36 for i in range(num_reqs):

37 if input_type == "image":

38 prompt = IMAGE_PROMPT

@@ -54,7 +55,7 @@ def _make_inputs(input_type: str, num_reqs: int) -> list[dict]:

55 return inputs

56

57

57 -def _assert_batch_invariant(llm: LLM, inputs: list[dict], sampling) -> None:

58 +def _assert_batch_invariant(llm: LLM, inputs: list[TextPrompt], sampling) -> None:

59 bs1 = [

60 _extract_step_logprobs(llm.generate([inp], sampling, use_tqdm=False)[0])

61 for inp in inputs

@@ -134,4 +135,4 @@ def _run_vlm_batch_invariance(

135 _assert_batch_invariant(llm, inputs, sampling)

136 finally:

137 with contextlib.suppress(Exception):

137 - llm.shutdown()

138 + llm.llm_engine.engine_core.shutdown()

tests/v1/determinism/test_cutlass_batch_invariance.py

+8 -6

@@ -221,13 +221,15 @@ def _make_cutlass_fp4_moe_batch_invariant_case(

221 in_dtype=dtype,

222 activation=activation,

223 )

224 + prepare_finalize = maybe_make_prepare_finalize(

225 + moe=moe_config,

226 + quant_config=quant_config,

227 + allow_new_interface=True,

228 + use_monolithic=False,

229 + )

230 + assert prepare_finalize is not None

231 kernel = mk.FusedMoEKernel(

225 - maybe_make_prepare_finalize(

226 - moe=moe_config,

227 - quant_config=quant_config,

228 - allow_new_interface=True,

229 - use_monolithic=False,

230 - ),

232 + prepare_finalize,

233 CutlassExpertsFp4(

234 moe_config=moe_config,

235 quant_config=quant_config,

tests/v1/determinism/test_nvfp4_batch_invariant.py

+1 -1

@@ -97,4 +97,4 @@ def test_dense_nvfp4_generation_is_deterministic_across_batch_sizes_e2e(backend)

97 finally:

98 if llm is not None:

99 with contextlib.suppress(Exception):

100 - llm.shutdown()

100 + llm.llm_engine.engine_core.shutdown()

tests/v1/determinism/test_rms_norm_batch_invariant.py

+2

@@ -524,6 +524,7 @@ def test_rms_norm_batch_invariance(dtype):

524 batch[4] = row[0]

525 out_batch = rms_norm_batch_invariant(batch, weight, eps=eps)

526

527 + assert isinstance(out_single, torch.Tensor) and isinstance(out_batch, torch.Tensor)

528 assert torch.equal(out_single[0], out_batch[4]), (

529 "rms_norm output for a row differs when batch context changes"

530 )

@@ -547,6 +548,7 @@ if __name__ == "__main__":

548

549 # Batch-invariant implementation

550 triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)

551 + assert isinstance(triton_output, torch.Tensor)

552

553 # Compare

554 max_diff = (triton_output - reference_output).abs().max().item()

tests/v1/executor/test_executor.py

+5 -5

@@ -85,12 +85,12 @@ def test_multiproc_executor_worker_termination_timeout(

85 monkeypatch.setattr(multiproc_executor_module.time, "sleep", clock.sleep)

86 executor = MultiprocExecutor.__new__(MultiprocExecutor)

87 proc = _FakeProcess(clock, exits_at=exits_at)

88 - executor._ensure_worker_termination([proc])

88 + executor._ensure_worker_termination([proc]) # type: ignore[list-item]

89 assert proc.terminate_called is expected_terminate

90

91

92 class CustomMultiprocExecutor(MultiprocExecutor):

93 - def collective_rpc(

93 + def collective_rpc( # type: ignore[override]

94 self,

95 method: str | Callable,

96 timeout: float | None = None,

@@ -98,7 +98,7 @@ class CustomMultiprocExecutor(MultiprocExecutor):

98 kwargs: dict | None = None,

99 non_block: bool = False,

100 unique_reply_rank: int | None = None,

101 - kv_output_aggregator: KVOutputAggregator = None,

101 + kv_output_aggregator: KVOutputAggregator | None = None,

102 ec_output_aggregator: ECOutputAggregator | None = None,

103 ) -> Any | list[Any] | Future[Any | list[Any]]:

104 # Drop marker to show that this was run

@@ -126,7 +126,7 @@ def test_custom_executor_type_checking():

126 model=MODEL,

127 gpu_memory_utilization=0.2,

128 max_model_len=8192,

129 - distributed_executor_backend=Mock,

129 + distributed_executor_backend=Mock, # type: ignore[arg-type]

130 )

131 LLMEngine.from_engine_args(engine_args)

132 with pytest.raises(ValueError):

@@ -134,7 +134,7 @@ def test_custom_executor_type_checking():

134 model=MODEL,

135 gpu_memory_utilization=0.2,

136 max_model_len=8192,

137 - distributed_executor_backend=Mock,

137 + distributed_executor_backend=Mock, # type: ignore[arg-type]

138 )

139 AsyncLLM.from_engine_args(engine_args)

140

tests/v1/simple_kv_offload/test_integration.py

+1

@@ -55,6 +55,7 @@ def _flush_gpu_cache(llm: LLM, sampling_params: SamplingParams, seed: int = 0):

55 block_size = cache_config.block_size

56 # Use 1.2x GPU capacity to give the lazy cursor enough scheduling steps

57 # to walk past all target blocks near the tail of the free queue.

58 + assert num_gpu_blocks is not None

59 total_tokens_needed = int(num_gpu_blocks * block_size * 1.5)

60

61 # Use token-id prompts so each filler is unique (no prefix sharing).

tests/v1/simple_kv_offload/test_kv_events.py

+5 -1

@@ -260,6 +260,7 @@ def test_block_stored_per_group_metadata_full_attention() -> None:

260 )

261 assert ev.block_hashes == [expected_hash_0]

262 assert ev.parent_block_hash is None

263 + assert req.prompt_token_ids is not None

264 assert ev.token_ids == req.prompt_token_ids[0:BLOCK_SIZE]

265

266 # Second block's parent_block_hash equals the first block's hash.

@@ -296,7 +297,7 @@ def test_eager_store_lora_metadata() -> None:

297 f"expected lora_name='test-lora', got {ev.lora_name!r}"

298 )

299 assert ev.extra_keys is not None, "expected extra_keys for lora request"

299 - assert any("test-lora" in keys for keys in ev.extra_keys), (

300 + assert any(keys is not None and "test-lora" in keys for keys in ev.extra_keys), (

301 f"expected 'test-lora' in extra_keys, got {ev.extra_keys}"

302 )

303 assert len(ev.token_ids) == BLOCK_SIZE

@@ -372,6 +373,7 @@ def test_finished_eager_store_emits_all_storage_events() -> None:

373 secondary_event_hash = maybe_convert_block_hash(get_block_hash(secondary_hash))

374 by_hash = {event.block_hashes[0]: event for event in stored}

375 assert set(by_hash) == {*primary_event_hashes, secondary_event_hash}

376 + assert req.prompt_token_ids is not None

377 assert (

378 by_hash[primary_event_hashes[0]].token_ids == req.prompt_token_ids[:BLOCK_SIZE]

379 )

@@ -451,6 +453,7 @@ def test_secondary_hash_block_stored_metadata() -> None:

453 by_hash = {event.block_hashes[0]: event for event in stored}

454 assert set(by_hash) == {primary_event_hash, secondary_event_hash}

455 assert by_hash[primary_event_hash].block_size == BLOCK_SIZE

456 + assert req.prompt_token_ids is not None

457 assert by_hash[primary_event_hash].token_ids == req.prompt_token_ids[:BLOCK_SIZE]

458 assert by_hash[secondary_event_hash].block_size == BLOCK_SIZE // 2

459 assert (

@@ -752,6 +755,7 @@ def test_mamba_align_skips_positional_event_metadata() -> None:

755 stored = [e for e in events if isinstance(e, BlockStored)]

756 assert len(stored) == 2

757 assert {ev.group_idx for ev in stored} == {0}

758 + assert req.prompt_token_ids is not None

759 for idx, ev in enumerate(stored):

760 assert ev.block_size == vbs

761 assert ev.token_ids == req.prompt_token_ids[idx * vbs : (idx + 1) * vbs]

tests/v1/simple_kv_offload/test_scheduler.py

+5 -4

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

19 SchedulerConfig,

20 VllmConfig,

21 )

22 +from vllm.config.cache import MambaCacheMode

23 from vllm.config.kv_events import KVEventsConfig

24 from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import (

25 SimpleCPUOffloadConnector,

@@ -1243,7 +1244,7 @@ def test_preemption_no_cpu_block_leak() -> None:

1244 block_hasher=req._block_hasher,

1245 )

1246 hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)

1246 - assert hit_tokens > 0

1247 + assert hit_tokens is not None and hit_tokens > 0

1248

1249 gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks)

1250 kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))

@@ -1338,7 +1339,7 @@ def test_inflight_finish_deferred_cleanup() -> None:

1339 block_hasher=req._block_hasher,

1340 )

1341 hit_tokens, _ = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)

1341 - assert hit_tokens > 0

1342 + assert hit_tokens is not None and hit_tokens > 0

1343

1344 gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks)

1345 kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))

@@ -1896,7 +1897,7 @@ def test_reset_pending_loads() -> None:

1897 block_hasher=req._block_hasher,

1898 )

1899 hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)

1899 - assert hit_tokens > 0

1900 + assert hit_tokens is not None and hit_tokens > 0

1901

1902 gpu_blocks2 = gpu_pool.get_new_blocks(num_blocks)

1903 kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))

@@ -2246,7 +2247,7 @@ def _make_hybrid_attention_mamba_scheduler(

2247 hash_block_size: int | None = None,

2248 dcp_world_size: int = 4,

2249 lazy: bool = False,

2249 - mamba_cache_mode: str = "align",

2250 + mamba_cache_mode: MambaCacheMode = "align",

2251 enable_kv_cache_events: bool = False,

2252 ) -> SchedulerFixture:

2253 """Build a scheduler for one attention group plus one Mamba group."""

tests/v1/simple_kv_offload/test_worker.py

+25 -18

@@ -36,6 +36,7 @@ from vllm.v1.core.kv_cache_utils import (

36 )

37 from vllm.v1.kv_cache_interface import (

38 FullAttentionSpec,

39 + KVCacheConfig,

40 KVCacheGroupSpec,

41 KVCacheLayout,

42 KVCacheSpec,

@@ -50,7 +51,10 @@ from vllm.v1.simple_kv_offload.cuda_mem_ops import (

51 build_params,

52 )

53 from vllm.v1.simple_kv_offload.disk_backend import DiskBackend

53 -from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadMetadata

54 +from vllm.v1.simple_kv_offload.metadata import (

55 + SimpleCPUOffloadMetadata,

56 + SimpleCPUOffloadWorkerMetadata,

57 +)

58 from vllm.v1.simple_kv_offload.worker import SimpleCPUOffloadWorker

59 from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector

60 from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin

@@ -64,6 +68,12 @@ ITERS = 30

68 SLEEP_CYCLES = 50_000_000

69

70

71 +def _make_worker(

72 + kv_cache_config: KVCacheConfig | None, cpu_capacity_bytes: int

73 +) -> SimpleCPUOffloadWorker:

74 + return SimpleCPUOffloadWorker(None, kv_cache_config, cpu_capacity_bytes) # type: ignore[arg-type]

75 +

76 +

77 def _make_backend() -> tuple[DmaCopyBackend, torch.Tensor, torch.Tensor]:

78 gpu = {"k": torch.zeros((NUM_BLOCKS, BLOCK_BYTES), dtype=torch.int8, device="cuda")}

79 cpu = {

@@ -87,7 +97,7 @@ def _make_backend() -> tuple[DmaCopyBackend, torch.Tensor, torch.Tensor]:

97 def test_no_forward_step_completes_cpu_store(monkeypatch, use_v2):

98 """A finished request's final store must drain without another model step."""

99 backend, gpu, cpu = _make_backend()

90 - worker = SimpleCPUOffloadWorker(None, None, cpu_capacity_bytes=0)

100 + worker = _make_worker(None, 0)

101 worker._backend = backend

102 connector = SimpleCPUOffloadConnector.__new__(SimpleCPUOffloadConnector)

103 connector.worker_handler = worker

@@ -105,13 +115,13 @@ def test_no_forward_step_completes_cpu_store(monkeypatch, use_v2):

115 active = ActiveKVConnector.__new__(ActiveKVConnector)

116 active.kv_connector = connector

117 active._disabled = False

108 - active._pending_load_start = False

118 + active._pending_load_kwargs = None

119 module = "vllm.v1.worker.gpu.kv_connector"

120 monkeypatch.setattr(f"{module}.is_forward_context_available", lambda: True)

121 monkeypatch.setattr(f"{module}.get_forward_context", lambda: None)

122 result = active.no_forward(output)

123 else:

114 - result = KVConnectorModelRunnerMixin.kv_connector_no_forward(output, None)

124 + result = KVConnectorModelRunnerMixin.kv_connector_no_forward(output, None) # type: ignore[arg-type]

125 completion = (

126 result.kv_connector_output.kv_connector_worker_meta

127 if result.kv_connector_output is not None

@@ -123,6 +133,7 @@ def test_no_forward_step_completes_cpu_store(monkeypatch, use_v2):

133 completion = connector.build_connector_worker_meta()

134 time.sleep(0.001)

135 assert completion is not None, "CPU store never completed on the empty step"

136 + assert isinstance(completion, SimpleCPUOffloadWorkerMetadata)

137 assert completion.completed_store_events == {0: 1}

138 assert torch.equal(cpu[0], gpu[0].cpu())

139 finally:

@@ -132,7 +143,7 @@ def test_no_forward_step_completes_cpu_store(monkeypatch, use_v2):

143 def test_deferred_store_waits_for_draft_forward(monkeypatch):

144 """Target completion must not submit the store before draft finalization."""

145 backend, gpu, cpu = _make_backend()

135 - worker = SimpleCPUOffloadWorker(None, None, cpu_capacity_bytes=0)

146 + worker = _make_worker(None, 0)

147 worker._backend = backend

148 connector = SimpleCPUOffloadConnector.__new__(SimpleCPUOffloadConnector)

149 connector.worker_handler = worker

@@ -267,11 +278,12 @@ class _RecordingBackend:

278

279 def test_transfer_hooks_pass_wait_event_for_store_only():

280 """wait_for_save gates stores on a compute-done event; start_load_kv does not."""

270 - worker = SimpleCPUOffloadWorker(

271 - vllm_config=None, kv_cache_config=None, cpu_capacity_bytes=0

281 + worker = _make_worker(

282 + kv_cache_config=None,

283 + cpu_capacity_bytes=0,

284 )

285 recording = _RecordingBackend()

274 - worker._backend = recording

286 + worker._backend = recording # type: ignore[assignment]

287 worker._connector_metadata = SimpleCPUOffloadMetadata(

288 load_event=0,

289 load_gpu_blocks=[0],

@@ -332,8 +344,7 @@ def test_register_shared_kv_cache_storage(monkeypatch, layout: KVCacheLayout):

344 )

345 ],

346 )

335 - worker = SimpleCPUOffloadWorker(

336 - vllm_config=None,

347 + worker = _make_worker(

348 kv_cache_config=cache_config,

349 cpu_capacity_bytes=raw.nbytes,

350 )

@@ -366,8 +377,7 @@ def test_register_kv_cache_storage_with_trailing_padding(monkeypatch):

377 cache_bytes = num_blocks * block_bytes

378 raw = torch.zeros(4096, dtype=torch.int8, device="cuda")

379 cache = raw[:cache_bytes].view(num_blocks, block_bytes)

369 - worker = SimpleCPUOffloadWorker(

370 - vllm_config=None,

380 + worker = _make_worker(

381 kv_cache_config=MagicMock(

382 num_blocks=num_blocks,

383 kv_cache_tensors=[

@@ -412,8 +422,7 @@ def test_register_separate_kv_head_groups(monkeypatch):

422 )

423 caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)

424 layer_names = [f"layer.{i}" for i in range(num_layers)]

415 - worker = SimpleCPUOffloadWorker(

416 - vllm_config=None,

425 + worker = _make_worker(

426 kv_cache_config=MagicMock(

427 num_blocks=num_blocks,

428 kv_cache_tensors=[

@@ -509,8 +518,7 @@ def test_register_mixed_page_sizes_in_one_cache_group(monkeypatch):

518 assert len(set(pages)) > 1, "the mixed page sizes are what this test covers"

519

520 kv_caches = allocate_kv_cache(kv_cache_config, torch.device("cuda"), layout)

512 - worker = SimpleCPUOffloadWorker(

513 - vllm_config=None,

521 + worker = _make_worker(

522 kv_cache_config=kv_cache_config,

523 cpu_capacity_bytes=sum(pages) * num_blocks,

524 )

@@ -567,8 +575,7 @@ def test_register_mixed_page_sizes_odd_block_counts(monkeypatch, rank_blocks):

575 assert kv_cache_config.num_blocks == rank_blocks

576

577 kv_caches = allocate_kv_cache(kv_cache_config, torch.device("cuda"), layout)

570 - worker = SimpleCPUOffloadWorker(

571 - vllm_config=None,

578 + worker = _make_worker(

579 kv_cache_config=kv_cache_config,

580 cpu_capacity_bytes=sum(pages) * rank_blocks,

581 )

tests/v1/structured_output/test_backend_guidance.py

+13 -8

@@ -91,7 +91,7 @@ def test_grammar_bitmask_with_specdec():

91 json='{"type": "object"}',

92 ),

93 )

94 - sampling_params.structured_outputs._backend = "guidance"

94 + sampling_params.structured_outputs._backend = "guidance" # type: ignore[union-attr]

95 sampling_params.update_from_generation_config({}, tokenizer.eos_token_id)

96

97 my_req_id = f"my_req_id_{i}"

@@ -107,19 +107,23 @@ def test_grammar_bitmask_with_specdec():

107 def grammar_bitmask(req: Request, tokens: list[int]) -> None:

108 structured_output_manager.grammar_bitmask(

109 requests={req.request_id: req},

110 - structured_output_request_ids={req.request_id: 0},

110 + structured_output_request_ids=[req.request_id],

111 scheduled_spec_decode_tokens={req.request_id: tokens},

112 )

113 # At this point, we rolled-back, so should not be terminated

114 - assert not req.structured_output_request.grammar.is_terminated()

114 + assert req.structured_output_request is not None

115 + grammar = req.structured_output_request.grammar

116 + assert grammar is not None and not isinstance(grammar, Exception)

117 + assert not grammar.is_terminated()

118

119 # The grammar might not yet be compiled, so we wait for it

120 + assert request.structured_output_request is not None

121 while not request.structured_output_request._check_grammar_completion():

122 continue

123

120 - assert request.structured_output_request.grammar.accept_tokens(

121 - request.request_id, prompt[:i]

122 - )

124 + grammar = request.structured_output_request.grammar

125 + assert grammar is not None and not isinstance(grammar, Exception)

126 + assert grammar.accept_tokens(request.request_id, prompt[:i])

127

128 grammar_bitmask(request, prompt[i:] + [tokenizer.eos_token_id])

129 grammar_bitmask(

@@ -155,7 +159,7 @@ def test_grammar_init_async_and_sync(async_grammar):

159 json='{"type": "object"}',

160 ),

161 )

158 - sampling_params.structured_outputs._backend = "guidance"

162 + sampling_params.structured_outputs._backend = "guidance" # type: ignore[union-attr]

163 sampling_params.update_from_generation_config({}, tokenizer.eos_token_id)

164

165 request = Request(

@@ -169,6 +173,7 @@ def test_grammar_init_async_and_sync(async_grammar):

173

174 # Check the internal _grammar type immediately after init

175 # Before _check_grammar_completion is called, async mode should have a Future

176 + assert request.structured_output_request is not None

177 raw_grammar = request.structured_output_request._grammar

178 if async_grammar:

179 assert isinstance(raw_grammar, Future), (

@@ -191,7 +196,7 @@ def test_grammar_init_async_and_sync(async_grammar):

196

197 # Verify grammar is properly initialized and functional

198 grammar = request.structured_output_request.grammar

194 - assert grammar is not None

199 + assert grammar is not None and not isinstance(grammar, Exception)

200 assert not grammar.is_terminated()

201

202 # Verify the grammar can accept valid tokens

tests/v1/structured_output/test_backend_xgrammar_lark.py

+6 -3

@@ -13,6 +13,7 @@ from mistral_common.protocol.instruct.tool_calls import Tool, ToolChoiceEnum

13 from transformers import AutoTokenizer

14

15 from vllm.config import StructuredOutputsConfig, VllmConfig

16 +from vllm.config.structured_outputs import StructuredOutputsBackend

17 from vllm.exceptions import VLLMValidationError

18 from vllm.sampling_params import SamplingParams, StructuredOutputsParams

19 from vllm.tokenizers.mistral import MistralTokenizer

@@ -41,13 +42,15 @@ EBNF_GRAMMAR = 'root ::= "id=" [0-9]+'

42

43

44 def _validate(

44 - grammar: str, backend: str = "xgrammar", tokenizer: object | None = None

45 + grammar: str,

46 + backend: StructuredOutputsBackend = "xgrammar",

47 + tokenizer: object | None = None,

48 ) -> StructuredOutputsParams:

49 params = SamplingParams(structured_outputs=StructuredOutputsParams(grammar=grammar))

50 params._validate_structured_outputs(

48 - _StubModelConfig(),

51 + _StubModelConfig(), # type: ignore[arg-type]

52 StructuredOutputsConfig(backend=backend),

50 - tokenizer=tokenizer or object(),

53 + tokenizer=tokenizer or object(), # type: ignore[arg-type]

54 )

55 assert params.structured_outputs is not None

56 return params.structured_outputs

tests/v1/structured_output/test_backend_xgrammar_plamo3.py

+3 -3

@@ -51,7 +51,7 @@ def make_vllm_config() -> VllmConfig:

51 config = VllmConfig(

52 structured_outputs_config=StructuredOutputsConfig(backend="xgrammar"),

53 )

54 - config.model_config = SimpleNamespace(

54 + config.model_config = SimpleNamespace( # type: ignore[assignment]

55 hf_config=SimpleNamespace(model_type="plamo3")

56 )

57 return config

@@ -73,7 +73,7 @@ def test_plamo3_python_tokenizer_compiles_and_accepts_structural_tag():

73 tokenizer = Plamo3Tokenizer(["<eos>", *call_tokens])

74 backend = XgrammarBackend(

75 make_vllm_config(),

76 - tokenizer=tokenizer,

76 + tokenizer=tokenizer, # type: ignore[arg-type]

77 vocab_size=len(tokenizer.get_vocab()),

78 )

79 tools = [

@@ -107,7 +107,7 @@ def test_plamo3_grammar_matches_decoded_bytes(value):

107 tokenizer = Plamo3Tokenizer(list(dict.fromkeys(["<eos>", value, *pieces])))

108 backend = XgrammarBackend(

109 make_vllm_config(),

110 - tokenizer=tokenizer,

110 + tokenizer=tokenizer, # type: ignore[arg-type]

111 vocab_size=len(tokenizer.get_vocab()),

112 )

113 for tokens in ([value], pieces):

tests/v1/structured_output/test_structured_output_manager.py

+16 -11

@@ -10,9 +10,11 @@ from transformers import AutoTokenizer

10 from vllm.config import DeviceConfig, StructuredOutputsConfig, VllmConfig

11 from vllm.config.model import ModelConfig

12 from vllm.config.speculative import SpeculativeConfig

13 +from vllm.config.structured_outputs import StructuredOutputsBackend

14 from vllm.sampling_params import SamplingParams, StructuredOutputsParams

15 from vllm.v1.request import Request

16 from vllm.v1.structured_output import StructuredOutputManager

17 +from vllm.v1.structured_output.backend_outlines import OutlinesGrammar

18

19 TOKENIZER = "gpt2"

20 THINK_END = "\n" # reasoning-end marker (single GPT-2 token)

@@ -84,7 +86,7 @@ def _wait_for_grammar(request: Request) -> None:

86

87 def _build_harness(

88 tokenizer,

87 - backend: str,

89 + backend: StructuredOutputsBackend,

90 prefix: str = "",

91 use_reasoner: bool = True,

92 reasoning_ended: bool | None = None,

@@ -106,13 +108,13 @@ def _build_harness(

108 )

109 manager = StructuredOutputManager(vllm_config)

110 if use_reasoner:

109 - manager.reasoner_cls = MockReasoner

111 + manager.reasoner_cls = MockReasoner # type: ignore[assignment]

112

113 sampling_params = SamplingParams(

114 structured_outputs=structured_outputs

115 or StructuredOutputsParams(json=JSON_SCHEMA)

116 )

115 - sampling_params.structured_outputs._backend = backend

117 + sampling_params.structured_outputs._backend = backend # type: ignore[union-attr]

118 sampling_params.update_from_generation_config({}, tokenizer.eos_token_id)

119

120 prompt_ids = tokenizer.encode(prefix) if prefix else []

@@ -130,9 +132,10 @@ def _build_harness(

132 structured_req = request.structured_output_request

133 assert request.prompt_token_ids is not None

134 assert structured_req is not None

133 - assert structured_req.grammar is not None

135 + grammar = structured_req.grammar

136 + assert grammar is not None and not isinstance(grammar, Exception)

137 if prompt_ids:

135 - assert structured_req.grammar.accept_tokens(

138 + assert grammar.accept_tokens(

139 request.request_id,

140 prompt_ids,

141 )

@@ -173,7 +176,7 @@ def _run_real_flow(

176 structured_req = request.structured_output_request

177 assert structured_req is not None

178 grammar = structured_req.grammar

176 - assert grammar is not None

179 + assert grammar is not None and not isinstance(grammar, Exception)

180 assert not grammar.is_terminated()

181 # `bitmask[i]` is the grammar state after the first `i` scheduled tokens.

182 # These tests commit the validated prefix, so the matching post-accept

@@ -335,7 +338,7 @@ FLOW_CASES = [

338 @pytest.mark.parametrize("case", FLOW_CASES)

339 def test_real_flow(

340 tokenizer,

338 - backend: str,

341 + backend: StructuredOutputsBackend,

342 case: FlowCase,

343 ):

344 if backend == "guidance" and case.xfail_guidance:

@@ -379,7 +382,7 @@ def test_real_flow(

382 )

383 def test_initial_constraint_activation(

384 tokenizer,

382 - backend: str,

385 + backend: StructuredOutputsBackend,

386 use_reasoner: bool,

387 reasoning_ended: bool | None,

388 enable_in_reasoning: bool,

@@ -431,7 +434,7 @@ def test_initial_constraint_activation(

434 )

435 def test_regex_flow(

436 tokenizer,

434 - backend: str,

437 + backend: StructuredOutputsBackend,

438 raw_drafts: tuple[str, ...],

439 expected_validated: tuple[str, ...] | None,

440 expected_row_pattern: str,

@@ -459,7 +462,9 @@ def test_regex_flow(

462

463

464 @pytest.mark.parametrize("backend", REGEX_BACKENDS)

462 -def test_rejected_draft_keeps_later_rows_constrained(tokenizer, backend: str):

465 +def test_rejected_draft_keeps_later_rows_constrained(

466 + tokenizer, backend: StructuredOutputsBackend

467 +):

468 """A draft the grammar rejects during bitmask fill must not unconstrain

469 the remaining rows or the bonus row."""

470 manager, request = _build_harness(

@@ -492,7 +497,7 @@ def test_outlines_termination(tokenizer):

497 structured_req = request.structured_output_request

498 assert structured_req is not None

499 grammar = structured_req.grammar

495 - assert grammar is not None

500 + assert isinstance(grammar, OutlinesGrammar)

501 one, eos = _to_token_ids(tokenizer, ("1", EOS))

502

503 assert not grammar.accept_tokens(request.request_id, [eos])

tests/v1/structured_output/test_validation.py

+11 -10

@@ -35,7 +35,7 @@ def test_structured_outputs_rejected_for_diffusion_models():

35 )

36 with pytest.raises(VLLMValidationError, match="not yet supported for diffusion"):

37 params._validate_structured_outputs(

38 - _StubModelConfig(is_diffusion=True),

38 + _StubModelConfig(is_diffusion=True), # type: ignore[arg-type]

39 StructuredOutputsConfig(),

40 tokenizer=None,

41 )

@@ -45,7 +45,7 @@ def test_plain_request_allowed_for_diffusion_models():

45 """Requests without structured outputs are unaffected by the guard."""

46 params = SamplingParams()

47 params._validate_structured_outputs(

48 - _StubModelConfig(is_diffusion=True),

48 + _StubModelConfig(is_diffusion=True), # type: ignore[arg-type]

49 StructuredOutputsConfig(),

50 tokenizer=None,

51 )

@@ -71,9 +71,9 @@ def test_degenerate_structured_outputs_rejected(structured_outputs, match):

71 params = SamplingParams(structured_outputs=structured_outputs)

72 with pytest.raises(VLLMValidationError, match=match):

73 params._validate_structured_outputs(

74 - _StubModelConfig(is_diffusion=False),

74 + _StubModelConfig(is_diffusion=False), # type: ignore[arg-type]

75 StructuredOutputsConfig(),

76 - tokenizer=object(),

76 + tokenizer=object(), # type: ignore[arg-type]

77 )

78

79

@@ -97,9 +97,9 @@ def test_regex_with_nul_byte_rejected(regex):

97 # (which would otherwise catch the error and fall back to another backend).

98 with pytest.raises(VLLMValidationError, match="NUL"):

99 params._validate_structured_outputs(

100 - _StubModelConfig(is_diffusion=False),

100 + _StubModelConfig(is_diffusion=False), # type: ignore[arg-type]

101 StructuredOutputsConfig(),

102 - tokenizer=object(),

102 + tokenizer=object(), # type: ignore[arg-type]

103 )

104

105 # The xgrammar backend also rejects it directly (defense in depth), before

@@ -133,9 +133,9 @@ def test_unsupported_grammar_is_a_client_error(backend, structured_outputs):

133 params = SamplingParams(structured_outputs=structured_outputs)

134 with pytest.raises(VLLMClientError):

135 params._validate_structured_outputs(

136 - _StubModelConfig(is_diffusion=False),

136 + _StubModelConfig(is_diffusion=False), # type: ignore[arg-type]

137 StructuredOutputsConfig(backend=backend),

138 - tokenizer=object(),

138 + tokenizer=object(), # type: ignore[arg-type]

139 )

140

141

@@ -180,8 +180,9 @@ def test_auto_backend_falls_back_on_unsupported_schema(schema, expected_backend)

180 """`auto` falls back on rejection, so it must catch what the validators raise."""

181 params = SamplingParams(structured_outputs=StructuredOutputsParams(json=schema))

182 params._validate_structured_outputs(

183 - _StubModelConfig(is_diffusion=False),

183 + _StubModelConfig(is_diffusion=False), # type: ignore[arg-type]

184 StructuredOutputsConfig(backend="auto"),

185 - tokenizer=object(),

185 + tokenizer=object(), # type: ignore[arg-type]

186 )

187 + assert params.structured_outputs is not None

188 assert params.structured_outputs._backend == expected_backend

tools/pre_commit/mypy.py

+15 -17

@@ -31,24 +31,37 @@ SILENT_GROUPS = [

31 "tests/compile/correctness_e2e",

32 "tests/compile/fullgraph",

33 "tests/compile/fusions_e2e",

34 + "tests/compile/passes",

35 "tests/config",

36 "tests/entrypoints/generate",

37 + "tests/entrypoints/serve",

38 "tests/entrypoints/tool_parsers",

39 + "tests/entrypoints/unit_tests",

40 "tests/entrypoints/weight_transfer",

41 "tests/kernels/core",

42 "tests/kernels/mamba",

43 + "tests/kernels/quantization",

44 + "tests/model_executor/layers",

45 "tests/models/language",

46 "tests/models/quantization",

47 "tests/plugins/bge_m3_sparse_plugin",

48 "tests/plugins/prithvi_io_processor_plugin",

49 "tests/plugins/vllm_add_dummy_platform",

50 "tests/plugins/vllm_add_dummy_stat_logger",

46 - "tests/plugins_tests/gguf",

47 - "tests/plugins_tests/lora_resolvers",

51 + "tests/plugins_tests",

52 + "tests/samplers",

53 "tests/spec_decode",

54 + "tests/tokenizers_",

55 + "tests/tool_use",

56 "tests/transformers_utils",

57 + "tests/utils_",

58 + "tests/v1/cudagraph",

59 + "tests/v1/determinism",

60 "tests/v1/distributed",

61 + "tests/v1/executor",

62 "tests/v1/shutdown",

63 + "tests/v1/simple_kv_offload",

64 + "tests/v1/structured_output",

65 ]

66

67 # After fixing errors resulting from changing follow_imports

@@ -57,57 +70,42 @@ SEPARATE_GROUPS = [

70 "tests",

71 "tests/benchmarks",

72 "tests/compile",

60 - "tests/compile/passes",

73 "tests/distributed",

74 "tests/entrypoints/anthropic",

75 "tests/entrypoints/llm",

76 "tests/entrypoints/multimodal",

77 "tests/entrypoints/openai",

78 "tests/entrypoints/pooling",

67 - "tests/entrypoints/serve",

79 "tests/entrypoints/speech_to_text",

69 - "tests/entrypoints/unit_tests",

80 "tests/kernels",

81 "tests/kernels/attention",

82 "tests/kernels/helion",

83 "tests/kernels/moe",

74 - "tests/kernels/quantization",

84 "tests/lora",

85 "tests/model_executor",

77 - "tests/model_executor/layers",

86 "tests/model_executor/model_loader",

87 "tests/models",

88 "tests/models/test_initialization.py",

89 "tests/models/multimodal",

90 "tests/multimodal",

91 "tests/parser",

84 - "tests/plugins_tests",

92 "tests/quantization",

93 "tests/reasoning",

94 "tests/renderers",

88 - "tests/samplers",

89 - "tests/tokenizers_",

95 "tests/tool_parsers",

91 - "tests/tool_use",

92 - "tests/utils_",

96 "tests/v1",

97 "tests/v1/attention",

98 "tests/v1/core",

96 - "tests/v1/cudagraph",

97 - "tests/v1/determinism",

99 "tests/v1/e2e",

100 "tests/v1/ec_connector",

101 "tests/v1/engine",

101 - "tests/v1/executor",

102 "tests/v1/kv_connector",

103 "tests/v1/kv_offload",

104 "tests/v1/logits_processors",

105 "tests/v1/metrics",

106 "tests/v1/sample",

107 - "tests/v1/simple_kv_offload",

107 "tests/v1/spec_decode",

108 "tests/v1/streaming_input",

110 - "tests/v1/structured_output",

109 "tests/v1/worker",

110 ]

111

[MyPy][2/N] Fix mypy errors in small tests/ dirs (#51043)