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 6, 23:00 UTC
fix bugs when token_classify & classify run concurrently (#36614)
Signed-off-by: augusto.yjh <augusto.yjh@antgroup.com>
Augusto Yao committed
Mar 11, 2026 at 11:16 UTC
b386bb3d7c871f380b96d0ec0f74c53ed4cadf62
1 file changed
+7
-4
vllm/model_executor/layers/pooler/tokwise/methods.py
+7
-4
@@ -47,10 +47,13 @@ class AllPool(TokenPoolingMethod):
47
pooling_metadata: PoolingMetadata,
48
) -> list[TokenPoolingMethodOutputItem]:
49
pooling_cursor = pooling_metadata.get_pooling_cursor()
50
- hidden_states_all = hidden_states.split(
51
- pooling_cursor.num_scheduled_tokens_cpu.tolist()
52
- )
53
- hidden_states_lst = [hidden_states_all[i] for i in pooling_cursor.index]
50
+ hidden_states_lst = [
51
+ hidden_states[first : last + 1]
52
+ for first, last in zip(
53
+ pooling_cursor.first_token_indices_gpu.tolist(),
54
+ pooling_cursor.last_token_indices_gpu.tolist(),
55
+ )
56
+ ]
57
58
if not self.enable_chunked_prefill:
59
return hidden_states_lst