Training-Inference Co-Design

"Training and inference used to be separate departments. In 2026 they share a Slack channel and an architectural budget."

Sched, Co-Design-Native AI Agent

The first decade of large-language-model engineering treated training and inference as separate concerns. Training was a one-time event you optimized for FLOPs efficiency; inference was a continuous workload you optimized for latency and throughput. By 2026 that separation has broken down. Sardana et al. (2024) showed that the classical Chinchilla "compute-optimal" scaling laws systematically under-train the model relative to where the inference-aware optimum sits: if you serve a model billions of times, it is worth spending more training compute to save inference compute later. This insight has filtered into every frontier-model decision since 2024.

This section closes Chapter 58 with the systems-level co-design questions this implies: how training and inference share hardware, how MoE routing affects both sides, how the speculative-decoding / draft-model pattern blurs the boundary, and where 2027 will arbitrate.

58.5.1 Inference-aware scaling laws

The Sardana et al. paper extends Chinchilla's compute-optimal frontier to account for amortized inference cost. The result is intuitive: if your model will be served for a long time at high volume, the optimal balance shifts toward "smaller model trained longer". For frontier closed-API models served to hundreds of millions of users, the inference-aware optimum is very different from the training-only optimum. The Sardana framework explains why we see the trend toward "dense models trained on aggressive token-to-parameter ratios" (Llama-4 8B on 30T+ tokens; SmolLM2 360M on FineWeb-Edu; Liquid LFM2.5-350M on 28T tokens, 80,000:1 ratio).

58.5.2 MoE: where co-design pays off the most

Mixture-of-experts models break the symmetry between training and inference. During training, all experts must be loaded and gradient-updated; during inference, only the active subset is touched per token. The economic story is therefore: pay for big training (all experts in memory), get cheap inference (sparse activation). DeepSeek V4 (671B parameters, ~37B active per token) and Qwen3-235B-A22B (128 experts, top-8 routing) are the 2026 reference points. The Friendli MoE comparison walks through the routing schemes.

# MoE top-k routing in PyTorch (the inference-time path).
# Demonstrates why only k expert FLOPs are paid per token even though
# N expert parameter sets live in memory during training.
import torch
import torch.nn.functional as F

batch, d_model, n_experts, k = 4, 32, 8, 2
x = torch.randn(batch, d_model)
gate = torch.nn.Linear(d_model, n_experts, bias=False)
experts = torch.nn.ModuleList(
    [torch.nn.Linear(d_model, d_model) for _ in range(n_experts)]
)

# 1) Router scores per token, top-k selection, renormalized weights.
scores = gate(x)                                # (batch, n_experts)
topk_w, topk_idx = scores.topk(k, dim=-1)       # (batch, k)
topk_w = F.softmax(topk_w, dim=-1)              # weights sum to 1 over the k

# 2) Dispatch each token to its selected experts and combine the outputs.
y = torch.zeros_like(x)
for j in range(k):
    for b in range(batch):
        e = topk_idx[b, j].item()
        y[b] += topk_w[b, j] * experts[e](x[b])

print("active experts per token :", topk_idx.tolist())
print("active FLOPs / total FLOPs:", f"{k}/{n_experts} = {k/n_experts:.2f}")

58.5.3 Speculative decoding and draft-model patterns

Speculative decoding (Leviathan et al., 2023; popularised in 2024-25 production) runs a small "draft" model that proposes several tokens at once, then verifies them with the target model in parallel. Practical 2-4x throughput gains at no quality cost. The pattern reaches its logical conclusion in Medusa and EAGLE, which co-train target and draft heads. The training cost is small; the inference savings are large. This is co-design done right.

58.5.4 Multi-stage inference pipelines

MIT HERMES (2025) formalises the multi-stage AI inference pipeline that 2026 production systems actually use: prefill on one fabric (often Blackwell GPU), token generation on another (often Groq LPU or Cerebras CS-3), retrieval and tool-use on a third (CPU + accelerators), and any vision / multimodal preprocessing on a fourth (specialized vision-encoder NPUs). HERMES's contribution is a scheduler that places each stage on its best-fit silicon and minimizes inter-stage copy cost.

58.5.5 Comparing co-design strategies