model-interfaces.md

API: the Model interfaces

This document is the normative contract between trellis.model and its consumers — the Trainer Core, which trains through these interfaces, and the Inference Engine, which drives the step-mode path. It states what a call takes, what it returns, and what is guaranteed about the relationship between the two modes. The implementation must match this document exactly; tests/test_model_docs.py checks that it does. Design rationale lives in docs/modules/backbone.md, state and cache semantics in docs/architecture/state-and-cache-contract.md, and neither is restated here.

Exported names

Name	Kind	Role
TrellisModel	module	the backbone: graph, loop engine, heads, audit
build_model	function	construct and audit from a TrellisConfig
build_preset_model	function	construct from a named preset
ModelSettings	dataclass	behaviour flags; parameter-count-neutral
RouterSettings	dataclass	nested, and the tie-break's name
AuxHeadSpec	dataclass	auxiliary-head widths and label vocabularies
TrellisOutput	dataclass	the single typed result object
LoopReadout	dataclass	one executed loop's deep-supervision readout
BankReadout	dataclass	one read point's slot indices and statistics
ModelState	dataclass	every state instance, plus the decode position
StateShape	dataclass	how many instances a graph needs
LoopEngine	module	the loop, the routing, the segmentation
CapacityRouter	module	the expert-choice router at boundaries j ≥ 2
RouterDecision	object	one boundary's mask, count and scores
PerLoopSpecialisation	module	per-loop adapters and norm gains, and the switch
TiedEmbeddingHead	module	the one 49,152 × 768 matrix, both directions
ExitHeads	module	one halting scorer per loop
AuxiliaryHeads	module	the four training-only heads
ModelParameterReport	dataclass	the audit
ModelError and subclasses	exceptions	every rejection this unit makes

TrellisModel.forward — the training and prefill path

out = model(token_ids, k=None, state=None, readout_mode=None,
            dfg_pairs=None, token_classes=None, plan=None, return_state=False)

Argument	Type	Meaning
token_ids	(batch, tokens) int64	ids in [0, vocab_size); a 1-D tensor is treated as one sequence
k	int in [1, K_max]	loops to execute; defaults to K_max
state	ModelState or None	a prefill continuation; a fresh bundle when None
readout_mode	"logits" or "hidden"	overrides the instance's setting for this call
dfg_pairs	(batch, n_pairs, 2) int64	candidate (def, use) positions for the DFG head
token_classes	(batch, tokens) int64	boundary classes for the A9 attention-bias arm
plan	sequence of (batch, tokens) bool	overrides the router at each boundary
return_state	bool	return the advanced state bundle

k outside [1, K_max] raises LoopEngineError. Ids outside the vocabulary raise ModelConfigError. A plan mask that revives a token the previous boundary dropped raises LoopEngineError under nested selection.

TrellisModel.step — the decode path

out = model.step(token_id, state, k=None, readout_mode=None,
                 dfg_pairs=None, token_classes=None, plan=None)

token_id is (batch, 1); a scalar or (batch,) is accepted and reshaped. state is required and is validated against model.state_shape() and the call's batch size; a mismatch raises ModelStateError. out.state is always populated.

step mutates nothing the caller did not pass in. ModelState is frozen and every advance returns a new bundle; the only in-place writes are into the preallocated global-KV storage inside the bundle the caller handed over, which is what preallocating it is for.

The routing plan, and why step needs one

Expert-choice selection keeps the top `⌈ρ_j·n⌉ tokens of a sequence, so a token's fate depends on the scores of the tokens after it. During training and prefill the whole sequence is present and this costs nothing. At decode time it is unavailable in principle. A router that pretended otherwise — ranking against a running quantile, or against the prefix only — would produce a different selection than the forward pass and silently break every equivalence below.

So step takes the decision rather than inventing one. The Inference Engine owns producing it, from the exit heads and the conformal thresholds, which is the causal signal that exists for the purpose. With no plan, every loop runs, which is the capacity schedule's own worst case and therefore always inside the budget.

TrellisOutput

Field	Type	Always present
logits	(batch, tokens, vocab)	yes
readouts	tuple of LoopReadout, one per executed loop	yes
exit_logits	(loops, batch, tokens)	when K_max > 1
aux	AuxiliaryOutputs, a dict keyed by head name	empty when the heads are absent
router_masks	tuple of (batch, tokens) bool, one per boundary	empty at K = 1
router_scores	tuple of (batch, tokens) float	alongside the masks
bank	tuple of BankReadout, one per read point	when the bank is enabled
loop_counts	(batch, tokens) int64	yes
executed_loops	int	yes
state	ModelState	from step, and from forward with return_state=True
readout_mode	"logits" or "hidden"	yes
readout_applicator	callable	yes

A LoopReadout carries hidden or logits according to the mode, and always carries applicator; as_logits() returns the logits either way. Both modes satisfy the same contract — the difference is only whether the K × n × 49,152 tensor is ever materialised. hidden is the default because at 4,096 tokens and four loops that tensor is 805M floats.

A BankReadout carries indices — the slot indices before the value gather, which is what the Inference Engine prefetches against — plus weights and the read point's utilization report.

exit_logits are reported and never acted on. Thresholding is the Inference Engine's decision, against thresholds the Training Program calibrates. Exits may never raise a token's depth above the capacity schedule.

Guarantees

Guarantee	Scope	Tolerance
Unrolled equivalence	every capacity 1.0, K ∈ {1,2,3,4}	bit-identical in fp32; 2e-2 relative in bf16
Router counts	every boundary	exactly ⌈ρ_j·n⌉, never a function of the scores
Nesting	default settings	selected set ⊆ previous boundary's set
Tie-break	equal scores	toward the lower position index, reproducible bit-for-bit
Skipped-token pass-through	every boundary	bit-identical; no KV write, no GDN state advance
Loop-invariant sliding cache	K = 1 vs K = 4	identical slots, fill and bytes
Global KV	any K	written once per token per global layer
Determinism	same input, seed and configuration	same routing, same K path, same logits
Auxiliary-head detachability	heads on vs off	logits, readouts, exit logits and routing bit-identical
Bank neutrality	deltas forced to zero	bit-identical to a bank-disabled model
Batch-composition invariance	routing and depth	independent of what shares the batch
Step vs forward	K = 1	1e-4 absolute in fp32
Step vs forward	any K, no sliding layer in the core	1e-4 absolute in fp32
Step vs forward	any K > 1, shipped configuration	does not hold — see below

Prefill and decode are not equivalent for K > 1

This is the one place where the specification asks for two things that cannot both be true, so it is stated here rather than reconciled away.

The sliding cache is loop-invariant: one entry per position, written at loop 1 and gate-updated at loops > 1. A position's cache entry therefore records how many loops that position has been through.

• In a full-sequence forward, the loops are the outer iteration. At loop j, every position's entry has been blended j − 1 times, so a token attends to a prefix that is exactly as refined as it is.
• In token-by-token decode, the tokens are the outer iteration. When token t runs its loop 1, every prefix position has already finished all of its loops, so token t attends to a prefix that is more refined than it is.

Those are different computations, and no ordering of the same operations makes them the same one. Reconciling them needs per-loop key and value storage — K entries per position instead of one — which is precisely what loop invariance forbids and what the 2.1 MB sliding-cache figure is a figure for.

What holds, and what the suite asserts. At K = 1 there is no gated update and decode reproduces prefill exactly. With the core's sliding layer replaced by a GDN layer — a configuration-only change that touches nothing else — decode reproduces prefill exactly at every K ∈ {1,2,3,4}, which is what identifies the loop-invariant sliding cache as the sole cause: the loop engine, the per-loop LoRA switching, the per-loop GDN state instancing, the routing, the readouts, the exit heads, the prelude, the coda and the tied head are all step/forward exact. For the shipped configuration at K > 1 the divergence is measured and pinned by a regression test, so that it cannot grow quietly and cannot silently disappear.

What a consumer should do about it. Use K = 1 for the speculative drafter, where exactness is what makes checkpoint-and-rollback sound. For K > 1, treat prefill and decode as the same model under two schedules rather than as two evaluations of one function — which is what an adaptive-depth model with a shared cache is.

The parameter audit

model.parameter_report() returns a ModelParameterReport computed at construction and compared component by component with trellis.config.budget.params.compute_parameters. Construction fails with ParameterAuditError on any mismatch outside the two named booking differences.

Row	Default configuration
prelude_untied	13,791,744
core_block_tied	44,579,328
coda_untied	13,791,744
per_loop_lora	2,064,384
routers_gates_norms	67,335
embeddings_tied	37,748,736
memory_bank	68,681,728
serving_total	180,724,999 — 180.72M
training_aux_heads	6,000,000
training_total	186,724,999 — 186.72M

Two rows are booking differences between the constructed model and Prompt 1's report, both named in RECONCILED_ROWS and neither rounded away:

• swa_kv_update_gate, 768. The state-and-cache contract and this unit's specification both place the sliding layer's w_g in the routers/gates/norms row; Prompt 1's decomposition of that row has no term for it. It is counted, and the budget side of the comparison is adjusted by exactly that amount with a note. It does not move the published totals.
• bank_sub_key_tables, 196,608. The product-key codebooks, which Prompt 1 books outside the 68.68M bank row deliberately and trellis.bank's own audit books the same way.

Other model interfaces

Method	Returns
initial_state(batch, capacity, device, dtype)	a zero ModelState; cache dtype defaults to the model's
state_shape()	StateShape — 22 GDN instances, 1 sliding cache, 2 global KV at the defaults
loop_capacity_schedule()	(1.0, 1.0, 0.5, 0.25)
expected_loops()	2.75
effective_depth(k)	4 + 6k
parameter_report(refresh=False)	the audit
strip_aux_heads()	a serving state dictionary; removals reported by stripped_keys()
load_serving_state_dict(sd, strict=True)	(load_result, surplus_keys)
all_layers / core_layers	the ten unique layers; the core appears once
engine.core_pass(...)	one loop's core forward — the activation-checkpointing unit
engine.segment_executor	substitute torch.utils.checkpoint.checkpoint(use_reentrant=False)

strip_aux_heads returns the dictionary rather than mutating the model, so a training process can write a serving checkpoint without disturbing the run still using the heads. The removed keys are reported through stripped_keys() rather than left for the caller to discover by diffing two key sets.

Change log

• Added the training-forward and step-mode contracts, the TrellisOutput field table, the routing-plan requirement for decode, the guarantee table, the parameter-audit table with its two named booking differences, and the prefill/decode incompatibility for K > 1.