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

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 12, 19:45 UTC

[CI] Fix mypy for `vllm/v1/worker` (#29037)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

Wentao Ye committed Nov 20, 2025 at 22:36 UTC 56669c1f293d5c53b6a19ddf2f78802fa9fff2c2

13 files changed +180 -104

tools/pre_commit/mypy.py

+1 -1

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

38 "vllm/usage",

39 "vllm/v1/core",

40 "vllm/v1/engine",

41 + "vllm/v1/worker",

42 ]

44 # After fixing errors resulting from changing follow_imports

@@ -62,7 +63,6 @@ SEPARATE_GROUPS = [

63 "vllm/v1/sample",

64 "vllm/v1/spec_decode",

65 "vllm/v1/structured_output",

65 - "vllm/v1/worker",

66 ]

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

vllm/model_executor/utils.py

+1 -1

@@ -10,7 +10,7 @@ import torch

10 from vllm.utils.torch_utils import is_torch_equal_or_newer

13 -def set_random_seed(seed: int) -> None:

13 +def set_random_seed(seed: int | None) -> None:

14 from vllm.platforms import current_platform

16 current_platform.seed_everything(seed)

vllm/multimodal/utils.py

+2 -2

@@ -3,7 +3,7 @@

4 import asyncio

5 import atexit

6 -from collections.abc import Iterable, Set

6 +from collections.abc import Generator, Set

7 from concurrent.futures import ThreadPoolExecutor

8 from itertools import groupby

9 from pathlib import Path

@@ -403,7 +403,7 @@ def group_mm_kwargs_by_modality(

403 pin_memory: bool = False,

404 merge_by_field_config: bool | None = None,

405 multimodal_cpu_fields: Set[str] = frozenset(),

406 -) -> Iterable[tuple[str, int, BatchedTensorInputs]]:

406 +) -> Generator[tuple[str, int, BatchedTensorInputs], None, None]:

407 """Group consecutive `MultiModalKwargsItem`s from `mm_kwargs` with the same

408 modality together into the same `MultiModalKwargs` instance.

409

vllm/v1/worker/cpu_worker.py

+8 -4

@@ -3,6 +3,7 @@

3 import os

4 import platform

5 from collections.abc import Callable

6 +from typing import Any

8 import torch

@@ -37,6 +38,9 @@ class CPUWorker(Worker):

39 self.parallel_config.disable_custom_all_reduce = True

41 + # Torch profiler. Enabled and configured through env vars:

42 + # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace

43 + self.profiler: Any | None = None

44 if envs.VLLM_TORCH_PROFILER_DIR:

45 torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR

46 worker_name = f"{vllm_config.instance_id}-rank-{self.rank}"

@@ -80,13 +84,13 @@ class CPUWorker(Worker):

84 self.local_omp_cpuid = "nobind"

85 else:

86 local_dp_rank = self.parallel_config.data_parallel_rank_local

83 - omp_cpuids = omp_cpuids.split("|")

87 + omp_cpuids_list = omp_cpuids.split("|")

88 if local_dp_rank is not None:

89 world_size = self.parallel_config.world_size

86 - omp_cpuids = omp_cpuids[

90 + omp_cpuids_list = omp_cpuids_list[

91 local_dp_rank * world_size : (local_dp_rank + 1) * world_size

92 ]

89 - self.local_omp_cpuid = omp_cpuids[self.rank]

93 + self.local_omp_cpuid = omp_cpuids_list[self.rank]

95 if self.local_omp_cpuid != "nobind":

96 ret = torch.ops._C_utils.init_cpu_threads_env(self.local_omp_cpuid)

@@ -120,7 +124,7 @@ class CPUWorker(Worker):

124 pass

125

126 def determine_available_memory(self) -> int:

123 - return self.cache_config.cpu_kvcache_space_bytes # type: ignore

127 + return self.cache_config.cpu_kvcache_space_bytes or 0

128

129 def compile_or_warm_up_model(self) -> None:

130 # Reset the seed to ensure that the random state is not affected by

vllm/v1/worker/gpu_model_runner.py

+82 -46

@@ -5,7 +5,7 @@ import gc

5 import itertools

6 import time

7 from collections import defaultdict

8 -from collections.abc import Iterator

8 +from collections.abc import Iterator, Sequence

9 from contextlib import contextmanager

10 from copy import copy, deepcopy

11 from functools import reduce

@@ -53,6 +53,7 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

53 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding

54 from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader

55 from vllm.model_executor.models.interfaces import (

56 + SupportsMRoPE,

57 SupportsMultiModal,

58 is_mixture_of_experts,

59 supports_eagle3,

@@ -126,6 +127,7 @@ from vllm.v1.outputs import (

127 )

128 from vllm.v1.pool.metadata import PoolingMetadata

129 from vllm.v1.sample.logits_processor import LogitsProcessors, build_logitsprocs

130 +from vllm.v1.sample.logits_processor.interface import LogitsProcessor

131 from vllm.v1.sample.metadata import SamplingMetadata

132 from vllm.v1.sample.rejection_sampler import RejectionSampler

133 from vllm.v1.sample.sampler import Sampler

@@ -404,7 +406,10 @@ class GPUModelRunner(

406 # solution, we initialize the input batch here, and re-initialize it

407 # in `initialize_kv_cache` if the block_sizes here is different from

408 # the block_sizes in the kv cache config.

407 - custom_logitsprocs = model_config.logits_processors

409 + logits_processors = model_config.logits_processors

410 + custom_logitsprocs: Sequence[str | type[LogitsProcessor]] = (

411 + tuple(logits_processors) if logits_processors is not None else ()

412 + )

413 self.input_batch = InputBatch(

414 max_num_reqs=self.max_num_reqs,

415 # We need to use the encoder length for encoder-decoer

@@ -959,9 +964,13 @@ class GPUModelRunner(

964 def _init_mrope_positions(self, req_state: CachedRequestState):

965 model = self.get_model()

966 assert supports_mrope(model), "M-RoPE support is not implemented."

967 + assert req_state.prompt_token_ids is not None, (

968 + "M-RoPE requires prompt_token_ids to be available."

969 + )

970 + mrope_model = cast(SupportsMRoPE, model)

971

972 req_state.mrope_positions, req_state.mrope_position_delta = (

964 - model.get_mrope_input_positions(

973 + mrope_model.get_mrope_input_positions(

974 req_state.prompt_token_ids,

975 req_state.mm_features,

976 )

@@ -1762,6 +1771,7 @@ class GPUModelRunner(

1771 dst_start = mrope_pos_ptr

1772 dst_end = mrope_pos_ptr + completion_part_len

1773

1774 + assert req.mrope_position_delta is not None

1775 MRotaryEmbedding.get_next_input_positions_tensor(

1776 out=self.mrope_positions.np,

1777 out_offset=dst_start,

@@ -1907,6 +1917,8 @@ class GPUModelRunner(

1917

1918 for mm_input_id in encoder_input_ids:

1919 mm_feature = req_state.mm_features[mm_input_id]

1920 + if mm_feature.data is None:

1921 + continue

1922 mm_hash = mm_feature.identifier

1923 mm_kwargs.append(mm_feature.data)

1924 mm_hashes_pos.append((mm_hash, mm_feature.mm_position))

@@ -1930,7 +1942,7 @@ class GPUModelRunner(

1942 # multimodal inputs. The proper solution should be reordering the

1943 # encoder outputs.

1944 model = cast(SupportsMultiModal, self.model)

1933 - encoder_outputs = []

1945 + encoder_outputs: list[torch.Tensor] = []

1946 for modality, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(

1947 mm_kwargs,

1948 device=self.device,

@@ -1938,7 +1950,7 @@ class GPUModelRunner(

1950 merge_by_field_config=model.merge_by_field_config,

1951 multimodal_cpu_fields=model.multimodal_cpu_fields,

1952 ):

1941 - curr_group_outputs = []

1953 + curr_group_outputs: list[torch.Tensor] = []

1954

1955 # EVS-related change.

1956 # (ekhvedchenia): Temporary hack to limit peak memory usage when

@@ -1980,7 +1992,7 @@ class GPUModelRunner(

1992 # 2. A list or tuple (length: num_items) of tensors,

1993 # each of shape (feature_size, hidden_size) in case the feature

1994 # size is dynamic depending on the input multimodal items.

1983 - curr_group_outputs = model.embed_multimodal(**mm_kwargs_group)

1995 + curr_group_outputs = model.embed_multimodal(**mm_kwargs_group) # type: ignore[assignment]

1996

1997 sanity_check_mm_encoder_outputs(

1998 curr_group_outputs,

@@ -2180,7 +2192,7 @@ class GPUModelRunner(

2192 def sync_and_slice_intermediate_tensors(

2193 self,

2194 num_tokens: int,

2183 - intermediate_tensors: IntermediateTensors,

2195 + intermediate_tensors: IntermediateTensors | None,

2196 sync_self: bool,

2197 ) -> IntermediateTensors:

2198 assert self.intermediate_tensors is not None

@@ -2397,6 +2409,7 @@ class GPUModelRunner(

2409 if is_first_rank:

2410 intermediate_tensors = None

2411 else:

2412 + assert intermediate_tensors is not None

2413 intermediate_tensors = self.sync_and_slice_intermediate_tensors(

2414 num_input_tokens, intermediate_tensors, True

2415 )

@@ -2765,14 +2778,14 @@ class GPUModelRunner(

2778 uniform_decode = (

2779 max_num_scheduled_tokens == self.uniform_decode_query_len

2780 ) and (num_scheduled_tokens == num_reqs * max_num_scheduled_tokens)

2768 - batch_descriptor = BatchDescriptor(

2781 + batch_desc = BatchDescriptor(

2782 num_tokens=num_input_tokens,

2783 uniform_decode=uniform_decode,

2784 has_lora=len(self.input_batch.lora_id_to_lora_request) > 0,

2785 )

2786 cudagraph_runtime_mode, batch_descriptor = (

2787 self.cudagraph_dispatcher.dispatch(

2775 - batch_descriptor,

2788 + batch_desc,

2789 use_cascade_attn=cascade_attn_prefix_lens is not None,

2790 )

2791 )

@@ -2856,15 +2869,15 @@ class GPUModelRunner(

2869 else:

2870 logits = self.model.compute_logits(sample_hidden_states)

2871

2859 - model_output_broadcast_data = {}

2872 + model_output_broadcast_data: dict[str, Any] = {}

2873 if logits is not None:

2874 model_output_broadcast_data["logits"] = logits.contiguous()

2875

2863 - model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(

2876 + broadcasted = get_pp_group().broadcast_tensor_dict(

2877 model_output_broadcast_data, src=len(get_pp_group().ranks) - 1

2878 )

2866 - assert model_output_broadcast_data is not None

2867 - logits = model_output_broadcast_data["logits"]

2879 + assert broadcasted is not None

2880 + logits = broadcasted["logits"]

2881

2882 self.execute_model_state = ExecuteModelState(

2883 scheduler_output,

@@ -2889,7 +2902,7 @@ class GPUModelRunner(

2902 if self.execute_model_state is None:

2903 # Nothing to do (PP non-final rank case), output isn't used.

2904 if not kv_connector_output:

2892 - return None # noqa

2905 + return None # type: ignore[return-value]

2906

2907 # In case of PP with kv transfer, we need to pass through the

2908 # kv_connector_output

@@ -2941,33 +2954,37 @@ class GPUModelRunner(

2954 spec_decode_common_attn_metadata,

2955 )

2956

2957 + spec_config = self.speculative_config

2958 use_padded_batch_for_eagle = (

2945 - self.speculative_config

2946 - and self.speculative_config.use_eagle()

2947 - and not self.speculative_config.disable_padded_drafter_batch

2959 + spec_config is not None

2960 + and spec_config.use_eagle()

2961 + and not spec_config.disable_padded_drafter_batch

2962 )

2963 effective_drafter_max_model_len = self.max_model_len

2964 if effective_drafter_max_model_len is None:

2965 effective_drafter_max_model_len = self.model_config.max_model_len

2966 if (

2953 - self.speculative_config

2954 - and self.speculative_config.draft_model_config is not None

2955 - and self.speculative_config.draft_model_config.max_model_len is not None

2967 + spec_config is not None

2968 + and spec_config.draft_model_config is not None

2969 + and spec_config.draft_model_config.max_model_len is not None

2970 ):

2971 effective_drafter_max_model_len = (

2958 - self.speculative_config.draft_model_config.max_model_len

2972 + spec_config.draft_model_config.max_model_len

2973 )

2974 input_fits_in_drafter = spec_decode_common_attn_metadata and (

2975 spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens

2976 <= effective_drafter_max_model_len

2977 )

2978 if use_padded_batch_for_eagle:

2979 + assert self.speculative_config is not None

2980 + assert isinstance(self.drafter, EagleProposer)

2981 sampled_token_ids = sampler_output.sampled_token_ids

2982 if input_fits_in_drafter:

2983 # EAGLE speculative decoding can use the GPU sampled tokens

2984 # as inputs, and does not need to wait for bookkeeping to finish.

2985 propose_draft_token_ids(sampled_token_ids)

2986 elif self.valid_sampled_token_count_event is not None:

2987 + assert spec_decode_common_attn_metadata is not None

2988 next_token_ids, valid_sampled_tokens_count = (

2989 self.drafter.prepare_next_token_ids_padded(

2990 spec_decode_common_attn_metadata,

@@ -3105,7 +3122,9 @@ class GPUModelRunner(

3122 common_attn_metadata: CommonAttentionMetadata,

3123 ) -> torch.Tensor | list[list[int]]:

3124 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

3108 - if self.speculative_config.method == "ngram":

3125 + spec_config = self.speculative_config

3126 + assert spec_config is not None

3127 + if spec_config.method == "ngram":

3128 assert isinstance(sampled_token_ids, list)

3129 assert isinstance(self.drafter, NgramProposer)

3130 draft_token_ids = self.drafter.propose(

@@ -3115,11 +3134,11 @@ class GPUModelRunner(

3134 self.input_batch.token_ids_cpu,

3135 self.input_batch.spec_decode_unsupported_reqs,

3136 )

3118 - elif self.speculative_config.method == "suffix":

3137 + elif spec_config.method == "suffix":

3138 assert isinstance(sampled_token_ids, list)

3139 assert isinstance(self.drafter, SuffixDecodingProposer)

3140 draft_token_ids = self.drafter.propose(self.input_batch, sampled_token_ids)

3122 - elif self.speculative_config.method == "medusa":

3141 + elif spec_config.method == "medusa":

3142 assert isinstance(sampled_token_ids, list)

3143 assert isinstance(self.drafter, MedusaProposer)

3144

@@ -3144,10 +3163,10 @@ class GPUModelRunner(

3163 target_hidden_states=hidden_states,

3164 sampling_metadata=sampling_metadata,

3165 )

3147 - elif self.speculative_config.use_eagle():

3166 + elif spec_config.use_eagle():

3167 assert isinstance(self.drafter, EagleProposer)

3168

3150 - if self.speculative_config.disable_padded_drafter_batch:

3169 + if spec_config.disable_padded_drafter_batch:

3170 # When padded-batch is disabled, the sampled_token_ids should be

3171 # the cpu-side list[list[int]] of valid sampled tokens for each

3172 # request, with invalid requests having empty lists.

@@ -3197,7 +3216,7 @@ class GPUModelRunner(

3216 else:

3217 target_hidden_states = hidden_states[:num_scheduled_tokens]

3218 else:

3200 - if self.speculative_config.disable_padded_drafter_batch:

3219 + if spec_config.disable_padded_drafter_batch:

3220 token_indices_to_sample = None

3221 common_attn_metadata, token_indices = self.drafter.prepare_inputs(

3222 common_attn_metadata,

@@ -3292,9 +3311,12 @@ class GPUModelRunner(

3311 and is_mixture_of_experts(self.drafter.model)

3312 and self.parallel_config.enable_eplb

3313 ):

3314 + spec_config = self.vllm_config.speculative_config

3315 + assert spec_config is not None

3316 + assert spec_config.draft_model_config is not None

3317 logger.info_once(

3318 "EPLB is enabled for drafter model %s.",

3297 - self.vllm_config.speculative_config.draft_model_config.model,

3319 + spec_config.draft_model_config.model,

3320 )

3321

3322 global_expert_load = (

@@ -3311,7 +3333,7 @@ class GPUModelRunner(

3333 self.eplb_state = EplbState(self.parallel_config, self.device)

3334 self.eplb_state.add_model(

3335 self.drafter.model,

3314 - self.vllm_config.speculative_config.draft_model_config,

3336 + spec_config.draft_model_config,

3337 global_expert_load,

3338 old_global_expert_indices,

3339 rank_mapping,

@@ -3346,9 +3368,11 @@ class GPUModelRunner(

3368 scope="local",

3369 )

3370 prepare_communication_buffer_for_model(self.model)

3371 + mm_config = self.model_config.multimodal_config

3372 self.is_multimodal_pruning_enabled = (

3373 supports_multimodal_pruning(self.get_model())

3351 - and self.model_config.multimodal_config.is_multimodal_pruning_enabled()

3374 + and mm_config is not None

3375 + and mm_config.is_multimodal_pruning_enabled()

3376 )

3377

3378 if is_mixture_of_experts(self.model) and self.parallel_config.enable_eplb:

@@ -3383,15 +3407,14 @@ class GPUModelRunner(

3407 # CudagraphWraper and CudagraphDispatcher of vllm.

3408

3409 # wrap the model with full cudagraph wrapper if needed.

3386 - if (

3387 - self.compilation_config.cudagraph_mode.has_full_cudagraphs()

3388 - and not self.parallel_config.enable_dbo

3389 - ):

3410 + cudagraph_mode = self.compilation_config.cudagraph_mode

3411 + assert cudagraph_mode is not None

3412 + if cudagraph_mode.has_full_cudagraphs() and not self.parallel_config.enable_dbo:

3413 self.model = CUDAGraphWrapper(

3414 self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL

3415 )

3416 elif self.parallel_config.enable_dbo:

3394 - if self.compilation_config.cudagraph_mode.has_full_cudagraphs():

3417 + if cudagraph_mode.has_full_cudagraphs():

3418 self.model = UBatchWrapper(

3419 self.model, self.vllm_config, CUDAGraphMode.FULL, self.device

3420 )

@@ -4071,7 +4094,8 @@ class GPUModelRunner(

4094 def profile_run(self) -> None:

4095 # Profile with multimodal encoder & encoder cache.

4096 if self.supports_mm_inputs:

4074 - if self.model_config.multimodal_config.skip_mm_profiling:

4097 + mm_config = self.model_config.multimodal_config

4098 + if mm_config is not None and mm_config.skip_mm_profiling:

4099 logger.info(

4100 "Skipping memory profiling for multimodal encoder and "

4101 "encoder cache."

@@ -4333,8 +4357,9 @@ class GPUModelRunner(

4357 def get_attn_backends_for_group(

4358 kv_cache_group_spec: KVCacheGroupSpec,

4359 ) -> tuple[dict[AttentionGroupKey, list[str]], set[type[AttentionBackend]]]:

4360 + layer_type = cast(type[Any], AttentionLayerBase)

4361 layers = get_layers_from_vllm_config(

4337 - self.vllm_config, AttentionLayerBase, kv_cache_group_spec.layer_names

4362 + self.vllm_config, layer_type, kv_cache_group_spec.layer_names

4363 )

4364 attn_backends = {}

4365 attn_backend_layers = defaultdict(list)

@@ -4349,7 +4374,7 @@ class GPUModelRunner(

4374 if layer_name in self.kv_sharing_fast_prefill_eligible_layers:

4375 attn_backend = create_fast_prefill_custom_backend(

4376 "FastPrefill",

4352 - attn_backend,

4377 + attn_backend, # type: ignore[arg-type]

4378 )

4379

4380 full_cls_name = attn_backend.full_cls_name()

@@ -4448,6 +4473,7 @@ class GPUModelRunner(

4473 min_cg_backend_name = attn_backend.__name__

4474 # Flexible resolve the cudagraph mode

4475 cudagraph_mode = self.compilation_config.cudagraph_mode

4476 + assert cudagraph_mode is not None

4477 # check cudagraph for mixed batch is supported

4478 if (

4479 cudagraph_mode.mixed_mode() == CUDAGraphMode.FULL

@@ -4562,12 +4588,17 @@ class GPUModelRunner(

4588 self.compilation_config.adjust_cudagraph_sizes_for_spec_decode(

4589 self.uniform_decode_query_len, self.parallel_config.tensor_parallel_size

4590 )

4565 - self.cudagraph_batch_sizes = self.compilation_config.cudagraph_capture_sizes

4591 + capture_sizes = self.compilation_config.cudagraph_capture_sizes

4592 + self.cudagraph_batch_sizes = (

4593 + capture_sizes if capture_sizes is not None else []

4594 + )

4595

4596 # Trigger cudagraph dispatching keys initialization after

4597 # resolved cudagraph mode.

4598 + cudagraph_mode = self.compilation_config.cudagraph_mode

4599 + assert cudagraph_mode is not None

4600 self.cudagraph_dispatcher.initialize_cudagraph_keys(

4570 - self.compilation_config.cudagraph_mode, self.uniform_decode_query_len

4601 + cudagraph_mode, self.uniform_decode_query_len

4602 )

4603

4604 def calculate_reorder_batch_threshold(self) -> None:

@@ -4579,7 +4610,7 @@ class GPUModelRunner(

4610 """

4611 min_none_high = lambda a, b: a if b is None else b if a is None else min(a, b)

4612

4582 - reorder_batch_thresholds = [

4613 + reorder_batch_thresholds: list[int | None] = [

4614 group.get_metadata_builder().reorder_batch_threshold

4615 for group in self._attn_group_iterator()

4616 ]

@@ -4588,7 +4619,7 @@ class GPUModelRunner(

4619 if len(reorder_batch_thresholds) == 0:

4620 self.reorder_batch_threshold = None

4621 return

4591 - self.reorder_batch_threshold = reduce(min_none_high, reorder_batch_thresholds)

4622 + self.reorder_batch_threshold = reduce(min_none_high, reorder_batch_thresholds) # type: ignore[assignment]

4623

4624 @staticmethod

4625 def select_common_block_size(

@@ -5048,12 +5079,16 @@ class GPUModelRunner(

5079 kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks)

5080

5081 if self.dcp_world_size > 1:

5051 - layers = get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase)

5082 + layer_type = cast(type[Any], AttentionLayerBase)

5083 + layers = get_layers_from_vllm_config(self.vllm_config, layer_type)

5084 for layer in layers.values():

5053 - assert layer.impl.need_to_return_lse_for_decode, (

5085 + layer_impl = getattr(layer, "impl", None)

5086 + if layer_impl is None:

5087 + continue

5088 + assert layer_impl.need_to_return_lse_for_decode, (

5089 "DCP requires attention impls to return"

5090 " the softmax lse for decode, but the impl "

5056 - f"{layer.impl.__class__.__name__} "

5091 + f"{layer_impl.__class__.__name__} "

5092 "does not return the softmax lse for decode."

5093 )

5094

@@ -5094,7 +5129,8 @@ class GPUModelRunner(

5129 if has_ec_transfer() and get_ec_transfer().is_producer:

5130 return {}

5131 kv_cache_spec: dict[str, KVCacheSpec] = {}

5097 - attn_layers = get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase)

5132 + layer_type = cast(type[Any], AttentionLayerBase)

5133 + attn_layers = get_layers_from_vllm_config(self.vllm_config, layer_type)

5134 for layer_name, attn_module in attn_layers.items():

5135 if isinstance(attn_module, Attention) and (

5136 kv_tgt_layer := attn_module.kv_sharing_target_layer_name

vllm/v1/worker/gpu_ubatch_wrapper.py

+13 -7

@@ -121,18 +121,24 @@ class UBatchWrapper:

121

122 @staticmethod

123 def _create_sm_control_context(vllm_config: VllmConfig):

124 - comm_sms = envs.VLLM_DBO_COMM_SMS

124 + comm_sms: int = envs.VLLM_DBO_COMM_SMS

125

126 set_comm_sms = lambda sms: None

127 if vllm_config.parallel_config.enable_expert_parallel:

128 # Currently only DeepEP highthroughput supports SM control so this

129 # only affects that case.

130 - all2all_manager = get_ep_group().device_communicator.all2all_manager

131 -

132 - if all2all_manager.max_sms_used() is not None:

133 - comm_sms = min(comm_sms, all2all_manager.max_sms_used())

134 -

135 - if comm_sms > 0:

130 + ep_group = get_ep_group()

131 + device_communicator = ep_group.device_communicator

132 + all2all_manager = None

133 + if device_communicator is not None:

134 + all2all_manager = device_communicator.all2all_manager

135 +

136 + if all2all_manager is not None:

137 + max_sms_used = all2all_manager.max_sms_used()

138 + if max_sms_used is not None:

139 + comm_sms = min(comm_sms, max_sms_used)

140 +

141 + if comm_sms > 0 and all2all_manager is not None:

142 set_comm_sms = lambda sms: all2all_manager.set_num_sms(sms)

143

144 # TODO(lucas): support other kernels besides DeepGEMM

vllm/v1/worker/gpu_worker.py

+33 -29

@@ -6,7 +6,7 @@ import gc

6 import os

7 from contextlib import AbstractContextManager, nullcontext

8 from types import NoneType

9 -from typing import TYPE_CHECKING, Any

9 +from typing import TYPE_CHECKING, Any, cast

11 import torch

12 import torch.distributed

@@ -87,8 +87,10 @@ class Worker(WorkerBase):

87 # Buffers saved before sleep

88 self._sleep_saved_buffers: dict[str, torch.Tensor] = {}

90 - # Torch profiler. Enabled and configured through env vars:

90 + # Torch/CUDA profiler. Enabled and configured through env vars:

91 # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace

92 + # VLLM_TORCH_CUDA_PROFILE=1

93 + self.profiler: Any | None = None

94 if envs.VLLM_TORCH_PROFILER_DIR:

95 worker_name = f"{vllm_config.instance_id}-rank-{self.rank}"

96 self.profiler = TorchProfilerWrapper(

@@ -146,17 +148,17 @@ class Worker(WorkerBase):

148 assert allocator.get_current_usage() == 0, (

149 "Sleep mode can only be used for one instance per process."

150 )

149 - context = allocator.use_memory_pool(tag=tag)

151 + return allocator.use_memory_pool(tag=tag)

152 else:

151 - context = nullcontext()

152 - return context

153 + return nullcontext()

154

155 def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None:

156 self.cache_config.num_gpu_blocks = num_gpu_blocks

157 self.cache_config.num_cpu_blocks = num_cpu_blocks

158

159 def init_device(self):

159 - if self.device_config.device.type == "cuda":

160 + device = self.device_config.device

161 + if isinstance(device, torch.device) and device.type == "cuda":

162 # This env var set by Ray causes exceptions with graph building.

163 os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)

164 if (

@@ -375,23 +377,21 @@ class Worker(WorkerBase):

377 from vllm.device_allocator.cumem import CuMemAllocator

378

379 allocator = CuMemAllocator.get_instance()

378 - context = allocator.use_memory_pool(tag="kv_cache")

380 + with allocator.use_memory_pool(tag="kv_cache"):

381 + self.model_runner.initialize_kv_cache(kv_cache_config)

382 else:

380 - context = nullcontext()

381 - with context:

383 self.model_runner.initialize_kv_cache(kv_cache_config)

384

385 def compile_or_warm_up_model(self) -> None:

386 # warm up sizes that are not in cudagraph capture sizes,

387 # but users still want to compile for better performance,

388 # e.g. for the max-num-batched token size in chunked prefill.

388 - warmup_sizes = self.vllm_config.compilation_config.compile_sizes.copy()

389 + compile_sizes = self.vllm_config.compilation_config.compile_sizes

390 + warmup_sizes = compile_sizes.copy() if compile_sizes is not None else []

391 if not self.model_config.enforce_eager:

390 - warmup_sizes = [

391 - x

392 - for x in warmup_sizes

393 - if x not in self.vllm_config.compilation_config.cudagraph_capture_sizes

394 - ]

392 + capture_sizes = self.vllm_config.compilation_config.cudagraph_capture_sizes

393 + if capture_sizes is not None:

394 + warmup_sizes = [x for x in warmup_sizes if x not in capture_sizes]

395 # We skip EPLB here since we don't want to record dummy metrics

396 for size in sorted(warmup_sizes, reverse=True):

397 logger.info("Compile and warming up model for size %d", size)

@@ -532,12 +532,12 @@ class Worker(WorkerBase):

532 )

533 }

534 if forward_pass and not get_pp_group().is_first_rank:

535 - intermediate_tensors = IntermediateTensors(

536 - get_pp_group().recv_tensor_dict(

537 - all_gather_group=get_tp_group(),

538 - all_gather_tensors=all_gather_tensors,

539 - )

535 + tensor_dict = get_pp_group().recv_tensor_dict(

536 + all_gather_group=get_tp_group(),

537 + all_gather_tensors=all_gather_tensors,

538 )

539 + assert tensor_dict is not None

540 + intermediate_tensors = IntermediateTensors(tensor_dict)

541

542 with self.annotate_profile(scheduler_output):

543 output = self.model_runner.execute_model(

@@ -605,7 +605,7 @@ class Worker(WorkerBase):

605 assert self.model_runner.eplb_state is not None

606 self.model_runner.eplb_state.rearrange(

607 execute_shuffle=True,

608 - global_expert_load=None,

608 + global_expert_loads=None,

609 rank_mapping=rank_mapping,

610 )

611 torch.cuda.synchronize()

@@ -661,7 +661,7 @@ class Worker(WorkerBase):

661

662 def _reconfigure_moe(

663 self, old_ep_size: int, new_ep_size: int

664 - ) -> torch.Tensor | None:

664 + ) -> list[torch.Tensor] | None:

665 """

666 Reconfigure MoE modules with provided reconfig_request

667

@@ -728,26 +728,29 @@ class Worker(WorkerBase):

728 num_local_physical_experts = num_local_experts

729 assert self.model_runner.eplb_state is not None

730 new_physical_experts = (

731 - self.model_runner.eplb_state.physical_to_logical_map.shape[1]

731 + self.model_runner.eplb_state.physical_to_logical_map.shape[1] # type: ignore[attr-defined]

732 )

733 parallel_config.eplb_config.num_redundant_experts = (

734 new_physical_experts

735 - - self.model_runner.eplb_state.logical_replica_count.shape[1]

735 + - self.model_runner.eplb_state.logical_replica_count.shape[1] # type: ignore[attr-defined]

736 )

737 global_expert_loads = None

738 else:

739 - num_local_physical_experts = torch.tensor(

739 + num_local_physical_experts_tensor = torch.tensor(

740 [num_local_experts], dtype=torch.int32, device="cpu"

741 )

742 torch.distributed.broadcast(

743 - num_local_physical_experts, group=get_ep_group().cpu_group, group_src=0

743 + num_local_physical_experts_tensor,

744 + group=get_ep_group().cpu_group,

745 + group_src=0,

746 )

745 - num_local_physical_experts = num_local_physical_experts.item()

747 + num_local_physical_experts = int(num_local_physical_experts_tensor.item())

748 new_physical_experts = num_local_physical_experts * new_ep_size

749 assert self.model_runner.eplb_state is not None

748 - global_expert_loads = self.model_runner.eplb_state.rearrange(

750 + global_expert_loads_any = self.model_runner.eplb_state.rearrange(

751 execute_shuffle=False

752 )

753 + global_expert_loads = cast(list[torch.Tensor], global_expert_loads_any)

754 parallel_config.eplb_config.num_redundant_experts = (

755 new_physical_experts - global_expert_loads[0].shape[1]

756 )

@@ -849,8 +852,9 @@ def init_worker_distributed_environment(

852 init_batch_invariance()

853 set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)

854

855 + init_method = distributed_init_method or "env://"

856 init_distributed_environment(

853 - parallel_config.world_size, rank, distributed_init_method, local_rank, backend

857 + parallel_config.world_size, rank, init_method, local_rank, backend

858 )

859

860 ensure_model_parallel_initialized(

vllm/v1/worker/kv_connector_model_runner_mixin.py

+1 -1

@@ -59,7 +59,7 @@ class KVConnectorModelRunnerMixin:

59 @staticmethod

60 def ensure_kv_transfer_shutdown() -> None:

61 # has_kv_transfer_group can be None during interpreter shutdown.

62 - if has_kv_transfer_group and has_kv_transfer_group():

62 + if has_kv_transfer_group and has_kv_transfer_group(): # type: ignore[truthy-function]

63 ensure_kv_transfer_shutdown()

65 @staticmethod

vllm/v1/worker/tpu_model_runner.py

+22 -6

@@ -572,7 +572,10 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

572 format. Layers that do not need KV cache are not included.

573 """

574

575 - layers = get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase)

575 + layers = get_layers_from_vllm_config(

576 + self.vllm_config,

577 + AttentionLayerBase, # type: ignore[type-abstract]

578 + )

579 block_size = self.vllm_config.cache_config.block_size

580 cache_dtype_str = self.vllm_config.cache_config.cache_dtype

581

@@ -725,7 +728,11 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

728 req_id = self.input_batch.req_ids[i]

729 assert req_id is not None

730 num_tokens = scheduler_output.num_scheduled_tokens[req_id]

728 - if not use_max_model_len and num_tokens > self.most_model_len:

731 + if (

732 + not use_max_model_len

733 + and self.most_model_len is not None

734 + and num_tokens > self.most_model_len

735 + ):

736 use_max_model_len = True

737 num_scheduled_tokens_per_req.append(num_tokens)

738 if use_max_model_len:

@@ -737,6 +744,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

744 else:

745 end_index = num_reqs

746 else:

747 + assert self.num_reqs_most_model_len is not None

748 if len(num_scheduled_tokens_per_req) > self.num_reqs_most_model_len:

749 num_scheduled_tokens_per_req = num_scheduled_tokens_per_req[

750 : self.num_reqs_most_model_len

@@ -829,6 +837,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

837 ].to(self.device)

838 seq_lens = self.seq_lens_cpu[: self.num_reqs_max_model_len].to(self.device)

839 else:

840 + assert self.num_reqs_most_model_len is not None

841 block_tables = self.block_table_cpu[

842 : self.num_reqs_most_model_len, : self.num_blocks_per_most_len_req

843 ]

@@ -931,6 +940,8 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

940

941 for mm_input_id in encoder_input_ids:

942 mm_feature = req_state.mm_features[mm_input_id]

943 + if mm_feature.data is None:

944 + continue

945 mm_hash = mm_feature.identifier

946 mm_kwargs.append(mm_feature.data)

947 mm_hashes_pos.append((mm_hash, mm_feature.mm_position))

@@ -1114,7 +1125,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1125 ) -> ModelRunnerOutput:

1126 if self.scheduler_output is None:

1127 # Nothing to do (PP non-final rank case), output isn't used.

1117 - return None # noqa

1128 + return None # type: ignore[return-value]

1129 scheduler_output = self.scheduler_output

1130 mm_embed_inputs = self.mm_embed_inputs

1131 self.scheduler_output = None

@@ -1696,7 +1707,8 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1707 ) -> None:

1708 # Profile with multimodal encoder & encoder cache.

1709 if self.supports_mm_inputs:

1699 - if self.model_config.multimodal_config.skip_mm_profiling:

1710 + mm_config = self.model_config.multimodal_config

1711 + if mm_config is not None and mm_config.skip_mm_profiling:

1712 logger.info(

1713 "Skipping memory profiling for multimodal encoder and "

1714 "encoder cache."

@@ -2166,5 +2178,9 @@ def replace_set_lora(model):

2178 if isinstance(module, BaseLayerWithLoRA):

2179 module._original_set_lora = module.set_lora

2180 module._original_reset_lora = module.reset_lora

2169 - module.set_lora = _tpu_set_lora.__get__(module, module.__class__)

2170 - module.reset_lora = _tpu_reset_lora.__get__(module, module.__class__)

2181 + module.set_lora = _tpu_set_lora.__get__( # type: ignore[method-assign]

2182 + module, module.__class__

2183 + )

2184 + module.reset_lora = _tpu_reset_lora.__get__( # type: ignore[method-assign]

2185 + module, module.__class__

2186 + )

vllm/v1/worker/tpu_worker.py

+2 -3

@@ -141,8 +141,7 @@ class TPUWorker:

141

142 # Set random seed.

143 set_random_seed(self.model_config.seed)

144 - if self.model_config.seed is not None:

145 - xm.set_rng_state(self.model_config.seed, self.device)

144 + xm.set_rng_state(self.model_config.seed, self.device)

145

146 # Increase the cache size limit, which is the maximum number of

147 # dynamo graphs that can be compiled.

@@ -332,7 +331,7 @@ class TPUWorker:

331 world_size=parallel_config.world_size,

332 rank=rank,

333 local_rank=local_rank,

335 - distributed_init_method=distributed_init_method,

334 + distributed_init_method=distributed_init_method or "env://",

335 backend=current_platform.dist_backend,

336 )

337 ensure_model_parallel_initialized(

vllm/v1/worker/utils.py

+5 -3

@@ -280,7 +280,7 @@ def bind_kv_cache(

280 kv_caches: dict[str, torch.Tensor],

281 forward_context: dict[str, "Attention"],

282 runner_kv_caches: list[torch.Tensor],

283 - num_attn_module: int | None = 1,

283 + num_attn_module: int = 1,

284 ) -> None:

285 """

286 Bind the allocated KV cache to both ModelRunner and forward context so

@@ -362,5 +362,7 @@ def is_residual_scattered_for_sp(

362 or vllm_config.compilation_config.use_inductor_graph_partition

363 ):

364 return True

365 -

366 - return num_input_tokens in vllm_config.compilation_config.compile_sizes

365 + compile_sizes = vllm_config.compilation_config.compile_sizes

366 + if compile_sizes is None:

367 + return False

368 + return num_input_tokens in compile_sizes

vllm/v1/worker/worker_base.py

@@ -315,10 +315,12 @@ class WorkerWrapperBase:

315

316 def initialize_from_config(self, kv_cache_configs: list[Any]) -> None:

317 kv_cache_config = kv_cache_configs[self.global_rank]

318 + assert self.vllm_config is not None

319 with set_current_vllm_config(self.vllm_config):

320 self.worker.initialize_from_config(kv_cache_config) # type: ignore

321

322 def init_device(self):

323 + assert self.vllm_config is not None

324 with set_current_vllm_config(self.vllm_config):

325 # To make vLLM config available during device initialization

326 self.worker.init_device() # type: ignore

vllm/v1/worker/xpu_worker.py

+8 -1

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 import os

4 +from typing import Any

6 import torch

7 import torch.distributed

@@ -37,6 +38,7 @@ class XPUWorker(Worker):

39 # Torch profiler. Enabled and configured through env vars:

40 # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace

41 + self.profiler: Any | None = None

42 if envs.VLLM_TORCH_PROFILER_DIR:

43 torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR

44 worker_name = f"{vllm_config.instance_id}-rank-{self.rank}"

@@ -148,7 +150,12 @@ class XPUWorker(Worker):

150 return int(available_kv_cache_memory)

151

152 def init_device(self):

151 - if self.device_config.device.type == "xpu" and current_platform.is_xpu():

153 + device = self.device_config.device

154 + if (

155 + isinstance(device, torch.device)

156 + and device.type == "xpu"

157 + and current_platform.is_xpu()

158 + ):

159 self.device = torch.device(f"xpu:{self.local_rank}")

160 current_platform.set_device(self.device)

161 current_platform.check_if_supports_dtype(self.model_config.dtype)