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

Pin scheduler metadata before asynchronous H2D copies (#35944)

Browse Source 
Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
This commit is contained in:
weireweire
2026-08-26 22:21:47 -07:00
committed by GitHub
co-authored by weireweire
parent 0f7b5b8b2a
commit ea48cb04cc
4 changed files with 49 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/forward_batch_info.py
+13 -7
View File
@@ -707,12 +707,15 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
self.original_global_num_tokens_cpu = batch.global_num_tokens
self.global_num_tokens_cpu = global_num_tokens
pin_memory = is_pin_memory_available(device)
self.global_num_tokens_gpu = torch.tensor(
global_num_tokens, dtype=torch.int64
global_num_tokens, dtype=torch.int64, pin_memory=pin_memory
).to(device, non_blocking=True)
self.global_num_tokens_for_logprob_cpu = global_num_tokens_for_logprob
self.global_num_tokens_for_logprob_gpu = torch.tensor(
global_num_tokens_for_logprob, dtype=torch.int64
global_num_tokens_for_logprob,
dtype=torch.int64,
pin_memory=pin_memory,
).to(device, non_blocking=True)
self.can_run_decode_cuda_graph = batch.can_run_decode_cuda_graph
@@ -850,9 +853,11 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
num_tokens = len(batch.input_ids) if batch.input_ids is not None else 0
if enable_num_token_non_padded():
ret.num_token_non_padded = torch.tensor(num_tokens, dtype=torch.int32).to(
device, non_blocking=True
)
ret.num_token_non_padded = torch.tensor(
num_tokens,
dtype=torch.int32,
pin_memory=is_pin_memory_available(device),
).to(device, non_blocking=True)
ret.num_token_non_padded_cpu = num_tokens
ret.init_mlp_sync_metadata(batch, device)
@@ -890,11 +895,12 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
if isinstance(extend_seq_lens, list):
# Main path: H2D from host lists; populate *_cpu mirrors.
assert isinstance(extend_prefix_lens, list)
pin_memory = is_pin_memory_available(device)
ret.extend_seq_lens = torch.tensor(
extend_seq_lens, dtype=torch.int32
extend_seq_lens, dtype=torch.int32, pin_memory=pin_memory
).to(device, non_blocking=True)
ret.extend_prefix_lens = torch.tensor(
extend_prefix_lens, dtype=torch.int32
extend_prefix_lens, dtype=torch.int32, pin_memory=pin_memory
).to(device, non_blocking=True)
ret.extend_prefix_lens_cpu = extend_prefix_lens
ret.extend_seq_lens_cpu = extend_seq_lens