From 686ef5067296a1d1ddaaa4587157d246453be8ff Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 28 May 2026 00:11:20 -0700 Subject: [PATCH] Group ScheduleBatch and ForwardBatch fields by data-flow role (#26022) --- python/sglang/srt/managers/schedule_batch.py | 201 ++++++++++-------- .../srt/model_executor/forward_batch_info.py | 201 ++++++++++-------- 2 files changed, 225 insertions(+), 177 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index ccb7c00a7..c301e0a02 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -86,7 +86,6 @@ from sglang.srt.mem_cache.common import ( ) from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.radix_cache import RadixKey -from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, @@ -1470,17 +1469,26 @@ def set_mamba_track_indices_from_reqs(batch): class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): """Store all information of a batch on the scheduler.""" - # Request, memory pool, and cache + # === Core: request list (ForwardBatch derives lora_ids / rids / grammars / positions from it) === reqs: List[Req] + + # === Global config and shared resources (engine-lifetime; identical across batches) === + # Memory pool and cache req_to_token_pool: ReqToTokenPool = None token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator = None tree_cache: BasePrefixCache = None - is_hybrid_swa: bool = False # Batch configs model_config: ModelConfig = None - forward_mode: ForwardMode = None enable_overlap: bool = False + + # Device + device: str = "cuda" + + # HiSparse (engine-level coordinator ref, same across batches) + hisparse_coordinator: Optional[HiSparseCoordinator] = None + + # === Batch-variant scheduler state (per-batch; not read by ForwardBatch) === # Tell whether the current running batch is full so that we can skip # the check of whether to prefill new requests. # This is an optimization to reduce the overhead of the prefill check. @@ -1490,22 +1498,62 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): chunked_req: Optional[Req] = None contains_last_prefill_chunk: bool = True - # Sampling info - sampling_info: SamplingBatchInfo = None + # For DP attention + inner_idle_batch: Optional[ScheduleBatch] = None + # Decode requests carried alongside a chunked-prefill batch + decoding_reqs: List[Req] = None + # For split prefill + split_index: int = 0 + split_prefill_finished: bool = False + split_forward_count: int = 1 + split_forward_batch: ForwardBatch = None + + # For logits and logprob post processing (ForwardBatch keeps its own copies) + temp_scaled_logprobs: bool = False + top_p_normalized_logprobs: bool = False + + # CPU mirror of req_pool_indices; schedule-path only (used in overlap_utils, + # not read by ForwardBatch), stale in spec draft window + req_pool_indices_cpu: torch.Tensor = None # shape: [b], int64 + + # Forward-pass metrics + fpm_start_time: float = 0.0 + + # Stream + has_stream: bool = False + + # Whether to return captured experts + return_routed_experts: bool = False + return_indexer_topk: bool = False + + # hicache pointer for synchronizing data loading from CPU to GPU + hicache_consumer_index: int = -1 + + # Metrics + dp_cooperation_info: Optional[DPCooperationInfo] = None + prefill_stats: Optional[PrefillStats] = None + forward_iter: Optional[int] = None + + # === GPU tensors crossing to ForwardBatch (clone targets for stream isolation) === # Batched arguments to model runner input_ids: torch.Tensor = None # shape: [b], int64 input_embeds: torch.Tensor = None # shape: [b, hidden_size], float32 + # Token replacement embeddings and absolute positions (optional). replace_embeds: Optional[torch.Tensor] = None replace_positions: Optional[torch.Tensor] = None + + # Read by ForwardBatch ngram embedding init ne_token_table: torch.Tensor = None + token_type_ids: torch.Tensor = None # shape: [b], int64 req_pool_indices: torch.Tensor = None # shape: [b], int64 seq_lens: torch.Tensor = None # shape: [b], int64 - seq_lens_cpu: torch.Tensor = None # shape: [b], int64 - # CPU mirror of req_pool_indices; schedule-path only, stale in spec draft window - req_pool_indices_cpu: torch.Tensor = None # shape: [b], int64 + + # The original sequence lengths, Qwen-1M related + orig_seq_lens: torch.Tensor = None # shape: [b], int32 + # The output locations of the KV cache out_cache_loc: torch.Tensor = None # shape: [b], int64 @@ -1518,47 +1566,31 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): mamba_cow_dst_indices: torch.Tensor = None mamba_clear_indices: torch.Tensor = None - # For multimodal inputs - multimodal_inputs: Optional[List] = None + # Encoder-decoder device tensors (host fields in the host metadata group) + encoder_lens: Optional[torch.Tensor] = None + encoder_out_cache_loc: Optional[torch.Tensor] = None - # The sum of all sequence lengths - seq_lens_sum: int = None - # The original sequence lengths, Qwen-1M related - orig_seq_lens: torch.Tensor = None # shape: [b], int32 + # It comes empty list if logprob is not required. + extend_input_logprob_token_ids: Optional[torch.Tensor] = None + + # === Config / flags crossing to ForwardBatch (by-value) === + forward_mode: ForwardMode = None + global_forward_mode: Optional[ForwardMode] = None # For DP attention - inner_idle_batch: Optional[ScheduleBatch] = None - global_num_tokens: Optional[List[int]] = None - global_num_tokens_for_logprob: Optional[List[int]] = None is_extend_in_batch: bool = False all_extend_in_batch: bool = False # plumbing for downstream forks (PR #19639) can_run_dp_cuda_graph: bool = False tbo_split_seq_index: Optional[int] = None - global_forward_mode: Optional[ForwardMode] = None # For processing logprobs return_logprob: bool = False - top_logprobs_nums: Optional[List[int]] = None - token_ids_logprobs: Optional[List[List[int]]] = None - # For logits and logprob post processing - temp_scaled_logprobs: bool = False - top_p_normalized_logprobs: bool = False + # Whether this batch is prefill-only (no token generation needed) + is_prefill_only: bool = False - # For extend and mixed chunekd prefill - prefix_lens: List[int] = None - extend_lens: List[int] = None - extend_num_tokens: Optional[int] = None - decoding_reqs: List[Req] = None - extend_logprob_start_lens: List[int] = None - # It comes empty list if logprob is not required. - extend_input_logprob_token_ids: Optional[torch.Tensor] = None - - # For encoder-decoder architectures - encoder_cached: Optional[List[bool]] = None - encoder_lens: Optional[torch.Tensor] = None - encoder_lens_cpu: Optional[List[int]] = None - encoder_out_cache_loc: Optional[torch.Tensor] = None + # Speculative decoding + spec_algorithm: SpeculativeAlgorithm = None # For matryoshka embeddings dimensions: Optional[list[int]] = None @@ -1566,61 +1598,57 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # Whether to return pooled hidden states (pre-head transformer output) return_pooled_hidden_states: bool = False - # For split prefill - split_index: int = 0 - split_prefill_finished: bool = False - split_forward_count: int = 1 - split_forward_batch: ForwardBatch = None - - # One-shot per-forward overrides; init_new consumes and resets. - seq_lens_cpu_cache: torch.Tensor = None - capture_hidden_mode: Optional[CaptureHiddenMode] = None - return_hidden_states_before_norm: bool = False - - # Forward-pass metrics - fpm_start_time: float = 0.0 - - # Stream - has_stream: bool = False + # Whether to return hidden states + return_hidden_states: bool = False # Has grammar has_grammar: bool = False - # Device - device: str = "cuda" - - # Speculative decoding - spec_algorithm: SpeculativeAlgorithm = None - # spec_info: Optional[SpecInput] = None - spec_info: Optional[SpecInput] = None - - # Whether to return hidden states - return_hidden_states: bool = False - - # Whether to return captured experts - return_routed_experts: bool = False - - return_indexer_topk: bool = False - - # Whether this batch is prefill-only (no token generation needed) - is_prefill_only: bool = False - - # Multi-item scoring delimiter indices (set during prepare_for_extend) - multi_item_delimiter_indices: Optional[List[torch.Tensor]] = None - - # hicache pointer for synchronizing data loading from CPU to GPU - hicache_consumer_index: int = -1 + # The sum of all sequence lengths + seq_lens_sum: int = None + extend_num_tokens: Optional[int] = None # Diffusion LLM dllm_config: Optional[DllmConfig] = None - # Metrics - dp_cooperation_info: Optional[DPCooperationInfo] = None - prefill_stats: Optional[PrefillStats] = None - forward_iter: Optional[int] = None + # === Host metadata crossing to ForwardBatch (CPU lists / mirrors) === + seq_lens_cpu: torch.Tensor = None # shape: [b], int64 - # HiSparse - hisparse_coordinator: Optional[HiSparseCoordinator] = None + # For multimodal inputs + multimodal_inputs: Optional[List] = None + + # For processing logprobs + top_logprobs_nums: Optional[List[int]] = None + token_ids_logprobs: Optional[List[List[int]]] = None + + # For encoder-decoder architectures + encoder_cached: Optional[List[bool]] = None + encoder_lens_cpu: Optional[List[int]] = None + + # Multi-item scoring delimiter indices (set during prepare_for_extend) + multi_item_delimiter_indices: Optional[List[torch.Tensor]] = None + + # For extend and mixed chunekd prefill + prefix_lens: List[int] = None + extend_lens: List[int] = None + extend_logprob_start_lens: List[int] = None + + # For DP attention + global_num_tokens: Optional[List[int]] = None + global_num_tokens_for_logprob: Optional[List[int]] = None + + # === Compound crossing to ForwardBatch (carry their own device tensors) === + # Sampling info + sampling_info: SamplingBatchInfo = None + + # Speculative decoding + # spec_info: Optional[SpecInput] = None + spec_info: Optional[SpecInput] = None + + # === One-shot per-forward overrides; init_new consumes and resets === + seq_lens_cpu_cache: torch.Tensor = None + capture_hidden_mode: Optional[CaptureHiddenMode] = None + return_hidden_states_before_norm: bool = False @classmethod def init_new( @@ -1637,16 +1665,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): ): return_logprob = any(req.return_logprob for req in reqs) - is_hybrid_swa = False - if isinstance(token_to_kv_pool_allocator, SWATokenToKVPoolAllocator): - is_hybrid_swa = True - batch = cls( reqs=reqs, req_to_token_pool=req_to_token_pool, token_to_kv_pool_allocator=token_to_kv_pool_allocator, tree_cache=tree_cache, - is_hybrid_swa=is_hybrid_swa, model_config=model_config, enable_overlap=enable_overlap, return_logprob=return_logprob, diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index f59b90ffe..6c0a278e4 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -274,6 +274,7 @@ class NgramEmbeddingInfo: class ForwardBatch(ForwardBatchDeepSeekMHAMixin): """Store all inputs of a forward pass.""" + # === Required core inputs (no default; input_ids / req_pool_indices / seq_lens / out_cache_loc are borrowed from ScheduleBatch) === # The forward mode forward_mode: ForwardMode # The batch size @@ -286,10 +287,13 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): seq_lens: torch.Tensor # The indices of output tokens in the token_to_kv_pool out_cache_loc: torch.Tensor - # The sum of all sequence lengths seq_lens_sum: int + # === Borrowed from ScheduleBatch: GPU tensors (cross-stream; clone targets for stream isolation) === + # FIXME(lsyin): these are currently aliased by reference from ScheduleBatch. Once + # they are cloned/relayed into FB-owned copies at the boundary, move them out of + # "Borrowed" into a dedicated "Forward-resolved snapshot" group. # The original sequence length without being chunked. Qwen-1M related. orig_seq_lens: Optional[torch.Tensor] = None @@ -304,21 +308,76 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): mamba_cow_dst_indices: Optional[torch.Tensor] = None mamba_clear_indices: Optional[torch.Tensor] = None - # Optional seq_lens on cpu + # For input embeddings + input_embeds: Optional[torch.Tensor] = None + # For token embedding overrides (sparse replacement at specific positions) + replace_embeds: Optional[torch.Tensor] = None + replace_positions: Optional[torch.Tensor] = None + + # For cross-encoder model + token_type_ids: Optional[torch.Tensor] = None + + # Encoder-decoder device tensors + encoder_lens: Optional[torch.Tensor] = None + encoder_out_cache_loc: Optional[torch.Tensor] = None + + # === Borrowed from ScheduleBatch: config / flags (by-value) === + # For logprob + return_logprob: bool = False + # Whether this batch is prefill-only (no token generation needed) + is_prefill_only: bool = False + spec_algorithm: SpeculativeAlgorithm = None + # For matryoshka embeddings + dimensions: Optional[list[int]] = None + # Whether to return pooled hidden states (pre-head transformer output) + return_pooled_hidden_states: bool = False + + # For DP attention + is_extend_in_batch: bool = False + # Mirrors ScheduleBatch.all_extend_in_batch; kept for downstream forks. + all_extend_in_batch: bool = False + can_run_dp_cuda_graph: bool = False + global_forward_mode: Optional[ForwardMode] = None + + # For two-batch overlap + tbo_split_seq_index: Optional[int] = None + + # === Borrowed from ScheduleBatch: host metadata (CPU lists / mirrors) === + # Optional seq_lens on cpu (CPU mirror of seq_lens) seq_lens_cpu: Optional[torch.Tensor] = None # For logprob - return_logprob: bool = False top_logprobs_nums: Optional[List[int]] = None token_ids_logprobs: Optional[List[List[int]]] = None - # For logits and logprobs post processing - next_token_logits_buffer: torch.Tensor = None - temp_scaled_logprobs: bool = False - temperature: torch.Tensor = None - top_p_normalized_logprobs: bool = False - top_p: torch.Tensor = None + # For multimodal + mm_inputs: Optional[List[MultimodalInputs]] = None + # Encoder-decoder host fields + encoder_cached: Optional[List[bool]] = None + encoder_lens_cpu: Optional[List[int]] = None + + # Pre-computed delimiter indices for multi-item scoring (CPU tensors, one per request) + multi_item_delimiter_indices: Optional[List[torch.Tensor]] = None + + # === Borrowed from ScheduleBatch: compound (carry their own device tensors) === + # Sampling info + sampling_info: SamplingBatchInfo = None + # Speculative decoding + spec_info: Optional[SpecInput] = None + + # === Derived from ScheduleBatch.reqs === + # For LoRA + lora_ids: Optional[List[str]] = None + # For dumper: request IDs for cross-step sequence tracking + rids: Optional[List[str]] = None + + # === Resolved from SB one-shot overrides (consumed + reset by init_new) === + capture_hidden_mode: CaptureHiddenMode = None + # For hidden states before normal + return_hidden_states_before_norm: bool = False + + # === Forward-derived (built in init_new on the forward stream; FB-owned) === # Position information positions: torch.Tensor = None @@ -332,6 +391,26 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): extend_logprob_start_lens_cpu: Optional[List[int]] = None extend_input_logprob_token_ids_gpu: Optional[torch.Tensor] = None + # For DP attention (MLP sync sizes) + original_global_num_tokens_cpu: Optional[List[int]] = None + global_num_tokens_cpu: Optional[List[int]] = None + global_num_tokens_gpu: Optional[torch.Tensor] = None + # Has to be None when cuda graph is captured. + global_num_tokens_for_logprob_cpu: Optional[List[int]] = None + global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] = None + + # For padding + num_token_non_padded: Optional[torch.Tensor] = None # scalar tensor + num_token_non_padded_cpu: int = None + + # === Runtime-filled (set during the forward pass / cuda graph / managers; not at construction) === + # For logits and logprobs post processing + next_token_logits_buffer: torch.Tensor = None + temp_scaled_logprobs: bool = False + temperature: torch.Tensor = None + top_p_normalized_logprobs: bool = False + top_p: torch.Tensor = None + # For split prefill # intermediate values for split prefill hidden_states: torch.Tensor = None @@ -341,39 +420,12 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): split_index: int = 0 # For multimodal - mm_inputs: Optional[List[MultimodalInputs]] = None + mm_input_embeds: Optional[torch.Tensor] = None - # Encoder-decoder - encoder_cached: Optional[List[bool]] = None - encoder_lens: Optional[torch.Tensor] = None - encoder_lens_cpu: Optional[List[int]] = None - encoder_out_cache_loc: Optional[torch.Tensor] = None + # Encoder-decoder cross-attention mask cross_attention_custom_mask: Optional[torch.Tensor] = None - # For LoRA - lora_ids: Optional[List[str]] = None - - # For input embeddings - input_embeds: Optional[torch.Tensor] = None - - # For token embedding overrides (sparse replacement at specific positions) - replace_embeds: Optional[torch.Tensor] = None - replace_positions: Optional[torch.Tensor] = None - - # For cross-encoder model - token_type_ids: Optional[torch.Tensor] = None - - # Sampling info - sampling_info: SamplingBatchInfo = None - - # For DP attention - original_global_num_tokens_cpu: Optional[List[int]] = None - global_num_tokens_cpu: Optional[List[int]] = None - global_num_tokens_gpu: Optional[torch.Tensor] = None - # Has to be None when cuda graph is captured. - global_num_tokens_for_logprob_cpu: Optional[List[int]] = None - global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] = None - # The padding mode for DP attention + # For DP attention (padding / local info) dp_padding_mode: Optional[DpPaddingMode] = None # for extend, local start pos and num tokens is different in logits processor # this will be computed in get_dp_local_info @@ -381,55 +433,23 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): dp_local_start_pos: Optional[torch.Tensor] = None # cached info at runtime dp_local_num_tokens: Optional[torch.Tensor] = None # cached info at runtime global_dp_buffer_len: Optional[int] = None - is_extend_in_batch: bool = False - # Mirrors ScheduleBatch.all_extend_in_batch; kept for downstream forks. - all_extend_in_batch: bool = False - can_run_dp_cuda_graph: bool = False - global_forward_mode: Optional[ForwardMode] = None - - # Whether this batch is prefill-only (no token generation needed) - is_prefill_only: bool = False - - # Pre-computed delimiter indices for multi-item scoring (CPU tensors, one per request) - multi_item_delimiter_indices: Optional[List[torch.Tensor]] = None - - # Speculative decoding - spec_info: Optional[SpecInput] = None - spec_algorithm: SpeculativeAlgorithm = None - mm_input_embeds: Optional[torch.Tensor] = None - capture_hidden_mode: CaptureHiddenMode = None # For padding padded_static_len: int = -1 # -1 if not padded - num_token_non_padded: Optional[torch.Tensor] = None # scalar tensor - num_token_non_padded_cpu: int = None # For Qwen2-VL mrope_positions: torch.Tensor = None # For two-batch overlap - tbo_split_seq_index: Optional[int] = None tbo_parent_token_range: Optional[Tuple[int, int]] = None tbo_padded_len: Optional[int] = None tbo_children: Optional[List[ForwardBatch]] = None - # For matryoshka embeddings - dimensions: Optional[list[int]] = None - attn_cp_metadata: Optional[ContextParallelMetadata] = None - # For hidden states before normal - return_hidden_states_before_norm: bool = False - - # Whether to return pooled hidden states (pre-head transformer output) - return_pooled_hidden_states: bool = False - # For ngram embedding ngram_embedding_info: Optional[NgramEmbeddingInfo] = None - # For dumper: request IDs for cross-step sequence tracking - rids: Optional[List[str]] = None - @classmethod def init_new( cls, @@ -494,49 +514,54 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): batch.seq_lens_sum = int(batch.seq_lens_cpu.sum()) ret = cls( + # Required core inputs forward_mode=batch.forward_mode, batch_size=len(batch.seq_lens), input_ids=batch.input_ids, req_pool_indices=batch.req_pool_indices, seq_lens=batch.seq_lens, out_cache_loc=batch.out_cache_loc, + seq_lens_sum=batch.seq_lens_sum, + # Inputs aliased by reference from ScheduleBatch + seq_lens_cpu=seq_lens_cpu, + orig_seq_lens=batch.orig_seq_lens, mamba_track_indices=batch.mamba_track_indices, mamba_track_mask=batch.mamba_track_mask, mamba_track_seqlens=batch.mamba_track_seqlens, mamba_cow_src_indices=batch.mamba_cow_src_indices, mamba_cow_dst_indices=batch.mamba_cow_dst_indices, mamba_clear_indices=batch.mamba_clear_indices, - mm_inputs=batch.multimodal_inputs, - encoder_cached=batch.encoder_cached, encoder_lens=batch.encoder_lens, - encoder_lens_cpu=batch.encoder_lens_cpu, encoder_out_cache_loc=batch.encoder_out_cache_loc, - seq_lens_sum=batch.seq_lens_sum, - seq_lens_cpu=seq_lens_cpu, - orig_seq_lens=batch.orig_seq_lens, + input_embeds=batch.input_embeds, + replace_embeds=batch.replace_embeds, + replace_positions=batch.replace_positions, + token_type_ids=batch.token_type_ids, + # Scalar config / flags return_logprob=batch.return_logprob, - top_logprobs_nums=batch.top_logprobs_nums, - token_ids_logprobs=batch.token_ids_logprobs, is_extend_in_batch=batch.is_extend_in_batch, all_extend_in_batch=batch.all_extend_in_batch, can_run_dp_cuda_graph=batch.can_run_dp_cuda_graph, global_forward_mode=batch.global_forward_mode, is_prefill_only=batch.is_prefill_only, - multi_item_delimiter_indices=batch.multi_item_delimiter_indices, - lora_ids=[req.lora_id for req in batch.reqs], - sampling_info=batch.sampling_info, spec_algorithm=batch.spec_algorithm, - spec_info=batch.spec_info, capture_hidden_mode=capture_hidden_mode, - input_embeds=batch.input_embeds, - replace_embeds=batch.replace_embeds, - replace_positions=batch.replace_positions, - token_type_ids=batch.token_type_ids, - tbo_split_seq_index=batch.tbo_split_seq_index, dimensions=batch.dimensions, return_pooled_hidden_states=batch.return_pooled_hidden_states, return_hidden_states_before_norm=return_hidden_states_before_norm, + tbo_split_seq_index=batch.tbo_split_seq_index, + # Host-side metadata + top_logprobs_nums=batch.top_logprobs_nums, + token_ids_logprobs=batch.token_ids_logprobs, + mm_inputs=batch.multimodal_inputs, + encoder_cached=batch.encoder_cached, + encoder_lens_cpu=batch.encoder_lens_cpu, + multi_item_delimiter_indices=batch.multi_item_delimiter_indices, + lora_ids=[req.lora_id for req in batch.reqs], rids=[req.rid for req in batch.reqs], + # Compound (carry their own device tensors) + sampling_info=batch.sampling_info, + spec_info=batch.spec_info, ) device = model_runner.device