nnterp · encyclopedia laguna · nnterp.families.laguna

family laguna

Laguna XSLaguna SLaguna M

Llama's pre-norm block with a softplus gate on each head's output before o_proj, and after the first blocks a sigmoid-routed mixture of experts beside a shared expert; on XS and S the blocks differ in query heads.

Per-block sizesQuery/key normsOutput gateSliding windowPartial rotaryMixture of expertsDense first blocksGrouped-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. The slider picks a block, each tick's color gives its block's shape, and the diagram redraws for the block picked.

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

identity layers[i].input + self_attn.attention_output + 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 LagunaForCausalLMlayoutwhat 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] LagunaDecoderLayerlayoutwhat 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 LagunaAttentionlayoutwhat 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

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 LagunaMLP / LagunaSparseMoeBlocklayoutwhat 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

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

inside router's forward: the output of F_linear_0
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, times routed_scaling_factor

inside the forward: the output of hidden_states_2
shared_expert_output Residual [batch seq hidden_size] The shared expert's output

the child module shared_experts's output

Printout

What print(model) shows on poolside/Laguna-XS-2.1: the native modules, standard names and every standard value with its layout.

LagunaForCausalLM(
  (lm_head): Linear(in_features=2048, out_features=100352, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(100352, 2048, padding_idx=9)
  (layers): ModuleList(
    (0): LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=6144, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=6144, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=48, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaMLP(
        (gate_proj): Linear(in_features=2048, out_features=8192, bias=False)
        (up_proj): Linear(in_features=2048, out_features=8192, bias=False)
        (down_proj): Linear(in_features=8192, out_features=2048, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
        (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
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (1-3): 3 x LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=8192, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=8192, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=64, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (4): LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=6144, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=6144, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=48, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (5-7): 3 x LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=8192, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=8192, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=64, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (8): LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=6144, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=6144, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=48, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (9-11): 3 x LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=8192, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=8192, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=64, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (12): LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=6144, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=6144, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=48, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (13-15): 3 x LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=8192, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=8192, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=64, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (16): LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=6144, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=6144, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=48, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (17-19): 3 x LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=8192, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=8192, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=64, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (20): LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=6144, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=6144, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=48, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (21-23): 3 x LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=8192, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=8192, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=64, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (24): LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=6144, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=6144, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=48, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (25-27): 3 x LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=8192, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=8192, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=64, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (28): LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=6144, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=6144, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=48, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (29-31): 3 x LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=8192, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=8192, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=64, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (32): LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=6144, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=6144, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=48, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (33-35): 3 x LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=8192, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=8192, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=64, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (36): LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=6144, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=6144, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=48, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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
    )
    (37-39): 3 x LagunaDecoderLayer(
      (self_attn): LagunaAttention(
        (q_proj): Linear(in_features=2048, out_features=8192, bias=False)
        (k_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (v_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (o_proj): Linear(in_features=8192, out_features=2048, bias=False)
        (q_norm): LagunaRMSNorm((128,), eps=1e-06)
        (k_norm): LagunaRMSNorm((128,), eps=1e-06)
        (g_proj): Linear(in_features=2048, out_features=64, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): LagunaSparseMoeBlock(
        (experts): LagunaExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): LagunaTopKRouter()
        (shared_experts): LagunaMLP(
          (gate_proj): Linear(in_features=2048, out_features=512, bias=False)
          (up_proj): Linear(in_features=2048, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
          (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
        )
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output, times routed_scaling_factor
        (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (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
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (post_attention_layernorm): LagunaRMSNorm((2048,), eps=1e-06)
      (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): LagunaRMSNorm((2048,), eps=1e-06)
  (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_heads48
model.num_kv_heads8
model.head_dim128
model.qk_head_dim128
model.vocab_size100352
model.intermediate_size8192
architecturesLagunaForCausalLM
num_hidden_layers40
hidden_size2048
intermediate_size8192
num_attention_heads48
num_key_value_heads8
head_dim128
vocab_size100352
max_position_embeddings262144
sliding_window512
hidden_actsilu
rms_norm_eps1e-06
tie_word_embeddingsFalse
dtypebfloat16
moe_intermediate_size512
num_experts256
num_experts_per_tok8
norm_topk_probTrue
attention_biasFalse
decoder_sparse_step1
shared_expert_intermediate_size512

The Notes

03 · what to know before you run

The block, in order

h   = x + self_attn(input_layernorm(x))       # q/k norms, a softplus gate
out = h + mlp(post_attention_layernorm(h))    # dense first, then shared + top_k of 256

Llama's pre-norm block and Llama's names. mlp_layer_types makes the first block dense on XS and S and the first three on M.1, and the rest mixtures. The block adds each sublayer's output to the stream and returns a tensor, so the identity is the plain sum on both block kinds:

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()

torch.testing.assert_close(x + attn + mlp, out)

Each block has its own number of query heads

num_attention_heads_per_layer sets each block's query heads. XS-2.1 and XS.2 have 48 on the full blocks and 64 on the sliding ones, S-2.1 has 48 and 72, and M.1 has 64 on every block. Every block has 8 key/value heads of 128 dimensions. The root's num_heads is the config's num_attention_heads, the full blocks' count; each block's own is on its attention:

model.num_heads                         # 48 on XS-2.1
model.layers[1].self_attn.num_heads     # 64: block 1 slides
model.layers[0].self_attn.num_heads     # 48: block 0 is full

attention_queries, the pattern and attention_head_outputs have the block's own head count, so the pattern is [batch, 64, seq, seq] on a sliding block of XS-2.1 and [batch, 48, seq, seq] on a full one: stack them per block, not across blocks. A key/value head serves 6 query heads on a full block and 8 (XS) or 9 (S) on a sliding one.

Sliding and full blocks differ in window and rotary

On XS and S, layer_types makes every fourth block full (0, 4, 8, ...) and the others sliding, with a window of 512 positions; every block of M.1 is full. A full block's rotary is YaRN with base 500,000, over the first half of each head (partial_rotary_factor 0.5) on XS and S and over the whole head on M.1. A sliding block's rotary is plain, with base 10,000, over the whole head. The YaRN factor is 32 on XS-2.1, 64 on XS.2 and M.1 and 128 on S-2.1. q_norm and k_norm are RMSNorms over each head (head_dim wide), applied before the rotary; attention_queries and attention_keys are read after both, so an edit to one head's slice of q_proj.output reaches that head only.

The head outputs are gated before o_proj

After the interface, the attention multiplies each head's output by softplus(g_proj(x)), where x is the attention's input. On XS-2.1, XS.2 and S-2.1 the gate has one value per head per token (gating is per-head or true); on M.1 one per channel (per-element). attention_head_outputs is the ungated interface output; o_proj.input is the gated one, and attention_output is o_proj of it:

attn = model.layers[1].self_attn
with model.trace(prompt):
    heads = attn.attention_head_outputs.save()  # [batch, seq, num_heads, head_dim], ungated
    g = attn.g_proj.output.save()               # [batch, seq, heads] when per-head
    gated = attn.o_proj.input.save()

gate = torch.nn.functional.softplus(g.float()).to(heads.dtype)
torch.testing.assert_close((heads * gate.unsqueeze(-1)).flatten(2), gated)

A head's term in attention_output is its gate times o_proj's slice applied to its output, so zeroing a head at attention_head_outputs removes the term, and the gate is a positive per-token weight on each head.

Load with eager for the attention interior

The attention interior needs attn_implementation="eager"; the default load runs sdpa. There is no softcap or sink, so the two loads compute the same attention up to rounding (1e-7 in the logits on the tiny checkpoint, in float32).

The mixture: a shared expert first, then sigmoid routing with a selection bias

Each mixture runs its shared expert (mlp.shared_experts, shared_expert_intermediate_size wide) first, on the same input as the router. router_logits is the router's projection in float32, read before the tanh cap that moe_router_logit_softcapping sets; that key is 0 (no cap) on every checkpoint. The router adds its selection bias, mlp.router.e_score_correction_bias, to the sigmoids to choose; expert_weights are the chosen experts' sigmoids without the bias, renormalized to sum to one. The experts' weighted sum is then multiplied by moe_routed_scaling_factor (2.5 on XS and S, 1.0 on M.1): that product is routed_output, and mlp_output is routed_output + shared_expert_output. Read shared_expert_output first in a trace:

moe = model.layers[1].mlp
with model.trace(prompt):
    shared = moe.shared_expert_output.save()    # first: it runs before the router
    logits = moe.router_logits.save()           # [batch, seq, num_experts], float32
    w = moe.expert_weights.save()               # [batch, seq, top_k]
    idx = moe.expert_indices.save()
    each = moe.expert_outputs.save()            # [batch, seq, top_k, hidden_size]
    routed = moe.routed_output.save()
    out = moe.mlp_output.save()

chosen = logits.sigmoid().gather(-1, idx)
scale = model.config.moe_routed_scaling_factor
torch.testing.assert_close(w, (chosen / chosen.sum(-1, keepdim=True)).to(w.dtype))
torch.testing.assert_close(each.sum(2) * scale, routed)
torch.testing.assert_close(routed + shared, out)

expert_outputs are the weighted outputs before the scale, so on XS and S a slot adds 2.5 times its expert_outputs entry to the stream. A written router_logits moves both the choice and the weights; a written bias moves only the choice. The shared expert is an Mlp whose own mlp_output is unavailable: what the block adds is the mixture's.

Ablating an expert touches only the tokens routed to it

Zeroing expert e's weight where expert_indices == e removes that slot's term and leaves the token's other weights as they were. mlp_output changes on exactly the tokens that chose e, and the shared expert's term stays:

moe, e = model.layers[1].mlp, 3
with model.trace(prompt):
    idx = moe.expert_indices.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()

changed = (ablated != clean).any(-1)                       # [batch, seq]
assert torch.equal(changed, (idx == e).any(-1))            # only the tokens that chose e

The readout and the tokenizer

logits is lm_head.output, with no cap or scale, and project_on_vocab on the last block's layer_output equals logits. The final norm is a plain RMSNorm whose gain is norm.weight. The tokenizer prepends id 2, 〈|EOS|〉, which is also an end token, so model.input_ids[:, 0] is 2; the configs' end tokens are ids 2 and 24 (</assistant>).

From the family module

04 · nnterp/families/laguna.py

Laguna (LagunaForCausalLM, poolside's Laguna).

Llama's tree and Llama's block: model.{embed_tokens, layers[i].{input_layernorm, self_attn, post_attention_layernorm, mlp}, norm} and lm_head, the residual added in the block, which returns a bare tensor. What differs is inside the sublayers:

  • Per-block head counts. num_attention_heads_per_layer gives each block its own number of query heads (48 on the full-attention blocks and 64 or 72 on the sliding-window ones in the released checkpoints); the key/value heads and head_dim are the same everywhere. layers[i].self_attn.num_heads is that block's count, read off the module; the root's num_heads is the config's top-level num_attention_heads.
  • Gated attention. q/k norms before the rotary (partial on the full blocks), the shared eager forward, then a softplus gate (g_proj, per head or per channel) on the interface's output before o_proj: attention_head_outputs are ungated and attention_output is the gated projection the block adds.
  • Mixture of experts. mlp_layer_types names the dense blocks (LagunaMLP, intermediate_size wide) and the sparse ones (LagunaSparseMoeBlock: routed experts moe_intermediate_size wide, times moe_routed_scaling_factor, plus a shared expert, as one tensor). The shared expert (mlp.shared_experts) is a LagunaMLP too, so it is an Mlp; its mlp_output is unavailable, since the block adds the mixture's sum.