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 pathThe 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§
- Decode
Dims - The model-scalar geometry the shared build path needs.
- Decode
Layer Inputs - The per-layer arguments of
DecodeBackbone::decode_apply_layer, bundled so the signature stays readable as families are added. - Decode
Token Host - Per-token host data, before it becomes either a baked Const or an upload.
- Mask
Plan - 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.
- Rope
Plan - Per-layer RoPE-BASE variation — the sibling axis to
MaskPlan, and the one Gemma3 forced.
Traits§
- Decode
Backbone - Everything the shared build path needs from one model family.
- Persistent
Decode Model - 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_capacityuses. On CPU the Storage wraps the host bytes; on GPU it performs the (tiny) H2D copy.