dew.nn.backbones.causal_transformer

The autoregressive transformer decoder every language model here trains.

Token embedding, rotary positions, pre-norm blocks of grouped-query causal attention and a gated MLP, a final RMSNorm, and an fp32 head. Attention goes through the one shared kernel path in dew.nn.attention, so a run picks reference/xla/cudnn/tpu the same way a diffusion run does, and decoding reuses the same fixed-size KV cache helpers.

Parameter names mirror the HF decoder layout - embed_tokens, layers_N.{input_layernorm, self_attn.{q,k,v,o}_proj, post_attention_layernorm, mlp.{gate,up,down}_proj}, norm, lm_head. A model family is supported only after its translator and same-weight reference parity test land. Gemma’s two extra norms are the exception: HF calls them post_attention_layernorm and post_feedforward_layernorm even though they normalize sublayer outputs, so here they are attention_output_norm and mlp_output_norm and the pre-norms keep their names. dew.interop.hf_decoders does that rename.

The block holds its token mixer in a slot: any module with the (x, decode=…, positions=…, segment_ids=…) -> x signature of CausalSelfAttention becomes self_attn without the block changing, which is where a linear-attention mixer goes.

Name	Summary
LayerKind	What the layers of one kind in the pattern do differently.
ResolvedKind	One kind of layer with the model’s defaults filled in.
LayerSpec	What one layer of the stack is, resolved: everything its block’s parameters and computation depend on that the layers do not share.
scan_groups	The stack as runs of layers, (first, count) each, in order.
group_name	The module name of a scanned run: layers_3_7 runs layers 3 through 7.
group_layers	The layers a stack module name runs: layers_3_7 as range(3, 8), layers_3 as range(3, 4), anything else as None.
INTERMEDIATES	The collection flax’s capture_intermediates fills.
STREAMS	Manifold-constrained hyper-connections’ [B, S, hc_mult, D] residual streams.
layer_outputs	The capture_intermediates filter that keeps every layer’s output.
layer_output	Layer index’s output out of the collection layer_outputs filled.
Mixture	The experts some layers route to, and how the router chooses.
GatedMLP	down_proj(act(gate_proj(x)) * up_proj(x)): swiglu is silu, geglu is the tanh approximation of gelu (HF’s gelu_pytorch_tanh) and geglu_exact the erf form (HF’s gelu, which Gemma’s released config names).
BlockWiring	How a block norms its two residuals, and whether it scales its output.
QKV_RESIDUALS
ATTENTION_RESIDUALS
MLP_RESIDUALS
RESIDUALS	The values a block names as it runs, each after the projection that produced it: kv_proj is latent attention’s fused kv_b_proj, context the attention kernel’s output before o_proj, and the MLP names cover the dense MLP and the routed experts alike.
RematPolicy	What a recomputed block keeps for its backward pass, and where.
REMAT_POLICIES	MaxText’s remat recipes (nnx_decoders.py, get_remat_policy), fastest and largest first.
remat_policy	value as the policy it names: a RematPolicy, a name in REMAT_POLICIES, a record of save/offload names, or None for no recomputation at all.
DecoderBlock	Pre-norm decoder block: token mixer, then feed-forward, both residual.
MTPBlock	One multi-token-prediction depth: the next depth’s hidden states.
Block	Layer index’s block under a module name: what the stack and its stages build their layers from, so one factory describes every view of them.
WRITTEN	The collections a decoder block writes: its decode cache, and what its router, its attention and its sparse indexer sow.
run_stack	Run the layers over x, one run at a time as groups says.
PipelineStage	The layers one pipeline stage holds, over one microbatch.
StackView	The layer stack’s variables as its loops read them.
DecoderBank	One decoder’s stored layer namespace below every variables collection.
CausalTransformer	Decoder-only transformer over token ids: [B, S] int32 -> [B, S, vocab] fp32.
gather_cache_rows	A decode cache reindexed on its batch axis, one gather per leaf.

LayerKind

dataclass source

class LayerKind(
    window: int | None = None,
    chunk: int | None = None,
    num_kv_heads: int | None = None,
    rope_theta: float | None = None,
    rope_scaling: RopeScaling | None = None,
    yarn: YarnScaling | None = None,
    head_dim: int | None = None,
    mixer: MixerBase | None = None,
)

What the layers of one kind in the pattern do differently.

The pattern names each layer’s kind, and this is what the kind means: a windowed kind is the “sliding attention” of the reference configs, and rope_theta and head_dim are the model’s unless this kind states its own. Rotary positions rotate every dimension of a windowed kind; Gemma 4 puts its partial rotary on the global layers and its sliding layers rotate whole.

mixer is this kind’s token mixer, a value from the mixers registry; None rides the model’s mixer. A hybrid stack names its per-layer mixers here, keyed by the names already in the pattern.

window: int | None

Keys a layer of this kind attends, its own included; None attends all.

chunk: int | None

Chunked local attention: a layer of this kind reads only the keys at or before each query whose position shares the query’s position // chunk, MaxText’s chunk_attn_window_size and Llama 4’s attention_chunk_size. None attends all; a kind sets a window or a chunk, not both.

num_kv_heads: int | None

This kind’s key/value head count; None takes the model’s. Gemma 4’s global layers keep fewer than its sliding ones (num_global_key_value_heads).

rope_scaling: RopeScaling | None

This kind’s llama3 ramp or its record; None rides the model’s.

yarn: YarnScaling | None

This kind’s YaRN ramp or its record; None rides the model’s. OLMo 3 scales its full-attention layers alone (configuration_olmo3.py:110-113), so a YaRN ramp is a kind’s as much as the model’s.

mixer: MixerBase | None

This kind’s mixer value or its record; None is the model’s mixer.

ResolvedKind

dataclass source

class ResolvedKind(
    window: int | None,
    chunk: int | None,
    num_kv_heads: int,
    rope_theta: float,
    rope_scaling: RopeScaling | None,
    yarn: YarnScaling | None,
    head_dim: int,
    mixer: MixerBase | None,
)

One kind of layer with the model’s defaults filled in.

LayerKind is what a config states, so a field it leaves to the model is None there. This is what the model resolved it to, so rope_theta and head_dim are numbers; only the window stays optional, because attending the whole sequence is what a kind without one does. mixer passes through: it needs no resolution, only the model’s default when unset.

LayerSpec

dataclass source

class LayerSpec(
    layer_type: str,
    kind: ResolvedKind,
    routed: bool,
    hash_routed: bool,
    width: int,
    sparsity: float,
    kv_shared: bool,
    provider: int | None,
    residual_site: ResidualSite | None,
)

What one layer of the stack is, resolved: everything its block’s parameters and computation depend on that the layers do not share.

Two layers with equal specs have parameters of the same shapes and run the same program, so a scan can run them as iterations of one body and a pipeline can run them at the same position of different stages. Everything the whole model sets (norms, the attention dials, per-layer inputs, AltUp) is the same for every layer and so is not repeated here.

routed: bool

The feed-forward routes to the mixture’s experts.

hash_routed: bool

The routed feed-forward selects its experts by the token table.

width: int

The dense feed-forward width, doubled on a sharing layer when the model asks.

sparsity: float

The gaussian top-k fraction on the feed-forward gate, 0 for none.

kv_shared: bool

The layer reads its keys and values from an earlier layer’s.

provider: int | None

The layer’s own index when a later layer reads its keys and values; such a layer runs unrolled, since what it stashes leaves the stack’s loop.

residual_site: ResidualSite | None

The layer’s place among Kimi K3’s blocks of attention residuals, None without them. It differs at every block boundary, so a scanned run never crosses one.

scan_groups

function source

def scan_groups(
    specs: Sequence[LayerSpec],
    bank_layers: int | None = None,
) -> tuple[tuple[int, int], ...]

The stack as runs of layers, (first, count) each, in order.

Consecutive layers with equal specs form one run, which a scan runs as iterations of one body; a layer with no equal neighbour is a run of one, which stays unrolled. The grouping is read off the specs, never written by hand, so a model’s pattern decides what scans.

bank_layers caps how many layers one run holds, which is how many a parameter bank stacks: a longer run splits into consecutive runs of at most that many layers. A host-resident bank is built and read one bank at a time, so the cap is what bounds the memory either costs.

group_name

function source

def group_name(first: int, count: int) -> str

The module name of a scanned run: layers_3_7 runs layers 3 through 7.

group_layers

function source

def group_layers(name: str) -> range | None

The layers a stack module name runs: layers_3_7 as range(3, 8), layers_3 as range(3, 4), anything else as None. The inverse of group_name, for a reader keyed by single layers that meets the run.

INTERMEDIATES

attribute source

INTERMEDIATES = 'intermediates'

The collection flax’s capture_intermediates fills.

STREAMS

attribute source

STREAMS = ('activation_batch', 'activation_length', None, 'activation_embed')

Manifold-constrained hyper-connections’ [B, S, hc_mult, D] residual streams.

layer_outputs

function source

def layer_outputs(module: nn.Module, method: str) -> bool

The capture_intermediates filter that keeps every layer’s output.

It matches the __call__ of the stack’s layers_N modules, a scanned run’s layers_3_7 and a pipeline stage’s layers_j included, and StackView.unstack lays what those sow out per layer, as the plain loop does; layer_output reads one layer back.

layer_output

function source

def layer_output(
    intermediates: Mapping[str, Mapping[str, Sequence[jax.Array]]],
    index: int,
) -> jax.Array

Layer index’s output out of the collection layer_outputs filled.

Mixture

dataclass source

class Mixture(
    experts: int,
    top_k: int = 2,
    layers: tuple[int, ...] | None = None,
    every: int | None = None,
    score_function: str = 'softmax',
    norm_topk_prob: bool = True,
    scaling: float = 1.0,
    groups: int = 1,
    groups_per_token: int = 1,
    group_score: str = 'top2',
    bias: bool = False,
    scale_inputs: bool = False,
    parallel: bool = False,
    expert_features: int | None = None,
    shared_features: int = 0,
    shared_gate: bool = False,
    implementation: str = 'auto',
    dispatch: str = 'global',
    capacity_factor: float | None = None,
    hash_layers: tuple[int, ...] | None = None,
    latent_features: int | None = None,
    latent_norm: bool = False,
)

The experts some layers route to, and how the router chooses.

experts is what the rest depends on, so they live together: a top_k, a cadence or a balancing bias says nothing about a model with no experts. layers names the sparse layers by index, or every makes every nth layer sparse counting from the end of the first group, the meaning of Qwen3-MoE’s decoder_sparse_step; neither makes every layer sparse, which is Mixtral.

The routing fields pass straight through to Router; that class documents each one.

parallel is Gemma 4’s placement (enable_moe_block). The experts run beside the dense feed-forward on the same residual, and the two are summed after a norm each, under Gemma4TextRouter. That router replaces the routing fields above, which are refused with it.

expert_features is the routed experts’ width, None for the model’s mlp_features; DeepSeek sizes its experts apart from its dense layers. shared_features is the width of the one dense gated MLP every token takes beside the routed experts, 0 for none. DeepseekV3MoE builds its n_shared_experts as a single MLP of that many times its expert width, so the product is the whole record of them. shared_gate multiplies that branch’s output by a learned scalar sigmoid per token, as Qwen3.5 MoE does.

implementation is the grouped matmul the experts run on, one of moe.grouped_matmul’s, the way attention_impl names an attention kernel. It changes which kernel computes the same contraction and nothing about the routing.

dispatch='exchange' is expert parallelism: each device trades its selected tokens with the expert shards that own them, in bounded all-to-all rounds, on an expert mesh axis larger than one that divides the expert count. The default 'global' sorts and gathers where the tokens are. Both share the projection precision and the differentiation contract.

capacity_factor drops slots past each sequence’s per-expert capacity, GShard’s and MaxText’s token dropping (moe.capacity_positions); None, the default, keeps every selected slot. Both dispatches drop the same slots on any placement, and the exchange then runs one round.

hash_layers names the sparse layers that route by DeepSeek V4’s fixed token table instead of the scores (DeepseekV4HashRouter). Their router holds tid2eid over the vocabulary in place of the balancing bias, and the block hands it the token ids.

latent_features is Kimi K3’s latent MoE: the routed experts run at that width between a down and an up projection, with the model’s RMSNorm on their weighted sum when latent_norm is set (SparseMLP). None runs them at the model width.

GatedMLP

Flax module source

class GatedMLP(
    activation: GatedActivation = 'swiglu',
    activation_sparsity: float = 0.0,
    swiglu_limit: float | None = None,
    init_std: float | None = None,
    output_init_std: float | None = None,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
)

down_proj(act(gate_proj(x)) * up_proj(x)): swiglu is silu, geglu is the tanh approximation of gelu (HF’s gelu_pytorch_tanh) and geglu_exact the erf form (HF’s gelu, which Gemma’s released config names). A Situ in place of the name is Kimi K3’s SiTU, which transforms both halves (dew.nn.moe.gated_product).

Bias-free, like the gated MLP of every open decoder this loads.

activation_sparsity is Gemma 3n’s gaussian top-k on the gate before its nonlinearity (dew.nn.gemma3n.gaussian_topk); 0 leaves the gate alone. swiglu_limit is the clamp GLM-5.3-Flash and DeepSeek V4 apply before the activation (Glm5NextTextMLP.forward, modeling_glm5_next.py:98-104): the gate capped at the limit from above and the up projection on both sides. None is the plain gated MLP.

GatedMLP.setup

def setup()

BlockWiring

dataclass source

class BlockWiring(
    pre_norms: bool = True,
    output_norms: bool = False,
    layer_scalar: Literal['frozen', 'trainable'] | None = None,
)

How a block norms its two residuals, and whether it scales its output.

pre_norms norms each sublayer’s input and output_norms its output. The input pair alone is the plain pre-norm block, both pairs Gemma’s sandwich block, and the output pair alone OLMo 3’s post-norm block (modeling_olmo3.py:249-266), where each sublayer reads the residual stream as it is and its output is normed before it is added. The output pair norms the sublayer outputs and not their inputs, so the input norms keep their names and their places, and a checkpoint without the output pair loads into the same tree minus two leaves per layer. layer_scalar selects the reference’s frozen or trainable output scalar. One wiring serves every layer, so it stays off the per-layer specs the scan groups by.

QKV_RESIDUALS

attribute source

QKV_RESIDUALS = ('q_proj', 'k_proj', 'v_proj', 'kv_proj')

ATTENTION_RESIDUALS

attribute source

ATTENTION_RESIDUALS = (*QKV_RESIDUALS, 'o_proj')

MLP_RESIDUALS

attribute source

MLP_RESIDUALS = ('gate_proj', 'up_proj', 'down_proj')

RESIDUALS

attribute source

RESIDUALS = (*ATTENTION_RESIDUALS, 'context', *MLP_RESIDUALS)

The values a block names as it runs, each after the projection that produced it: kv_proj is latent attention’s fused kv_b_proj, context the attention kernel’s output before o_proj, and the MLP names cover the dense MLP and the routed experts alike. A remat policy picks from these; a name off the list is a typo that would otherwise recompute silently.

RematPolicy

dataclass source

class RematPolicy(save: tuple[str, ...] = (), offload: tuple[str, ...] = ())

What a recomputed block keeps for its backward pass, and where.

A block under remat saves its inputs and recomputes its forward when the backward pass asks. save names the residuals (RESIDUALS) it keeps in device memory instead, offload the ones it moves to pinned host memory after the forward pass and fetches back for the backward; everything else is recomputed. Both empty is MaxText’s full, which recomputes the whole block. REMAT_POLICIES holds MaxText’s named recipes under this decoder’s names: MaxText’s query_proj/key_proj/value_proj/ out_proj are q_proj/k_proj/v_proj/o_proj and its mlpwi_0/ mlpwi_1/mlpwo are gate_proj/up_proj/down_proj. Its fused qkv_proj/mlpwi name projections this decoder does not fuse, and its quantization names AQT intermediates Qwix does not produce. A config gives a policy by name or as a record of the two lists.

RematPolicy.checkpoint_policy

def checkpoint_policy()

The policy nn.remat runs the block under; None recomputes everything.

REMAT_POLICIES

attribute source

REMAT_POLICIES: Mapping[str, RematPolicy]

MaxText’s remat recipes (nnx_decoders.py, get_remat_policy), fastest and largest first. full keeps nothing but the block’s inputs.

remat_policy

function source

def remat_policy(
    value: RematPolicy | str | Mapping[str, Sequence[str]] | None,
) -> RematPolicy | None

value as the policy it names: a RematPolicy, a name in REMAT_POLICIES, a record of save/offload names, or None for no recomputation at all. A config’s record arrives here untyped, so a value of another kind is refused rather than passed on.

DecoderBlock

Flax module source

class DecoderBlock(
    norm_eps: float = 1e-05,
    scale_offset: bool = False,
    scale_after_cast: bool = False,
    per_layer_input_dim: int = 0,
    gate_activation: GatedActivation = 'swiglu',
    parallel: Callable[..., nn.Module] | None = None,
    altup: AltUp | None = None,
    laurel_rank: int | None = None,
    hyper_connections: HyperConnections | None = None,
    hash_routed: bool = False,
    residual_multiplier: float = 1.0,
    residual_site: ResidualSite | None = None,
    routed: bool = False,
    dropout_rate: float = 0.0,
    remat: RematPolicy | None = None,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
)

Pre-norm decoder block: token mixer, then feed-forward, both residual.

mixer and feedforward are factories taking only a name. What mixer builds lands in the tree as self_attn and has to accept (x, decode=…, positions=…, segment_ids=…), the last two None outside a packed batch. What feedforward builds lands there as mlp and takes the normalized states alone, which is the one call GatedMLP and moe.SparseMLP share; a hash_routed block hands it the token ids too, which the metadata carries down the stack for DeepSeek V4’s hash router. A feedforward of None is a block of the mixer alone, norm, mixer, residual, which is Mamba-2’s (Mamba2Block, modeling_mamba2.py:608-632): no post_attention_layernorm, no mlp, no output norm for either.

wiring places the block’s norms: the input pair alone is the plain pre-norm block, both pairs Gemma’s sandwich block, and the output pair alone OLMo 3’s post-norm block (modeling_olmo3.py:249-266), where each sublayer reads the residual stream as it is and its output is normed before it is added. The output pair norms the sublayer outputs and not their inputs, so the input norms keep their names and their places, and a checkpoint without the output pair loads into the same tree minus two leaves per layer.

kv_store threads one dict down the layer stack so a KV-sharing mixer reads its provider’s keys and values; a mixer without a kv_store keyword fails loudly when a run shares. per_layer_input is the layer’s slice of LayerInputs: its input signal for the per-layer residual, and on a routed block the replayed experts its router uses (dew.nn.moe.Routes), None when the model reads neither.

altup makes the block take and return Gemma 3n’s stack of residual copies, [num_inputs, B, S, D]: it predicts the copies, runs on the active prediction, corrects every copy by what it computed, and adds the per-layer residual to the copies past the first and not to its own output (modeling_gemma3n.py, Gemma3nTextDecoderLayer.forward). laurel_rank adds the LAuReL block over the attention’s normed input, averaged with the attention residual over sqrt(2).

hyper_connections makes the block take and return the mHC stack of residual streams, [B, S, hc_mult, D]: each sublayer reads the collapse its site’s mapping chooses and writes back into every stream over the Sinkhorn-mixed residual (dew.nn.hyper_connections), the plain pre-norm block otherwise (modeling_glm5_next.py:1293-1327).

residual_site makes the block take and return Kimi K3’s depth state, [B, S, blocks + 1, D]: the finished blocks and the partial sum (dew.nn.attention_residuals). Each sublayer reads the softmax mixture of the finished blocks and the partial its site holds, as a plain pre-norm block reads the residual, and adds its output to the partial.

gate_activation: GatedActivation

The per-layer residual’s gated product, which Gemma 3n/4 share with the feed-forward’s own (modeling_gemma4.py, Gemma4TextDecoderLayer).

parallel: Callable[..., nn.Module] | None

A branch summed with the feed-forward’s output before its output norm, called with the residual and that output (Gemma 4’s routed experts).

DecoderBlock.setup

def setup()

MTPBlock

Flax module source

class MTPBlock(
    hyper_connections: HyperConnections | None = None,
    norm_eps: float = 1e-05,
    scale_offset: bool = False,
    scale_after_cast: bool = False,
    dropout_rate: float = 0.0,
    remat: RematPolicy | None = None,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
)

One multi-token-prediction depth: the next depth’s hidden states.

The depth norms the token embeddings with enorm and the previous hidden states with hnorm, projects the pair concatenated in that order back to the model width, and runs one decoder block over it. That composition is what the released MTP weights were trained for, which the engines state (vLLM deepseek_mtp.py, glm4_moe_mtp.py and qwen3_5_mtp.py). Training shifts complete sequences through this block; prediction steps may use an independently allocated KV cache.

MTPBlock.setup

def setup()

MTPBlock.states

def states(
    hidden,
    embeds,
    train: bool = False,
    positions=None,
    segment_ids=None,
    attention_metadata=None,
    decode: bool = False,
    prediction_phase: PredictionPhase = 'ordinary',
)

The normalized head input and the state a subsequent prediction reads.

Block

attribute source

Block = Callable[[int, str], DecoderBlock]

Layer index’s block under a module name: what the stack and its stages build their layers from, so one factory describes every view of them.

WRITTEN

attribute source

WRITTEN = ('cache', 'router', 'qk', INDEXER_COLLECTION)

The collections a decoder block writes: its decode cache, and what its router, its attention and its sparse indexer sow. A run whose parameters are fetched is applied in a scope of its own, so these are the names whose values the loop has to carry back out to the scope that asked for them.

run_stack

function source

def run_stack(
    layers: Sequence[DecoderBlock],
    block: Block,
    specs: Sequence[LayerSpec],
    groups: Sequence[tuple[int, int]],
    x,
    *,
    train: bool,
    decode: bool,
    positions,
    segment_ids,
    kv_store,
    per_layer_input,
    attention_metadata=None,
    banked: bool = False,
)

Run the layers over x, one run at a time as groups says.

A run of one layer is layers[first], called as the plain loop calls it. A longer run is one block, named for its range, under flax’s scan. Its variables carry a leading layer axis that the view outside stacks and unstacks, and each iteration reads its own slice of the per-layer inputs.

A run’s layers all share or all own their keys and values. Sharing layers read the store their providers filled before the run, a constant the loop closes over. Owning layers would write into it from inside the loop, where a Python dict cannot follow, so they get no store. Nothing reads what they would have written, because a provider is always a run of one.

banked says the store already holds each run’s parameters as one array (dew.inference.banks) rather than as the layers the view stacked. Those runs, and any run the layout left in host memory, are read one layer at a time in _prefetched_run. At most two layers of one run are in device memory then, whatever the depth, and nothing is fetched that nothing computes with. Staging crosses the runs’ boundaries, so a stack of single layers, of unequal runs, or of both is pipelined the same way. No run’s fetch can be hoisted above the layer before it, because it is issued inside that layer’s scan iteration or ordered after it by the carry it lands in. A resident bank and a host-resident one are read by the same loop and give the same values.

Training uses the native Linen scan instead. map_variables stages one row under remat, so the backward pass refetches the original pinned bank rather than retaining a device copy of every layer. That leaves training no duplicate weight bank and no prefetch carry; inference keeps its prefetch.

PipelineStage

Flax module source

class PipelineStage()

The layers one pipeline stage holds, over one microbatch.

The pipeline vmaps this module over the stage axis, so every stage runs it over its own slice of the stacked layer weights and its variables carry a leading stage axis that the view outside stacks and unstacks. Layer j of a stage is layers_j here and layers_{stage * count + j} in the stored tree; specs and groups are stage 0’s, which every stage repeats.

PipelineStage.setup

def setup()

StackView

dataclass source

class StackView(
    groups: tuple[tuple[int, int], ...],
    stages: int = 1,
    microbatches: int = 1,
    broadcast: tuple[str, ...] = (),
    banked: tuple[str, ...] = (),
)

The layer stack’s variables as its loops read them.

Outside, every collection holds one subtree per layer, layers_N, which is the tree a checkpoint stores and a Hugging Face loader fills. Inside a scanned run the same leaves are stacked along a leading layer axis under the run’s name (layers_3_7), and inside a pipeline every stage’s copy of a position is stacked along a leading stage axis under stages. stack builds the inside from the outside and unstack the outside from the inside, so what a run reads, sows and caches lands leaf for leaf where the plain loop puts it.

banked names the collections a store already holds the inside way, one array per run with the layer axis in it. Those the view leaves alone in both directions: the bank a run scans is the one array the store holds, with no copy of it under either name and no per-layer mirror beside it. The stored identity is still layers_N: bank_names says which bank a run’s layers are in, and unstack on a banked store outside a scope is what a save or an export reads, one layer’s slice of one bank at a time.

groups are the runs of one stage (of the whole stack without a pipeline). A collection that entered the pipeline’s loop keeps [stage, ...] leaves; one the loop created (what the routers sow) keeps [iteration, stage, microbatch, ...] leaves, of which the real iterations of each stage are its microbatches in order.

StackView.bank_names

def bank_names() -> list[str]

Every run’s stored name, in order: what a banked store’s keys are.

StackView.stack

def stack(variables: Mapping[str, Mapping]) -> dict

StackView.unstack

def unstack(variables: Mapping[str, Mapping]) -> dict

DecoderBank

dataclass source

class DecoderBank(namespace: tuple[str, ...], view: StackView, scanned: bool = True)

One decoder’s stored layer namespace below every variables collection.

A container prefixes the namespace and leaves the view untouched: layer groups, module names and RNG streams remain the decoder’s. Shared scopes declare one site even when several methods read the same parameters.

scanned: bool

Whether the stack runs its banks under scan, which is what sequences a host-resident bank’s fetches one row at a time. A plain loop declares the same banks, but nothing orders their fetches, so the compiler hoists every layer’s copy to the front and the whole stack lands on the device at once; a host layout refuses it (dew.training.execution.resident).

CausalTransformer

Flax module source

class CausalTransformer(
    emb_features: int = 512,
    num_layers: int = 8,
    num_heads: int = 8,
    num_kv_heads: int | None = None,
    head_dim: int | None = None,
    mlp: GatedActivation = 'swiglu',
    mlp_features: int | tuple[int, ...] | None = None,
    max_seq_len: int = 2048,
    rope_theta: float = 10000.0,
    rope_scaling: RopeScaling | None = None,
    partial_rotary_factor: float | None = None,
    partial_rotary_type: str = 'proportional',
    layer_types: tuple[str, ...] | None = None,
    kinds: Mapping[str, LayerKind] | None = None,
    norm_eps: float = 1e-05,
    scale_offset: bool = False,
    scale_after_cast: bool = False,
    sandwich_norms: bool = False,
    pre_norms: bool = True,
    qk_norm: bool = True,
    qk_norm_scope: str = 'head',
    v_norm: bool = False,
    attention_k_eq_v: bool = False,
    layer_scalar: Literal['frozen', 'trainable'] | None = None,
    attention_bias: bool = False,
    o_proj_bias: bool | None = None,
    attention_scale: float | None = None,
    attention_sinks: bool = False,
    yarn: YarnScaling | None = None,
    attn_logit_softcap: float | None = None,
    output_gate: bool = False,
    embedding_scale: bool = False,
    embedding_multiplier: float = 1.0,
    residual_multiplier: float = 1.0,
    logits_scaling: float = 1.0,
    initializer_range: float | None = None,
    depth_scaled_init: bool = False,
    final_logit_softcap: float | None = None,
    tie_embeddings: bool = True,
    embedding_zero_ids: tuple[int, ...] = (),
    dropout_rate: float = 0.0,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    force_fp32_for_softmax: bool = True,
    attention_impl: str = 'auto',
    kv_cache: KVCache = KVCache(),
    mixture: Mixture | None = None,
    use_double_wide_mlp: bool = False,
    causal: bool = True,
    per_layer_input_dim: int | None = None,
    per_layer_input_vocab: int | None = None,
    num_kv_shared_layers: int = 0,
    kv_shared_layers: tuple[int, ...] | None = None,
    mixer: MixerBase | None = None,
    num_nextn_predict_layers: int = 0,
    index_share_for_mtp_iteration: bool = False,
    mtp_layer_type: str | None = None,
    mtp_hyper_connections: HyperConnections | None = None,
    altup: AltUp | None = None,
    laurel_rank: int | None = None,
    hyper_connections: HyperConnections | None = None,
    attention_residuals: AttentionResiduals | None = None,
    swiglu_limit: float | None = None,
    activation_sparsity_pattern: tuple[float, ...] | None = None,
    mask_token_id: int | None = None,
    scan_layers: bool = False,
    bank_layers: int | None = None,
    remat: RematPolicy | None = None,
)

Decoder-only transformer over token ids: [B, S] int32 -> [B, S, vocab] fp32.

The defaults train a model from scratch: multi-head attention, swiglu, tied embeddings, no softcap. Every field an open decoder varies is a field here, so loading Qwen3 or Gemma3 is a field mapping and not a subclass. The field comments below name which family sets each one.

layer_types is the pattern, one kind per layer, and kinds says what a kind does: its window, and its own rope base or head dim. Deriving the pattern from a checkpoint’s config belongs to that translation, not here; this takes the tuple.

mixture turns the feed-forward of some layers into moe.SparseMLP, routing each token to a few of its experts. None is a dense model. The LM objective’s balance_rate is what moves a mixture’s balancing bias.

causal=False turns every layer into full attention with no cache, which is the encoder a masked diffusion language model denoises with. The parameter tree is the same either way.

per_layer_input_dim turns on Gemma 3n/4 per-layer input embeddings: an extra table, read per layer and added to that layer’s input through its own gate. None leaves the tree unchanged.

num_kv_shared_layers makes that many trailing layers reuse an earlier layer’s keys and values instead of projecting their own, which is Gemma 3n/4 cross-layer KV sharing. use_double_wide_mlp, which widens the sharing layers’ MLP, needs it. kv_shared_layers names the sharing layers one by one instead, for a pattern that is not a trailing run: GLM’s IndexShare puts one after every indexer layer but the first three. Either spelling resolves to the same plan. A sharing layer reads what the last earlier non-sharing layer of its own kind stashed, and what a layer stashes is its mixer’s own: keys and values for attention, the indexer’s selection for MLA.

altup carries Gemma 3n’s copies of the residual stream (dew.nn.gemma3n). The embeddings enter the layers as a stack, each block predicts the copies, runs on the active one and corrects them all, and the copies come back through their own projections to a mean the final norm reads. laurel_rank adds the LAuReL block to every layer and activation_sparsity_pattern the gaussian top-k on each layer’s gate. A tuple mlp_features gives each layer its own width; a width of 0 is a layer without a feed-forward, which is Mamba-2’s block of the mixer alone.

attention_residuals replaces the running residual with Kimi K3’s softmax over finished blocks of layers (dew.nn.attention_residuals): the embeddings enter as the first partial sum, every sublayer reads a mixture over depth, and a model-level site mixes the blocks once more before the final norm. mlp may be a Situ in place of an activation name: Kimi K3’s SiTU, which every feed-forward then shares (dew.nn.moe.Situ).

partial_rotary_factor rotates that fraction of an unwindowed kind’s head dims and passes the rest through; a windowed kind rotates whole. partial_rotary_type names which published convention the fraction follows, because the two rotate different angles: dew.nn.rope.rotary_freqs documents both. Interleaved mRoPE (Qwen3.5’s mrope_section) is this same rotation for text. With one position per token the three grids’ angles are equal and the interleave reads the same value from each, so text-only input reduces to this partial rope exactly; image-grid positions are not modelled.

mixer names the per-layer token mixer as a value from the mixers registry, one frozen dataclass per kind under the reference’s own field names. None is grouped-query causal attention. A non-standard kind reads its own record and ignores the GQA projection geometry the context still carries; those fields stay validated, so a translation fills them with consistent values.

num_nextn_predict_layers stacks that many multi-token-prediction depths after the final norm, each an MTPBlock. Depth d pairs the previous depth’s state at position p with the embedding of the token at p + d and scores what follows p + d (arXiv 2412.19437, section 2.2), so each depth is one position shorter than the last.

scan_layers runs each run of consecutive like layers as iterations of one body under flax’s scan, and the layers between such runs unrolled. Layers are alike when they share a parameter shape and a computation, which is read off the resolved layers and never configured. A body compiles once however many layers it runs, so compile time stops growing with depth. The variables tree is the unscanned one leaf for leaf: init always runs the plain loop, and the scan reads and writes a stacked view of the same leaves (StackView). A stage axis above one on the mesh runs the stack as a pipeline over that axis (_pipeline), whether or not the layers scan.

embedding_multiplier: float

Token embeddings times this before the first layer: muP’s m_emb in lm-engine, GraniteMoeHybrid’s embedding_multiplier.

residual_multiplier: float

Every sublayer output times this before it joins the residual stream: lm-engine’s m_residual, GraniteMoeHybrid’s residual_multiplier.

logits_scaling: float

The logits divided by this: lm-engine’s m_width (lm_logits * (1 / m_width)), GraniteMoeHybrid’s logits_scaling. The final states carry the division, in fp32, so every head that contracts them with head_weight scores the same logits __call__ returns.

initializer_range: float | None

lm-engine’s initialisation: the embedding table (and an untied head) drawn from N(0, initializer_range^2), every hidden matrix (attention and Mamba-2 projections, conv taps, router, experts, dense MLPs) from N(0, (initializer_range / sqrt(logits_scaling))^2), biases zero and norms one. With logits_scaling as m_width that is lm-engine’s init_method="mup", and with it 1 its "normal" (init_utils.py at 45b6b57b). None keeps each module’s own initializer.

depth_scaled_init: bool

With initializer_range, the projections back into the residual stream (o_proj, out_proj, down_proj) divide their std by sqrt(2 * num_layers), lm-engine’s use_depth_scaled_init.

embedding_zero_ids: tuple[int, ...]

Placeholder ids looked up as token zero, without changing labels (modeling_kimi_k25.py:686-690, the text-only wrapper path).

kv_cache: KVCache

How attention layers store the decode cache: dense or paged, full or quantized (dew.nn.kv_cache). The parameters do not depend on it.

mtp_layer_type: str | None

An explicit prediction-layer kind, which need not occur in the trunk.

mtp_hyper_connections: HyperConnections | None

None gives prediction depths plain residuals and normalized trunk inputs. A stream depth explicitly opts in, independently of the trunk’s residuals.

bank_layers: int | None

The most layers one scanned run holds, which is how many its parameter bank stacks. A longer run of like layers splits into consecutive runs of at most this many, each its own bank under its own name; None puts a whole run in one bank. Only scan_layers reads it, and the split is what bounds the memory that building a host-resident bank and reading it back cost, so a deep stack offloaded to the host sets it.

remat: RematPolicy | None

Recompute each block in the backward pass, keeping its inputs, any K/V supplied to later layers and the residuals the policy names. A name from REMAT_POLICIES or a record of save/offload residual names arrives from a config; None recomputes nothing. Init and cached decode follow the direct block path; stored parameters have the same layout.

init_stds: tuple[float | None, float | None]

The normal std of the hidden matrices and of the projections back into the residual stream, or (None, None) for the modules’ own initializers (initializer_range).

mlp_widths: tuple[int, ...]

Each layer’s dense feed-forward width.

hidden_features: int

The model’s one feed-forward width, which the routed experts fall back to and the prediction depths take; a model whose layers differ has none.

hash_layers: set

The sparse layers routing by the mixture’s token table.

sparse_layers: tuple[int, ...]

The layers whose feed-forward routes to experts.

sharing_layers: tuple[int, ...]

The layers that read another layer’s stash, in order: the trailing num_kv_shared_layers or the ones kv_shared_layers names.

kv_sharing: dict

Sharing layer index to the provider it reads, both of one layer type.

A sharing layer owns no K/V (no indexer, for MLA) and reads the last earlier non-sharing layer of its own type (modeling_gemma4.py, Gemma4TextAttention; modeling_glm_moe_dsa.py:739-748 carries the last full layer’s top-k forward). Empty unless sharing is on.

bank_sites: tuple[DecoderBank, ...]

This decoder’s stored stack: scanned runs, or one layer per bank.

groups already says which is which. A scanned stack declares its runs and a plain loop declares singletons, and run_stack fetches a run of one the same way it fetches a longer one; scanned says whether those fetches are sequenced, which a host layout requires.

CausalTransformer.kind_of

def kind_of(layer_type: str) -> ResolvedKind

What the layers of layer_type do, the model’s defaults included.

CausalTransformer.mixer_context

def mixer_context(kind: ResolvedKind, layer_type: str, kv_shared: bool) -> MixerContext

One layer’s mixer geometry: the kind’s resolved values as a context.

head_dim, rope_theta, window and the two rotary ramps already carry the layer kind’s overrides; a windowed kind rotates every dimension, so the partial rotary belongs to the kinds that attend the whole sequence, where Gemma 4 puts it. A kind builds its DecoderBlock factory from this and its own record; setup chooses the mixer there and nowhere else.

CausalTransformer.layer_kinds

def layer_kinds(types: Sequence[str]) -> dict[str, ResolvedKind]

Resolve every kind the pattern names, and the prediction depths’.

Raises when the pattern is not one kind per layer, or when kinds describes a kind no layer has.

CausalTransformer.refuse_unbuildable_mup

def refuse_unbuildable_mup(kinds: Mapping[str, ResolvedKind])

Raise for lm-engine’s multipliers or init on a model that cannot carry them.

CausalTransformer.refuse_unbuildable_fields

def refuse_unbuildable_fields(kinds: Mapping[str, ResolvedKind])

Raise for a field, or a pair of fields, this model cannot build.

Each check names the field the caller set and what a model without it looks like, so a translated config says which entry to fix.

CausalTransformer.feedforward_factories

def feedforward_factories()

Build the three feed-forward factories a layer can be given.

Returns the dense gated MLP, the routed experts, and Gemma 4’s parallel branch, each a partial the block calls with a name. A model with no mixture has only the first.

CausalTransformer.setup

def setup()

Build the embeddings, the layers, the prediction depths and the head.

specs is one LayerSpec per layer, block builds a layer from its index and name, layers holds them all, and groups says which consecutive runs of them scan together. mtp holds the prediction depths and norm the final norm, with lm_head beside it when the embeddings are not tied.

CausalTransformer.states_and_logits

def states_and_logits(tokens, **kwargs={})

The prediction input states and logits from one forward.

A speculative decoder verifies with both: the logits give the target distribution and the states seed the next block’s prediction depths: normalized states normally, uncollapsed residual streams for V4.

CausalTransformer.states_and_logits_at

def states_and_logits_at(tokens, slots, **kwargs={})

The prediction input states, and the logits of one slot per row.

A prefill scores the position the first draw reads, slots, and no other: the head over every prompt position is the largest array the forward allocates, [rows, width, vocab], and a decoder keeps one row of it. Gathering the state before the head leaves the head [rows, features] of work.

CausalTransformer.mtp_hidden_states

def mtp_hidden_states(
    hidden,
    tokens,
    train: bool = False,
    positions=None,
    segment_ids=None,
    input_embeddings=None,
    embedding_positions=None,
    attention_mask=None,
    image_groups=None,
    rotary_positions=None,
)

One final-normed state array per shifted prediction depth.

A depth combines the preceding hidden state and the next token’s embedding, including any media replacement. Its positions are the next token’s positions. Both ends must be valid and belong to the same packed document; padded intermediates cannot become keys.

CausalTransformer.mtp_logits

def mtp_logits(
    hidden,
    tokens,
    train: bool = False,
    positions=None,
    segment_ids=None,
    input_embeddings=None,
    embedding_positions=None,
    attention_mask=None,
    image_groups=None,
    rotary_positions=None,
)

The shared language head over each prediction depth’s hidden states.

CausalTransformer.mtp_step

def mtp_step(
    hidden,
    tokens,
    *,
    depth: int = 0,
    positions=None,
    input_embeddings=None,
    attention_mask=None,
    rotary_positions=None,
    decode: bool = False,
    prediction_phase: PredictionPhase = 'ordinary',
)

One unshifted prediction step, optionally appending its own KV cache.

Call init_mtp_cache before cached steps. The hidden input is the target model’s preceding state; tokens or input_embeddings supply the candidate next token, as in vLLM’s Qwen3_5MultiTokenPredictor. Returns the step’s logits and its own hidden state, which the next step of a chained draft consumes in place of the target’s. With index_share_for_mtp_iteration, cached extend publishes index selections and draft reuses them; ordinary always recomputes. Uncached training never carries a selection between queries.

CausalTransformer.token_embeddings

def token_embeddings(tokens)

The embeddings a prediction depth pairs with tokens.

mtp_hidden_states reads the unscaled table, so this is that lookup and nothing else: a drawn token is text, and media never reaches it.

CausalTransformer.scaled_embeddings

def scaled_embeddings(x)

Token embeddings x as the first layer reads them: times sqrt(emb_features) when embedding_scale is set, as Gemma scales them, then times embedding_multiplier, as lm-engine and GraniteMoeHybrid scale them. The decoder’s forward, the multimodal wrapper and the Qwen-Image conditioner all scale here.

Gemma casts embed_scale to the embedding weight dtype (modeling_gemma3.py:117). The token lookup holds that table in fp32 and returns the compute dtype, so the factor keeps its fp32 value and only the product rounds with the activations. A factor rounded to bf16 would be 34.0 at hidden 1152, where sqrt(1152) is 33.94112549695428. lm-engine multiplies the looked-up states (hidden_states * m_emb, mixins/dense/base.py at 45b6b57b) in fp32 opmath, which scaled keeps.

CausalTransformer.init_mtp_cache

def init_mtp_cache(batch_size: int)

Allocate prediction-layer caches independently of the trunk cache.

CausalTransformer.hidden_states

def hidden_states(tokens, **kwargs={})

The final normalized states, excluding the vocabulary projection.

CausalTransformer.hidden_and_mtp_inputs

def hidden_and_mtp_inputs(
    tokens,
    train: bool = False,
    decode: bool = False,
    positions=None,
    segment_ids=None,
    input_embeddings=None,
    embedding_positions=None,
    attention_mask=None,
    image_groups=None,
    rotary_positions=None,
    attention_pairwise_mask=None,
    attention_key_positions=None,
    routed_experts=None,
    routed=None,
)

The final normalized states and the prediction depth’s input.

V4’s depth reads the raw residual streams before the collapse head and final norm (official inference/model.py MTPBlock.forward at b5968e9); ordinary depths read the final normalized states.

A packed batch passes per-document positions and segment_ids through to the layers, where RoPE and the mask read them.

A caller that fuses another encoder’s outputs passes them as input_embeddings with their token positions in embedding_positions: both or neither, and the values replace the scaled token embeddings before the layers read them.

attention_pairwise_mask is an explicit boolean [B, queries, keys] visibility mask for ordinary attention mixers. Optional attention_key_positions supplies logical [B, keys] coordinates; local layers apply their configured window to those coordinates. These are call-local cached-read metadata, not sliceable token fields.

routed_experts replays a rollout engine’s routing: [B, S, layers, top_k] expert ids indexed by decoder layer, dense layers included (the layout vLLM’s routed_experts and SGLang’s meta_info.routed_experts return), with routed, [B, S], marking the tokens the record covers (None for all). Every sparse layer’s router selects its slice (dew.nn.moe.Routes).

CausalTransformer.stack

def stack(
    x,
    *,
    train: bool,
    decode: bool,
    positions,
    segment_ids,
    per_layer_input,
    attention_metadata=None,
)

The layer stack over x: the plain loop, the scanned runs, or the pipeline over the mesh’s stages.

The plain loop is what init always runs, so the variables tree is the one it creates whatever the model is asked to do afterwards. With scan_layers, or a stage axis above one on the mesh in context, the stack runs under StackView: the same leaves, stacked along the loops’ axes while the loops run and unstacked on the way out. The loops carry the residual stream in one dtype, so it enters them in the dtype it settles in (residual_dtype).

A store built bank by bank (dew.inference.banks) already holds the parameters the way the runs read them, one array per run, so the view leaves that collection alone in both directions and the run scans the one array the store holds.

CausalTransformer.banked_collections

def banked_collections() -> tuple[str, ...]

The collections whose layer subtrees the store holds as banks.

A store built per layer holds layers_0; one built bank by bank holds the name of each run of more than one layer, layers_0_15. A run of one is the same tree either way, so it says nothing about which of the two a store is and neither form has to be converted for it. A collection may hold a run’s bank beside that run’s layers when the two hold different leaves: a host layout keeps the leaves every row froze as the bank and the rest per layer, and StackView.stack stacks the rows into the bank. Such a collection is not banked, since its rows still stack. A bank and a row holding the same leaf, or some of the banks and not the others with no rows, is refused: which of the two the run would read is not a question this answers by guessing.

CausalTransformer.residual_dtype

def residual_dtype(
    x,
    *,
    train: bool,
    decode: bool,
    positions,
    segment_ids,
    per_layer_input,
    attention_metadata=None,
) -> jnp.dtype

The dtype the residual stream settles in: x’s promoted with what the first layer returns for it.

A scan carries one dtype from its first iteration to its last, and every pipeline stage takes the dtype the stage before it returns. Under a bf16 policy the dense feed-forward returns fp32, so the plain loop’s stream is fp32 from the first layer on and the loops take it in fp32 from the start; at fp32 nothing changes. The layer runs abstractly, in a scope of its own, so it writes no cache and sows nothing here, and its RNG streams take placeholder keys, since a shape needs a key of each name and no value.

Only the shapes and dtypes matter, so a banked store answers this from the first row of its first bank, without fetching it: the layer the plain loop would read is not in that store under a name of its own.

CausalTransformer.first_layer_shapes

def first_layer_shapes() -> dict

The stack’s first layer’s variables as shape/dtype structs.

A store holding one subtree per layer answers with layer 0’s; one holding banks answers with the first row of the first bank’s, the layer axis dropped; one holding both (banked_collections) answers with layer 0’s leaves completed by the bank’s. Nothing is read, fetched or sliced: this is what an abstract run of one layer needs and no more, and a bank’s first row is a shape, not a copy.

CausalTransformer.stage_layers

def stage_layers(stages: int) -> int

Layers per stage when the stack splits into stages, or why it cannot.

Every stage runs one program over its own layers, so the stages have to be the same length and layer j of every stage the same kind of layer as layer j of the first: same kind, same feed-forward, same width, the same say in keys and values. A pattern that does not repeat every num_layers / stages layers is refused with the first pair of layers that differ and what differs between them.

CausalTransformer.per_layer_inputs

def per_layer_inputs(tokens, inputs_embeds)

Every layer’s input signal [B, S, L, P] (Gemma 3n/4 PLE).

The token-identity component is the packed table’s row for each token, scaled like the main embedding; the context component is the input embeddings projected down, scaled and normed. Their sum over sqrt(2) is what each layer’s gate multiplies in (modeling_gemma4.py, get_per_layer_inputs/project_per_layer_inputs).

CausalTransformer.head_weight

def head_weight(params)

The [D, vocab] head matrix in its stored dtype, as the forward contracts it (_logits: fp32 accumulation over the stored operand).

params is the parameter tree the forward runs under, so this is a plain read: a tied head is the embedding table transposed, an untied one is lm_head’s kernel, which is [D, vocab] already. The Gemma embedding scale multiplies the input embeddings only, so it has no place here. A vocabulary-sized fp32 copy is what a loss over this head would hold for its backward, so none is made here.

CausalTransformer.head_table

def head_table(params)

The head matrix as the tree stores it, and whether its rows are the vocabulary: the [vocab, D] embedding table and True for a tied head, lm_head’s [D, vocab] kernel and False otherwise.

No operation sits between the parameter and this value, so a loss that keeps the head for its backward (chunked_cross_entropy with vocab_major) keeps the parameter itself rather than a transposed copy of it, which for a tied head is a vocabulary-sized array. head_weight is this value as [D, vocab].

CausalTransformer.init_cache

def init_cache(batch_size: int)

Allocate a zeroed decode cache for batch_size sequences.

cache = model.apply(params, batch_size, method=CausalTransformer.init_cache, mutable=[‘cache’])[1][‘cache’]

The forward pass this runs is a single dummy token whose keys are never written: allocation happens on the first decode-mode call, the write on the ones after it.

gather_cache_rows

function source

def gather_cache_rows(cache, rows)

A decode cache reindexed on its batch axis, one gather per leaf.

Outside the layer stack a cache holds one subtree per layer, and every leaf a decode step writes carries its batch on axis zero: dense keys and values with their cached validity and cursor, a gated delta net’s convolution and recurrent state, latent attention’s compressed cache, cached image groups, and a multimodal model’s next position. The scanned stack’s layer axis exists only inside run_stack; StackView removes it before the cache crosses apply, so axis zero is the row here whatever the stack did.

rows is any index array: repeats duplicate a row’s whole decode state, a permutation reparents rows, and a shorter or longer array changes the row count. Beam branching and speculative rollback are both this operation. Nothing else in the tree depends on the row order, so the gathered cache decodes exactly as the rows it came from.

A paged cache (dew.nn.kv_cache) keeps its keys in a pool the rows share through their page tables, so gathered rows would write into each other’s pages; it is refused.