pad customized_info for mixed output batches (#29842)

This commit is contained in:
Feng Yao
2026-07-02 12:07:58 +08:00
committed by GitHub
parent 4605d4d94b
commit e6f6a353bf
2 changed files with 111 additions and 4 deletions
@@ -499,11 +499,23 @@ class _GenerationStreamAccumulator:
req.indexer_topk if req.return_indexer_topk else None
)
current_output_len = len(self.output_ids[-1])
if req.customized_info is not None:
for k, v in req.customized_info.items():
if k not in self.customized_info:
self.customized_info[k] = []
self.customized_info[k].append(v[send_token_offset : len(output_ids_)])
for key, req_values in req.customized_info.items():
if key not in self.customized_info:
self.customized_info[key] = [
[None] * len(prev_output_ids)
for prev_output_ids in self.output_ids[:-1]
]
self.customized_info[key].append(
[None] * current_output_len
if req_values is None
else req_values[send_token_offset : len(output_ids_)]
)
for per_request_values in self.customized_info.values():
if len(per_request_values) < len(self.output_ids):
per_request_values.append([None] * current_output_len)
def to_payload(
self, *, dp_rank: int, is_idle_batch: bool