The composition this graph was built for: its KV scheme and paging are its caches' (makeCache).
ReadonlycacheReadonlycacheOne cache per layer, the lego the composition's KV scheme names for it
(bf16, bf16-first affine, per listed layer under a per-layer scheme, or
TurboQuant on full-attention layers, read deferred as Gemma 4 reads it), a
ring of the sliding window on sliding layers, all sharing one mask memo;
under the composition's paging the full-attention layers hold pages of
what the scheme names (maybePageKv, growing from one block). The encoder
prefill fills them, the decoder canvas pass reads them.
ReadonlycanvasReadonlyconfigReadonlydenoisingReadonlydenoisingReadonlydenoisingReadonlyembedUnscaled token embeddings, for media prompts that splice embeddings.
ReadonlyembedReadonlyencoderReadonlyfinalReadonlyimageReadonlylayersReadonlyloraReadonlymakeThe graph binds compiled GeGLU and softcap closures, matching the reference's mx.compile'd operations.
ReadonlyprefixReadonlyselfReadonlyslidingReadonlysoftcapReadonlyvisionDedicated DiffusionGemma SigLIP vision tower (image-text-to-text).
ReadonlyweightsDenoising rows interleave through their own grouped method; images enter as the method's pixel prefill input.
The denoising objective: LoRA on the decoder blocks (the encoder reuses them, weight-tied).
Decoder canvas pass over an on-device canvas: embed canvas -> 30 layers, each reading its cache through the block read (the cache builds the reference _make_decoder_masks sliding mask and keeps the newest sliding_window - 1 encoder positions) -> final norm. Returns hidden [1, canvasLength, hidden]. The D2 engine drives this per denoising step (canvas stays on-device).
One denoising-step decoder pass -> softcapped logits [1, canvas, vocab].
Dequantized embedding table (once per call), for self-conditioning soft
embeddings — matches optiq's _diffusion_soft_embedding_weight.
Releases the rope frequency tables its layers built; the weights are its owner's.
Encoder prefill: run all layers over the prompt through the window read, filling the cache. Discards the hidden.
Diagnostic: run the encoder, invoking onLayer(i, h) after each layer
(h is the raw layer output, pre-final-norm) and onEmbed(h) on the
embedding. Returns the final pre-norm hidden. Fills cache.
Image-text-to-text encoder prefill: merged embeddings through all layers,
each reading through the encoder window (causal over the whole prompt, no
sliding window, image positions both ways; reference _make_encoder_masks +
_vision_block_overlay, one mask for every layer), filling cache. The
canvas decoder then reads it like the text path.
Re-prefill the encoder over an accepted canvas block (continuation), so the next block's decoder sees it as context. Appends to the existing cache.
D1 entry point: prompt + canvas ids -> softcapped logits [1, canvas, vocab]. Builds a fresh cache, prefills the encoder, runs one decoder pass.
Training forward over an on-device canvas (D5 LoRA): prompt + canvas array -> softcapped logits [1, L, vocab], differentiable w.r.t. mounted LoRA.
LoRA-able decoder linears (the optiq diffusion DEFAULT_LORA_KEYS: attention q/k/v/o + dense MLP gate/up/down per layer; full layers have no v_proj).
media pixels: the pixel input a denoising method consumes; null when the checkpoint has none.
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).