Training and Inference Systems › Latency, Memory, and Energy Optimization
World Models HandbookAll chapters Training and Inference Systems › Latency, Memory, and Energy OptimizationOne section · Sections in this chapter5 / 8 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
We now turn from correct model execution to three resource questions: response time, memory and energy. Their tradeoffs depend on the workload. Batching can change amortized compute cost while adding admission waiting; quantization can change storage and arithmetic while also changing predictions. Neither gives a fixed energy or latency factor without measurement. Choose settings against the deployment's actual quality and resource constraints, which differ between a data-center training job and a battery-powered robot.
Latency: Decomposition and Budgets
Choose the timing boundary before decomposing latency. For an illustrative serial control loop, from the start of sensor acquisition to completion of the commanded actuator response, use disjoint stages:
where:
- : acquisition of the sensor data used for this decision.
- : all waiting for admission, batching and resources along the measured path, excluding active work counted elsewhere.
- : movement of inputs, outputs and commands between components, excluding their encoding and queue waiting.
- : input preparation such as resizing and normalization; learned encoder computation is counted under the model instead.
- : the model computation, including any learned observation encoder and dynamics calls used for this decision.
- : output interpretation or uncertainty calculations not already in the model.
- : serialization/deserialization of messages, not network transit.
- : remaining planner/controller work, excluding model calls already counted above.
- : command processing and actuator response within the chosen endpoint.
This sum assumes serial stages without overlap. Pipelined capture, asynchronous transfers or overlapping computation require critical-path timestamp accounting, not addition of every raw service duration. Count an operation once even if software names put it in two components. Optimizing model time still reduces the serial sum, but may yield only a small relative improvement when other terms dominate. The local benchmarks below cover a much narrower model/bookkeeping boundary, not this complete loop.
In a synchronous one-decision-per-period design at 30 Hz, the period is about 33.3 ms, and the scheduled work needs margin inside it. Asynchronous control may use different observation-age and stability requirements. A browser's display interval and input-to-display latency are distinct targets: 60 Hz means about 16.7 ms between frames, not a universal 100 ms frame budget. Offline imagination may prioritize throughput while still having completion deadlines, memory limits or feedback-staleness constraints. Choose the metric and its boundary for the actual consumer.
Concrete techniques:
- Precision strategy. A tensor stored in fp16 or bf16 uses half the scalar bytes of fp32; that does not halve the entire footprint or guarantee a throughput factor. Long rollouts can amplify numerical perturbations in sensitive transitions, while contractions can attenuate them. Evaluate precision at one-step, rollout and decision levels on the target hardware. fp32 accumulation or a higher-precision reference may help diagnose an error, but periodic reference recomputation has its own cost and is not a generic cure. Quantization sensitivity depends on the operation and dynamics, not a universal transformer-versus-recurrence ordering.
- Kernel fusion. Compatible fused kernels can reduce launches and intermediate memory traffic while preserving the intended computation. Fusing execution is not the same as algebraically folding input-dependent normalization into a fixed linear weight matrix. Measure launch and memory costs on the target runtime, and validate numerical differences introduced by a replacement kernel.
- Sequence-length reduction. Dense full-prefix attention has quadratic arithmetic in context length; a cached one-token query still performs work linear in retained context, as discussed earlier. A fixed window or fixed-size recurrent state can bound per-step work at fixed model size. Chunking full-context attention bounds temporary working memory, not all pairwise arithmetic. Truncating context or changing architecture requires prediction and decision-quality validation.
- Speculative decoding. Exact language-model sampling with a draft model requires a specified acceptance and correction procedure preserving the target distribution, as in Leviathan, Kalman and Matias. A fast approximate world-model transition checked against another model or incoming sensors is a different heuristic; it does not automatically inherit that exactness. Define the property being checked, the correction or fallback and the remaining approximation error. Bounded approximations can be intentional controller designs, but a timing gain alone does not establish their acceptability.
Memory: Anatomy of a Serving Footprint
Four useful tensor categories contribute memory pressure, before additional runtime overhead:
- Weights. Parameters and any additional network copies actually loaded for inference. Their fraction of the footprint depends on model size and the other terms, not batch size alone.
- Activations. Transient inference computation buffers; their sizes depend on the operation, batch and context. Ordinary inference does not retain backward activations.
- Session state. Full-context KV storage grows with retained tokens and concurrent sessions. A fixed-width recurrent state grows with concurrent sessions but not with rollout duration unless extra history is retained.
- Input/output buffers. These hold observations, actions, latent representations and predictions passed between encoders, dynamics models, decoders and policies. Multimodal inputs and outputs can make these buffers substantial.
Managing session state can be important, but profile the dominant term rather than assume it. For an explicitly disjoint tensor accounting, add weights, inference activations, retained session state and I/O, then account separately for runtime, allocator and workspace memory:
where:
- : bytes of loaded parameters and any required additional network copies
- : transient tensors for the current inference computation
- : retained state for each session, with growth determined by its recurrent or attention representation
- : input and output buffers for encoders, decoders, and policies
- : other allocated or reserved memory not assigned to the preceding terms, including runtime and workspaces
Strategies include:
- Eviction policies. An evicted session can explicitly expire, with the next call returning an error, or it can be recovered by a defined protocol. Recovery may need a saved boundary state or observation context, subsequent actions, matching model/configuration and stochastic state; actions alone generally do not reconstruct it. LRU is only a selection policy, not a proof of recovery or request ordering. Silently replacing lost causal state with a default is the correctness hazard.
- State compression. Quantization or a low-rank representation can reduce payload bytes, but is generally lossy; omitted information cannot simply be recomputed from the compressed state. Under a quantified contraction with the same future inputs, an initial-state perturbation can decay, but ongoing compression injects new error and control constraints may still be sensitive to it. Measure reconstruction, rollout and decision effects before accepting the memory tradeoff.
- Sliding-window attention with action replay. When old observations become irrelevant, drop their KV entries but track the drop in a per-session ledger. Any subsequent request that references the dropped range is refused or reconstructed. This makes the memory/latency tradeoff explicit and auditable.
Length-aware grouping can change padding and allocation patterns. In an allocator requiring contiguous blocks, unusable free fragments can prevent an allocation even when total free bytes would suffice. Paged caches, preallocation and allocator policy can change that behavior. Monitor allocated, reserved and usable memory under changing session lengths; this CPU tutorial neither measures GPU fragmentation nor predicts an OOM time.
Energy: The Third Axis
Energy needs its own measurement. Battery capacity limits total energy, while electrical and thermal constraints limit power; data centers can have hard power and cooling limits as well as cost budgets. For a declared measurement boundary, energy is , where power is in watts and elapsed time is in seconds, giving joules. Runtime alone does not determine it. Account for the actual model lifecycle:
- Training and adaptation may repeat. Model updates, imagined transitions and deployment queries all contribute; which dominates depends on their counts and measured costs, not whether the learner is called PPO or imagination-based.
- Rollout length multiplies serial work. At a constant 10 ms per step, 10,000 sequential steps take 100 seconds before other overhead. Both horizon and per-step implementation are design choices to evaluate against decision quality. Time is not yet energy without power measurement.
- Batching can amortize fixed overhead, but can also change active power, memory, waiting and unused work. Establish joules per useful inference from measured power and time, and test mobile deadlines and thermal limits rather than assume batching is unambiguously beneficial.
The common optimizations are:
- Precision choices, assessed for numerical quality and measured energy on the target hardware.
- Kernel fusion and reduced memory traffic, assessed with power instrumentation rather than treating bytes moved as an energy meter.
- Avoiding redundant evaluations. Sampling-based planners may score many candidates. Valid repeated-prefix reuse can avoid computation when lookup and storage costs are lower; independent stochastic samples must retain their distinct sample identities. Pruning can reduce work but also change planning quality.
Energy is observable during development using a suitable meter or validated device counters and a declared scope. Measure idle and active intervals consistently, and include the components relevant to the decision, not only a convenient accelerator kernel. Energy per inference and per useful action answer different questions from average power; record the metrics the deployment budget actually constrains.
A Small Measurements Example
We measure a small CPU workload at several vectorized batch sizes. Each row times a different number of independent states advanced by one forward pass. Dividing that time by batch size gives amortized compute time per state, not end-to-end request latency. These timings vary with the CPU, runtime, thread settings and system load; they are not GPU throughput or energy measurements. Energy requires measured power integrated over time, as in the MLCommons power measurement rules.
@torch.inference_mode()
def measure_batch(batch_size, steps=40, latent_dim=16, action_dim=2):
z = torch.zeros(batch_size, latent_dim)
a = torch.zeros(batch_size, action_dim)
times = []
for _ in range(3):
model(z, a) # warm-up
for _ in range(steps):
t0 = time.perf_counter()
z = model(z, a)
times.append(time.perf_counter() - t0)
times = np.array(times)
return float(np.median(times))
batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128]
latencies = np.array([measure_batch(bs) for bs in batch_sizes])
per_session = latencies / batch_sizesbatch | step ms | per-session ms
1 | 0.015 | 0.0148
2 | 0.015 | 0.0074
4 | 0.016 | 0.0039
8 | 0.019 | 0.0024
16 | 0.020 | 0.0012
32 | 0.023 | 0.0007
64 | 0.024 | 0.0004
128 | 0.029 | 0.0002Compare the actual table: batching may reduce amortized compute time even when a whole batch takes longer. The experiment does not simulate an arrival queue or establish a universal optimal batch size. A deployment must jointly measure throughput and end-to-end latency at its own request sizes and load.
Sequential Scaling Across Rollout Horizon
The next experiment isolates history-buffer bookkeeping during a recursive rollout. Both models use only the current state and action, so trimming the stored history cannot change predictions here. A transformer that attends to its history is different: truncating its cache may change the prediction distribution.
@torch.inference_mode()
def measure_recursive(steps=400, latent_dim=16, action_dim=2):
z = torch.zeros(1, latent_dim)
a = torch.zeros(1, action_dim)
history = torch.zeros(0, latent_dim)
step_times = []
for _ in range(steps):
t0 = time.perf_counter()
# naive recursive core
z = model(z, a)
history = torch.cat([history, z], dim=0)
step_times.append(time.perf_counter() - t0)
return np.array(step_times)
recursive_times = measure_recursive(steps=400)
print(
f"Growing-history median per-step time: {np.median(recursive_times) * 1e3:.3f} ms"
)Now the same rollout with a rolling history capped at a small window, so the working set stays bounded.
@torch.inference_mode()
def measure_windowed(steps=400, latent_dim=16, action_dim=2, cap=32):
if (
isinstance(cap, bool)
or not isinstance(cap, (int, np.integer))
or cap <= 0
):
raise ValueError("cap must be a positive integer")
z = torch.zeros(1, latent_dim)
a = torch.zeros(1, action_dim)
history = torch.zeros(0, latent_dim)
step_times = []
for _ in range(steps):
t0 = time.perf_counter()
z = model(z, a)
history = torch.cat([history, z], dim=0)
if history.shape[0] > cap:
history = history[-cap:].clone()
step_times.append(time.perf_counter() - t0)
return np.array(step_times)
windowed_times = measure_windowed(steps=400, cap=32)
print(
f"Windowed-history median per-step time: {np.median(windowed_times) * 1e3:.3f} ms"
)With fixed state dimension, repeatedly concatenating a growing history copies a number of elements proportional to , giving quadratic total buffer-copy work. A fixed window of width bounds each copy, giving copy work, linear in horizon when is fixed. These are bookkeeping operation counts, not a guarantee that noisy CPU timings increase at every step. A preallocated buffer can avoid repeated copies without truncating history. Neither implementation measures attention or demonstrates a KV-cache speedup.

Read the observed variability rather than assuming either curve rises or stays flat. The analytical copy-count difference can exist even when tiny noisy timing measurements do not reveal it. Long-context attention has its own computation and memory costs, which must be measured separately.
The batch-size plot presents the amortized forward timings recorded earlier without assuming a knee occurs in the tested range.

Use this plot as a local compute diagnostic, not a serving configuration prescription. A production batching policy adds queue waiting, arrival rates, padding, memory limits and deadlines to the measured model cost. There need not be a single visible knee, and per-request latency does not universally grow linearly with batch size. Triton's batching documentation makes the deliberate queue-delay tradeoff explicit.