[PD] Support structured output (#6560)

This commit is contained in:
Byron Hsu
2025-05-23 21:49:00 -07:00
committed by GitHub
parent ed0c3035cd
commit 2d831c6ef9
6 changed files with 106 additions and 13 deletions

View File

@@ -101,6 +101,9 @@ class ScheduleBatchDisaggregationDecodeMixin:
for req in self.reqs:
self.output_ids.append(req.output_ids[-1])
self.tree_cache.cache_unfinished_req(req)
if req.grammar is not None:
req.grammar.accept_token(req.output_ids[-1])
req.grammar.finished = req.finished()
self.output_ids = torch.tensor(self.output_ids, device=self.device)
# Simulate the eagle run. We add mock data to hidden states for the