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 7, 09:00 UTC
[Perf] [Hybrid] Copy num_accepted_tokens in non-blocking way when not using prefix caching (#35442)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Thomas Parnell committed
Mar 3, 2026 at 13:15 UTC
ad9d09e2b8a601b50d07c76fb8736c2bbda2d6fb
1 file changed
+12
-6
vllm/v1/worker/gpu_model_runner.py
+12
-6
@@ -1191,13 +1191,14 @@ class GPUModelRunner(
1191
return
1192
1193
# Find the number of accepted tokens for each sequence.
1194
- num_accepted_tokens = (
1194
+ num_reqs = output_token_ids.size(0)
1195
+ self.num_accepted_tokens.gpu[:num_reqs] = (
1196
(
1197
torch.cat(
1198
[
1199
output_token_ids,
1200
torch.full(
1200
- (output_token_ids.size(0), 1),
1201
+ (num_reqs, 1),
1202
-1,
1203
device=output_token_ids.device,
1204
),
@@ -1208,12 +1209,13 @@ class GPUModelRunner(
1209
)
1210
.int()
1211
.argmax(-1)
1211
- .cpu()
1212
- .numpy()
1212
)
1214
- for i, num_tokens in enumerate(num_accepted_tokens):
1215
- self.input_batch.num_accepted_tokens_cpu[i] = num_tokens
1213
if self.cache_config.mamba_cache_mode == "align":
1214
+ for i, num_tokens in enumerate(
1215
+ self.num_accepted_tokens.gpu[:num_reqs].cpu().numpy()
1216
+ ):
1217
+ self.input_batch.num_accepted_tokens_cpu[i] = num_tokens
1218
+
1219
mamba_utils.postprocess_mamba(
1220
scheduler_output,
1221
self.kv_cache_config,
@@ -1224,6 +1226,10 @@ class GPUModelRunner(
1226
self.model.get_mamba_state_copy_func(),
1227
self._get_mamba_copy_bufs(),
1228
)
1229
+ else:
1230
+ self.input_batch.num_accepted_tokens_cpu_tensor[:num_reqs].copy_(
1231
+ self.num_accepted_tokens.gpu[:num_reqs], non_blocking=True
1232
+ )
1233
1234
def _update_streaming_request(
1235
self, req_id: str, new_req_data: NewRequestData