Practice/xAI/Dynamic Batch Inference

Dynamic Batch Inference

CodingMust

Problem Overview

Design and implement a Dynamic Batch Inference Engine that efficiently processes multiple generation requests by batching them together. This is a simplified version of what production LLM inference engines (like vLLM, TensorRT-LLM, or TGI) do to serve models like Grok.

You are given a simulated language model interface that generates the next token for a batch of sequences. Your task is to implement a BatchInferenceEngine that:

• Accepts generation requests with callbacks
• Batches requests together for efficient inference
• Handles dynamic batching: as some sequences complete, fill empty slots with waiting requests
• Supports termination conditions: max_tokens limit and stop_token

Example

` model = SimulatedLLM() engine = BatchInferenceEngine(model, batch_size=4, stop_token=0)

results = []

Submit requests

engine.submit_request([1, 2, 3], max_tokens=5, callback=lambda seq: results.append(seq)) engine.submit_request([10, 20], max_tokens=3, callback=lambda seq: results.append(seq)) engine.submit_request([100], max_tokens=10, callback=lambda seq: results.append(seq))

Run until all complete

engine.run()

Results contain completed sequences (prompt + generated tokens)

Each sequence ends either at max_tokens or when stop_token is generated

`

Requirements

• Maintain a queue of pending requests
• Fill batch slots dynamically as sequences complete
• Track which batch slot corresponds to which request
• Handle termination conditions correctly (stop token or max tokens)
• Invoke callbacks when sequences complete

Clarification Questions to Ask

• Should we support streaming (partial results) or only final results?
• Can requests be submitted while run() is executing, or only before?
• How should we handle empty prompts?
• Should we support per-request stop sequences (multi-token patterns)?
• Is there a priority system for requests, or is it FIFO?

Part 1: Basic Batch Processing

Implement the core batching logic without dynamic slot filling. All requests submitted before run() are processed together.

Example

` model = SimulatedLLM() engine = BatchInferenceEngine(model, batch_size=2, stop_token=0)

results = [] engine.submit_request([1, 2], max_tokens=3, callback=lambda s: results.append(s)) engine.submit_request([10, 20], max_tokens=3, callback=lambda s: results.append(s)) engine.run()

Both requests processed in a single batch

results contains two sequences

`

Requirements

• Process up to batch_size requests in parallel
• Generate one token per sequence per iteration
• Track tokens generated per sequence
• Stop when a sequence reaches max_tokens or generates stop_token

Part 2: Dynamic Slot Filling

Extend the implementation to fill empty slots with waiting requests as sequences complete.

Example

` engine = BatchInferenceEngine(model, batch_size=2, stop_token=0)

Submit 4 requests (batch_size=2)

engine.submit_request([1], max_tokens=1, callback=...) # Finishes after 1 step engine.submit_request([2], max_tokens=3, callback=...) # Finishes after 3 steps engine.submit_request([3], max_tokens=2, callback=...) # Queued initially engine.submit_request([4], max_tokens=1, callback=...) # Queued initially

engine.run()

Request [1] finishes first, slot freed

Request [3] immediately fills the empty slot

All 4 requests complete efficiently

`

Requirements

• Maintain a queue of pending requests
• After each generation step, check for completed sequences
• Fill empty slots with requests from the queue
• Ensure correct mapping between batch positions and sequences