Unify forward output datastructure (#11124)

This commit is contained in:
Liangsheng Yin
2025-10-03 00:28:57 +08:00
committed by GitHub
parent 3511b37099
commit 458611de77
12 changed files with 180 additions and 135 deletions

View File

@@ -860,10 +860,6 @@ class Req:
)
# Batch id
bid = 0
@dataclasses.dataclass
class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
"""Store all information of a batch on the scheduler."""
@@ -1829,10 +1825,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
seq_lens_cpu_cache if seq_lens_cpu_cache is not None else self.seq_lens_cpu
)
global bid
bid += 1
return ModelWorkerBatch(
bid=bid,
forward_mode=self.forward_mode,
input_ids=self.input_ids,
req_pool_indices=self.req_pool_indices,
@@ -1952,8 +1945,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
@dataclasses.dataclass
class ModelWorkerBatch:
# The batch id
bid: int
# The forward mode
forward_mode: ForwardMode
# The input ids