From b409c230336eb2134f2aff45f6f57560cacd7b8d Mon Sep 17 00:00:00 2001 From: yewentao256 Date: Sun, 12 Apr 2026 16:37:05 +0000 Subject: [PATCH] address comments Signed-off-by: yewentao256 --- vllm/model_executor/layers/pooler/tokwise/heads.py | 2 +- vllm/model_executor/layers/pooler/tokwise/methods.py | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/pooler/tokwise/heads.py b/vllm/model_executor/layers/pooler/tokwise/heads.py index 0d8d0926ad9..9e8000b3206 100644 --- a/vllm/model_executor/layers/pooler/tokwise/heads.py +++ b/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 diff --git a/vllm/model_executor/layers/pooler/tokwise/methods.py b/vllm/model_executor/layers/pooler/tokwise/methods.py index b3f3fe07880..97d4fdc00fb 100644 --- a/vllm/model_executor/layers/pooler/tokwise/methods.py +++ b/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,