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: ```python 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. ```python 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.""" ``` ```hint Per-slot state List exactly what you need to know about each slot at every step, and what has to happen to a slot in the step where its request finishes. ``` ```hint Order inside one step Within one decoding step, decide when you check for completion and when you admit waiting requests, so that no slot sits empty while work is queued. ``` ### 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? ### What a Strong Answer Covers - An explicit mapping from slot to request, a token buffer per request, and a first-in, first-out waiting queue - Completion checks for both the stop token and `max_tokens`, applied to each sequence after every step - Freed slots refilled before the next model call, so batching is continuous rather than static - Model outputs mapped back to the right requests, whether the batch is compacted or padded - Correct termination, with no request lost or served twice - A comparison with static batching in the number of model calls, and the cost of each call - Tests with uneven output lengths, immediate stops and more requests than slots ### Follow-up Questions - A real model keeps a key/value cache for each sequence instead of re-reading the full prefix. What must happen to that cache when a slot is freed and reused? - Requests now arrive continuously while the loop is running. How do you admit them, and how do you keep one very long request from hurting everyone else's latency? - The model requires a fixed batch shape. How do you represent empty slots, and how do you make sure their outputs are ignored? - How would you stream each request's tokens to its caller as they are generated, and support cancelling a request midway?