minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[sgl]add pin_mem to avoid cpu->gpu copy sync point (#19590)

Browse Source
This commit is contained in:
Bi Xue
2026-03-02 21:08:31 +08:00
committed by GitHub
parent b3718982a1
commit eaf18ebe8d
1 changed files with 10 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+10 -7
View File
@@ -1449,21 +1449,22 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
r.token_type_ids for r in reqs if r.token_type_ids is not None 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( 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) ).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 self.device, non_blocking=True
) )
seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int64) seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int64)
orig_seq_lens_tensor = torch.tensor(orig_seq_lens, dtype=torch.int32).to( orig_seq_lens_tensor = torch.tensor(
self.device, non_blocking=True orig_seq_lens, dtype=torch.int32, pin_memory=_pin
) ).to(self.device, non_blocking=True)
token_type_ids_tensor = None token_type_ids_tensor = None
if len(token_type_ids) > 0: if len(token_type_ids) > 0:
token_type_ids_tensor = torch.tensor( 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) ).to(self.device, non_blocking=True)
# Set batch fields needed by alloc_for_extend # Set batch fields needed by alloc_for_extend
@@ -1601,7 +1602,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.orig_seq_lens = orig_seq_lens_tensor self.orig_seq_lens = orig_seq_lens_tensor
self.out_cache_loc = out_cache_loc self.out_cache_loc = out_cache_loc
self.input_embeds = ( 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 if input_embeds
else None else None
) )