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 5, 17:15 UTC
[Perf] Optimize mean pooling using chunks and index_add, 5.9% E2E throughput improvement (#38559)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Mar 31, 2026 at 23:54 UTC
7b01d97a22c977aaad2f38fa03cd11742d2893e8
1 file changed
+33
-9
vllm/model_executor/layers/pooler/seqwise/methods.py
+33
-9
@@ -14,6 +14,8 @@ from vllm.v1.pool.metadata import PoolingMetadata
14
15
SequencePoolingMethodOutput: TypeAlias = torch.Tensor | list[torch.Tensor]
16
17
+_MEAN_POOL_ACCUMULATION_CHUNK_BYTES = 16 * 1024 * 1024 # 16MB
18
+
19
20
class SequencePoolingMethod(nn.Module, ABC):
21
def get_supported_tasks(self) -> Set[PoolingTask]:
@@ -67,19 +69,41 @@ class MeanPool(SequencePoolingMethod):
69
)
70
71
prompt_lens = pooling_cursor.prompt_lens_cpu.to(
70
- hidden_states.device, non_blocking=True
72
+ hidden_states.device, dtype=torch.int64, non_blocking=True
73
+ )
74
+
75
+ num_seqs = prompt_lens.numel()
76
+ hidden_size = hidden_states.shape[-1]
77
+
78
+ if num_seqs == 0:
79
+ # early return for empty batch
80
+ return hidden_states.new_empty((0, hidden_size), dtype=torch.float32)
81
+
82
+ # eg. [2, 1, 3] -> [0, 0, 1, 2, 2, 2]
83
+ segment_ids = torch.repeat_interleave(
84
+ torch.arange(num_seqs, device=hidden_states.device, dtype=torch.long),
85
+ prompt_lens,
86
+ )
87
+ segment_sums = torch.zeros(
88
+ (num_seqs, hidden_size),
89
+ dtype=torch.float32,
90
+ device=hidden_states.device,
91
)
92
73
- # Use float32 for torch.cumsum in MeanPool,
74
- # otherwise precision will be lost significantly.
75
- cumsum = torch.cumsum(hidden_states, dim=0, dtype=torch.float32)
93
+ bytes_per_token = hidden_size * torch.finfo(torch.float32).bits // 8
94
+ chunk_size = max(1, _MEAN_POOL_ACCUMULATION_CHUNK_BYTES // bytes_per_token)
95
77
- start_indices = pooling_cursor.first_token_indices_gpu
78
- end_indices = pooling_cursor.last_token_indices_gpu
96
+ # iterate over the batch in chunks
97
+ for start in range(0, hidden_states.shape[0], chunk_size):
98
+ end = min(start + chunk_size, hidden_states.shape[0])
99
+ # using index_add_ to accumulate for each segment
100
+ segment_sums.index_add_(
101
+ 0,
102
+ segment_ids[start:end],
103
+ hidden_states[start:end].to(dtype=torch.float32),
104
+ )
105
80
- return (
81
- cumsum[end_indices] - cumsum[start_indices] + hidden_states[start_indices]
82
- ) / prompt_lens.unsqueeze(1)
106
+ return segment_sums / prompt_lens.unsqueeze(1)
107
108
109
def get_seq_pooling_method(pooling_type: SequencePoolingType | str):