Warm-start mid-loop

build_resume_handle + resume_generate let you prefill a prompt once at some loop depth L_used, over-allocate a KV cache with capacity for a larger max_loops, and then generate at any new_L in [L_used, max_loops] without re-running the first L_used loop iterations.

This is the right pattern for interactive UIs where the user might ask "try harder on that" — you route the follow-up to resume_generate(handle, new_L) instead of a fresh generate.

It's also the right pattern for adaptive-quality offline pipelines that score the same prompt at multiple depths (e.g. for uncertainty calibration) — run once at L_max and reuse the intermediate state for all shallower depths.

Model.build_resume_handle(prompt, L_used, max_loops=32, cache_window=0)

Runs prefill once at L_used loop iterations and returns a ResumeHandle. The KV cache inside the handle is over-allocated to hold state for up to max_loops iterations, so later resume_generate calls extend in place with no reallocation.

handle = model.build_resume_handle(prompt, L_used=8, max_loops=32, cache_window=0)

Parameters

Name	Type	Default	Description
prompt	np.ndarray[int32]	—	Prompt token ids, shape [seq_len]. Must satisfy seq_len + future max_tokens <= model.max_seq_len.
L_used	int	8	Loop depth to prefill at. Must satisfy 0 <= L_used <= max_loops.
max_loops	int	32	Upper bound on any new_L passed to resume_generate on this handle. The cache is over-allocated to this size so no reallocation is needed at resume time. VRAM cost scales linearly with max_loops; set aggressively if you plan to resume many times at various depths.
cache_window	int	0	Sliding-window cap; 0 means full cache.

Returns

tinyloop_py.ResumeHandle — opaque object.

Raises

• RuntimeError — model has been closed; OOM during cache over-allocation; L_used > max_loops; seq_len > max_seq_len.
• ValueError — L_used < 0; max_loops < 1; cache_window < 0.

Memory cost

The handle holds a RuntimeKVCache with n_pre_blocks + max_loops layers, sized for seq_len tokens (or cache_window if non-zero). Only the first n_pre_blocks + L_used layers are populated after build; the remaining (max_loops - L_used) layers are allocated but empty until a resume_generate extends into them.

Example VRAM cost (FP16 KV, D=2048):

seq_len	max_loops	Cache bytes
128	32	34 MB
512	32	136 MB
1024	32	272 MB
1024	64	528 MB

Enable TINYLOOP_KV_H_MODE=1 TINYLOOP_KV_INT4=1 to cut this by ~78 %; see KV cache modes.

ResumeHandle methods

Method	Returns	Description
handle.seq_len()	int	Prompt length.
handle.loops_used()	int	The L_used the handle was built at.
handle.max_loops()	int	The upper bound for new_L.

The handle is ref-counted by pybind11; GPU memory is released when the last Python reference drops.

Model.resume_generate(handle, new_L, max_tokens=128, ...)

Extends the handle's loop state to new_L and decodes. The first handle.loops_used() iterations are not re-run.

tokens = model.resume_generate(handle, new_L=16, max_tokens=64, temperature=0.0)

Parameters

Name	Type	Default	Description
handle	ResumeHandle	—	Returned by build_resume_handle.
new_L	int	—	Loop depth to resume at. Must satisfy handle.loops_used() <= new_L <= handle.max_loops(). When new_L == handle.loops_used() the extend step is a no-op and only decode runs.
max_tokens	int	128	Maximum decode length.
temperature	float	0.0	As Model.generate.
top_k	int	50	As Model.generate.
top_p	float	1.0	As Model.generate.
cache_window	int	0	Sliding-window cap. Overrides the handle's setting for this call only.
repetition_penalty	float	1.0	As Model.generate.

Returns

List[int] — full token stream (original prompt + decoded tokens).

Raises

• RuntimeError — model has been closed; handle belongs to a different (now closed) Model; OOM during cache clone; new_L out of [loops_used(), max_loops()] range.
• ValueError — max_tokens < 1; negative sampling kwargs.

What actually happens

1. Clone cache. The handle's cache is cloned to a working cache sized for n_pre_blocks + new_L layers (so multiple resume_generate calls on the same handle don't race on its state).
2. Restore residual. The snapshot of the residual stream going into iter L_used is copied into buf.main.
3. Extend loop. run_loop_iters_range(model, seq_len, L_used, new_L, cache) runs iters [L_used, new_L) in place on the restored residual. The first L_used iterations are not re-run.
4. Synthesise PrefixCache. A temporary PrefixCache is built over the now-extended working cache.
5. Decode. generate_from_prefix_cache does the actual token generation.

The implementation is ~150 lines in src/generate.cpp — the heavy lifting is the internal run_loop_iters_range primitive that is also the subject of the kernel-level parity test.

Typical usage

import numpy as np
import tinyloop_py as tl

with tl.Model("model.tinyloop", max_seq_len=2048) as model:
    prompt = np.array([...], dtype=np.int32)

    # Initial L=8 generation.
    handle = model.build_resume_handle(prompt, L_used=8, max_loops=32)
    short_reply = model.resume_generate(handle, new_L=8, max_tokens=32)

    # Same prompt, follow-up at deeper L without re-prefilling iters 1-8.
    deeper_reply = model.resume_generate(handle, new_L=16, max_tokens=64)

    # Any new_L up to max_loops is cheap on this handle.
    deepest_reply = model.resume_generate(handle, new_L=32, max_tokens=64)

Measured savings

H100, 407M GPTQ-INT4, prefix=128, decode=16, L_mid=4, L_new=8:

config	wall time	vs fresh
fresh generate(loops=L_new)	329.0 ms	—
build_resume_handle(L_used=L_mid, max_loops=32)	6.4 ms	one-time
resume_generate(new_L=L_new)	222.1 ms	−32.5 %

Break-even point: the 6.4 ms build cost pays for itself after the first resume_generate call that would otherwise have cost fresh − resume = 106.9 ms. So one follow-up at a different new_L makes warm-start strictly faster than calling generate twice. Any additional follow-up is pure profit.

Scale with max_loops:

• The build cost is roughly linear in seq_len × L_used × D. At seq_len=128, L_used=4, D=2048 it is 6.4 ms; at seq_len=512 proportionally larger.
• The resume cost is roughly linear in seq_len × (new_L - L_used) × D plus the per-token decode cost.
• The break-even point stays at N=1 as long as fresh − resume > build, which is the normal case when L_used is a meaningful fraction of new_L.

Correctness invariant

Under greedy decoding (temperature=0.0), resume_generate(build_resume_handle(prompt, L_used=k), new_L=L) produces the same token stream as model.generate(prompt, loops=L). Verified by:

• tests/test_warmstart_parity.cu — kernel-level bit-exact residual match across all cache layers and all 32 768 FP32 residual elements.
• tests/test_warmstart_python.py — Python-level token-stream match over the full decode on 407M GPTQ-INT4.

Non-greedy (temperature > 0) is not bit-exact with fresh generate — the sampling PRNG is not shared between the two paths. Token streams will diverge even with the same input. Tracked as an open follow-up.

Composition with other APIs

With	Supported	Notes
generate	Independent	generate does not touch a handle; resume_generate does not touch generate's cache.
Prefix cache	Not wired	In principle a resume handle could be built off a prefix cache; not implemented yet. Open follow-up.
Speculative decode	Independent	Both use the shared-weights property but along different axes. Orthogonal.
KV cache modes	✓ all modes	Resume handle inherits the process-level KV mode. Bit-exact parity guaranteed under FP16 modes; lossy modes' parity is up to the mode's noise bound.
cache_window > 0	✓	The handle records the window at build time; override per-call if needed. Ring-buffer wrap is handled correctly in all h-mode paths.

Architectural note — why this is weight-sharing-exclusive

A standard deep transformer has different weights per layer, so "resume at layer k" would mean picking up from layer k's output and applying layer (k+1)'s different weights — and the residual-stream state from the original generation's layer k has no meaning for the new layer. You cannot just "run more layers"; the layers above k don't exist in the context of that residual.

TinyLoop's looped block shares weights across all iterations, so the iter-k residual is a valid starting point for iter (k+1) of the same shared block. Applying the block (new_L - L_used) more times is exactly what the original generate(loops=new_L) would have done; the only difference is that we start from a snapshot rather than re-running iters [0, L_used).

This is distinct from prefix cache (same prompt across many requests at the same L) and from speculative decode (same weights at two depths within one generation). Warm-start reuses state across separate generations that differ in L.

Common patterns

Uncertainty escalation pipeline

# Score at L=8. If high uncertainty, escalate to L=16 using warm-start.
handle = model.build_resume_handle(prompt, L_used=8, max_loops=16)
shallow = model.resume_generate(handle, new_L=8, max_tokens=1, temperature=0.0)

# Check top-1 logit margin from a separate score_with_uncertainty call;
# if below threshold, deepen without repeating the first 8 iters.
logits_hi, kl = model.score_with_uncertainty(prompt, L_lo=8, L_hi=16)
if kl > 0.5:
    deep = model.resume_generate(handle, new_L=16, max_tokens=32)

"Try harder" UI button

# User hits submit → answer at L=8
handle = model.build_resume_handle(user_prompt, L_used=8, max_loops=32)
answer = model.resume_generate(handle, new_L=8, max_tokens=128)

# User hits "try harder" → reuse handle, go deeper
better = model.resume_generate(handle, new_L=16, max_tokens=128)

# User hits "really hard" → reuse same handle, go deepest
best = model.resume_generate(handle, new_L=32, max_tokens=128)

The handle is built once; each click reuses it. Build cost (6.4 ms on 407M) is paid on the first click only; every subsequent click is ~32.5 % faster than a fresh generate.