Group ScheduleBatch and ForwardBatch fields by data-flow role (#26022)

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
+    # For DP attention
-    sampling_info: SamplingBatchInfo = None
+    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
+    # The original sequence lengths, Qwen-1M related
-    req_pool_indices_cpu: torch.Tensor = None  # shape: [b], int64
+    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
+    # Encoder-decoder device tensors (host fields in the host metadata group)
-    multimodal_inputs: Optional[List] = None
+    encoder_lens: Optional[torch.Tensor] = None
+    encoder_out_cache_loc: Optional[torch.Tensor] = None
 
-    # The sum of all sequence lengths
+    # It comes empty list if logprob is not required.
-    seq_lens_sum: int = None
+    extend_input_logprob_token_ids: Optional[torch.Tensor] = None
-    # The original sequence lengths, Qwen-1M related
+
-    orig_seq_lens: torch.Tensor = None  # shape: [b], int32
+    # === 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
+    # Whether this batch is prefill-only (no token generation needed)
-    temp_scaled_logprobs: bool = False
+    is_prefill_only: bool = False
-    top_p_normalized_logprobs: bool = False
 
-    # For extend and mixed chunekd prefill
+    # Speculative decoding
-    prefix_lens: List[int] = None
+    spec_algorithm: SpeculativeAlgorithm = 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
 
     # 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
+    # Whether to return hidden states
-    split_index: int = 0
+    return_hidden_states: bool = False
-    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
 
     # Has grammar
     has_grammar: bool = False
 
-    # Device
+    # The sum of all sequence lengths
-    device: str = "cuda"
+    seq_lens_sum: int = None
-
+    extend_num_tokens: Optional[int] = None
-    # 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
 
     # Diffusion LLM
     dllm_config: Optional[DllmConfig] = None
 
-    # Metrics
+    # === Host metadata crossing to ForwardBatch (CPU lists / mirrors) ===
-    dp_cooperation_info: Optional[DPCooperationInfo] = None
+    seq_lens_cpu: torch.Tensor = None  # shape: [b], int64
-    prefill_stats: Optional[PrefillStats] = None
-    forward_iter: Optional[int] = None
 
-    # HiSparse
+    # For multimodal inputs
-    hisparse_coordinator: Optional[HiSparseCoordinator] = None
+    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,

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
+    # For multimodal
-    next_token_logits_buffer: torch.Tensor = None
+    mm_inputs: Optional[List[MultimodalInputs]] = None
-    temp_scaled_logprobs: bool = False
-    temperature: torch.Tensor = None
-    top_p_normalized_logprobs: bool = False
-    top_p: torch.Tensor = 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-decoder cross-attention mask
-    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
     cross_attention_custom_mask: Optional[torch.Tensor] = None
 
-    # For LoRA
+    # For DP attention (padding / local info)
-    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
     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,
+            input_embeds=batch.input_embeds,
-            seq_lens_cpu=seq_lens_cpu,
+            replace_embeds=batch.replace_embeds,
-            orig_seq_lens=batch.orig_seq_lens,
+            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