[core] step 2: drop seq_lens sentinel; SB maintains GPU as seq_lens_cpu mirror (#26020)

python/sglang/srt/managers/overlap_utils.py

@@ -42,6 +42,17 @@ class FutureIndices:
 
 
 class FutureMap:
+    """Cross-iter relay buffer for values the next iter's schedule cannot
+    compute locally (e.g. spec_v2 seq_lens after accept_lens, sampled tokens).
+
+    Forward stream publishes into a buf; next iter's schedule pulls lazily.
+    Schedule-deterministic values (e.g. non-spec seq_lens via +1) stay
+    maintained by SB directly and do not need the relay.
+
+    SB.seq_lens GPU is always a faithful seq_lens_cpu mirror; forward path
+    treats it as read-only, spec mutations land on forward_batch.seq_lens.
+    """
+
     def __init__(
         self,
         device: torch.device,
@@ -89,10 +100,9 @@ class FutureMap:
             )
 
     def resolve_future(self, batch: ScheduleBatch):
-        if batch.forward_mode.is_decode():
-            batch.seq_lens = self.new_seq_lens_buf[batch.req_pool_indices]
-            torch._assert_async((batch.seq_lens > 0).all())
-
+        # seq_lens is already real on entry (SB +1 for non-spec;
+        # resolve_seq_lens_cpu pulled from buf for spec_v2). Only resolve
+        # input_ids tokens / spec extras here.
         if self.spec_algo.is_none():
             _resolve_future_token_ids(batch.input_ids, self.output_tokens_buf)
         else:
@@ -113,18 +123,26 @@ class FutureMap:
         if spec_need_hidden_states():
             draft_input.hidden_states = self.hidden_states_buf[indices]
 
-    def invalidate(self, batch: ScheduleBatch, future_indices: FutureIndices) -> None:
-        sentinel = -future_indices.indices
-        batch.input_ids = sentinel
-        batch.seq_lens = sentinel
+    def set_input_ids_sentinel(
+        self, batch: ScheduleBatch, future_indices: FutureIndices
+    ) -> None:
+        # Sentinel for the decode portion so mixed batches can cat extend
+        # (positive real tokens) + decode (negative sentinels) into one
+        # input_ids; resolve_future translates negatives via output_tokens_buf.
+        batch.input_ids = -future_indices.indices
 
     def resolve_seq_lens_cpu(self, batch: ScheduleBatch) -> None:
+        # Lazy pull from new_seq_lens_buf for spec_v2 (accept_lens not known to
+        # schedule). Write into both CPU and GPU so SB.seq_lens stays a faithful
+        # seq_lens_cpu mirror.
         fi = batch.spec_info.future_indices if batch.spec_info is not None else None
         if fi is None:
             return
         if self.publish_ready is not None:
             self.publish_ready.wait()
-        batch.seq_lens_cpu = self.new_seq_lens_buf[fi.indices].cpu()
+        new_seq_lens = self.new_seq_lens_buf[fi.indices]
+        batch.seq_lens = new_seq_lens
+        batch.seq_lens_cpu = new_seq_lens.cpu()
         batch.seq_lens_sum = int(batch.seq_lens_cpu.sum())
 
     def publish(

python/sglang/srt/managers/schedule_batch.py

@@ -2161,14 +2161,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             req.fill_ids = req.origin_input_ids + req.output_ids
             req.set_extend_input_len(1)
 
-        if running_batch.enable_overlap:
-            # running_batch.seq_lens (GPU) is the FutureMap sentinel between iters;
-            # restore from CPU shadow before merge so MIXED's seq_lens has real values.
-            # (resolve_future only restores for is_decode(), not is_mixed().)
-            running_batch.seq_lens = running_batch.seq_lens_cpu.to(
-                running_batch.device, non_blocking=True
-            )
-
         input_ids = torch.cat([self.input_ids, running_batch.input_ids])
         out_cache_loc = torch.cat([self.out_cache_loc, running_batch.out_cache_loc])
 
@@ -2419,13 +2411,13 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             req.kv_committed_len += 1
             req.kv_allocated_len += 1
 
-        # Update seq_lens after allocation
         if self.enable_overlap:
-            # Overlap: GPU seq_lens restored by resolve_future from FutureMap buf.
+            # New-tensor avoids racing model_worker_batch refs queued for
+            # overlap forward.
+            self.seq_lens = self.seq_lens + 1
             self.seq_lens_cpu = self.seq_lens_cpu + 1
             self.orig_seq_lens = self.orig_seq_lens + 1
         else:
-            # A faster in-place version
             self.seq_lens.add_(1)
             self.seq_lens_cpu.add_(1)
             self.orig_seq_lens.add_(1)

python/sglang/srt/managers/scheduler.py

@@ -2890,7 +2890,7 @@ class Scheduler(
                         else:
                             batch_result.future_indices = future_indices
 
-                self.future_map.invalidate(batch, future_indices)
+                self.future_map.set_input_ids_sentinel(batch, future_indices)
 
                 if batch.is_spec_v2:
                     batch.spec_info = batch_result.next_draft_input