7ef5873752 · tduncan/Nvidia-Optimization-Stack-VLLM

tools/pre_commit/mypy.py

+1 -1

@@ -49,6 +49,7 @@ FILES = [

49 "vllm/v1/metrics",

50 "vllm/v1/pool",

51 "vllm/v1/sample",

52 + "vllm/v1/structured_output",

53 "vllm/v1/worker",

54 ]

55

@@ -64,7 +65,6 @@ SEPARATE_GROUPS = [

65 # v1 related

66 "vllm/v1/kv_offload",

67 "vllm/v1/spec_decode",

67 - "vllm/v1/structured_output",

68 ]

69

70 # TODO(woosuk): Include the code from Megatron and HuggingFace.

vllm/reasoning/abs_reasoning_parsers.py

+2 -2

@@ -51,7 +51,7 @@ class ReasoningParser:

51 return self.model_tokenizer.get_vocab()

52

53 @abstractmethod

54 - def is_reasoning_end(self, input_ids: list[int]) -> bool:

54 + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:

55 """

56 Check if the reasoning content ends in the input_ids.

57

@@ -68,7 +68,7 @@ class ReasoningParser:

68 """

69

70 def is_reasoning_end_streaming(

71 - self, input_ids: list[int], delta_ids: list[int]

71 + self, input_ids: Sequence[int], delta_ids: Sequence[int]

72 ) -> bool:

73 """

74 Check if the reasoning content ends in the input_ids on a

vllm/reasoning/basic_parsers.py

+2 -2

@@ -65,7 +65,7 @@ class BaseThinkingReasoningParser(ReasoningParser):

65 "think start/end tokens in the tokenizer!"

66 )

67

68 - def is_reasoning_end(self, input_ids: list[int]) -> bool:

68 + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:

69 start_token_id = self.start_token_id

70 end_token_id = self.end_token_id

71

@@ -77,7 +77,7 @@ class BaseThinkingReasoningParser(ReasoningParser):

77 return False

78

79 def is_reasoning_end_streaming(

80 - self, input_ids: list[int], delta_ids: list[int]

80 + self, input_ids: Sequence[int], delta_ids: Sequence[int]

81 ) -> bool:

82 end_token_id = self.end_token_id

83 return end_token_id in delta_ids

vllm/reasoning/deepseek_v3_reasoning_parser.py

+1 -1

@@ -41,7 +41,7 @@ class DeepSeekV3ReasoningParser(ReasoningParser):

41 return self._parser.is_reasoning_end(input_ids)

42

43 def is_reasoning_end_streaming(

44 - self, input_ids: list[int], delta_ids: list[int]

44 + self, input_ids: Sequence[int], delta_ids: Sequence[int]

45 ) -> bool:

46 return self._parser.is_reasoning_end_streaming(input_ids, delta_ids)

47

vllm/reasoning/gptoss_reasoning_parser.py

+1 -1

@@ -78,7 +78,7 @@ class GptOssReasoningParser(ReasoningParser):

78 self.reasoning_end_token_ids_suffix = self.model_tokenizer.encode("<|message|>")

79 self.reasoning_max_num_between_tokens = 20

80

81 - def is_reasoning_end(self, input_ids: list[int]) -> bool:

81 + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:

82 end_token_ids_prefix = self.reasoning_end_token_ids_prefix

83 end_token_ids_suffix = self.reasoning_end_token_ids_suffix

84 assert len(end_token_ids_prefix) > 0, "reasoning_end_token_ids_prefix is empty"

vllm/reasoning/holo2_reasoning_parser.py

+1 -1

@@ -61,7 +61,7 @@ class Holo2ReasoningParser(ReasoningParser):

61 return self._parser.is_reasoning_end(input_ids)

62

63 def is_reasoning_end_streaming(

64 - self, input_ids: list[int], delta_ids: list[int]

64 + self, input_ids: Sequence[int], delta_ids: Sequence[int]

65 ) -> bool:

66 return self._parser.is_reasoning_end_streaming(input_ids, delta_ids)

67

vllm/reasoning/hunyuan_a13b_reasoning_parser.py

+1 -1

@@ -79,7 +79,7 @@ class HunyuanA13BReasoningParser(ReasoningParser):

79 self.token_buffer = []

80 self.text_buffer = ""

81

82 - def is_reasoning_end(self, input_ids: list[int]) -> bool:

82 + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:

83 return self.current_state == "response"

84

85 def extract_content_ids(self, input_ids: list[int]) -> list[int]:

vllm/reasoning/identity_reasoning_parser.py

+2 -2

@@ -31,12 +31,12 @@ class IdentityReasoningParser(ReasoningParser):

31 "constructor during construction."

32 )

33

34 - def is_reasoning_end(self, input_ids: list[int]) -> bool:

34 + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:

35 # Always return True, since we never treat reasoning specially

36 return True

37

38 def is_reasoning_end_streaming(

39 - self, input_ids: list[int], delta_ids: list[int]

39 + self, input_ids: Sequence[int], delta_ids: Sequence[int]

40 ) -> bool:

41 return True

42

vllm/reasoning/minimax_m2_reasoning_parser.py

+1 -1

@@ -88,7 +88,7 @@ class MiniMaxM2AppendThinkReasoningParser(ReasoningParser):

88 super().__init__(tokenizer, *args, **kwargs)

89 self.end_token_id = self.vocab.get("</think>")

90

91 - def is_reasoning_end(self, input_ids: list[int]) -> bool:

91 + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:

92 end_token_id = self.end_token_id

93 return any(input_id == end_token_id for input_id in reversed(input_ids))

94

vllm/reasoning/mistral_reasoning_parser.py

+2 -1

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 +from collections.abc import Sequence

5 from functools import cached_property

6

7 from vllm.entrypoints.openai.chat_completion.protocol import (

@@ -65,7 +66,7 @@ class MistralReasoningParser(BaseThinkingReasoningParser):

66

67 return SpecialTokens.end_think

68

68 - def is_reasoning_end(self, input_ids: list[int]) -> bool:

69 + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:

70 has_eot_token = False

71

72 for id in input_ids[::-1]:

vllm/reasoning/olmo3_reasoning_parser.py

+1 -1

@@ -242,7 +242,7 @@ class Olmo3ReasoningParser(ReasoningParser):

242 think_start=self.think_start, think_end=self.think_end

243 )

244

245 - def is_reasoning_end(self, input_ids: list[int]) -> bool:

245 + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:

246 text = self.model_tokenizer.decode(input_ids)

247 return self.think_end in text

248

vllm/reasoning/step3_reasoning_parser.py

+2 -2

@@ -100,11 +100,11 @@ class Step3ReasoningParser(ReasoningParser):

100

101 return reasoning, content

102

103 - def is_reasoning_end(self, input_ids: list[int]) -> bool:

103 + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:

104 return self.think_end_token_id in input_ids

105

106 def is_reasoning_end_streaming(

107 - self, input_ids: list[int], delta_ids: list[int]

107 + self, input_ids: Sequence[int], delta_ids: Sequence[int]

108 ) -> bool:

109 end_token_id = self.think_end_token_id

110 return end_token_id in delta_ids

vllm/v1/attention/backends/fa_utils.py

+6 -4

@@ -1,6 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3

4 +from typing import Any

5 +

6 from vllm.logger import init_logger

7 from vllm.platforms import current_platform

8

@@ -17,15 +19,15 @@ elif current_platform.is_xpu():

19 from vllm._ipex_ops import ipex_ops

20

21 reshape_and_cache_flash = ipex_ops.reshape_and_cache_flash

20 - flash_attn_varlen_func = ipex_ops.flash_attn_varlen_func

21 - get_scheduler_metadata = ipex_ops.get_scheduler_metadata

22 + flash_attn_varlen_func = ipex_ops.flash_attn_varlen_func # type: ignore[assignment]

23 + get_scheduler_metadata = ipex_ops.get_scheduler_metadata # type: ignore[assignment]

24

25 elif current_platform.is_rocm():

26 try:

25 - from flash_attn import flash_attn_varlen_func # noqa: F401

27 + from flash_attn import flash_attn_varlen_func # type: ignore[no-redef]

28 except ImportError:

29

28 - def flash_attn_varlen_func(*args, **kwargs):

30 + def flash_attn_varlen_func(*args: Any, **kwargs: Any) -> Any: # type: ignore[no-redef,misc]

31 raise ImportError(

32 "ROCm platform requires upstream flash-attn "

33 "to be installed. Please install flash-attn first."

vllm/v1/attention/backends/mla/aiter_triton_mla.py

+1 -1

@@ -49,7 +49,7 @@ class AiterTritonMLAImpl(AiterMLAImpl):

49 def _flash_attn_varlen_diff_headdims(

50 self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs

51 ):

52 - result = self.flash_attn_varlen_func(

52 + result = self.flash_attn_varlen_func( # type: ignore[call-arg]

53 q,

54 k,

55 v,

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

+1 -1

@@ -230,7 +230,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

230 def _flash_attn_varlen_diff_headdims(

231 self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs

232 ):

233 - output = self.flash_attn_varlen_func(

233 + output = self.flash_attn_varlen_func( # type: ignore[call-arg]

234 q=q,

235 k=k,

236 v=v,

vllm/v1/structured_output/__init__.py

+3 -2

@@ -294,7 +294,7 @@ class StructuredOutputManager:

294 assert request.structured_output_request is not None

295 if request.structured_output_request.reasoning_ended is None:

296 request.structured_output_request.reasoning_ended = (

297 - self.reasoner.is_reasoning_end(request.prompt_token_ids)

297 + self.reasoner.is_reasoning_end(request.prompt_token_ids or [])

298 )

299 return request.structured_output_request.reasoning_ended

300 return True

@@ -323,8 +323,9 @@ class StructuredOutputManager:

323

324 # Check if reasoning ends in *this* step

325 delta_from = request.num_computed_tokens - request.num_output_placeholders

326 + all_token_ids = request.all_token_ids

327 if self.reasoner.is_reasoning_end_streaming(

327 - request.all_token_ids, request.all_token_ids[delta_from:]

328 + all_token_ids, all_token_ids[delta_from:]

329 ):

330 # Reasoning just ended, so we shouldn't advance til

331 # next pass

vllm/v1/structured_output/backend_guidance.py

+3

@@ -284,6 +284,9 @@ def serialize_guidance_grammar(

284 def validate_guidance_grammar(

285 sampling_params: SamplingParams, tokenizer: llguidance.LLTokenizer | None = None

286 ) -> None:

287 + # if structured output is not enabled, there is nothing to validate

288 + if sampling_params.structured_outputs is None:

289 + return

290 tp, grm = get_structured_output_key(sampling_params.structured_outputs)

291 guidance_grm = serialize_guidance_grammar(tp, grm)

292 err = llguidance.LLMatcher.validate_grammar(guidance_grm, tokenizer)

vllm/v1/structured_output/backend_xgrammar.py

+1 -1

@@ -69,7 +69,7 @@ class XgrammarBackend(StructuredOutputBackend):

69 if idx < vocab_size:

70 encoded_vocab[idx] = token

71 stop_token_ids = [self.tokenizer.eos_token_id]

72 - backend_str = self.tokenizer.tokenizer.backend_tokenizer.to_str()

72 + backend_str = self.tokenizer.tokenizer.backend_tokenizer.to_str() # type: ignore[attr-defined]

73 metadata = xgr.TokenizerInfo._detect_metadata_from_hf(backend_str)

74 tokenizer_info = xgr.TokenizerInfo(

75 encoded_vocab=encoded_vocab,

[CI] Fix mypy for `vllm/v1/structured_output` (#32722)