From e02f7ca482c18e628a1ed303b49c3dff279d53a1 Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Fri, 12 Jun 2026 20:30:17 -0700 Subject: [PATCH] [perf] remove several h2d sync (#28076) --- python/sglang/srt/mem_cache/common.py | 6 +-- .../srt/sampling/sampling_batch_info.py | 38 ++++++++++++------- 2 files changed, 28 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 358d14e36..f52e668df 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -22,7 +22,7 @@ from sglang.srt.mem_cache.triton_ops.common import ( ) from sglang.srt.server_args import ServerArgs, get_global_server_args from sglang.srt.utils import is_hip, support_triton -from sglang.srt.utils.common import ceil_align +from sglang.srt.utils.common import ceil_align, is_pin_memory_available _is_hip = is_hip() @@ -121,9 +121,9 @@ def write_cache_indices( if support_triton(get_global_server_args().attention_backend): prefix_pointers = torch.tensor( [t.data_ptr() for t in prefix_tensors], - device=req_to_token_pool.device, dtype=torch.uint64, - ) + pin_memory=is_pin_memory_available(req_to_token_pool.device), + ).to(req_to_token_pool.device, non_blocking=True) # TODO: some tensors can be reused for ForwardBatchInfo (e.g., extend_lens, cumsum_start) write_req_to_token_pool_triton[(req_pool_indices_tensor.shape[0],)]( req_to_token_pool.req_to_token, diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index 492bb7dae..cfb22d841 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/python/sglang/srt/sampling/sampling_batch_info.py @@ -11,6 +11,7 @@ from sglang.srt.sampling.custom_logit_processor import CustomLogitProcessor from sglang.srt.sampling.penaltylib.repetition_penalty import apply_scaling_penalties from sglang.srt.sampling.sampling_params import TOP_K_ALL from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils.common import is_pin_memory_available if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import ScheduleBatch @@ -79,20 +80,31 @@ class SamplingBatchInfo: reqs = batch.reqs device = batch.device - temperatures = torch.tensor( - [r.sampling_params.temperature for r in reqs], - dtype=torch.float, - device=device, - ).view(-1, 1) + _pin = is_pin_memory_available(device) + temperatures = ( + torch.tensor( + [r.sampling_params.temperature for r in reqs], + dtype=torch.float, + pin_memory=_pin, + ) + .to(device, non_blocking=True) + .view(-1, 1) + ) top_ps = torch.tensor( - [r.sampling_params.top_p for r in reqs], dtype=torch.float, device=device - ) + [r.sampling_params.top_p for r in reqs], + dtype=torch.float, + pin_memory=_pin, + ).to(device, non_blocking=True) top_ks = torch.tensor( - [r.sampling_params.top_k for r in reqs], dtype=torch.int32, device=device - ) + [r.sampling_params.top_k for r in reqs], + dtype=torch.int32, + pin_memory=_pin, + ).to(device, non_blocking=True) min_ps = torch.tensor( - [r.sampling_params.min_p for r in reqs], dtype=torch.float, device=device - ) + [r.sampling_params.min_p for r in reqs], + dtype=torch.float, + pin_memory=_pin, + ).to(device, non_blocking=True) sampling_seed = ( torch.tensor( [ @@ -104,8 +116,8 @@ class SamplingBatchInfo: for r in reqs ], dtype=torch.int64, - device=device, - ) + pin_memory=_pin, + ).to(device, non_blocking=True) if enable_deterministic else None )