[core] step 2: drop seq_lens sentinel; SB maintains GPU as seq_lens_cpu mirror (#26020)
This commit is contained in:
@@ -183,15 +183,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
|
|||||||
future_map.stash(spec_info.future_indices, spec_info)
|
future_map.stash(spec_info.future_indices, spec_info)
|
||||||
self.spec_info = spec_info
|
self.spec_info = spec_info
|
||||||
else:
|
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
|
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)
|
|
||||||
|
|||||||
@@ -42,6 +42,17 @@ class FutureIndices:
|
|||||||
|
|
||||||
|
|
||||||
class FutureMap:
|
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__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
device: torch.device,
|
device: torch.device,
|
||||||
@@ -89,10 +100,9 @@ class FutureMap:
|
|||||||
)
|
)
|
||||||
|
|
||||||
def resolve_future(self, batch: ScheduleBatch):
|
def resolve_future(self, batch: ScheduleBatch):
|
||||||
if batch.forward_mode.is_decode():
|
# seq_lens is already real on entry (SB +1 for non-spec;
|
||||||
batch.seq_lens = self.new_seq_lens_buf[batch.req_pool_indices]
|
# resolve_seq_lens_cpu pulled from buf for spec_v2). Only resolve
|
||||||
torch._assert_async((batch.seq_lens > 0).all())
|
# input_ids tokens / spec extras here.
|
||||||
|
|
||||||
if self.spec_algo.is_none():
|
if self.spec_algo.is_none():
|
||||||
_resolve_future_token_ids(batch.input_ids, self.output_tokens_buf)
|
_resolve_future_token_ids(batch.input_ids, self.output_tokens_buf)
|
||||||
else:
|
else:
|
||||||
@@ -113,18 +123,26 @@ class FutureMap:
|
|||||||
if spec_need_hidden_states():
|
if spec_need_hidden_states():
|
||||||
draft_input.hidden_states = self.hidden_states_buf[indices]
|
draft_input.hidden_states = self.hidden_states_buf[indices]
|
||||||
|
|
||||||
def invalidate(self, batch: ScheduleBatch, future_indices: FutureIndices) -> None:
|
def set_input_ids_sentinel(
|
||||||
sentinel = -future_indices.indices
|
self, batch: ScheduleBatch, future_indices: FutureIndices
|
||||||
batch.input_ids = sentinel
|
) -> None:
|
||||||
batch.seq_lens = sentinel
|
# 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:
|
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
|
fi = batch.spec_info.future_indices if batch.spec_info is not None else None
|
||||||
if fi is None:
|
if fi is None:
|
||||||
return
|
return
|
||||||
if self.publish_ready is not None:
|
if self.publish_ready is not None:
|
||||||
self.publish_ready.wait()
|
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())
|
batch.seq_lens_sum = int(batch.seq_lens_cpu.sum())
|
||||||
|
|
||||||
def publish(
|
def publish(
|
||||||
|
|||||||
@@ -2161,14 +2161,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
|||||||
req.fill_ids = req.origin_input_ids + req.output_ids
|
req.fill_ids = req.origin_input_ids + req.output_ids
|
||||||
req.set_extend_input_len(1)
|
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])
|
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])
|
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_committed_len += 1
|
||||||
req.kv_allocated_len += 1
|
req.kv_allocated_len += 1
|
||||||
|
|
||||||
# Update seq_lens after allocation
|
|
||||||
if self.enable_overlap:
|
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.seq_lens_cpu = self.seq_lens_cpu + 1
|
||||||
self.orig_seq_lens = self.orig_seq_lens + 1
|
self.orig_seq_lens = self.orig_seq_lens + 1
|
||||||
else:
|
else:
|
||||||
# A faster in-place version
|
|
||||||
self.seq_lens.add_(1)
|
self.seq_lens.add_(1)
|
||||||
self.seq_lens_cpu.add_(1)
|
self.seq_lens_cpu.add_(1)
|
||||||
self.orig_seq_lens.add_(1)
|
self.orig_seq_lens.add_(1)
|
||||||
|
|||||||
@@ -2890,7 +2890,7 @@ class Scheduler(
|
|||||||
else:
|
else:
|
||||||
batch_result.future_indices = future_indices
|
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:
|
if batch.is_spec_v2:
|
||||||
batch.spec_info = batch_result.next_draft_input
|
batch.spec_info = batch_result.next_draft_input
|
||||||
|
|||||||
@@ -531,12 +531,7 @@ def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
|
|||||||
|
|
||||||
batch.maybe_evict_swa()
|
batch.maybe_evict_swa()
|
||||||
|
|
||||||
if batch.enable_overlap:
|
seq_lens_gpu = batch.seq_lens
|
||||||
# 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
|
|
||||||
bs = seq_lens_gpu.shape[0]
|
bs = seq_lens_gpu.shape[0]
|
||||||
|
|
||||||
if batch.tree_cache.page_size == 1:
|
if batch.tree_cache.page_size == 1:
|
||||||
|
|||||||
@@ -226,9 +226,6 @@ class EagleDraftInputV2Mixin:
|
|||||||
|
|
||||||
batch.spec_info = self
|
batch.spec_info = self
|
||||||
batch.input_ids = predict
|
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.extend_lens = [num_draft_tokens for _ in range(len(batch.seq_lens))]
|
||||||
batch.prefix_lens = seq_lens_cpu_.tolist()
|
batch.prefix_lens = seq_lens_cpu_.tolist()
|
||||||
batch.extend_num_tokens = extend_num_tokens
|
batch.extend_num_tokens = extend_num_tokens
|
||||||
@@ -244,6 +241,11 @@ class EagleDraftInputV2Mixin:
|
|||||||
)
|
)
|
||||||
batch.capture_hidden_mode = capture_mode
|
batch.capture_hidden_mode = capture_mode
|
||||||
forward_batch = ForwardBatch.init_new(batch, draft_model_runner)
|
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)
|
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:
|
if not batch.forward_mode.is_idle() and not can_cuda_graph:
|
||||||
draft_model_runner.attn_backend.init_forward_metadata(forward_batch)
|
draft_model_runner.attn_backend.init_forward_metadata(forward_batch)
|
||||||
|
|||||||
Reference in New Issue
Block a user