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 14, 05:30 UTC
[Disagg][Perf] Use CUDA event sync instead of blocking `tolist` to avoid unintentional copy ops blocking across different CUDA streams, improving disagg TTIT/TTFT (#22760)
Signed-off-by: Zijing Liu <liuzijing2014@gmail.com> Signed-off-by: Zijing Liu <liuzijing2014@users.noreply.github.com>
Zijing Liu committed
Aug 25, 2025 at 21:06 UTC
b395b3b0a3166d17c75e74f4eaf0ff4b15f2554f
1 file changed
+23
-2
vllm/v1/worker/gpu_model_runner.py
+23
-2
@@ -316,6 +316,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
316
# Cached outputs.
317
self._draft_token_ids: Optional[Union[list[list[int]],
318
torch.Tensor]] = None
319
+ self.transfer_event = torch.cuda.Event()
320
+ self.sampled_token_ids_pinned_cpu = torch.empty(
321
+ (self.max_model_len, 1),
322
+ dtype=torch.int64,
323
+ device="cpu",
324
+ pin_memory=True)
325
326
def _make_buffer(self, *args, dtype: torch.dtype) -> CpuGpuBuffer:
327
return CpuGpuBuffer(*args,
@@ -1691,7 +1697,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1697
max_gen_len = sampled_token_ids.shape[-1]
1698
if max_gen_len == 1:
1699
# No spec decode tokens.
1694
- valid_sampled_token_ids = sampled_token_ids.tolist()
1700
+ valid_sampled_token_ids = self._to_list(sampled_token_ids)
1701
else:
1702
# Includes spec decode tokens.
1703
valid_sampled_token_ids = self.rejection_sampler.parse_output(
@@ -2219,7 +2225,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2225
- CUDAGraphMode.PIECEWISE: Piecewise cudagraph.
2226
- CUDAGraphMode.FULL: Full cudagraph, attention metadata is
2227
needed.
2222
- force_attention: If True, always create attention metadata. Used to
2228
+ force_attention: If True, always create attention metadata. Used to
2229
warm up attention backend when mode is NONE.
2230
uniform_decode: If True, the batch is a uniform decode batch.
2231
skip_eplb: If True, skip EPLB state update.
@@ -3233,3 +3239,18 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3239
mamba_type=mamba_module.mamba_type)
3240
3241
return kv_cache_spec
3242
+
3243
+ def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]:
3244
+ # This is a short term mitigation for issue mentioned in
3245
+ # https://github.com/vllm-project/vllm/issues/22754.
3246
+ # `tolist` would trigger a cuda wise stream sync, which
3247
+ # would block other copy ops from other cuda streams.
3248
+ # A cuda event sync would avoid such a situation. Since
3249
+ # this is in the critical path of every single model
3250
+ # forward loop, this has caused perf issue for a disagg
3251
+ # setup.
3252
+ pinned = self.sampled_token_ids_pinned_cpu[:sampled_token_ids.shape[0]]
3253
+ pinned.copy_(sampled_token_ids, non_blocking=True)
3254
+ self.transfer_event.record()
3255
+ self.transfer_event.synchronize()
3256
+ return pinned.tolist()