diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 93208639f..27eb699b2 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1621,6 +1621,9 @@ class SchedulerDisaggregationDecodeMixin: if self._engine_paused: continue + # WAR barrier: this iter's schedule writes to shared GPU buffers wait for prev forward's reads. + self.schedule_stream.wait_stream(self.forward_stream) + # Get the next batch to run batch = self.get_next_disagg_decode_batch_to_run() self.cur_batch = batch diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 93dee0cc2..ca2237861 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -436,6 +436,9 @@ class SchedulerDisaggregationPrefillMixin: if self._engine_paused: continue + # WAR barrier on shared GPU buffers (req_to_token_pool / SWA mapping). + self.schedule_stream.wait_stream(self.forward_stream) + # Get the next batch to run batch = self.get_next_disagg_prefill_batch_to_run() self.cur_batch = batch diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 0dd42f08e..2c474c1ba 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -105,6 +105,16 @@ class FutureMap: self.new_seq_lens_buf = torch.empty( (self.req_pool_size,), dtype=torch.int64, device=self.device ) + # Pinned host copy of new_seq_lens_buf + private stream for fwd-prepare + # D2H pulls (gated only on publish, off the schedule stream). + if _is_cuda or _is_hip: + self.new_seq_lens_cpu_pinned = torch.empty( + (self.req_pool_size,), dtype=torch.int64, pin_memory=True + ) + self.fwd_prepare_d2h_stream = torch.get_device_module(self.device).Stream() + else: + self.new_seq_lens_cpu_pinned = None + self.fwd_prepare_d2h_stream = None if self.spec_algo.is_some(): self._forward_buf_initialized = False @@ -194,9 +204,20 @@ class FutureMap: return if self.publish_ready is not None: self.publish_ready.wait() - new_seq_lens = self.new_seq_lens_buf[fi] - batch.seq_lens = new_seq_lens - batch.seq_lens_cpu = new_seq_lens.cpu() + batch.seq_lens = self.new_seq_lens_buf[fi] + + if self.fwd_prepare_d2h_stream is None or self.publish_ready is None: + batch.seq_lens_cpu = batch.seq_lens.cpu() # bootstrap / non-CUDA + batch.seq_lens_sum = int(batch.seq_lens_cpu.sum()) + return + + # seq_lens_cpu off the schedule stream: D2H the relay buf on a private + # stream (gated on publish), host-select via req_pool_indices_cpu. + self.fwd_prepare_d2h_stream.wait_event(self.publish_ready) + with torch.get_device_module(self.device).stream(self.fwd_prepare_d2h_stream): + self.new_seq_lens_cpu_pinned.copy_(self.new_seq_lens_buf, non_blocking=True) + self.fwd_prepare_d2h_stream.synchronize() + batch.seq_lens_cpu = self.new_seq_lens_cpu_pinned[batch.req_pool_indices_cpu] batch.seq_lens_sum = int(batch.seq_lens_cpu.sum()) def publish(self, future_indices: torch.Tensor, new_seq_lens: torch.Tensor) -> None: diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 2824e562f..868a28d43 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1408,6 +1408,9 @@ class Scheduler( if self._engine_paused: continue + # WAR barrier: this iter's schedule writes to shared GPU buffers wait for prev forward's reads. + self.schedule_stream.wait_stream(self.forward_stream) + # Get the next batch to run batch = self.get_next_batch_to_run() self.cur_batch = batch diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index 390e2e78d..a4b5064a3 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -142,21 +142,23 @@ class EagleDraftInputV2Mixin: cur_kv_lens_cpu = torch.tensor(cur_kv_lens, dtype=torch.int32, device="cpu") nxt_kv_lens_cpu = torch.tensor(nxt_kv_lens, dtype=torch.int32, device="cpu") + # non_blocking H2D: a blocking .to() syncs the schedule stream, which the WAR + # barrier has chained to the prev forward -> host stalls a full forward. + cur_kv_lens_device = cur_kv_lens_cpu.to(device=batch.device, non_blocking=True) + nxt_kv_lens_device = nxt_kv_lens_cpu.to(device=batch.device, non_blocking=True) if page_size == 1: out_cache_loc = alloc_token_slots(batch.tree_cache, num_needed_tokens) else: - cur_kv_lens = cur_kv_lens_cpu.to(device=batch.device) - nxt_kv_lens = nxt_kv_lens_cpu.to(device=batch.device) last_loc = get_last_loc( batch.req_to_token_pool.req_to_token, batch.req_pool_indices, - cur_kv_lens, + cur_kv_lens_device, ) out_cache_loc = alloc_paged_token_slots_extend( batch.tree_cache, - cur_kv_lens, + cur_kv_lens_device, cur_kv_lens_cpu, - nxt_kv_lens, + nxt_kv_lens_device, nxt_kv_lens_cpu, last_loc, num_needed_tokens, @@ -165,8 +167,8 @@ class EagleDraftInputV2Mixin: assign_req_to_token_pool_func( batch.req_pool_indices, batch.req_to_token_pool.req_to_token, - cur_kv_lens_cpu.to(device=batch.device), - nxt_kv_lens_cpu.to(device=batch.device), + cur_kv_lens_device, + nxt_kv_lens_device, out_cache_loc, bs, )