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 1, 20:00 UTC
[Misc] Replaced asserts with proper exceptions to improve UX for pooling (#44593)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Jun 6, 2026 at 00:57 UTC
eafbb06331caf580023b8da6871063499a37b403
7 files changed
+44
-16
tests/model_executor/layers/test_pooler_methods.py
+2
-2
@@ -119,7 +119,7 @@ class TestCLSPool:
119
hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
120
metadata = _make_metadata([3], num_scheduled_tokens=[2])
121
pooler = CLSPool()
122
- with pytest.raises(AssertionError, match="partial prefill"):
122
+ with pytest.raises(RuntimeError, match="partial prefill"):
123
pooler(hidden, metadata)
124
125
@@ -202,7 +202,7 @@ class TestMeanPool:
202
hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32)
203
metadata = _make_metadata([3], num_scheduled_tokens=[2])
204
pooler = MeanPool()
205
- with pytest.raises(AssertionError, match="partial prefill"):
205
+ with pytest.raises(RuntimeError, match="partial prefill"):
206
pooler(hidden, metadata)
207
208
def test_chunked_accumulation(self):
vllm/config/pooler.py
+10
-2
@@ -138,11 +138,19 @@ class PoolerConfig:
138
raise NotImplementedError(pooling_type)
139
140
def get_seq_pooling_type(self) -> SequencePoolingType:
141
- assert self.seq_pooling_type is not None, "Should be resolved by ModelConfig"
141
+ if self.seq_pooling_type is None:
142
+ raise ValueError(
143
+ "seq_pooling_type is not set; it should be resolved by"
144
+ " ModelConfig before calling get_seq_pooling_type()"
145
+ )
146
return self.seq_pooling_type
147
148
def get_tok_pooling_type(self) -> TokenPoolingType:
145
- assert self.tok_pooling_type is not None, "Should be resolved by ModelConfig"
149
+ if self.tok_pooling_type is None:
150
+ raise ValueError(
151
+ "tok_pooling_type is not set; it should be resolved by"
152
+ " ModelConfig before calling get_tok_pooling_type()"
153
+ )
154
return self.tok_pooling_type
155
156
def compute_hash(self) -> str:
vllm/model_executor/layers/pooler/seqwise/heads.py
+15
-3
@@ -52,7 +52,11 @@ class EmbeddingPoolerHead(SequencePoolerHead):
52
pooling_metadata: PoolingMetadata,
53
) -> SequencePoolerHeadOutput:
54
pooling_params = pooling_metadata.pooling_params
55
- assert len(pooled_data) == len(pooling_params)
55
+ if len(pooled_data) != len(pooling_params):
56
+ raise ValueError(
57
+ f"pooled_data length ({len(pooled_data)}) does not match "
58
+ f"pooling_params length ({len(pooling_params)})"
59
+ )
60
61
if isinstance(pooled_data, list):
62
pooled_data = torch.stack(pooled_data)
@@ -72,7 +76,11 @@ class EmbeddingPoolerHead(SequencePoolerHead):
76
dimensions_list = [pooling_param.dimensions for pooling_param in pooling_params]
77
if any(d is not None for d in dimensions_list):
78
# change the output dimension
75
- assert len(embeddings) == len(dimensions_list)
79
+ if len(embeddings) != len(dimensions_list):
80
+ raise ValueError(
81
+ f"embeddings length ({len(embeddings)}) does not match "
82
+ f"dimensions_list length ({len(dimensions_list)})"
83
+ )
84
if len(set(dimensions_list)) == 1 and not isinstance(embeddings, list):
85
# if all dimensions are the same
86
d = dimensions_list[0]
@@ -125,7 +133,11 @@ class ClassifierPoolerHead(SequencePoolerHead):
133
pooling_metadata: PoolingMetadata,
134
) -> SequencePoolerHeadOutput:
135
pooling_params = pooling_metadata.pooling_params
128
- assert len(pooled_data) == len(pooling_params)
136
+ if len(pooled_data) != len(pooling_params):
137
+ raise ValueError(
138
+ f"pooled_data length ({len(pooled_data)}) does not match "
139
+ f"pooling_params length ({len(pooling_params)})"
140
+ )
141
142
if isinstance(pooled_data, list):
143
pooled_data = torch.stack(pooled_data)
vllm/model_executor/layers/pooler/seqwise/methods.py
+4
-6
@@ -40,9 +40,8 @@ class CLSPool(SequencePoolingMethod):
40
pooling_metadata: PoolingMetadata,
41
) -> SequencePoolingMethodOutput:
42
pooling_cursor = pooling_metadata.get_pooling_cursor()
43
- assert not pooling_cursor.is_partial_prefill(), (
44
- "partial prefill not supported with CLS pooling"
45
- )
43
+ if pooling_cursor.is_partial_prefill():
44
+ raise RuntimeError("partial prefill is not supported with CLS pooling")
45
46
return hidden_states[pooling_cursor.first_token_indices_gpu]
47
@@ -64,9 +63,8 @@ class MeanPool(SequencePoolingMethod):
63
pooling_metadata: PoolingMetadata,
64
) -> SequencePoolingMethodOutput:
65
pooling_cursor = pooling_metadata.get_pooling_cursor()
67
- assert not pooling_cursor.is_partial_prefill(), (
68
- "partial prefill not supported with MEAN pooling"
69
- )
66
+ if pooling_cursor.is_partial_prefill():
67
+ raise RuntimeError("partial prefill is not supported with MEAN pooling")
68
69
prompt_lens_cpu = pooling_cursor.prompt_lens_cpu
70
num_seqs = prompt_lens_cpu.numel()
vllm/model_executor/layers/pooler/seqwise/poolers.py
+4
-1
@@ -115,7 +115,10 @@ def pooler_for_classify(
115
116
vllm_config = get_current_vllm_config()
117
model_config = vllm_config.model_config
118
- assert model_config.pooler_config is not None
118
+ if model_config.pooler_config is None:
119
+ raise ValueError(
120
+ "model_config.pooler_config must be set for classification pooling"
121
+ )
122
head = ClassifierPoolerHead(
123
head_dtype=model_config.head_dtype,
124
classifier=classifier,
vllm/model_executor/layers/pooler/tokwise/heads.py
+5
-1
@@ -36,7 +36,11 @@ class TokenPoolerHead(nn.Module, ABC):
36
pooling_metadata: PoolingMetadata,
37
) -> list[TokenPoolerHeadOutputItem]:
38
pooling_params = pooling_metadata.pooling_params
39
- assert len(pooled_data) == len(pooling_params)
39
+ if len(pooled_data) != len(pooling_params):
40
+ raise ValueError(
41
+ f"pooled_data length ({len(pooled_data)}) does not match "
42
+ f"pooling_params length ({len(pooling_params)})"
43
+ )
44
45
return [self.forward_chunk(d, p) for d, p in zip(pooled_data, pooling_params)]
46
vllm/model_executor/layers/pooler/tokwise/poolers.py
+4
-1
@@ -124,7 +124,10 @@ def pooler_for_token_classify(
124
125
vllm_config = get_current_vllm_config()
126
model_config = vllm_config.model_config
127
- assert model_config.pooler_config is not None
127
+ if model_config.pooler_config is None:
128
+ raise ValueError(
129
+ "model_config.pooler_config must be set for token classification pooling"
130
+ )
131
head = TokenClassifierPoolerHead(
132
head_dtype=model_config.head_dtype,
133
classifier=classifier,