[Misc] Replaced asserts with proper exceptions to improve UX for pooling (#44593)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>

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):

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:

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)

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()

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,

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)]
 

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,