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

[Perf] Overlap result D2H copy with the next forward step (#29075)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-24 01:34:02 -07:00
committed by GitHub
parent e04ed05193
commit 8e1988b746
4 changed files with 60 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/eplb/expert_distribution.py
+4 -2
View File
@@ -52,8 +52,10 @@ _OutputMode = Literal["file", "object"]
class ExpertDistributionMetrics:
eplb_balancedness: torch.Tensor
def copy_to_cpu(self):
self.eplb_balancedness = self.eplb_balancedness.to("cpu", non_blocking=True)
def map_device_tensors(self, fn):
# Device-tensor fields only; caller injects the copy+safety primitive
# (see GenerationBatchResult.copy_to_cpu).
self.eplb_balancedness = fn(self.eplb_balancedness)
class ExpertDistributionRecorder(ABC):
python/sglang/srt/managers/scheduler.py
+9 -4
View File
@@ -3253,10 +3253,15 @@ class Scheduler(
else batch_result.next_token_ids
)
self.future_map.stash(future_indices, stash_payload)
batch_result.copy_to_cpu(
return_logprob=batch.return_logprob,
return_hidden_states=batch.return_hidden_states,
)
# Result D2H on copy_stream overlaps the next forward
# instead of serializing on forward_stream; it's a leaf
# gated by copy_done, so nothing on forward_stream waits.
self.copy_stream.wait_stream(self.forward_stream)
with self.copy_stream_ctx:
batch_result.copy_to_cpu(
return_logprob=batch.return_logprob,
return_hidden_states=batch.return_hidden_states,
)
else:
batch_result.future_indices = future_indices
python/sglang/srt/managers/utils.py
+40 -27
View File
@@ -22,6 +22,19 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def _async_d2h(t: torch.Tensor) -> torch.Tensor:
"""Async D2H copy for overlap scheduling. On CUDA the dest is pinned (a D2H
to pageable host memory blocks the caller until done) and record_stream keeps
the source alive until the copy stream drains, so the caching allocator can't
recycle it early. Non-CUDA falls back to a plain copy."""
if not t.is_cuda:
return t.to("cpu", non_blocking=True)
cpu_t = torch.empty(t.shape, dtype=t.dtype, pin_memory=True)
cpu_t.copy_(t, non_blocking=True)
t.record_stream(torch.cuda.current_stream(t.device))
return cpu_t
@dataclasses.dataclass
class GenerationBatchResult:
logits_output: Optional[LogitsProcessorOutput] = None
@@ -72,6 +85,7 @@ class GenerationBatchResult:
fpm_start_event: Optional[torch.cuda.Event] = None
fpm_end_event: Optional[torch.cuda.Event] = None
@torch.profiler.record_function("copy_result_to_cpu")
def copy_to_cpu(self, return_logprob: bool, return_hidden_states: bool = True):
"""Copy tensors to CPU in overlap scheduling.
Only the tensors which are needed for processing results are copied,
@@ -79,45 +93,47 @@ class GenerationBatchResult:
"""
if return_logprob:
if self.logits_output.next_token_logprobs is not None:
self.logits_output.next_token_logprobs = (
self.logits_output.next_token_logprobs.to("cpu", non_blocking=True)
self.logits_output.next_token_logprobs = _async_d2h(
self.logits_output.next_token_logprobs
)
if self.logits_output.input_token_logprobs is not None:
self.logits_output.input_token_logprobs = (
self.logits_output.input_token_logprobs.to("cpu", non_blocking=True)
self.logits_output.input_token_logprobs = _async_d2h(
self.logits_output.input_token_logprobs
)
if self.logits_output.next_token_top_logprobs_val is not None:
self.logits_output.next_token_top_logprobs_val = [
v.to("cpu", non_blocking=True) if torch.is_tensor(v) else v
_async_d2h(v) if torch.is_tensor(v) else v
for v in self.logits_output.next_token_top_logprobs_val
]
if self.logits_output.next_token_top_logprobs_idx is not None:
self.logits_output.next_token_top_logprobs_idx = [
x.to("cpu", non_blocking=True) if torch.is_tensor(x) else x
_async_d2h(x) if torch.is_tensor(x) else x
for x in self.logits_output.next_token_top_logprobs_idx
]
if self.logits_output.next_token_token_ids_logprobs_val is not None:
self.logits_output.next_token_token_ids_logprobs_val = [
v.to("cpu", non_blocking=True) if torch.is_tensor(v) else v
_async_d2h(v) if torch.is_tensor(v) else v
for v in self.logits_output.next_token_token_ids_logprobs_val
]
if return_hidden_states and self.logits_output.hidden_states is not None:
self.logits_output.hidden_states = self.logits_output.hidden_states.to(
"cpu", non_blocking=True
self.logits_output.hidden_states = _async_d2h(
self.logits_output.hidden_states
)
self.next_token_ids = self.next_token_ids.to("cpu", non_blocking=True)
self.next_token_ids = _async_d2h(self.next_token_ids)
if self.accept_lens is not None:
self.accept_lens = self.accept_lens.to("cpu", non_blocking=True)
self.accept_lens = _async_d2h(self.accept_lens)
if self.routed_experts_output is not None:
self.routed_experts_output.copy_to_cpu()
if self.indexer_topk_output is not None:
self.indexer_topk_output.copy_to_cpu()
if (x := self.expert_distribution_metrics) is not None:
x.copy_to_cpu()
# Sub-objects only declare their device fields; the single copy+safety
# primitive (_async_d2h: pinned D2H + record_stream) is injected here so
# all device->host copying and lifetime safety lives in one place.
for holder in (
self.routed_experts_output,
self.indexer_topk_output,
self.expert_distribution_metrics,
):
if holder is not None:
holder.map_device_tensors(_async_d2h)
self.copy_done.record()
@@ -243,30 +259,27 @@ class EmbeddingBatchResult:
def can_run_cuda_graph(self) -> bool:
return False
@torch.profiler.record_function("copy_embedding_to_cpu")
def copy_to_cpu(self):
"""Copy embeddings and pooled hidden states to CPU for overlap scheduling."""
if isinstance(self.embeddings, torch.Tensor):
self.copy_done = torch.get_device_module(self.embeddings.device).Event()
self.embeddings = self.embeddings.to("cpu", non_blocking=True)
self.embeddings = _async_d2h(self.embeddings)
else:
assert isinstance(self.embeddings, list)
if len(self.embeddings) == 0:
return
self.copy_done = torch.get_device_module(self.embeddings[0].device).Event()
self.embeddings = [
emb.to("cpu", non_blocking=True) for emb in self.embeddings
]
self.embeddings = [_async_d2h(emb) for emb in self.embeddings]
if self.pooled_hidden_states is not None:
if isinstance(self.pooled_hidden_states, list):
self.pooled_hidden_states = [
t.to("cpu", non_blocking=True) for t in self.pooled_hidden_states
_async_d2h(t) for t in self.pooled_hidden_states
]
else:
self.pooled_hidden_states = self.pooled_hidden_states.to(
"cpu", non_blocking=True
)
self.pooled_hidden_states = _async_d2h(self.pooled_hidden_states)
self.copy_done.record()
python/sglang/srt/state_capturer/base.py
+7 -5
View File
@@ -79,17 +79,19 @@ class BaseHostCache:
@dataclasses.dataclass
class TopkCaptureOutput:
"""Holds GPU tensors captured during forward for overlap scheduling.
Call copy_to_cpu() inside forward stream (before copy_done.record()),
then finalize() after copy_done.synchronize().
map_device_tensors() D2H-copies them before copy_done.record() (may run on
the dedicated result-copy stream); finalize() runs after copy_done.synchronize().
"""
out_cache_loc: torch.Tensor
topk: torch.Tensor
host_cache: BaseHostCache
def copy_to_cpu(self):
self.out_cache_loc = self.out_cache_loc.to("cpu", non_blocking=True)
self.topk = self.topk.to("cpu", non_blocking=True)
def map_device_tensors(self, fn):
# Device-tensor fields only; caller injects the copy+safety primitive
# (see GenerationBatchResult.copy_to_cpu).
self.out_cache_loc = fn(self.out_cache_loc)
self.topk = fn(self.topk)
def finalize(self):
self.host_cache.buffer[self.out_cache_loc] = self.topk