nnterp · encyclopedia zaya · nnterp.families.zaya

family zaya

ZAYA1

A pre-norm block that scales and shifts, per channel, both the sublayer's output and the stream at every add; its queries and keys are mixed over tokens by a causal convolution, and its mixture runs one expert per token, chosen beside a skip class by a router whose state carries from block to block.

Scaled residual addsEmbedding multiplierPartial rotaryConvolved queries and keysSliding windowMixture of expertsUnnormalized routingGrouped-query attention

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card.

embed_tokenstoken_embeddings
→
layers× 40
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
i = 0

identity h = (layers[i].input + a.residual_bias) * a.residual_scale + self_attn.attention_output; (h + m.residual_bias) * m.residual_scale + mlp.mlp_output == layer_output

The API

02 · support()

Values, by host

⚠︎ marks a value this checkpoint has only under some load: hover it for the condition.

model ZayaForCausalLMlayoutwhat it isread at
logits Logits [batch seq vocab_size] The logits the model returns, after anything it does past lm_head (a softcap, a scale)

the module's own output
token_embeddings Residual [batch seq hidden_size] The embedding module's output; layers[0].input is what enters the first block

the child module embed_tokens's output
final_norm_output Residual [batch seq hidden_size] The residual stream after the final norm, what lm_head reads

the child module norm's output
lm_head_output Logits [batch seq vocab_size] The head's logits before anything the model does past lm_head (a softcap, a scale)

the child module lm_head's output
next_token_probs NextTokenProbs [batch vocab_size] The next-token distribution at the last position; derived, read-only

the module's own output
input_ids Tokens [batch seq] The token ids the model was called with

the module's own input
attention_mask Tokens [batch seq] The attention mask the model was called with; zeros are padding

the module's own input
input_size scalar [batch, seq] of the current call; read-only

the module's own input
model.layers[i] ZayaDecoderLayerlayoutwhat it isread at
layer_input Residual [batch seq hidden_size] The residual stream entering the block

the module's own input
attention_norm_output Residual [batch seq hidden_size] The normed stream entering the attention

the child module input_layernorm's output
layer_mid Residual [batch seq hidden_size] The residual stream after the attention add

the child module post_attention_layernorm's input
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP

the child module post_attention_layernorm's output
mlp_input Residual [batch seq hidden_size] What the MLP is called with

the child module mlp's input
layer_output Residual [batch seq hidden_size] The residual stream leaving the block, a tensor even when the block returns a tuple

the module's own output
model.layers[i].self_attn ZayaAttentionlayoutwhat it isread at
⚠︎attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the arguments of attention_interface_1, element 1
⚠︎attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the arguments of attention_interface_1, element 2
⚠︎attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the arguments of attention_interface_1, element 3
⚠︎attention_scores Pattern [batch num_heads query key] The attention scores entering the softmax, masked

inside the forward: the input of nn_functional_softmax_0 inside attention_interface_1
attention_premix HeadOutputs [batch seq num_heads head_dim] Every head's output side by side, as the output projection takes it

the child module o_proj's input
attention_output Residual [batch seq hidden_size] What the attention adds to the residual stream: its output plus the merge's hidden_states_bias, times its hidden_states_scale

the module's own output
⚠︎attention_probabilities Pattern [batch num_heads query key] The attention pattern the values are mixed with

inside the forward: the output of nn_functional_dropout_0 inside attention_interface_1
⚠︎attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of attention_interface_1, element 0
model.layers[i].mlp ZayaSparseMoeBlocklayoutwhat it isread at
⚠︎mlp_activation Neurons [batch seq intermediate_size] The activation function's output

the child module act_fn's output
⚠︎mlp_neurons Neurons [batch seq intermediate_size] What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise

the child module down_proj's input
mlp_output Residual [batch seq hidden_size] What the MLP adds to the residual stream: its output plus the merge's hidden_states_bias, times its hidden_states_scale

the module's own output
router_logits RouterLogits [batch seq num_experts] The router's logits, one per expert and a last one for skip, before the scoring

the child module router.router_mlp's output
expert_weights ExpertWeights [batch seq top_k] The weight each selected expert's output is scaled by

the child module experts's inputs, element 2
expert_indices ExpertIndices [batch seq top_k] The expert each routing slot sends the token to

the child module experts's inputs, element 1
expert_outputs ExpertOutputs [batch seq top_k hidden_size] Each routing slot's weighted expert output

inside experts's forward: the output of weighted_out_view_0 inside experts_forward_1
routed_output Residual [batch seq hidden_size] The routed experts' combined output

the child module experts's output

Printout

What print(model) shows on Zyphra/ZAYA1-8B: the native modules, standard names and every standard value with its layout.

ZayaForCausalLM(
  (lm_head): Linear(in_features=2048, out_features=262272, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(262272, 2048, padding_idx=0)
  (layers): ModuleList(
    (0-39): 40 x ZayaDecoderLayer(
      (self_attn): ZayaAttention(
        (o_proj): Linear(in_features=1024, out_features=2048, bias=False)
        (qkv_proj): ZayaCCAProjection(
          (q_proj): Linear(in_features=2048, out_features=1024, bias=False)
          (k_proj): Linear(in_features=2048, out_features=256, bias=False)
          (v_proj_current): Linear(in_features=2048, out_features=128, bias=False)
          (v_proj_delayed): Linear(in_features=2048, out_features=128, bias=False)
          (conv_qk_depthwise): Conv1d(1280, 1280, kernel_size=(2,), stride=(1,), groups=1280)
          (conv_qk_grouped): Conv1d(1280, 1280, kernel_size=(2,), stride=(1,), groups=10)
        )
        (qk_norm): ZayaQKNorm()
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: its output plus the merge's hidden_states_bias, times its hidden_states_scale
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): ZayaSparseMoeBlock(
        (router/gate): ZayaRouter(
          (down_proj): Linear(in_features=2048, out_features=256, bias=True)
          (router_mlp): ZayaRouterMLP(
            (norm): ZayaRMSNorm((256,), eps=1e-05)
            (fc1): Linear(in_features=256, out_features=256, bias=True)
            (fc2): Linear(in_features=256, out_features=256, bias=True)
            (out_proj): Linear(in_features=256, out_features=17, bias=False)
            (act_fn): GELU(approximate='none')
          )
        )
        (experts): ZayaExperts(
          (act_fn): SiLUActivation()
        )
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: its output plus the merge's hidden_states_bias, times its hidden_states_scale
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert and a last one for skip, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
      )
      (input_layernorm): ZayaRMSNorm((2048,), eps=1e-05)
      (post_attention_layernorm): ZayaRMSNorm((2048,), eps=1e-05)
      (post_attention_residual_scale): ZayaResidualScaling()
      (post_mlp_residual_scale): ZayaResidualScaling()
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
  )
  (norm): ZayaRMSNorm((2048,), eps=1e-05)
  (logits) -> Logits [batch seq vocab_size]: The logits the model returns, after anything it does past lm_head (a softcap, a scale)
  (token_embeddings) -> Residual [batch seq hidden_size]: The embedding module's output; layers[0].input is what enters the first block
  (final_norm_output) -> Residual [batch seq hidden_size]: The residual stream after the final norm, what lm_head reads
  (lm_head_output) -> Logits [batch seq vocab_size]: The head's logits before anything the model does past lm_head (a softcap, a scale)
  (next_token_probs) -> NextTokenProbs [batch vocab_size]: The next-token distribution at the last position; derived, read-only
  (input_ids) -> Tokens [batch seq]: The token ids the model was called with
  (attention_mask) -> Tokens [batch seq]: The attention mask the model was called with; zeros are padding
  (input_size): [batch, seq] of the current call; read-only
)

Config

model.num_layers40
model.hidden_size2048
model.num_heads8
model.num_kv_heads2
model.head_dim128
model.qk_head_dim128
model.vocab_size262272
model.intermediate_size2048
architecturesZayaForCausalLM
num_hidden_layers40
hidden_size2048
num_attention_heads8
num_key_value_heads2
head_dim128
vocab_size262272
max_position_embeddings131072
hidden_actsilu
rms_norm_eps1e-05
tie_word_embeddingsTrue
dtypebfloat16
partial_rotary_factor0.5
moe_intermediate_size2048
num_experts16
num_experts_per_tok1
attention_biasFalse

The Notes

03 · what to know before you run

The block, in order

h   = post_attention_residual_scale(self_attn(input_layernorm(x)), x)
out = post_mlp_residual_scale(mlp(post_attention_layernorm(h)), h)

merge(o, r) = (o + hidden_states_bias) * hidden_states_scale
            + (r + residual_bias) * residual_scale

Each merge is a ZayaResidualScaling module on the block with four [hidden_size] parameters. The block returns (hidden_states, prev_router_hidden_states); layer_output is the first element.

The contributions are the merges' scaled terms

attention_output is (self_attn.output[0] + hidden_states_bias) * hidden_states_scale with the first merge's parameters, and mlp_output the same over the mixture's output with the second's: a computed copy, divided back on assignment. The merge also rescales the stream, so the identity carries both merges' residual_scale and residual_bias. It is exact in float32 on the pinned tiny checkpoint with its merges moved off their initial values:

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn = model.layers[1].self_attn.attention_output.save()
    mlp = model.layers[1].mlp.mlp_output.save()
    out = model.layers[1].layer_output.save()

a = model.layers[1].post_attention_residual_scale._module
m = model.layers[1].post_mlp_residual_scale._module
h = (x + a.residual_bias) * a.residual_scale + attn
torch.testing.assert_close((h + m.residual_bias) * m.residual_scale + mlp, out)

A write to a contribution is divided back by hidden_states_scale over the whole tensor, so the positions it does not touch come back changed by rounding; load in float32 for small edits. Zeroing attention_output removes the term and leaves (x + residual_bias) * residual_scale, not x.

The merges shrink the stream

On ZAYA1-8B post_attention_residual_scale.residual_scale averages 0.52 to 2.03 per block and post_mlp_residual_scale.residual_scale 0.046 to 1.19, and the product of all 80 residual_scale vectors has a median of 6.8e-6 over the channels (at most 0.034). A term added early therefore reaches the last stream multiplied by that product, channel by channel; a direct attribution to the logits has to carry it. The hidden_states_scale that sizes each contribution reaches 3.8 on the attention merges and 13.2 on the mixture's. The biases are at most 1.15 in size and average zero.

The stream is float32 whatever the load

The model casts the embedding to float32 after its shift and scale, and each merge adds a float32 stream, so layers[i].input and layer_output are float32 under a bfloat16 load, while token_embeddings, attention_output, mlp_output and logits keep the load's dtype. input_layernorm, post_attention_layernorm and the final norm cast the stream to their weight's dtype before they norm it.

Queries and keys are convolved, normed to √head_dim and the keys scaled

qkv_proj computes the queries and keys with q_proj and k_proj and mixes them over the current and two previous tokens with two causal convolutions (cca_time0 and cca_time1 are 2 on ZAYA1-8B). It then adds a residual: each query head gets the mean of its own projection and its key head's projection, and each key head gets the mean of those query residuals over the query heads it serves. Half of each value head is the current token's (v_proj_current) and half the previous token's (v_proj_delayed). qk_norm then L2-normalizes every query and key head to length √head_dim and multiplies the keys by temp, one per key/value head; rotary turns the first half of each head (partial_rotary_factor 0.5). attention_queries and attention_keys are read after all of it, at the attention interface. temp is 0.8 to 12.9 on ZAYA1-8B; on the pinned tiny checkpoint it is zero, which zeroes every key, so the suite's copy sets it to one. ZAYA1-74B-preview attends over a 4097-token window on its hybrid_sliding blocks, every other block.

One expert per token, chosen beside a skip class

The router (gate, aliased router) projects the block's normed stream down to router_hidden_size, adds the previous block's router state times router_states_scale (not on block 0), and scores num_experts + 1 classes with a small MLP; router_logits is that MLP's output, and its last column is skip. The choice is the softmax plus balancing_biases (a buffer); the weight is the softmax's own entry for the chosen class, not renormalized, so expert_weights is below one. A slot that picks skip runs no expert: its weight is 0 and its index 0.

moe = model.layers[1].mlp
with model.trace(prompt):
    logits = moe.router_logits.save()     # [batch, seq, num_experts + 1]
    w = moe.expert_weights.save()         # [batch, seq, top_k], top_k = 1
    idx = moe.expert_indices.save()

probs = logits.float().softmax(-1)
torch.testing.assert_close(probs.gather(-1, idx).to(w.dtype), w)

On ZAYA1-8B and ZAYA1-74B-preview balancing_biases is −1.0 on skip and above zero on at least one expert of every block, so in this forward the skip class never wins a slot: its softmax mass lowers the chosen expert's weight instead. The routing state carried to the next block means an edit to one block's routing input reaches every later block's router.

Ablating an expert leaves the merge's bias

With one slot per token, zeroing expert e's weight removes the token's whole routed output, and mlp_output there becomes hidden_states_bias * hidden_states_scale, not zero. Count usage on slots with a weight, since a skipped slot reads as expert 0:

moe, e = model.layers[1].mlp, 3
with model.trace(prompt):
    idx = moe.expert_indices.save()
    w = moe.expert_weights.save()
    clean = moe.mlp_output.save()
with model.trace(prompt):
    moe.expert_weights = moe.expert_weights.masked_fill(moe.expert_indices == e, 0)
    ablated = moe.mlp_output.save()

chose = ((idx == e) & (w > 0)).any(-1)                 # [batch, seq]
assert torch.equal((ablated != clean).any(-1), chose)

There is no shared expert, so routed_output is the mixture's own output, before the merge.

The readout and the tokenizer

logits is lm_head.output, with no cap or scale. The tokenizer is a GemmaTokenizer that prepends <bos> (id 2). On ZAYA1-8B it also appends <|im_end|> (id 106) to every text, so a plain prompt's last position is <|im_end|>, not the prompt's last word; ZAYA1-74B-preview's tokenizer appends nothing. To read the prediction after the prompt on ZAYA1-8B, tokenize with add_special_tokens=False and prepend <bos> yourself, or read position -2.

From the family module

04 · nnterp/families/zaya.py

ZAYA (ZayaForCausalLM, Zyphra's ZAYA1).

Llama's containers over a block with residual scaling modules: after each sublayer a ZayaResidualScaling merges the sublayer's output o into the stream r with learned per-channel scales and biases::

merge(o, r) = (o + hidden_states_bias) * hidden_states_scale + (r + residual_bias) * residual_scale

h = post_attention_residual_scale(self_attn(input_layernorm(x)), x)
out = post_mlp_residual_scale(mlp(post_attention_layernorm(h)), h)

What each sublayer adds is its scaled term, (o + hidden_states_bias) * hidden_states_scale, so attention_output and mlp_output are that product: a computed copy of the module's output, inverted on assignment, with a transform carrying an in-place edit back, as on Granite. The stream itself is also rescaled, so the contribution identity is layer_output == (h + residual_bias) * residual_scale + mlp_output with h == (input + residual_bias) * residual_scale + attention_output (each merge's own parameters, layers[i].post_*_residual_scale._module, which each sublayer reads through its parent as merge), the plain one only at their initial values (scales one, biases zero). The model also scales and shifts the embedding module's output before the first block.

The attention runs the shared eager forward; its queries and keys come from a compressed convolutional projection (qkv_proj, a causal convolution over q/k and a delayed half of the values) and qk_norm, so there are no q_proj / k_proj on the module. The block returns (hidden_states, prev_router_hidden_states): the mixture's router carries a state from one block to the next. The mixture (ZayaSparseMoeBlock) is on every block and returns (hidden_states, router_state); its experts are moe_intermediate_size wide, which the root's intermediate_size reports (the config has no dense width). qk_norm L2-normalizes queries and keys and multiplies the keys by a learned per-head temp; at its initial value of zero every key is zero and the pattern uniform.

The mixture is a Moe; its router, gate (aliased router), mixes the previous block's router state into its input and scores num_experts + 1 classes with a small MLP (router_mlp), whose output is router_logits: the last column is skip. A slot that picks it runs no expert: its weight is 0 and its index 0, an alias of expert 0, so usage counts mask expert_weights == 0. The router picks on its softmax plus balancing_biases; on ZAYA1-8B and ZAYA1-74B-preview (num_experts_per_tok 1) the skip class's bias is -1 and some expert's is positive on every block, so skip is never chosen there; the tiny picks it.