Training and Inference Systems › Rollout Serving, Caching, and Batching

One deployment option is a service performing autoregressive rollouts. A controller may submit a request every 20 milliseconds, while an offline evaluator or imagination-based learner may evaluate thousands of independent candidates. These are illustrative workloads. They can use separate stacks; if they share capacity, serving must account for both interactive deadlines and batch throughput rather than assume one policy suits both.

For autoregressive models, two properties drive serving design:

• Causal state. A recurrent transition can consume a fixed-size sufficient state and the current action. A full-context causal transformer can instead retain a growing token history or KV cache. State must have an owner and identity, but that owner can be the client, server or local controller; model statefulness does not require server-held sessions.
• Decode granularity. Some models emit one token or a small block per call; others predict a joint trajectory or video. Small operations can expose dispatch overhead, whereas a large encoder or denoiser can dominate computation. Measure the actual boundary and choose batching against throughput and deadlines.

Two possible API designs are:

Stateless request/response (POST /rollout). A request supplies enough state or observation context, actions and inference configuration to produce its result. Any compatible worker can handle that self-contained request, which simplifies routing but does not remove model-placement, workload or deadline constraints. It can serve offline or interactive workloads. A caller holding a sufficient recurrent state can send that compact state and one action rather than full history. A transformer API that requires retransmitting a long context can have substantial transfer and recomputation costs.

Stateful session (POST /session/{id}/step). The server owns the current state and any cache; each ordered call advances that session. This can avoid repeatedly transferring context for interactive clients. Define session lifetimes, capacity, model-version binding, step ordering, retries and routing to the owner or a correct state-transfer mechanism. An eviction must return an explicit expiry or use a valid recovery protocol, not silently continue from a default state.

A system can offer both designs and route requests by declared workload or latency class. That is an architecture option, not a requirement or a measured prevalence claim. Triton's sequence batcher illustrates one stateful implementation: requests in a sequence are routed to the same model instance.

Caching in a world-model server comes in three categories, each solving a different problem.

Key-value (KV) cache. For causal transformer decoding with one new token per call and fixed layer count and width, recomputing a dense-attention prefix of length $t$ costs $O(t^2)$ attention arithmetic per call. Reusing past keys and values leaves the new query attending to $t$ cached positions, costing $O(t)$ attention arithmetic. Over $T$ generated tokens, those attention totals are $O(T^3)$ versus $O(T^2)$, not quadratic versus linear. Cached decoding remains increasingly expensive with unbounded context. A fixed window of $W$ positions gives $O(TW)$ attention arithmetic, but changes available context unless the model was defined for that window. Token count is not environment-step count when each observation contains several tokens; projections and other layers add their own costs. KV caching avoids recomputing old keys and values, not the new query's attention to retained history.

For one session with ordinary multi-head attention, where the number of KV heads equals the query-head count, cache tensor storage grows linearly with cached token count. Ignoring allocator overhead, its size is

$$M_{\text{KV}} = 2 \cdot L \cdot H \cdot d_h \cdot S \cdot 2 \text{ bytes}$$

where:

• $L$: the number of transformer layers
• $H$: the number of KV heads per layer, equal to query heads in this MHA example; grouped-query or multi-query attention has fewer KV heads
• $d_h$: the head dimension
• $S$: cached token positions, not necessarily environment time steps when a frame has multiple tokens
• the leading $2$: one key tensor and one value tensor per layer/head
• the trailing $2$: bytes per fp16 scalar

With $L=32$, $H=16$, $d_h=64$ and $S=8192$ token positions, the cache tensors occupy $1,073,741,824$ bytes: exactly 1 GiB or about 1.074 GB decimal per session. Multiply this tensor count by concurrent sessions before adding allocator overhead. Sliding-window attention and cache eviction may reduce stored history but require explicit semantics for lost context; they are not automatically equivalent to full-context inference. Cache layout and length-aware scheduling affect actual allocation behavior.

Result caching. Memoization helps when requests repeat a complete, deterministic computation, but the benefit depends on hit rate and lookup cost rather than a universal speedup range. Include model/checkpoint version, preprocessing and inference configuration, initial state or observation history, exact action prefix, horizon and numeric execution policy in the key. For seeded sampling, include the RNG algorithm and explicit stream state or a stable keyed counter; a seed alone does not identify a draw after arbitrary consumption. Reusing one cached stochastic result is not equivalent to requesting independent samples. Either disable memoization for fresh-sample requests or define a reproducible sample identity in the API and key. Invalidate on model or configuration changes, and never share session state across independent particles merely because their visible inputs match.

Retrieved-context caching. As an illustrative design, a map-conditioned model could retrieve a region's features before predicting local dynamics. Reuse is valid only when region content, query, preprocessing and model-dependent representation remain compatible. A region identifier alone is not enough after the map or representation changes. Whether caching helps depends on hit rate, retrieval cost, lookup overhead and invalidation; this tutorial does not benchmark a hierarchical retrieval system.

Batching groups independent requests into a forward pass to amortize overhead. Iteration-level or continuous batching can admit new sequences as others finish, as in Orca. Rollout sessions can be at different indices with different retained lengths, so ragged serving needs correct cache ownership, offsets and eligible-token masks. In ordinary attention with a separate batch row per session, the batch dimension isolates rows; omitting a padding mask does not by itself read another row. Packed sequences within one token axis additionally need boundary masking, and causal masking excludes future tokens. Masks cannot repair a cache routed to the wrong session.

The main tradeoffs of batching:

• Larger batches can amortize fixed dispatch overhead, but extra queue waiting, work and memory can increase response time. Their measured latency need not grow monotonically with size. A batching wait that exceeds a controller's remaining slack misses its deadline.
• Smaller batches can be appropriate for low demand or tight deadlines. Their utilization depends on operation sizes and hardware, not batch size alone.
• Length-aware grouping reduces padding in a representation padded to its longest member. It is not chunked prefill: that technique splits the processing of one input context into smaller pieces. Packed or paged layouts have different waste and allocation behavior.

A scheduler may deliberately wait for more requests up to a bounded deadline, or dispatch immediately. Choose that bound from arrivals, service times and available slack; there is no general few-microsecond setting. The packet-buffer analogy is useful for thinking about delay versus utilization, but does not establish correctness. Triton exposes an optional maximum queue delay for dynamic batching; a deployment still needs end-to-end measurements.

The next block builds a single-process, single-threaded in-memory server with state per session and FIFO eviction by creation order. It requires positive integer capacities and valid tensor shapes, copies externally supplied state and returned results, and raises KeyError for missing or evicted IDs. It does not implement LRU, timeouts, batching, locks, authenticated requests or idempotent retry/step ordering. History is bookkeeping only, not an attention KV cache; the fixed model consumes current state and action.