[core] WAR barrier for overlap schedule buffer writes, without fwd occupancy cost (#26380)

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

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

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 = self.new_seq_lens_buf[fi]
-        batch.seq_lens = new_seq_lens
+
-        batch.seq_lens_cpu = new_seq_lens.cpu()
+        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:

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

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),
+            cur_kv_lens_device,
-            nxt_kv_lens_cpu.to(device=batch.device),
+            nxt_kv_lens_device,
             out_cache_loc,
             bs,
         )