From eaf18ebe8d461e88caa99244d9ee652419aa68dc Mon Sep 17 00:00:00 2001 From: Bi Xue Date: Mon, 2 Mar 2026 05:08:31 -0800 Subject: [PATCH] [sgl]add pin_mem to avoid cpu->gpu copy sync point (#19590) --- python/sglang/srt/managers/schedule_batch.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 6013069d3..3902a4caa 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1449,21 +1449,22 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): r.token_type_ids for r in reqs if r.token_type_ids is not None ] + _pin = self.device != "cpu" and torch.cuda.is_available() input_ids_tensor = torch.tensor( - list(chain.from_iterable(input_ids)), dtype=torch.int64 + list(chain.from_iterable(input_ids)), dtype=torch.int64, pin_memory=_pin ).to(self.device, non_blocking=True) - seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int64).to( + seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int64, pin_memory=_pin).to( self.device, non_blocking=True ) seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int64) - orig_seq_lens_tensor = torch.tensor(orig_seq_lens, dtype=torch.int32).to( - self.device, non_blocking=True - ) + orig_seq_lens_tensor = torch.tensor( + orig_seq_lens, dtype=torch.int32, pin_memory=_pin + ).to(self.device, non_blocking=True) token_type_ids_tensor = None if len(token_type_ids) > 0: token_type_ids_tensor = torch.tensor( - sum(token_type_ids, []), dtype=torch.int64 + sum(token_type_ids, []), dtype=torch.int64, pin_memory=_pin ).to(self.device, non_blocking=True) # Set batch fields needed by alloc_for_extend @@ -1601,7 +1602,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.orig_seq_lens = orig_seq_lens_tensor self.out_cache_loc = out_cache_loc self.input_embeds = ( - torch.tensor(input_embeds).to(self.device, non_blocking=True) + torch.tensor(input_embeds, pin_memory=_pin).to( + self.device, non_blocking=True + ) if input_embeds else None )