Speculative decoding for distributed-inference — design note

Status: design + reference prototype (not wired into the production decode path). Origin: assessment of NextLat (arXiv 2511.05963, code, NeurIPS 2025).

TL;DR — what we are and are NOT taking from NextLat

NextLat is a training-time technique. It adds an auxiliary latent-dynamics MLP that, during pretraining, learns to predict the next hidden state h_{t+1} from (h_t, embed(next_token)) under a stop-gradient self-predictive loss (smooth-L1 on the latent + a KL keeping the predicted latent's token distribution close to the real one). The MLP is discarded at inference. Its two payoffs are (a) latents that provably converge toward belief states (compressed sufficient statistics of history), and (b) up to 3.3× inference speedup via variable-length self-speculative decoding — the trained latent head cheaply rolls latents forward to draft several tokens, which the full model then verifies.

We cannot adopt the NextLat method directly. This fleet is inference-only over frozen, quantized GGUF weights (tinyllama / mistral / phi3 / gemma4). We have no pretraining loop, so we cannot bake in the latent-dynamics objective. Adopting "NextLat" literally would require finetuning every served model — that contradicts how the fleet is built.

What transfers is the speedup's underlying primitive — and we already own the hard part of it. The training-free, codebase-native win is conformal-gated layer early-exit: skip transformer layers per token whenever the certified 2·τ ≤ γ gate (tukey_bowl.rs + conformal_gate.rs) guarantees the argmax can't change. That gate exists and is Lean-proven; it just isn't pointed at the text decode yet. This serves the inference-squeeze roadmap directly (tinyllama ~1000× too slow on the Worker lane; mistral 0.58 tok/s under the liquid-KV-save regression). NextLat's variable-length self-signal falls out for free from the per-token decision margin. Heavier token-level draft/verify (the literal 3.3× path) is a second, optional axis.

Where it slots into the current decode path

The autoregressive loop is single-token:

• Pipeline.generate (src/pipeline.ts:1352) — native single-station fast path calls entryStation.generateTokens (Rust); otherwise falls through to prefill + generateFromPrefill.
• Pipeline.generateFromPrefill (src/pipeline.ts:1734) — the host autoregressive loop: one chooseToken → one forward per step via entryStation.decodeNext (/decode-next, single token) or the slower chooseTokenThroughStations.
• Pipeline.chooseToken (src/pipeline.ts:775) — logits → token (argmax default, injectable config.sample).

Two integration tiers, in priority order:

1. Rust fat-station (generateTokens) — the real win. Most deployments are a single GPU fat-station, so the native loop is where tokens/sec actually live. A draft source + batched verify forward (γ positions in one pass) belongs here. This is roadmap-sized (new Rust path + speculative KV handling), not in this PR.
2. TS host loop (generateFromPrefill) — reference + multi-station. The host loop is the slow fallback, but it is where a clean, unit-tested reference implementation of the accept/verify math lives, and it is the only place the multi-station (trisplit) deployments decode. The prototype (src/speculative-sampling.ts) is transport-agnostic so it can drive either tier.

Two distinct axes — keep them separate

NextLat's trained latent head is unavailable to us, but it actually bundles two mechanisms that we should pursue independently:

1. Layer early-exit (depth shortcut for ONE token) — the primary, codebase-native path. For each decode step, run transformer layers until the induced token distribution has provably stabilized, then exit early and skip the remaining layers. This is what "fast-pathing out of layers" means. If the exit is governed by a certified gate, there is no draft/verify loop and no KV rollback — the certificate is the guarantee. We already own the certified gate (next section).
2. Token-level self-speculation (draft several tokens, verify in one pass) — the heavier, NextLat-3.3× path. Shallow layers draft x_1..x_γ, full depth verifies. Needs speculative-KV rollback. The lossless accept/verify math lives in src/speculative-sampling.ts (reference). Pursue only after early-exit lands.

Draft sources for axis 2 (training-free)

• Self-speculation via layer-skip — first k layers draft, full depth verifies. No second model, reuses the same KV. Best fit for the single fat-station.
• Small draft model — e.g. TinyLlama drafts for Mistral-7B (shared tokenizer family). Higher acceptance, costs a second resident model.
• N-gram / prompt-lookup draft — for templated output (JSON, code, tool calls). Free; strong on structured spans.

Axis 1 reuses the gate we already have: 2·τ ≤ γ

tukey_bowl.rs ("tao bowl") proves the energy bound: if the true logits have a strict argmax with runner-up margin γ, and a draft differs from truth by < τ in sup-norm with 2·τ ≤ γ, the argmax cannot flip. The Tukey bisquare tuning constant is exactly c = γ/2. conformal_gate.rs makes it truth-free at serve time: calibrate the residual once, take the conformal q-quantile Q_q, admit iff 2·Q_q ≤ γ, and the serve-time flip ("lie") rate is certified ≤ 1 − q with zero ground truth at inference. The Lean twin is Gnosis/ConformalCoverageAdmissibility.lean + EnergyMarginBridgeAdmissibility.

Today this gate governs quantization-precision drafters (dequant tile masking), not layer early-exit, and not the LLM text decode. Re-pointing it is the whole job:

conformal_admit(cal_k, q, γ) then yields, per step, the earliest layer k whose 2·Q_q(k) ≤ γ. Variable length falls out for free: γ varies per token, so high-margin ("easy") tokens exit very early and low-margin ("hard") tokens run full depth — the training-free analogue of NextLat's latent-rollout self-signal. The γ/2 factor is not a knob to tune; it is the proven bound. The real levers are the calibration coverage q (flip-rate budget) and the per-layer calibration set.

Missing pieces for axis 1 (small):

• Apply lm_head at an intermediate layer. The exit head (final RMSNorm + lm_head
  • softcap) already exists — model.rs:835, amplituhedron_tau.rs:577; it just needs to run after the layer-k RMSNorm instead of only at the last layer.
• Offline per-(layer-k, step-bucket) residual calibration to build each cal_k.
• No KV rollback, no verify pass — the conformal certificate replaces verification.

The acceptance math (standard speculative sampling)

Draft proposes x_1..x_γ with probs q_i. Target computes p_i(· | prefix·x_{<i}) for i = 1..γ in one batched forward, plus p_{γ+1} for the bonus token. Then:

• For i = 1..γ: accept x_i with prob min(1, p_i(x_i) / q_i(x_i)).
• On first rejection at i: resample from the normalized residual max(0, p_i − q_i), discard x_i..x_γ, stop. (Lossless — output distribution is exactly the target's.)
• If all γ accepted: sample the bonus token from p_{γ+1}.

Variable-length drafting (the NextLat-flavored refinement): γ is not fixed. The draft keeps proposing while its own confidence stays above a floor (top-prob or inverse-entropy gate) and stops early when it gets unsure — long drafts on easy spans, short on hard ones. The accept loop above is unchanged; only the proposal length adapts. This is the training-free analogue of NextLat's latent-rollout self-signal.

Cost / correctness constraints specific to this codebase

• KV-cache rollback. Rejected draft positions must be evicted from kv_cache.rs (the per-token save_kv writes are position-indexed). The verify forward writes KV for all γ speculative positions; on rejection at i, positions i..γ must be rolled back. This is the main piece of new state machinery.
• Batched verify forward. The host has lmHead over a single tail residual; a true verify needs logits at γ positions from one forward. The Rust models already have forward_range_chunk (batched prefill) — verify is the same primitive run over the speculative span.
• Terminal-prosody bias (chooseToken) must apply to the committed token only, not to speculative draft scoring, or the lossless guarantee breaks.
• Amplituhedron replay cache is orthogonal (prefix replay), but the speculative span should not poison the replay cache until committed.

Roadmap tie-in

This is a benchmark-gated item for the inference-squeeze loop. Earn prod via benchmark, per [[inference-squeeze-roadmap]]. Suggested order:

1. Axis 1 (early-exit, primary). Add an intermediate-layer lm_head apply to one Rust model; calibrate per-(layer-k, step) residuals offline; reuse conformal_admit to pick the exit layer per token. Benchmark mean exit-depth and tokens/sec at a fixed certified flip-rate (e.g. q=0.99 → ≤1%). No verify, no KV rollback. This is the path that moves fleet tokens/sec.
2. Axis 2 (token-level, optional). The reference accept/verify module (src/speculative-sampling.ts) + tests already exist for the host loop. Wire it behind a flag with a layer-skip self-draft only if axis-1 headroom is exhausted; this one needs speculative-KV rollback in the Rust generateTokens path.

NextLat's training-side benefit (belief-state compression) is being tracked separately as a gnosis-math formalization, not an inference change.