Group ScheduleBatch and ForwardBatch fields by data-flow role (#26022)
This commit is contained in:
Liangsheng Yin
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user