[perf] remove several h2d sync (#28076)

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,

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
         )