diff --git a/tests/model_executor/layers/test_pooler_methods.py b/tests/model_executor/layers/test_pooler_methods.py index cb8533cacb..28b2fc7a78 100644 --- a/tests/model_executor/layers/test_pooler_methods.py +++ b/tests/model_executor/layers/test_pooler_methods.py @@ -119,7 +119,7 @@ class TestCLSPool: hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) metadata = _make_metadata([3], num_scheduled_tokens=[2]) pooler = CLSPool() - with pytest.raises(AssertionError, match="partial prefill"): + with pytest.raises(RuntimeError, match="partial prefill"): pooler(hidden, metadata) @@ -202,7 +202,7 @@ class TestMeanPool: hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32) metadata = _make_metadata([3], num_scheduled_tokens=[2]) pooler = MeanPool() - with pytest.raises(AssertionError, match="partial prefill"): + with pytest.raises(RuntimeError, match="partial prefill"): pooler(hidden, metadata) def test_chunked_accumulation(self): diff --git a/vllm/config/pooler.py b/vllm/config/pooler.py index b52b0abd1d..37bcc1b223 100644 --- a/vllm/config/pooler.py +++ b/vllm/config/pooler.py @@ -138,11 +138,19 @@ class PoolerConfig: raise NotImplementedError(pooling_type) def get_seq_pooling_type(self) -> SequencePoolingType: - assert self.seq_pooling_type is not None, "Should be resolved by ModelConfig" + if self.seq_pooling_type is None: + raise ValueError( + "seq_pooling_type is not set; it should be resolved by" + " ModelConfig before calling get_seq_pooling_type()" + ) return self.seq_pooling_type def get_tok_pooling_type(self) -> TokenPoolingType: - assert self.tok_pooling_type is not None, "Should be resolved by ModelConfig" + if self.tok_pooling_type is None: + raise ValueError( + "tok_pooling_type is not set; it should be resolved by" + " ModelConfig before calling get_tok_pooling_type()" + ) return self.tok_pooling_type def compute_hash(self) -> str: diff --git a/vllm/model_executor/layers/pooler/seqwise/heads.py b/vllm/model_executor/layers/pooler/seqwise/heads.py index 2424d84107..efeda47a7d 100644 --- a/vllm/model_executor/layers/pooler/seqwise/heads.py +++ b/vllm/model_executor/layers/pooler/seqwise/heads.py @@ -52,7 +52,11 @@ class EmbeddingPoolerHead(SequencePoolerHead): pooling_metadata: PoolingMetadata, ) -> SequencePoolerHeadOutput: pooling_params = pooling_metadata.pooling_params - assert len(pooled_data) == len(pooling_params) + if len(pooled_data) != len(pooling_params): + raise ValueError( + f"pooled_data length ({len(pooled_data)}) does not match " + f"pooling_params length ({len(pooling_params)})" + ) if isinstance(pooled_data, list): pooled_data = torch.stack(pooled_data) @@ -72,7 +76,11 @@ class EmbeddingPoolerHead(SequencePoolerHead): dimensions_list = [pooling_param.dimensions for pooling_param in pooling_params] if any(d is not None for d in dimensions_list): # change the output dimension - assert len(embeddings) == len(dimensions_list) + if len(embeddings) != len(dimensions_list): + raise ValueError( + f"embeddings length ({len(embeddings)}) does not match " + f"dimensions_list length ({len(dimensions_list)})" + ) if len(set(dimensions_list)) == 1 and not isinstance(embeddings, list): # if all dimensions are the same d = dimensions_list[0] @@ -125,7 +133,11 @@ class ClassifierPoolerHead(SequencePoolerHead): pooling_metadata: PoolingMetadata, ) -> SequencePoolerHeadOutput: pooling_params = pooling_metadata.pooling_params - assert len(pooled_data) == len(pooling_params) + if len(pooled_data) != len(pooling_params): + raise ValueError( + f"pooled_data length ({len(pooled_data)}) does not match " + f"pooling_params length ({len(pooling_params)})" + ) if isinstance(pooled_data, list): pooled_data = torch.stack(pooled_data) diff --git a/vllm/model_executor/layers/pooler/seqwise/methods.py b/vllm/model_executor/layers/pooler/seqwise/methods.py index d99216fc10..06dddde7de 100644 --- a/vllm/model_executor/layers/pooler/seqwise/methods.py +++ b/vllm/model_executor/layers/pooler/seqwise/methods.py @@ -40,9 +40,8 @@ class CLSPool(SequencePoolingMethod): pooling_metadata: PoolingMetadata, ) -> SequencePoolingMethodOutput: pooling_cursor = pooling_metadata.get_pooling_cursor() - assert not pooling_cursor.is_partial_prefill(), ( - "partial prefill not supported with CLS pooling" - ) + if pooling_cursor.is_partial_prefill(): + raise RuntimeError("partial prefill is not supported with CLS pooling") return hidden_states[pooling_cursor.first_token_indices_gpu] @@ -64,9 +63,8 @@ class MeanPool(SequencePoolingMethod): pooling_metadata: PoolingMetadata, ) -> SequencePoolingMethodOutput: pooling_cursor = pooling_metadata.get_pooling_cursor() - assert not pooling_cursor.is_partial_prefill(), ( - "partial prefill not supported with MEAN pooling" - ) + if pooling_cursor.is_partial_prefill(): + raise RuntimeError("partial prefill is not supported with MEAN pooling") prompt_lens_cpu = pooling_cursor.prompt_lens_cpu num_seqs = prompt_lens_cpu.numel() diff --git a/vllm/model_executor/layers/pooler/seqwise/poolers.py b/vllm/model_executor/layers/pooler/seqwise/poolers.py index 81e19001b7..2894bc45c6 100644 --- a/vllm/model_executor/layers/pooler/seqwise/poolers.py +++ b/vllm/model_executor/layers/pooler/seqwise/poolers.py @@ -115,7 +115,10 @@ def pooler_for_classify( vllm_config = get_current_vllm_config() model_config = vllm_config.model_config - assert model_config.pooler_config is not None + if model_config.pooler_config is None: + raise ValueError( + "model_config.pooler_config must be set for classification pooling" + ) head = ClassifierPoolerHead( head_dtype=model_config.head_dtype, classifier=classifier, diff --git a/vllm/model_executor/layers/pooler/tokwise/heads.py b/vllm/model_executor/layers/pooler/tokwise/heads.py index d9f41132c0..4f4580ae12 100644 --- a/vllm/model_executor/layers/pooler/tokwise/heads.py +++ b/vllm/model_executor/layers/pooler/tokwise/heads.py @@ -36,7 +36,11 @@ class TokenPoolerHead(nn.Module, ABC): pooling_metadata: PoolingMetadata, ) -> list[TokenPoolerHeadOutputItem]: pooling_params = pooling_metadata.pooling_params - assert len(pooled_data) == len(pooling_params) + if len(pooled_data) != len(pooling_params): + raise ValueError( + f"pooled_data length ({len(pooled_data)}) does not match " + f"pooling_params length ({len(pooling_params)})" + ) return [self.forward_chunk(d, p) for d, p in zip(pooled_data, pooling_params)] diff --git a/vllm/model_executor/layers/pooler/tokwise/poolers.py b/vllm/model_executor/layers/pooler/tokwise/poolers.py index 0b66097e38..ec6624f431 100644 --- a/vllm/model_executor/layers/pooler/tokwise/poolers.py +++ b/vllm/model_executor/layers/pooler/tokwise/poolers.py @@ -124,7 +124,10 @@ def pooler_for_token_classify( vllm_config = get_current_vllm_config() model_config = vllm_config.model_config - assert model_config.pooler_config is not None + if model_config.pooler_config is None: + raise ValueError( + "model_config.pooler_config must be set for token classification pooling" + ) head = TokenClassifierPoolerHead( head_dtype=model_config.head_dtype, classifier=classifier,