Dynamic Batching Decode Loop for a Mock Language Model with Slot Refilling

You are given a mock language model with a single method. Each call takes a batch of token sequences, where every entry is the full token prefix generated so far for one request, and returns the next token for each entry:

class MockModel:
 def next_tokens(self, batch: list[list[int]]) -> list[int]:
 """batch[i] is the token prefix of one sequence.
 Returns one next token id per entry, in the same order as the batch."""

Implement the generation loop that serves a list of requests with dynamic batching. The batch has a fixed number of slots, and each slot holds one in-progress request. A request finishes when the model generates the stop token or when it reaches its max_tokens limit, so requests finish at different steps and slots empty out over time. When a request finishes, its slot must be refilled from the queue of waiting requests, so the batch stays as full as possible until every request is done.

from dataclasses import dataclass

@dataclass
class Request:
 request_id: str
 prompt: list[int]
 max_tokens: int

def generate(model: MockModel, requests: list[Request], num_slots: int, stop_token: int) -> dict[str, list[int]]:
 """Returns the generated tokens of every request, keyed by request_id."""

Constraints and Clarifications

• One model call may receive at most num_slots sequences.
• The mock model keeps no state between calls: every call receives each sequence's full token prefix.
• Waiting requests are admitted in the order in which they appear in requests.
• model.next_tokens returns exactly one token per input sequence, in input order.

Clarifying Questions

• Does max_tokens limit only the newly generated tokens, or does it include the prompt?
• Should the returned tokens include the stop token when it is generated?
• Must every model call receive exactly num_slots entries (a fixed batch shape, padded when slots are empty), or may it receive only the occupied slots?
• Is the full list of requests known up front, or can requests arrive while generation is running?
• How should the loop treat a max_tokens of zero, an empty prompt, or a model call that fails?