ReadonlyappendKnown tokens: this graph declares no committed spans, so they read as a window.
Readonlycachelayer idx → cache index (donors only; -1 for sharers).
ReadonlycacheCaches for donor layers only (sharers attend their donors' reads), all
sharing one mask memo, so each forward builds each mask once. Sliding
layers compose a ring of the sliding window, whose reads own the window
mask, and stay bf16 under TurboQuant. The full-attention caches' TurboQuant
read is deferred (deferredTurboQuantKernels): values stay rotated through
the fused SDPA and the read un-rotates the output.
ReadonlyconfigReadonlydecodeOne new token per sequence: ids [sequences, 1], so rows are sequences.
ReadonlyembedUnscaled token embeddings, for media prompts that splice embeddings.
ReadonlyembeddingPrepared embeddings without explicit positions, indexed by total rows.
ReadonlyembedReadonlyfinalSet by the trainer around its training forwards to checkpoint each layer
(null = off). While set, the scoring and training forward
(sequenceHidden) runs the training forward (forwardTrainingHidden's),
whose donors hand their sharers keys and values as checkpoint tensors.
ReadonlylayersReadonlylogitsThe one output head at every width.
ReadonlyloraActive LoRA adapters for the current generation (see LoraState).
ReadonlymediaReadonlymixedMixed rows in one forward, declared by a composition that serves them.
ReadonlymixedMixed groups that retain their own projection geometry. The scheduler routes each group independently when a verify group is present.
ReadonlymultiReadonlymultiReadonlynumDonor count: layers [0, numDonors) own caches; the rest share.
ReadonlyperReadonlyperReadonlyperReadonlyperReadonlyperReadonlyprefixWeight-name prefix ("language_model.model" or "model").
Readonlypreviouslayer idx → donor layer idx whose fetched KV it consumes (self for donors).
ReadonlypromptOne whole chunk of a prompt: tokens per sequence equal the composition's
prefill chunk size; rows are sequences times that size. Token ids only: a
prompt with prepared media enters through the MlxPromptInput that
bindMediaInput returns. The GPU is this graph's one unit.
ReadonlypromptOptional chunk bindings when a file distinguishes chunks from prompt tails.
ReadonlysequenceFull-sequence scoring operations, indexed by total token rows.
ReadonlytraceReadonlytraceReadonlyverifyA speculative window: each sequence's last accepted token and its drafted
tokens, right-padded by the caller to the round's depth, ids
[sequences, depth + 1], so rows are sequences times (depth + 1). The graph
appends all depth + 1 positions; dropping the rejected suffix is the
caller's speculative transaction, begun before this call.
ReadonlyweightsTotal weight-shard bytes (for the conditional wired-limit scope).
ReadonlywindowPrepared image/audio embeddings, layer taps for drafts, and a compiled decode step, except through experts: shapeless replay retraces their growing windows.
Donor layer indices that KV-sharing sharers reuse (e4b: {22,23} — the last donor of each attention type; empty for dense models). Their K/V must survive across segment boundaries (saved as detached boundaries).
Segmented backward with KV-shared donor threading, prefix-shared ORPO, gradient checkpointing, and the quantized tied head. Flash training attention is refused: e4b crashes natively on it at long sequences.
media embeddings/embeddings+positions: bind one request's prepared
embeddings into an input whose prefill is indivisible.
The compiled decode step declared above: the step records this graph's own forward, so it shares the graph's weights and constants.
ProtectedcomputeRelease model-owned graph constants after all executions and compiled runners have stopped. Weight tensors and hidden taps are borrowed.
Ports over this graph's live caches that draft sources may consume.
e2b/e4b training entry: the SCALED input embedding (the first segment boundary) plus the per-layer-input tensor [B,L,nLayers,width] (a detached boundary, sliced per layer inside runLayerRange). perLayer is null for dense models. Caller owns both.
ProtectedembedProtectedforwardThe decoder over h0 [B, L, hidden] (consumed), every donor layer's
attention reading its cache for phase and every KV-sharing layer
attending its donor's read. ids [B, L] feed the per-layer inputs
(e2b/e4b). capture receives each layer's residual output after the layer
runs (see MlxLayerCapture); the caller hands it the final hidden.
Layer i's input [B, L, width] (owned) from the [B, L, nLayers, width]
tensor compute returned (borrowed). A compiled-decode trace slices at an
array start, which replays at every step.
The training forward over stateless training caches (TrainingCache,
PrefixSharedCache), one per donor layer: ids [1, T] to the final-norm
hidden [1, T, hidden]. Every attention reads through its training cache's
window read (a KV-sharing layer through its donor's, with the donor's keys
and values), and a prefix-shared layout ropes block-wise; under gradCkpt
each layer is checkpointed with its donor keys and values as checkpoint
tensors.
LoRA-mountable linears, keyed by weight-file module path: optiq mount.py's 7 target suffixes PLUS the e2b/e4b per-layer-input projections — mlx-lm's trainer targets those on e4b and optiq's mount silently drops their trained weights (we apply every adapter weight we can map). Expert pools stay non-targets (LoRASwitchLinear is future work).
Computed model constants must retain their constructor arithmetic when a restored request enters compiled decode without a prefill first.
media: the lazy encoders this graph loads from the checkpoint at modelDir.
Retire the compiled step before the graph unloads: its recorded graph borrows the graph's weights and constants. Safe to call when none exists.
Run decoder layers [aIdx, bIdx) on hidden h, returning the residual
stream after the last layer in the range (NO finalNorm) plus the K/V of
any reused-donor layers in the range (donorKvOut, for later segments).
The input h is NEVER disposed (the caller / vjp owns the boundary leaf);
intra-range hidden intermediates and per-layer slices ARE disposed.
Non-reused donors' K/V are disposed; reused donors' K/V are returned
(caller threads + disposes). NOT gradient-checkpointed — the segmentation
IS the memory strategy.
the pass's training caches, one per donor layer.
the [B,L,nLayers,width] tensor (caller-owned) or null.
K/V of reused donors from EARLIER segments, keyed by donor layer index (caller-owned; not disposed here).
Readonlytrace
The graph's operations and construction facts. Pooled embedding support is declared by its operation alone. Other operations accompany the promises in
graphCapabilities; composition checks their pairing when it binds (declaredGraph).