Remove dead ScheduleBatch fields and avoid inplace seq_lens bump (#30669)
This commit is contained in:
@@ -723,7 +723,6 @@ class TboForwardBatchPreparer:
|
||||
for key in [
|
||||
"forward_mode",
|
||||
"is_extend_in_batch",
|
||||
"all_extend_in_batch",
|
||||
"return_logprob",
|
||||
"can_run_dp_cuda_graph",
|
||||
"can_run_dp_breakable_cuda_graph",
|
||||
|
||||
@@ -1892,7 +1892,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
||||
|
||||
# For DP attention
|
||||
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
|
||||
can_run_dp_breakable_cuda_graph: bool = False
|
||||
tbo_split_seq_index: Optional[int] = None
|
||||
@@ -1956,7 +1955,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
||||
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
|
||||
|
||||
@@ -2811,16 +2809,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
||||
req.kv_committed_len += 1
|
||||
req.kv_allocated_len += 1
|
||||
|
||||
if self.enable_overlap:
|
||||
# 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:
|
||||
self.seq_lens.add_(1)
|
||||
self.seq_lens_cpu.add_(1)
|
||||
self.orig_seq_lens.add_(1)
|
||||
# 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
|
||||
# Sum is recomputed lazily by ForwardBatch.init_new.
|
||||
self.seq_lens_sum = None
|
||||
|
||||
@@ -3013,7 +3006,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
||||
can_run_dp_cuda_graph=self.can_run_dp_cuda_graph,
|
||||
can_run_dp_breakable_cuda_graph=self.can_run_dp_breakable_cuda_graph,
|
||||
is_extend_in_batch=self.is_extend_in_batch,
|
||||
all_extend_in_batch=self.all_extend_in_batch,
|
||||
is_prefill_only=self.is_prefill_only,
|
||||
seq_lens_cpu=self.seq_lens_cpu,
|
||||
enable_overlap=self.enable_overlap,
|
||||
|
||||
@@ -396,8 +396,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
|
||||
|
||||
# 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
|
||||
can_run_dp_breakable_cuda_graph: bool = False
|
||||
global_forward_mode: Optional[ForwardMode] = None
|
||||
@@ -632,8 +630,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
|
||||
# for the contract.
|
||||
capture_hidden_mode = batch.capture_hidden_mode
|
||||
batch.capture_hidden_mode = None
|
||||
seq_lens_cpu_cache = batch.seq_lens_cpu_cache
|
||||
batch.seq_lens_cpu_cache = None
|
||||
return_hidden_states_before_norm = batch.return_hidden_states_before_norm
|
||||
batch.return_hidden_states_before_norm = False
|
||||
|
||||
@@ -668,18 +664,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
|
||||
# ScheduleBatch.sampling_info is already swapped to the forward-only
|
||||
# copy by Scheduler.run_batch under overlap mode (see save/restore
|
||||
# block there). Use it directly.
|
||||
if seq_lens_cpu_cache is not None:
|
||||
# Stale-cache guard: shape must match current GPU seq_lens. Mismatch
|
||||
# means caller forgot to refresh the override after batch size
|
||||
# changed (e.g. filter/merge_batch); using a stale cache would
|
||||
# propagate wrong CPU mirror to downstream DP / cudagraph logic.
|
||||
assert seq_lens_cpu_cache.shape == batch.seq_lens.shape, (
|
||||
f"seq_lens_cpu_cache shape {seq_lens_cpu_cache.shape} != "
|
||||
f"seq_lens {batch.seq_lens.shape}; stale override on batch?"
|
||||
)
|
||||
seq_lens_cpu = seq_lens_cpu_cache
|
||||
else:
|
||||
seq_lens_cpu = batch.seq_lens_cpu
|
||||
seq_lens_cpu = batch.seq_lens_cpu
|
||||
|
||||
if batch.seq_lens_sum is None and seq_lens_cpu is not None:
|
||||
batch.seq_lens_sum = int(seq_lens_cpu.sum())
|
||||
@@ -711,7 +696,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
|
||||
# Scalar config / flags
|
||||
return_logprob=batch.return_logprob,
|
||||
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,
|
||||
can_run_dp_breakable_cuda_graph=batch.can_run_dp_breakable_cuda_graph,
|
||||
global_forward_mode=batch.global_forward_mode,
|
||||
|
||||
Reference in New Issue
Block a user