address comments

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

vllm/model_executor/layers/pooler/tokwise/heads.py

@@ -119,7 +119,7 @@ class TokenEmbeddingPoolerHead(TokenPoolerHead):
         return embeddings
 
     def _has_uniform_postprocess(self, pooling_params: list[PoolingParams]) -> bool:
-        # check if we can apply the same postprocess to all tokens in the batch
+        """Return whether all pooling params share the same postprocess."""
         if not pooling_params:
             return True
 

vllm/model_executor/layers/pooler/tokwise/methods.py

@@ -38,6 +38,13 @@ class RaggedTokenBatch:
         return self.cu_lengths_cpu.shape[0] - 1
 
     def with_values(self, values: torch.Tensor) -> "RaggedTokenBatch":
+        expected_num_values = int(self.cu_lengths_cpu[-1])
+        if values.ndim == 0 or values.shape[0] != expected_num_values:
+            raise ValueError(
+                "values must preserve the flattened token dimension: "
+                f"{values.shape[0] if values.ndim > 0 else 0} "
+                f"!= {expected_num_values}."
+            )
         return RaggedTokenBatch(
             values=values,
             cu_lengths_cpu=self.cu_lengths_cpu,