Mock-LLM Inference Engine — Dynamic Batching

Implement a simplified inference-engine loop on top of a `model(batch) → next_tokens` API. The bar is dynamic batching: as sequences finish (max-tokens or stop-token), free their slots and fill them with new requests from a waiting queue, while keeping the slot↔request mapping correct so results never get crossed. Requirements Given a synchronous mock model API: Each call takes a batch of currently-running token prefixes and returns the next token (or next tokens) for each row. You receive a stream of requests; each request has a token budget (max_tokens) and a stop token / stop sequence. Implement a decoding loop that: Initial fill. Pull from the waiting queue until the running batch is full (up to batch_size). Step. Call model(batch) once per iteration, append the produced tokens to each row, and check termination per row (max_tokens reached or stop token observed). Completion. When a row finishes, deliver its result via the per-request callback / future, and free the slot. Refill. After every step, pull new requests from the waiting queue to fill any empty slots. When the queue is exhausted, run the remaining batch even if it is not full. Maintain a stable slot_id → request_id (or slot_id → sequence_id) mapping so that as sequences enter and leave the batch, no result is delivered to the wrong request. Canonical OA signatures: class SimulatedLLM: def generate_next_tokens(self, batch_prefixes: list[list[int]]) -> list[int]: ... # one next-token per row class BatchInferenceEngine: def __init__(self, model: SimulatedLLM, batch_size: int, stop_token: int): ... def submit_request(self, prompt_tokens: list[int], max_tokens: int, callback: Callable[[list[int]], None]) -> None: ... # enqueue; callback gets the full token list on finish def run(self) -> None: ... # loop until the queue is empty and all slots are free Two follow-up variants the spec layers on: # Fixed-size batch: pad empty slots and pass a mask instead of a variable-length batch. def generate_next_tokens_padded(self, batch_prefixes: list[list[int]], active_mask: list[bool]) -> list[int]: ... # Multi-token stop: terminate a row when its tail matches any stop sequence. def submit_request(self, prompt_tokens, max_tokens, callback, stop_sequences: list[list[int]] | None = None): ... Notes The interviewer cares more about correctness of the slot mapping than about real GPU semantics — the model is mocked. A common bug: copying the entire batch tensor each step instead of in-place updating row indices, which makes the slot mapping ambiguous after a refill. Use a dict[slot_id] = sequence_state and rebuild the batch view from it each iteration. Stop-sequence detection (vs. single stop token) is a likely follow-up; track a per-row tail buffer of the stop-sequence length. The terminating token is appended to the row before the termination check, so a finished sequence includes its stop token — with stop_token=0, a row that emits 5, 5, 0 is delivered as [...prompt, 5, 5, 0]. Confirm this convention rather than silently trimming the stop token. The technique is canonically called continuous batching or iteration-level batching: rather than waiting for every sequence in the batch to finish (static batching), the scheduler returns control after each decode step, evicts finished sequences, and immediately admits new ones from the waiting queue. Public benchmarks credit this with ~8× throughput vs. naive batching; PagedAttention adds another large multiplier (≈23× total at saturation) by storing the KV cache in fixed-size pages so newly admitted sequences do not require contiguous-buffer reallocation. Be ready to discuss continuous batching vs. static batching tradeoffs, and how production engines (vLLM, TGI) implement paged KV to make refill cheap. The round is one of three back-to-back 45-minute rounds (two coding + one research talk); pace yourself. Preparation Sketch a Scheduler class with running: dict[int, Seq], waiting: deque[Request], step() and tick() methods. Practice the slot-mapping invariant explicitly: after every refill, assert that set(running.keys()) == set(active slots). Internalize the iteration-level scheduling loop: step the model, sweep the running set for terminations, refill from the waiting queue, repeat. You do not need PagedAttention details, but you need the why. Be able to verbally extend the design to handle prefill vs. decode separation (the next-step follow-up if you finish early).