Skip to content

Commit

Permalink
fix(Jetstream Pt): batch returned in prefill initialized to None
Browse files Browse the repository at this point in the history
This is required when there are no more tokens generated after prefill.
  • Loading branch information
tengomucho committed Sep 16, 2024
1 parent d18b65f commit abf5f5a
Showing 1 changed file with 1 addition and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -476,6 +476,7 @@ def prefill(self, batch: Batch) -> Tuple[List[Generation], CachedBatch]:
self.slots.append(slot)
len_active_slots += 1

batch = None
if len_active_slots > 0:
# Whatever initial batch these requests came from, we always return all pending requests in a single batch
request_ids = [slot.request_id for slot in self.slots if slot.state == Slot.State.READY]
Expand Down

0 comments on commit abf5f5a

Please sign in to comment.