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 4, 12:30 UTC
[Perf] Optimize `AllPool.forward` by slicing first, 51% faster in the method level benchmark (#41163)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Apr 29, 2026 at 11:11 UTC
39a7f4f4e2635012ead0ad127970d7b6778890af
1 file changed
+11
-7
vllm/model_executor/layers/pooler/tokwise/methods.py
+11
-7
@@ -47,13 +47,17 @@ class AllPool(TokenPoolingMethod):
47
pooling_metadata: PoolingMetadata,
48
) -> list[TokenPoolingMethodOutputItem]:
49
pooling_cursor = pooling_metadata.get_pooling_cursor()
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
- ]
50
+ split_sizes = pooling_cursor.num_scheduled_tokens_cpu.tolist()
51
+ if split_sizes:
52
+ # DispatchPooler passes the full hidden_states tensor.
53
+ # slice out the subgroup once, then split it by
54
+ # per-request token counts
55
+ group_start = int(pooling_cursor.first_token_indices_gpu[0].item())
56
+ group_end = int(pooling_cursor.last_token_indices_gpu[-1].item()) + 1
57
+ hidden_states_group = hidden_states[group_start:group_end]
58
+ hidden_states_lst = list(hidden_states_group.split(split_sizes))
59
+ else:
60
+ hidden_states_lst = []
61
62
if not self.enable_chunked_prefill:
63
return hidden_states_lst