From c4b6b5ea1e0e7f9cd94638cd985082023d279fe5 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Fri, 22 May 2026 00:12:51 -0700 Subject: [PATCH] [core] step 2: drop seq_lens sentinel; SB maintains GPU as `seq_lens_cpu` mirror (#26020) --- .../decode_schedule_batch_mixin.py | 14 ++------ python/sglang/srt/managers/overlap_utils.py | 36 ++++++++++++++----- python/sglang/srt/managers/schedule_batch.py | 14 ++------ python/sglang/srt/managers/scheduler.py | 2 +- python/sglang/srt/mem_cache/common.py | 7 +--- .../sglang/srt/speculative/eagle_info_v2.py | 8 +++-- 6 files changed, 40 insertions(+), 41 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 92fd55ab7..0516e5a70 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -183,15 +183,7 @@ class ScheduleBatchDisaggregationDecodeMixin: future_map.stash(spec_info.future_indices, spec_info) self.spec_info = spec_info else: - # Non-spec: input_ids feeds the next decode forward directly. + # Non-spec: positive last token feeds decode directly. No FutureMap + # bootstrap needed (SB self-maintains seq_lens; resolve_future is + # a no-op on positive input_ids). self.input_ids = last_tokens_tensor - if self.enable_overlap: - from sglang.srt.managers.overlap_utils import FutureIndices - - future_indices = FutureIndices(indices=self.req_pool_indices) - # Bootstrap FutureMap so the first DECODE after PREBUILT can - # resolve_future from buf. Non-spec convention: batch.seq_lens - # at decode forward INCLUDES this iter's new token, so publish - # current + 1. - future_map.publish(future_indices, self.seq_lens + 1) - future_map.stash(future_indices, last_tokens_tensor) diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 1458191d2..52784d4c6 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/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( diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 9d266a221..e4f28e617 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 701a8eab2..621941f3a 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index ca3733e08..94d31c6c3 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -531,12 +531,7 @@ def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor: batch.maybe_evict_swa() - if batch.enable_overlap: - # batch.seq_lens (GPU) is a sentinel between iters (FutureMap.invalidate); - # materialize from CPU shadow for the allocator. Tensor stays local. - seq_lens_gpu = batch.seq_lens_cpu.to(batch.device, non_blocking=True) - else: - seq_lens_gpu = batch.seq_lens + seq_lens_gpu = batch.seq_lens bs = seq_lens_gpu.shape[0] if batch.tree_cache.page_size == 1: diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index 52d4ad4aa..a7c507297 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -226,9 +226,6 @@ class EagleDraftInputV2Mixin: batch.spec_info = self batch.input_ids = predict - batch.seq_lens = batch.seq_lens + num_draft_tokens - batch.seq_lens_cpu = batch.seq_lens_cpu + num_draft_tokens - batch.seq_lens_sum = int(batch.seq_lens_cpu.sum()) batch.extend_lens = [num_draft_tokens for _ in range(len(batch.seq_lens))] batch.prefix_lens = seq_lens_cpu_.tolist() batch.extend_num_tokens = extend_num_tokens @@ -244,6 +241,11 @@ class EagleDraftInputV2Mixin: ) batch.capture_hidden_mode = capture_mode forward_batch = ForwardBatch.init_new(batch, draft_model_runner) + # Forward sees post-write length (draft extend writes num_draft_tokens + # slots); mutation stays on forward_batch to preserve SB.seq_lens. + forward_batch.seq_lens = forward_batch.seq_lens + num_draft_tokens + forward_batch.seq_lens_cpu = forward_batch.seq_lens_cpu + num_draft_tokens + forward_batch.seq_lens_sum = int(forward_batch.seq_lens_cpu.sum()) can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run(forward_batch) if not batch.forward_mode.is_idle() and not can_cuda_graph: draft_model_runner.attn_backend.init_forward_metadata(forward_batch)