From ea48cb04cc33233e01674e218102746931f72128 Mon Sep 17 00:00:00 2001 From: weireweire Date: Thu, 27 Aug 2026 13:21:47 +0800 Subject: [PATCH] Pin scheduler metadata before asynchronous H2D copies (#35944) Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com> --- python/sglang/srt/mem_cache/allocation.py | 13 ++++++++--- .../srt/model_executor/forward_batch_info.py | 20 ++++++++++------ .../dspark_components/dspark_draft.py | 23 +++++++++++-------- .../dspark_components/dspark_worker_v2.py | 16 +++++++++---- 4 files changed, 49 insertions(+), 23 deletions(-) diff --git a/python/sglang/srt/mem_cache/allocation.py b/python/sglang/srt/mem_cache/allocation.py index af5038ec1..45ff96aa1 100644 --- a/python/sglang/srt/mem_cache/allocation.py +++ b/python/sglang/srt/mem_cache/allocation.py @@ -302,8 +302,13 @@ def alloc_for_extend( ] # Create tensors for allocation - prefix_lens_cpu = torch.tensor(batch.prefix_lens, dtype=torch.int64) - extend_lens_cpu = torch.tensor(batch.extend_lens, dtype=torch.int64) + pin_memory = is_pin_memory_available(batch.device) + prefix_lens_cpu = torch.tensor( + batch.prefix_lens, dtype=torch.int64, pin_memory=pin_memory + ) + extend_lens_cpu = torch.tensor( + batch.extend_lens, dtype=torch.int64, pin_memory=pin_memory + ) prefix_lens_device = prefix_lens_cpu.to(batch.device, non_blocking=True) extend_lens_device = extend_lens_cpu.to(batch.device, non_blocking=True) @@ -311,7 +316,9 @@ def alloc_for_extend( req_pool_indices = alloc_req_slots( batch.req_to_token_pool, batch.reqs, batch.tree_cache ) - req_pool_indices_cpu = torch.tensor(req_pool_indices, dtype=torch.int64) + req_pool_indices_cpu = torch.tensor( + req_pool_indices, dtype=torch.int64, pin_memory=pin_memory + ) req_pool_indices_device = req_pool_indices_cpu.to(batch.device, non_blocking=True) # Allocate KV cache (throws exception on failure) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index b2515aaaa..7c504656f 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -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 diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index e0c2bca89..3217393f9 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft.py @@ -28,6 +28,7 @@ from sglang.srt.speculative.spec_info import ( spec_scale_global_num_tokens, ) from sglang.srt.speculative.spec_utils import draft_tp_context +from sglang.srt.utils.common import is_pin_memory_available from sglang.srt.utils.invariants import Bucket, Invariant, NotNaN, expect logger = logging.getLogger(__name__) @@ -55,7 +56,11 @@ def _make_num_token_non_padded( ) -> Optional[torch.Tensor]: if not enable_num_token_non_padded(): return None - return torch.tensor(num_tokens, dtype=torch.int32).to(device, non_blocking=True) + return torch.tensor( + num_tokens, + dtype=torch.int32, + pin_memory=is_pin_memory_available(device), + ).to(device, non_blocking=True) class DraftBlockResult(msgspec.Struct, frozen=True): @@ -467,16 +472,16 @@ class DraftBlockProposer: device = self.draft_model_runner.device forward_batch.original_global_num_tokens_cpu = batch.global_num_tokens num_tokens = forward_batch.input_ids.numel() - if enable_num_token_non_padded(): - forward_batch.num_token_non_padded = torch.tensor( - num_tokens, dtype=torch.int32, device=device - ) + num_token_non_padded = _make_num_token_non_padded(num_tokens, device) + if num_token_non_padded is not None: + forward_batch.num_token_non_padded = num_token_non_padded forward_batch.num_token_non_padded_cpu = num_tokens forward_batch.global_num_tokens_cpu = gnt forward_batch.global_num_tokens_for_logprob_cpu = gnt_logprob - forward_batch.global_num_tokens_gpu = torch.tensor(gnt, dtype=torch.int64).to( - device, non_blocking=True - ) + pin_memory = is_pin_memory_available(device) + forward_batch.global_num_tokens_gpu = torch.tensor( + gnt, dtype=torch.int64, pin_memory=pin_memory + ).to(device, non_blocking=True) forward_batch.global_num_tokens_for_logprob_gpu = torch.tensor( - gnt_logprob, dtype=torch.int64 + gnt_logprob, dtype=torch.int64, pin_memory=pin_memory ).to(device, non_blocking=True) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 08c8d7896..4fad38915 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -74,7 +74,12 @@ from sglang.srt.speculative.spec_utils import ( draft_tp_context, prepare_mamba_track_for_verify, ) -from sglang.srt.utils import get_available_gpu_memory, is_cuda, is_npu +from sglang.srt.utils import ( + get_available_gpu_memory, + is_cuda, + is_npu, + is_pin_memory_available, +) logger = logging.getLogger(__name__) @@ -477,10 +482,13 @@ class DSparkWorkerV2(BaseSpecWorker): # Must inject before prefill returns: the scheduler may update radix # afterward, invalidating out_cache_loc. device = next_token_ids.device - ctx_lens = torch.tensor(batch.extend_lens, dtype=torch.int32, device=device) + pin_memory = is_pin_memory_available(device) + ctx_lens = torch.tensor( + batch.extend_lens, dtype=torch.int32, pin_memory=pin_memory + ).to(device, non_blocking=True) draft_seq_lens = torch.tensor( - batch.prefix_lens, dtype=torch.int32, device=device - ) + batch.prefix_lens, dtype=torch.int32, pin_memory=pin_memory + ).to(device, non_blocking=True) positions, _ = compute_position( self.model_runner.prefill_attention_backend_str, draft_seq_lens,