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, 10:45 UTC
[Perf] Optimize token_embed for pooling models, 1.0% token throughput improvement (#37347)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Mar 18, 2026 at 22:27 UTC
e37ff5b5c8c1ca692f4a83df3ba4c564a4bf1df6
2 files changed
+24
-6
vllm/v1/pool/metadata.py
+20
-5
@@ -101,6 +101,7 @@ class PoolingMetadata:
101
num_scheduled_tokens_np: np.ndarray,
102
seq_lens_cpu: torch.Tensor,
103
device: torch.device,
104
+ query_start_loc_gpu: torch.Tensor | None = None,
105
):
106
n_seq = len(num_scheduled_tokens_np)
107
prompt_lens = self.prompt_lens
@@ -109,11 +110,25 @@ class PoolingMetadata:
110
111
index = list(range(n_seq))
112
num_scheduled_tokens_cpu = torch.from_numpy(num_scheduled_tokens_np)
112
- cumsum = torch.zeros(
113
- n_seq + 1, dtype=torch.int64, pin_memory=pin_memory, device="cpu"
114
- )
115
- torch.cumsum(num_scheduled_tokens_cpu, dim=0, out=cumsum[1:])
116
- cumsum = cumsum.to(device, non_blocking=True)
113
+ if query_start_loc_gpu is None:
114
+ cumsum = torch.zeros(
115
+ n_seq + 1, dtype=torch.int64, pin_memory=pin_memory, device="cpu"
116
+ )
117
+ torch.cumsum(num_scheduled_tokens_cpu, dim=0, out=cumsum[1:])
118
+ cumsum = cumsum.to(device, non_blocking=True)
119
+ else:
120
+ if query_start_loc_gpu.shape[0] != n_seq + 1:
121
+ raise ValueError(
122
+ "query_start_loc_gpu length does not match "
123
+ f"the number of sequences: {query_start_loc_gpu.shape[0]} "
124
+ f"!= {n_seq + 1}."
125
+ )
126
+ if query_start_loc_gpu.device != device:
127
+ raise ValueError(
128
+ "query_start_loc_gpu must be on the same device as the "
129
+ f"hidden states: {query_start_loc_gpu.device} != {device}."
130
+ )
131
+ cumsum = query_start_loc_gpu
132
self.pooling_cursor = PoolingCursor(
133
index=index,
134
first_token_indices_gpu=cumsum[:n_seq],
vllm/v1/worker/gpu_model_runner.py
+4
-1
@@ -2928,7 +2928,10 @@ class GPUModelRunner(
2928
2929
pooling_metadata = self.input_batch.get_pooling_metadata()
2930
pooling_metadata.build_pooling_cursor(
2931
- num_scheduled_tokens_np, seq_lens_cpu, device=hidden_states.device
2931
+ num_scheduled_tokens_np,
2932
+ seq_lens_cpu,
2933
+ device=hidden_states.device,
2934
+ query_start_loc_gpu=self.query_start_loc.gpu[: num_reqs + 1],
2935
)
2936
2937
model = cast(VllmModelForPooling, self.model)