@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a3205beffb

Fork 0

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

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

Synced Nov 8, 05:00 UTC

[CI] Enable mypy coverage for individual excluded files (#34292)

Signed-off-by: Lucas Kabela <lucaskabela@meta.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Lucas Kabela committed Feb 16, 2026 at 07:34 UTC a3205beffb6b3d2923fd9ad8e1ef8b4fd5f7ed29

9 files changed +60 -27

tools/pre_commit/mypy.py

-4

@@ -36,7 +36,6 @@ SEPARATE_GROUPS = [

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

38 EXCLUDE = [

39 - "vllm/engine/arg_utils.py",

39 "vllm/model_executor/parallel_utils",

40 "vllm/model_executor/models",

41 "vllm/model_executor/layers/fla/ops",

@@ -49,9 +48,6 @@ EXCLUDE = [

48 "vllm/profiler",

49 "vllm/reasoning",

50 "vllm/tool_parser",

52 - "vllm/v1/cudagraph_dispatcher.py",

53 - "vllm/outputs.py",

54 - "vllm/logger.py",

51 ]

vllm/config/cache.py

+1 -1

@@ -39,7 +39,7 @@ KVOffloadingBackend = Literal["native", "lmcache"]

39 class CacheConfig:

40 """Configuration for the KV cache."""

42 - block_size: SkipValidation[BlockSize] = None # type: ignore

42 + block_size: SkipValidation[BlockSize] = None # type: ignore[assignment]

43 """Size of a contiguous cache block in number of tokens. On CUDA devices,

44 only block sizes up to 32 are supported.

vllm/config/parallel.py

+1 -1

@@ -182,7 +182,7 @@ class ParallelConfig:

182 threshold, microbatching will be used. Otherwise, the request will be

183 processed in a single batch."""

184

185 - disable_nccl_for_dp_synchronization: bool = Field(default=None)

185 + disable_nccl_for_dp_synchronization: bool | None = Field(default=None)

186 """Forces the dp synchronization logic in vllm/v1/worker/dp_utils.py

187 to use Gloo instead of NCCL for its all reduce.

188

vllm/config/scheduler.py

+2 -2

@@ -115,7 +115,7 @@ class SchedulerConfig:

115

116 # scheduler class or path. "vllm.v1.core.sched.scheduler.Scheduler"

117 # (default) or "mod.custom_class".

118 - scheduler_cls: str | type[object] = Field(default=None)

118 + scheduler_cls: str | type[object] | None = Field(default=None)

119 """The scheduler class to use. "vllm.v1.core.sched.scheduler.Scheduler" is

120 the default scheduler. Can be a class directly or the path to a class of

121 form "mod.custom_class"."""

@@ -128,7 +128,7 @@ class SchedulerConfig:

128 and starting configuration.

129 """

130

131 - async_scheduling: bool = Field(default=None)

131 + async_scheduling: bool | None = Field(default=None)

132 """If set to False, disable async scheduling. Async scheduling helps to

133 avoid gaps in GPU utilization, leading to better latency and throughput.

134 """

vllm/config/utils.py

+2 -2

@@ -10,7 +10,7 @@ import json

10 import pathlib

11 import textwrap

12 from collections.abc import Callable, Mapping, Sequence, Set

13 -from dataclasses import MISSING, Field, field, fields, is_dataclass

13 +from dataclasses import MISSING, field, fields, is_dataclass

14 from itertools import pairwise

15 from typing import TYPE_CHECKING, Any, Protocol, TypeVar, cast

@@ -66,7 +66,7 @@ def config(

66 return decorator(cls)

69 -def get_field(cls: ConfigType, name: str) -> Field:

69 +def get_field(cls: ConfigType, name: str) -> Any:

70 """Get the default factory field of a dataclass by name. Used for getting

71 default factory fields in `EngineArgs`."""

72 if not is_dataclass(cls):

vllm/engine/arg_utils.py

+36 -13

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

67 PrefixCachingHashAlgo,

68 )

69 from vllm.config.device import Device

70 +from vllm.config.lora import MaxLoRARanks

71 from vllm.config.model import (

72 ConvertOption,

73 HfOverrides,

@@ -77,7 +78,12 @@ from vllm.config.model import (

78 )

79 from vllm.config.multimodal import MMCacheType, MMEncoderTPMode

80 from vllm.config.observability import DetailedTraceModules

80 -from vllm.config.parallel import DistributedExecutorBackend, ExpertPlacementStrategy

81 +from vllm.config.parallel import (

82 + All2AllBackend,

83 + DataParallelBackend,

84 + DistributedExecutorBackend,

85 + ExpertPlacementStrategy,

86 +)

87 from vllm.config.scheduler import SchedulerPolicy

88 from vllm.config.utils import get_field

89 from vllm.config.vllm import OptimizationLevel

@@ -257,7 +263,7 @@ def _compute_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]:

263 # VllmConfig's Fields have default_factory set to config classes.

264 # These could emit logs on init, which would be confusing.

265 with suppress_logging():

260 - default = default.default_factory()

266 + default = default.default_factory() # type: ignore[call-arg]

267 elif field.default_factory is not MISSING:

268 default = field.default_factory()

269

@@ -373,7 +379,7 @@ class EngineArgs:

379 dtype: ModelDType = ModelConfig.dtype

380 kv_cache_dtype: CacheDType = CacheConfig.cache_dtype

381 seed: int = ModelConfig.seed

376 - max_model_len: int | None = ModelConfig.max_model_len

382 + max_model_len: int = ModelConfig.max_model_len

383 cudagraph_capture_sizes: list[int] | None = (

384 CompilationConfig.cudagraph_capture_sizes

385 )

@@ -405,9 +411,9 @@ class EngineArgs:

411 data_parallel_rpc_port: int | None = None

412 data_parallel_hybrid_lb: bool = False

413 data_parallel_external_lb: bool = False

408 - data_parallel_backend: str = ParallelConfig.data_parallel_backend

414 + data_parallel_backend: DataParallelBackend = ParallelConfig.data_parallel_backend

415 enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel

410 - all2all_backend: str = ParallelConfig.all2all_backend

416 + all2all_backend: All2AllBackend = ParallelConfig.all2all_backend

417 enable_dbo: bool = ParallelConfig.enable_dbo

418 ubatch_size: int = ParallelConfig.ubatch_size

419 dbo_decode_token_threshold: int = ParallelConfig.dbo_decode_token_threshold

@@ -425,7 +431,7 @@ class EngineArgs:

431 max_parallel_loading_workers: int | None = (

432 ParallelConfig.max_parallel_loading_workers

433 )

428 - block_size: BlockSize | None = CacheConfig.block_size

434 + block_size: BlockSize = CacheConfig.block_size

435 enable_prefix_caching: bool | None = None

436 prefix_caching_hash_algo: PrefixCachingHashAlgo = (

437 CacheConfig.prefix_caching_hash_algo

@@ -451,7 +457,7 @@ class EngineArgs:

457 hf_token: bool | str | None = ModelConfig.hf_token

458 hf_overrides: HfOverrides = get_field(ModelConfig, "hf_overrides")

459 tokenizer_revision: str | None = ModelConfig.tokenizer_revision

454 - quantization: QuantizationMethods | None = ModelConfig.quantization

460 + quantization: QuantizationMethods | str | None = ModelConfig.quantization

461 allow_deprecated_quantization: bool = ModelConfig.allow_deprecated_quantization

462 enforce_eager: bool = ModelConfig.enforce_eager

463 disable_custom_all_reduce: bool = ParallelConfig.disable_custom_all_reduce

@@ -479,11 +485,11 @@ class EngineArgs:

485 )

486 io_processor_plugin: str | None = None

487 skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling

482 - video_pruning_rate: float = MultiModalConfig.video_pruning_rate

488 + video_pruning_rate: float | None = MultiModalConfig.video_pruning_rate

489 # LoRA fields

490 enable_lora: bool = False

491 max_loras: int = LoRAConfig.max_loras

486 - max_lora_rank: int = LoRAConfig.max_lora_rank

492 + max_lora_rank: MaxLoRARanks = LoRAConfig.max_lora_rank

493 default_mm_loras: dict[str, str] | None = LoRAConfig.default_mm_loras

494 fully_sharded_loras: bool = LoRAConfig.fully_sharded_loras

495 max_cpu_loras: int | None = LoRAConfig.max_cpu_loras

@@ -557,7 +563,7 @@ class EngineArgs:

563 ModelConfig, "override_generation_config"

564 )

565 model_impl: str = ModelConfig.model_impl

560 - override_attention_dtype: str = ModelConfig.override_attention_dtype

566 + override_attention_dtype: str | None = ModelConfig.override_attention_dtype

567 attention_backend: AttentionBackendEnum | None = AttentionConfig.backend

568

569 calculate_kv_scales: bool = CacheConfig.calculate_kv_scales

@@ -569,7 +575,7 @@ class EngineArgs:

575 additional_config: dict[str, Any] = get_field(VllmConfig, "additional_config")

576

577 use_tqdm_on_load: bool = LoadConfig.use_tqdm_on_load

572 - pt_load_map_location: str = LoadConfig.pt_load_map_location

578 + pt_load_map_location: str | dict[str, str] = LoadConfig.pt_load_map_location

579

580 logits_processors: list[str | type[LogitsProcessor]] | None = (

581 ModelConfig.logits_processors

@@ -1280,7 +1286,7 @@ class EngineArgs:

1286 hf_config_path=self.hf_config_path,

1287 runner=self.runner,

1288 convert=self.convert,

1283 - tokenizer=self.tokenizer,

1289 + tokenizer=self.tokenizer, # type: ignore[arg-type]

1290 tokenizer_mode=self.tokenizer_mode,

1291 trust_remote_code=self.trust_remote_code,

1292 allowed_local_media_path=self.allowed_local_media_path,

@@ -1445,12 +1451,16 @@ class EngineArgs:

1451 self.kv_cache_dtype, model_config

1452 )

1453

1454 + assert self.enable_prefix_caching is not None, (

1455 + "enable_prefix_caching must be set by this point"

1456 + )

1457 +

1458 cache_config = CacheConfig(

1459 block_size=self.block_size,

1460 gpu_memory_utilization=self.gpu_memory_utilization,

1461 kv_cache_memory_bytes=self.kv_cache_memory_bytes,

1462 swap_space=self.swap_space,

1453 - cache_dtype=resolved_cache_dtype,

1463 + cache_dtype=resolved_cache_dtype, # type: ignore[arg-type]

1464 is_attention_free=model_config.is_attention_free,

1465 num_gpu_blocks_override=self.num_gpu_blocks_override,

1466 sliding_window=sliding_window,

@@ -1676,6 +1686,16 @@ class EngineArgs:

1686 target_parallel_config=parallel_config,

1687 )

1688

1689 + assert self.max_num_batched_tokens is not None, (

1690 + "max_num_batched_tokens must be set by this point"

1691 + )

1692 + assert self.max_num_seqs is not None, "max_num_seqs must be set by this point"

1693 + assert self.enable_chunked_prefill is not None, (

1694 + "enable_chunked_prefill must be set by this point"

1695 + )

1696 + assert model_config.max_model_len is not None, (

1697 + "max_model_len must be set by this point"

1698 + )

1699 scheduler_config = SchedulerConfig(

1700 runner_type=model_config.runner_type,

1701 max_num_batched_tokens=self.max_num_batched_tokens,

@@ -2043,6 +2063,9 @@ class EngineArgs:

2063 )

2064

2065 if orig_max_num_batched_tokens is None:

2066 + assert model_config.max_model_len is not None, (

2067 + "max_model_len must be set by this point"

2068 + )

2069 if not self.enable_chunked_prefill:

2070 # If max_model_len is too short, use the default for higher throughput.

2071 self.max_num_batched_tokens = max(

vllm/logger.py

+4 -3

@@ -38,7 +38,7 @@ def _use_color() -> bool:

38 return False

41 -DEFAULT_LOGGING_CONFIG = {

41 +DEFAULT_LOGGING_CONFIG: dict[str, dict[str, Any] | Any] = {

42 "formatters": {

43 "vllm": {

44 "class": "vllm.logging_utils.NewLineFormatter",

@@ -157,7 +157,7 @@ _METHODS_TO_PATCH = {

157

158

159 def _configure_vllm_root_logger() -> None:

160 - logging_config = dict[str, dict[str, Any] | Any]()

160 + logging_config: dict[str, dict[str, Any] | Any] = {}

161

162 if not envs.VLLM_CONFIGURE_LOGGING and envs.VLLM_LOGGING_CONFIG_PATH:

163 raise RuntimeError(

@@ -225,7 +225,8 @@ def suppress_logging(level: int = logging.INFO) -> Generator[None, Any, None]:

225 logging.disable(current_level)

226

227

228 -def current_formatter_type(lgr: Logger) -> Literal["color", "newline", None]:

228 +def current_formatter_type(logger: Logger) -> Literal["color", "newline", None]:

229 + lgr: Logger | None = logger

230 while lgr is not None:

231 if lgr.handlers and len(lgr.handlers) == 1 and lgr.handlers[0].name == "vllm":

232 formatter = lgr.handlers[0].formatter

vllm/outputs.py

+1 -1

@@ -162,7 +162,7 @@ class RequestOutput:

162 completion.token_ids.extend(next_completion.token_ids)

163 if next_completion.logprobs:

164 assert completion.logprobs is not None

165 - completion.logprobs.extend(next_completion.logprobs)

165 + completion.logprobs.extend(next_completion.logprobs) # type: ignore[arg-type]

166 completion.cumulative_logprob = (

167 next_completion.cumulative_logprob

168 )

vllm/v1/cudagraph_dispatcher.py

+13

@@ -71,6 +71,9 @@ class CudagraphDispatcher:

71 """Pre-compute the mapping from batch size to padded graph size."""

72 max_size = self.compilation_config.max_cudagraph_capture_size

73 capture_sizes = self.compilation_config.cudagraph_capture_sizes

74 + assert capture_sizes is not None, (

75 + "Cudagraph capture sizes must be set when cudagraphs are enabled."

76 + )

77 self._bs_to_padded_graph_size: list[int] = [0] * (max_size + 1)

78 for end, start in zip(

79 capture_sizes + [max_size + 1],

@@ -89,6 +92,7 @@ class CudagraphDispatcher:

92 and self.cudagraph_mode != CUDAGraphMode.NONE

93 ):

94 for size in self.compilation_config.compile_sizes:

95 + size = int(size)

96 if size <= self.compilation_config.max_cudagraph_capture_size:

97 padded = self._bs_to_padded_graph_size[size]

98 if padded != size:

@@ -178,6 +182,9 @@ class CudagraphDispatcher:

182 # guarantee all keys would be used. For example, if we allow lazy

183 # capturing in future PR, some keys may never be triggered.

184 if cudagraph_mode.mixed_mode() != CUDAGraphMode.NONE:

185 + assert self.compilation_config.cudagraph_capture_sizes is not None, (

186 + "Cudagraph capture sizes must be set when mixed mode is enabled."

187 + )

188 for bs, num_active_loras in product(

189 self.compilation_config.cudagraph_capture_sizes, lora_cases

190 ):

@@ -200,6 +207,9 @@ class CudagraphDispatcher:

207 uniform_decode_query_len

208 * self.vllm_config.scheduler_config.max_num_seqs

209 )

210 + assert self.compilation_config.cudagraph_capture_sizes is not None, (

211 + "Cudagraph capture sizes must be set when full mode is enabled."

212 + )

213 cudagraph_capture_sizes_for_decode = [

214 x

215 for x in self.compilation_config.cudagraph_capture_sizes

@@ -262,6 +272,9 @@ class CudagraphDispatcher:

272 else:

273 # When not specializing, graphs are captured only with max_loras + 1,

274 # so we must use max_loras + 1 for dispatch to find a matching graph.

275 + assert self.vllm_config.lora_config is not None, (

276 + "LoRA config must be set when has_lora is True."

277 + )

278 effective_num_active_loras = self.vllm_config.lora_config.max_loras + 1

279

280 batch_desc = self._create_padded_batch_descriptor(