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 2, 23:15 UTC
[Misc] Added missing return type annotations to improve mypy and IDE tooling (#43383)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
May 23, 2026 at 00:28 UTC
3a1c0621512f6a58f14a19abfc106f7c11863681
4 files changed
+14
-11
vllm/model_executor/layers/pooler/activations.py
+3
-3
@@ -64,7 +64,7 @@ def resolve_classifier_act_fn(
64
model_config: ModelConfig,
65
static_num_labels: bool = True,
66
act_fn: "PoolerActivation | None" = None,
67
-):
67
+) -> "PoolerActivation":
68
if act_fn is None:
69
return get_act_fn(model_config.hf_config, static_num_labels)
70
@@ -78,7 +78,7 @@ _T = TypeVar("_T", torch.Tensor, list[torch.Tensor])
78
79
class PoolerActivation(nn.Module, ABC):
80
@staticmethod
81
- def wraps(module: nn.Module):
81
+ def wraps(module: nn.Module) -> "PoolerActivation":
82
if isinstance(module, nn.Identity):
83
return PoolerIdentity()
84
if isinstance(module, (nn.Sigmoid, nn.Softmax)):
@@ -140,7 +140,7 @@ class PoolerClassify(PoolerActivation):
140
141
142
class LambdaPoolerActivation(PoolerActivation):
143
- def __init__(self, fn: Callable[[torch.Tensor], torch.Tensor]):
143
+ def __init__(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> None:
144
super().__init__()
145
146
self.fn = fn
vllm/model_executor/layers/pooler/seqwise/methods.py
+3
-1
@@ -108,7 +108,9 @@ class MeanPool(SequencePoolingMethod):
108
return segment_sums / prompt_lens.unsqueeze(1)
109
110
111
-def get_seq_pooling_method(pooling_type: SequencePoolingType | str):
111
+def get_seq_pooling_method(
112
+ pooling_type: SequencePoolingType | str,
113
+) -> SequencePoolingMethod:
114
if pooling_type == "CLS":
115
return CLSPool()
116
if pooling_type == "LAST":
vllm/model_executor/layers/pooler/special.py
+1
@@ -102,6 +102,7 @@ class DispatchPooler(Pooler):
102
# portion of the batch. Token offset is computed from the CPU
103
# `num_scheduled_tokens_cpu` to avoid a GPU->CPU sync.
104
group_cursor = group_metadata.pooling_cursor
105
+ assert group_cursor is not None
106
num_group_tokens = int(group_cursor.num_scheduled_tokens_cpu.sum())
107
group_hidden_states = hidden_states[
108
token_offset : token_offset + num_group_tokens
vllm/v1/pool/metadata.py
+7
-7
@@ -20,7 +20,7 @@ class PoolingCursor:
20
seq_lens_cpu: torch.Tensor
21
num_scheduled_tokens_cpu: torch.Tensor
22
23
- def __getitem__(self, indices: slice):
23
+ def __getitem__(self, indices: slice) -> "PoolingCursor":
24
return PoolingCursor(
25
first_token_indices_gpu=self.first_token_indices_gpu[indices],
26
last_token_indices_gpu=self.last_token_indices_gpu[indices],
@@ -29,19 +29,19 @@ class PoolingCursor:
29
num_scheduled_tokens_cpu=self.num_scheduled_tokens_cpu[indices],
30
)
31
32
- def is_partial_prefill(self):
32
+ def is_partial_prefill(self) -> bool:
33
return not torch.all(self.prompt_lens_cpu == self.num_scheduled_tokens_cpu)
34
35
- def is_finished(self):
35
+ def is_finished(self) -> torch.Tensor:
36
return self.prompt_lens_cpu == self.seq_lens_cpu
37
38
39
class PoolingStates:
40
- def __init__(self):
40
+ def __init__(self) -> None:
41
# for chunked prefill with ALL pooling
42
self.hidden_states_cache: list[torch.Tensor] = []
43
44
- def clean(self):
44
+ def clean(self) -> None:
45
self.hidden_states_cache.clear()
46
47
@@ -72,7 +72,7 @@ class PoolingMetadata:
72
73
self.tasks = tasks
74
75
- def __getitem__(self, indices: slice):
75
+ def __getitem__(self, indices: slice) -> "PoolingMetadata":
76
return PoolingMetadata(
77
prompt_lens=self.prompt_lens[indices],
78
prompt_token_ids=None
@@ -121,7 +121,7 @@ class PoolingMetadata:
121
seq_lens_cpu: torch.Tensor,
122
device: torch.device,
123
query_start_loc_gpu: torch.Tensor | None = None,
124
- ):
124
+ ) -> None:
125
n_seq = len(num_scheduled_tokens_np)
126
prompt_lens = self.prompt_lens
127