Skip to main content

Module persistent_decode

Module persistent_decode 

Source
Expand description

The shared persistent-decode rebind driver (GAP-029 increment 2b).

§Why this module exists

rebind_and_realize_prebuilt — the per-token hot path that reuses a held, already-optimized decode graph — was hand-copied per model. LlamaModel and PhiModel each carried a 48-line copy, structurally identical, and DeepSeek2Model carries a third that has diverged (rebind_..._mla).

Two hand-maintained copies of a hot path are a reproduction mechanism, not two implementations. GAP-029 adds up to eight more model families; porting them against the copied shape would have produced ten copies and kept the generator that produces copy eleven. So the driver is collapsed here before any new family is added.

§Why the two existing impls were ported first

This refactor is behaviour-preserving, which means it has no born-red state: “the tests still pass” is also exactly what a no-op produces. The only thing that makes it verifiable is an oracle that predates it — LlamaModel and PhiModel already work, so correctness is “their behaviour is unchanged”. A newly added family cannot serve that role: it has no prior behaviour to be unchanged from, so its test can only assert what the new code already does.

Measured oracle (not assumed): 8 Phi decode tests + the parallel Llama set, including phi_generate_loop_persistent_byte_exact_and_plans_once, which asserts byte-exact logits and plans-once and therefore already walks the multi-token rebind path this driver implements.

§The path divergence this driver must PRESERVE, not unify

Measured 2026-08-12 on CPU/F32 — and stated nowhere else in the tree:

llama: offset_node.is_some() == true   → device-offset path (Op::WriteSliceDoff)
phi:   offset_node.is_some() == false  → SymEnv path

The KV-write offset rides a device buffer for one and a symbol for the other. Forcing either model onto the other’s path would be a behaviour change wearing a refactor’s clothes, so the divergence lives behind PersistentDecodeModel::build_decode_token_data: each model keeps building its own token data exactly as before, and this driver never inspects which path it got.

Structs§

DecodeDims
The model-scalar geometry the shared build path needs.
DecodeLayerInputs
The per-layer arguments of DecodeBackbone::decode_apply_layer, bundled so the signature stays readable as families are added.
DecodeTokenHost
Per-token host data, before it becomes either a baked Const or an upload.
MaskPlan
Per-layer attention-mask variation — the axis GAP-029 measured across increment 3’s families and the earlier six-axis check could not see, every one of its axes being model-scalar.
RopePlan
Per-layer RoPE-BASE variation — the sibling axis to MaskPlan, and the one Gemma3 forced.

Traits§

DecodeBackbone
Everything the shared build path needs from one model family.
PersistentDecodeModel
The per-model variation points of one persistent-decode rebind step.

Functions§

build_and_realize_first_decode_token
D2 build — mint the held graph with re-bindable data Consts, optimize it ONCE, populate session, and return the first token’s logits.
build_decode_mask_variants
Build the [n_variants, 1, seq, max_seq_len] decode mask, variants concatenated along the leading axis.
compute_decode_token_host
Compute one decode step’s host-side data — the RoPE tables and the stacked mask. Pure host math, no graph and no upload, shared by the build path (which bakes or uploads it) and every family’s rebind hook (which uploads it).
drop_decode_session
Drop a held session, releasing its stable node bindings from ctx.
forward_with_kv_context
D1 — build the decode graph fresh, realize it, discard it.
forward_with_kv_context_persistent
The persistent decode entry point, shared by every family on this path.
rebind_and_realize_prebuilt
One persistent-decode step on an already-built session: recompute the per-token data, rebind, realize the held graph, advance cache state.
upload_decode_token_data
Upload one step’s host data to device-resident Arcs — the same upload path KvCache::with_capacity uses. On CPU the Storage wraps the host bytes; on GPU it performs the (tiny) H2D copy.