chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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 05: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

17
python/sglang/srt/managers/schedule_batch.py
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
]
_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
)