nnterp · encyclopedia granitemoehybrid · nnterp.families.granitemoehybrid

family granitemoehybrid

Granite 4.0

Granite's scaled block with a Mamba-2 mixer on most blocks and attention without rotary on a few (rotary attention on every block of the non-H models), and a feed-forward of a shared MLP plus, on the mixture checkpoints, routed experts, added as one scaled sum.

Scaled residual addsHybridState space (Mamba-2)Blocks without rotaryMixture of expertsEmbedding multiplierScaled logitsGrouped-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
→
logitsscaled
i = 0

identity layers[i].input + linear_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 GraniteMoeHybridForCausalLMlayoutwhat 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] GraniteMoeHybridDecoderLayerlayoutwhat 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].mlp GraniteMoeHybridMLPlayoutwhat 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 routed experts plus the shared expert, times residual_multiplier

inside the forward, at the operation that _feed_forward_sum finds on each call
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

the child module experts's output
shared_expert_output Residual [batch seq hidden_size] The shared expert's output, unscaled

the module's own output
model.layers[i].linear_attn GraniteMoeHybridMambaLayerlayoutwhat it isread at
attention_output Residual [batch seq hidden_size] What the Mamba-2 mixer adds to the residual stream: its output times residual_multiplier

the module's own output
⚠︎state scalar Unavailable: the chunk scan computes every token's state in one tensor per call, not one occurrence per token to walk with tracer.iter; read states, or state_after(t), after nnterp.chunk_per_token(model)

nowhere: the family does not serve it
⚠︎states States [batch seq heads key_dim value_dim] The state after every token of this call; needs nnterp.chunk_per_token(model)

inside the forward, at the operation that _states finds on each call
⚠︎attention_queries SSDQueries [batch seq groups state_dim] C, the queries reading the state

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument C
⚠︎attention_keys SSDKeys [batch seq groups state_dim] B, the keys writing into the state

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument B
⚠︎attention_values SSDValues [batch seq heads head_dim] x, the values written into the state

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument hidden_states
⚠︎betas Gates [batch seq heads] dt, the per-token write strength into the state

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument dt
⚠︎decays Gates [batch seq heads] A * dt, the per-token log decay of the state

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument dt
⚠︎state_input State | None [batch heads key_dim value_dim] The state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument state
⚠︎attention_head_outputs SSDHeadOutputs [batch seq heads head_dim] y, the per-head outputs before the gated norm and the output projection

inside the forward: the output of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), the element that call returns it in
⚠︎state_output State [batch heads key_dim value_dim] The state leaving the layer

inside the forward, at the operation that _state_output finds on each call
⚠︎set_state_after scalar Unavailable: the chunk scan computes every boundary state in one cumulative step from the initial state, so a state written at token t does not flow into later tokens' states; assign state_input to change where a call starts

nowhere: the family does not serve it
model.layers[i].self_attn GraniteMoeHybridAttentionlayoutwhat it isread at
⚠︎attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

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

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

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

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

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

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

Printout

What print(model) shows on ibm-granite/granite-4.0-h-tiny-base: the native modules, standard names and every standard value with its layout.

GraniteMoeHybridForCausalLM(
  (lm_head): Linear(in_features=1536, out_features=100352, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(100352, 1536, padding_idx=100256)
  (layers): ModuleList(
    (0-4): 5 x GraniteMoeHybridDecoderLayer(
      (input_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (post_attention_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (block_sparse_moe): GraniteMoeHybridMoE(
        (router): GraniteMoeHybridTopKRouter()
        (experts): GraniteMoeHybridExperts(
          (act_fn): SiLUActivation()
        )
      )
      (mlp/shared_mlp): GraniteMoeHybridMLP(
        (act_fn/activation): SiLUActivation()
        (input_linear): Linear(in_features=1536, out_features=2048, bias=False)
        (down_proj/output_linear): Linear(in_features=1024, out_features=1536, bias=False)
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output, unscaled
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the routed experts plus the shared expert, times residual_multiplier
        (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
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (linear_attn/mamba): GraniteMoeHybridMambaLayer(
        (act): SiLUActivation()
        (conv1d): Conv1d(3328, 3328, kernel_size=(4,), stride=(1,), padding=(3,), groups=3328)
        (in_proj): Linear(in_features=1536, out_features=6448, bias=False)
        (norm): GraniteMoeHybridRMSNormGated()
        (out_proj): Linear(in_features=3072, out_features=1536, bias=False)
        (attention_output) -> Residual [batch seq hidden_size]: What the Mamba-2 mixer adds to the residual stream: its output times residual_multiplier
        (attention_queries) -> SSDQueries [batch seq groups state_dim]: C, the queries reading the state
        (attention_keys) -> SSDKeys [batch seq groups state_dim]: B, the keys writing into the state
        (attention_values) -> SSDValues [batch seq heads head_dim]: x, the values written into the state
        (betas) -> Gates [batch seq heads]: dt, the per-token write strength into the state
        (decays) -> Gates [batch seq heads]: A * dt, the per-token log decay of the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> SSDHeadOutputs [batch seq heads head_dim]: y, the per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The state leaving the layer
        (states) -> States [batch seq heads key_dim value_dim]: The state after every token of this call; needs nnterp.chunk_per_token(model)
        (state): Unavailable: the chunk scan computes every token's state in one tensor per call, not one occurrence per token to walk with tracer.iter; read states, or state_after(t), after nnterp.chunk_per_token(model)
        (set_state_after): Unavailable: the chunk scan computes every boundary state in one cumulative step from the initial state, so a state written at token t does not flow into later tokens' states; assign state_input to change where a call starts
      )
      (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): GraniteMoeHybridDecoderLayer(
      (input_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (post_attention_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (block_sparse_moe): GraniteMoeHybridMoE(
        (router): GraniteMoeHybridTopKRouter()
        (experts): GraniteMoeHybridExperts(
          (act_fn): SiLUActivation()
        )
      )
      (mlp/shared_mlp): GraniteMoeHybridMLP(
        (act_fn/activation): SiLUActivation()
        (input_linear): Linear(in_features=1536, out_features=2048, bias=False)
        (down_proj/output_linear): Linear(in_features=1024, out_features=1536, bias=False)
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output, unscaled
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the routed experts plus the shared expert, times residual_multiplier
        (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
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (self_attn): GraniteMoeHybridAttention(
        (q_proj): Linear(in_features=1536, out_features=1536, bias=False)
        (k_proj): Linear(in_features=1536, out_features=512, bias=False)
        (v_proj): Linear(in_features=1536, out_features=512, bias=False)
        (o_proj): Linear(in_features=1536, out_features=1536, bias=False)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: its output times residual_multiplier
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (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
    )
    (6-14): 9 x GraniteMoeHybridDecoderLayer(
      (input_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (post_attention_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (block_sparse_moe): GraniteMoeHybridMoE(
        (router): GraniteMoeHybridTopKRouter()
        (experts): GraniteMoeHybridExperts(
          (act_fn): SiLUActivation()
        )
      )
      (mlp/shared_mlp): GraniteMoeHybridMLP(
        (act_fn/activation): SiLUActivation()
        (input_linear): Linear(in_features=1536, out_features=2048, bias=False)
        (down_proj/output_linear): Linear(in_features=1024, out_features=1536, bias=False)
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output, unscaled
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the routed experts plus the shared expert, times residual_multiplier
        (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
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (linear_attn/mamba): GraniteMoeHybridMambaLayer(
        (act): SiLUActivation()
        (conv1d): Conv1d(3328, 3328, kernel_size=(4,), stride=(1,), padding=(3,), groups=3328)
        (in_proj): Linear(in_features=1536, out_features=6448, bias=False)
        (norm): GraniteMoeHybridRMSNormGated()
        (out_proj): Linear(in_features=3072, out_features=1536, bias=False)
        (attention_output) -> Residual [batch seq hidden_size]: What the Mamba-2 mixer adds to the residual stream: its output times residual_multiplier
        (attention_queries) -> SSDQueries [batch seq groups state_dim]: C, the queries reading the state
        (attention_keys) -> SSDKeys [batch seq groups state_dim]: B, the keys writing into the state
        (attention_values) -> SSDValues [batch seq heads head_dim]: x, the values written into the state
        (betas) -> Gates [batch seq heads]: dt, the per-token write strength into the state
        (decays) -> Gates [batch seq heads]: A * dt, the per-token log decay of the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> SSDHeadOutputs [batch seq heads head_dim]: y, the per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The state leaving the layer
        (states) -> States [batch seq heads key_dim value_dim]: The state after every token of this call; needs nnterp.chunk_per_token(model)
        (state): Unavailable: the chunk scan computes every token's state in one tensor per call, not one occurrence per token to walk with tracer.iter; read states, or state_after(t), after nnterp.chunk_per_token(model)
        (set_state_after): Unavailable: the chunk scan computes every boundary state in one cumulative step from the initial state, so a state written at token t does not flow into later tokens' states; assign state_input to change where a call starts
      )
      (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
    )
    (15): GraniteMoeHybridDecoderLayer(
      (input_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (post_attention_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (block_sparse_moe): GraniteMoeHybridMoE(
        (router): GraniteMoeHybridTopKRouter()
        (experts): GraniteMoeHybridExperts(
          (act_fn): SiLUActivation()
        )
      )
      (mlp/shared_mlp): GraniteMoeHybridMLP(
        (act_fn/activation): SiLUActivation()
        (input_linear): Linear(in_features=1536, out_features=2048, bias=False)
        (down_proj/output_linear): Linear(in_features=1024, out_features=1536, bias=False)
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output, unscaled
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the routed experts plus the shared expert, times residual_multiplier
        (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
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (self_attn): GraniteMoeHybridAttention(
        (q_proj): Linear(in_features=1536, out_features=1536, bias=False)
        (k_proj): Linear(in_features=1536, out_features=512, bias=False)
        (v_proj): Linear(in_features=1536, out_features=512, bias=False)
        (o_proj): Linear(in_features=1536, out_features=1536, bias=False)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: its output times residual_multiplier
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (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-24): 9 x GraniteMoeHybridDecoderLayer(
      (input_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (post_attention_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (block_sparse_moe): GraniteMoeHybridMoE(
        (router): GraniteMoeHybridTopKRouter()
        (experts): GraniteMoeHybridExperts(
          (act_fn): SiLUActivation()
        )
      )
      (mlp/shared_mlp): GraniteMoeHybridMLP(
        (act_fn/activation): SiLUActivation()
        (input_linear): Linear(in_features=1536, out_features=2048, bias=False)
        (down_proj/output_linear): Linear(in_features=1024, out_features=1536, bias=False)
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output, unscaled
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the routed experts plus the shared expert, times residual_multiplier
        (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
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (linear_attn/mamba): GraniteMoeHybridMambaLayer(
        (act): SiLUActivation()
        (conv1d): Conv1d(3328, 3328, kernel_size=(4,), stride=(1,), padding=(3,), groups=3328)
        (in_proj): Linear(in_features=1536, out_features=6448, bias=False)
        (norm): GraniteMoeHybridRMSNormGated()
        (out_proj): Linear(in_features=3072, out_features=1536, bias=False)
        (attention_output) -> Residual [batch seq hidden_size]: What the Mamba-2 mixer adds to the residual stream: its output times residual_multiplier
        (attention_queries) -> SSDQueries [batch seq groups state_dim]: C, the queries reading the state
        (attention_keys) -> SSDKeys [batch seq groups state_dim]: B, the keys writing into the state
        (attention_values) -> SSDValues [batch seq heads head_dim]: x, the values written into the state
        (betas) -> Gates [batch seq heads]: dt, the per-token write strength into the state
        (decays) -> Gates [batch seq heads]: A * dt, the per-token log decay of the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> SSDHeadOutputs [batch seq heads head_dim]: y, the per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The state leaving the layer
        (states) -> States [batch seq heads key_dim value_dim]: The state after every token of this call; needs nnterp.chunk_per_token(model)
        (state): Unavailable: the chunk scan computes every token's state in one tensor per call, not one occurrence per token to walk with tracer.iter; read states, or state_after(t), after nnterp.chunk_per_token(model)
        (set_state_after): Unavailable: the chunk scan computes every boundary state in one cumulative step from the initial state, so a state written at token t does not flow into later tokens' states; assign state_input to change where a call starts
      )
      (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): GraniteMoeHybridDecoderLayer(
      (input_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (post_attention_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (block_sparse_moe): GraniteMoeHybridMoE(
        (router): GraniteMoeHybridTopKRouter()
        (experts): GraniteMoeHybridExperts(
          (act_fn): SiLUActivation()
        )
      )
      (mlp/shared_mlp): GraniteMoeHybridMLP(
        (act_fn/activation): SiLUActivation()
        (input_linear): Linear(in_features=1536, out_features=2048, bias=False)
        (down_proj/output_linear): Linear(in_features=1024, out_features=1536, bias=False)
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output, unscaled
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the routed experts plus the shared expert, times residual_multiplier
        (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
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (self_attn): GraniteMoeHybridAttention(
        (q_proj): Linear(in_features=1536, out_features=1536, bias=False)
        (k_proj): Linear(in_features=1536, out_features=512, bias=False)
        (v_proj): Linear(in_features=1536, out_features=512, bias=False)
        (o_proj): Linear(in_features=1536, out_features=1536, bias=False)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: its output times residual_multiplier
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (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
    )
    (26-34): 9 x GraniteMoeHybridDecoderLayer(
      (input_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (post_attention_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (block_sparse_moe): GraniteMoeHybridMoE(
        (router): GraniteMoeHybridTopKRouter()
        (experts): GraniteMoeHybridExperts(
          (act_fn): SiLUActivation()
        )
      )
      (mlp/shared_mlp): GraniteMoeHybridMLP(
        (act_fn/activation): SiLUActivation()
        (input_linear): Linear(in_features=1536, out_features=2048, bias=False)
        (down_proj/output_linear): Linear(in_features=1024, out_features=1536, bias=False)
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output, unscaled
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the routed experts plus the shared expert, times residual_multiplier
        (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
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (linear_attn/mamba): GraniteMoeHybridMambaLayer(
        (act): SiLUActivation()
        (conv1d): Conv1d(3328, 3328, kernel_size=(4,), stride=(1,), padding=(3,), groups=3328)
        (in_proj): Linear(in_features=1536, out_features=6448, bias=False)
        (norm): GraniteMoeHybridRMSNormGated()
        (out_proj): Linear(in_features=3072, out_features=1536, bias=False)
        (attention_output) -> Residual [batch seq hidden_size]: What the Mamba-2 mixer adds to the residual stream: its output times residual_multiplier
        (attention_queries) -> SSDQueries [batch seq groups state_dim]: C, the queries reading the state
        (attention_keys) -> SSDKeys [batch seq groups state_dim]: B, the keys writing into the state
        (attention_values) -> SSDValues [batch seq heads head_dim]: x, the values written into the state
        (betas) -> Gates [batch seq heads]: dt, the per-token write strength into the state
        (decays) -> Gates [batch seq heads]: A * dt, the per-token log decay of the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> SSDHeadOutputs [batch seq heads head_dim]: y, the per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The state leaving the layer
        (states) -> States [batch seq heads key_dim value_dim]: The state after every token of this call; needs nnterp.chunk_per_token(model)
        (state): Unavailable: the chunk scan computes every token's state in one tensor per call, not one occurrence per token to walk with tracer.iter; read states, or state_after(t), after nnterp.chunk_per_token(model)
        (set_state_after): Unavailable: the chunk scan computes every boundary state in one cumulative step from the initial state, so a state written at token t does not flow into later tokens' states; assign state_input to change where a call starts
      )
      (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
    )
    (35): GraniteMoeHybridDecoderLayer(
      (input_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (post_attention_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (block_sparse_moe): GraniteMoeHybridMoE(
        (router): GraniteMoeHybridTopKRouter()
        (experts): GraniteMoeHybridExperts(
          (act_fn): SiLUActivation()
        )
      )
      (mlp/shared_mlp): GraniteMoeHybridMLP(
        (act_fn/activation): SiLUActivation()
        (input_linear): Linear(in_features=1536, out_features=2048, bias=False)
        (down_proj/output_linear): Linear(in_features=1024, out_features=1536, bias=False)
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output, unscaled
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the routed experts plus the shared expert, times residual_multiplier
        (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
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (self_attn): GraniteMoeHybridAttention(
        (q_proj): Linear(in_features=1536, out_features=1536, bias=False)
        (k_proj): Linear(in_features=1536, out_features=512, bias=False)
        (v_proj): Linear(in_features=1536, out_features=512, bias=False)
        (o_proj): Linear(in_features=1536, out_features=1536, bias=False)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: its output times residual_multiplier
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (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-39): 4 x GraniteMoeHybridDecoderLayer(
      (input_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (post_attention_layernorm): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
      (block_sparse_moe): GraniteMoeHybridMoE(
        (router): GraniteMoeHybridTopKRouter()
        (experts): GraniteMoeHybridExperts(
          (act_fn): SiLUActivation()
        )
      )
      (mlp/shared_mlp): GraniteMoeHybridMLP(
        (act_fn/activation): SiLUActivation()
        (input_linear): Linear(in_features=1536, out_features=2048, bias=False)
        (down_proj/output_linear): Linear(in_features=1024, out_features=1536, bias=False)
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output, unscaled
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the routed experts plus the shared expert, times residual_multiplier
        (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
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
      )
      (linear_attn/mamba): GraniteMoeHybridMambaLayer(
        (act): SiLUActivation()
        (conv1d): Conv1d(3328, 3328, kernel_size=(4,), stride=(1,), padding=(3,), groups=3328)
        (in_proj): Linear(in_features=1536, out_features=6448, bias=False)
        (norm): GraniteMoeHybridRMSNormGated()
        (out_proj): Linear(in_features=3072, out_features=1536, bias=False)
        (attention_output) -> Residual [batch seq hidden_size]: What the Mamba-2 mixer adds to the residual stream: its output times residual_multiplier
        (attention_queries) -> SSDQueries [batch seq groups state_dim]: C, the queries reading the state
        (attention_keys) -> SSDKeys [batch seq groups state_dim]: B, the keys writing into the state
        (attention_values) -> SSDValues [batch seq heads head_dim]: x, the values written into the state
        (betas) -> Gates [batch seq heads]: dt, the per-token write strength into the state
        (decays) -> Gates [batch seq heads]: A * dt, the per-token log decay of the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> SSDHeadOutputs [batch seq heads head_dim]: y, the per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The state leaving the layer
        (states) -> States [batch seq heads key_dim value_dim]: The state after every token of this call; needs nnterp.chunk_per_token(model)
        (state): Unavailable: the chunk scan computes every token's state in one tensor per call, not one occurrence per token to walk with tracer.iter; read states, or state_after(t), after nnterp.chunk_per_token(model)
        (set_state_after): Unavailable: the chunk scan computes every boundary state in one cumulative step from the initial state, so a state written at token t does not flow into later tokens' states; assign state_input to change where a call starts
      )
      (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): GraniteMoeHybridRMSNorm((1536,), eps=1e-05)
  (logits) -> Logits [batch seq vocab_size]: The logits the model returns, after anything it does past lm_head (a softcap, a scale)
  (token_embeddings) -> Residual [batch seq hidden_size]: The embedding module's output; layers[0].input is what enters the first block
  (final_norm_output) -> Residual [batch seq hidden_size]: The residual stream after the final norm, what lm_head reads
  (lm_head_output) -> Logits [batch seq vocab_size]: The head's logits before anything the model does past lm_head (a softcap, a scale)
  (next_token_probs) -> NextTokenProbs [batch vocab_size]: The next-token distribution at the last position; derived, read-only
  (input_ids) -> Tokens [batch seq]: The token ids the model was called with
  (attention_mask) -> Tokens [batch seq]: The attention mask the model was called with; zeros are padding
  (input_size): [batch, seq] of the current call; read-only
)

Config

model.num_layers40
model.hidden_size1536
model.num_heads12
model.num_kv_heads4
model.head_dim128
model.qk_head_dim128
model.vocab_size100352
model.intermediate_size1024
architecturesGraniteMoeHybridForCausalLM
num_hidden_layers40
hidden_size1536
intermediate_size512
num_attention_heads12
num_key_value_heads4
vocab_size100352
max_position_embeddings131072
hidden_actsilu
rms_norm_eps1e-05
tie_word_embeddingsTrue
dtypebfloat16
embedding_multiplier12
residual_multiplier0.22
attention_multiplier0.0078125
logits_scaling6
num_experts_per_tok6
num_local_experts64
mamba_n_heads48
mamba_d_head64
mamba_d_state128
mamba_expand2
mamba_chunk_size256
mamba_n_groups1
attention_biasFalse

The Notes

03 · what to know before you run

The block, in order

x      = embed_tokens(ids) * embedding_multiplier
h      = x + mixer(input_layernorm(x)) * residual_multiplier       # mixer: Mamba-2 or attention
n      = post_attention_layernorm(h)
out    = h + (block_sparse_moe(n) + shared_mlp(n)) * residual_multiplier
logits = lm_head(norm(out)) / logits_scaling

Every block has the same class and holds one mixer, chosen by config.layer_types: the Mamba-2 mixer mamba (served as linear_attn) or self_attn. The H models (H-350M, H-1B, H-Micro, H-Tiny, H-Small) have 40 blocks with attention on blocks 5, 15, 25 and 35 (H-350M: 32 blocks, attention on 10, 13, 17 and 27) and a Mamba-2 mixer on the rest; the non-H 350M, 1B and Micro have attention on every block. The feed-forward is the shared MLP, shared_mlp (served as mlp). Where num_local_experts is above 0, the routed experts in block_sparse_moe run beside it: H-Tiny (64 experts, top 6), H-Small (72, top 10) and the Tiny previews (62, top 6). On the dense checkpoints shared_mlp(n) alone is the feed-forward. Pick blocks outside the trace:

kinds = model.config.layer_types
ssm = [i for i, t in enumerate(kinds) if t == "linear_attention"]
attn_blocks = [i for i, t in enumerate(kinds) if t == "full_attention"]

The contributions are scaled; the mixture's values are not

Both mixers' attention_output are the module's output times residual_multiplier, and mlp_output is the block's binding of the feed-forward's sum times it, so the plain identity holds, bit for bit in bfloat16 on H-1B on both block kinds. mlp is the shared MLP and hosts the mixture's values, read before the multiplier: routed_output is layers[i].block_sparse_moe.output and shared_expert_output is mlp.output. The routed experts run first.

moe = model.layers[1].mlp                    # a mixture checkpoint
with model.trace(prompt):
    x = model.layers[1].input.save()
    attn = model.layers[1].self_attn.attention_output.save()
    routed = moe.routed_output.save()
    shared = moe.shared_expert_output.save()
    mlp = moe.mlp_output.save()
    out = model.layers[1].layer_output.save()

m = model.config.residual_multiplier         # 0.22
torch.testing.assert_close((routed + shared) * m, mlp)
torch.testing.assert_close(x + attn + mlp, out)

mlp.output is the shared MLP's output alone, also on a dense checkpoint, where mlp_output == mlp.output * m. A write to mlp_output lands on the sum: mlp_output[:] = 0 removes the routed and the shared experts together. intermediate_size and mlp.intermediate_size are the shared MLP's width, shared_intermediate_size; the routed experts are config.intermediate_size wide (512 on H-Tiny). The router takes the top num_experts_per_tok of its logits and softmaxes those alone, GraniteMoE's routing.

Route the kernels before the first trace

With mamba_ssm installed, transformers runs the Mamba-2 scan in its CUDA kernels, and every linear_attn value but attention_output is unavailable until nnterp.route_kernels(model.family, "torch") is called before the first trace. On a CPU the unrouted model does not run at all. states and state_after also need nnterp.chunk_per_token(model): the scan keeps the state every mamba_chunk_size tokens, 256 on every released checkpoint.

import nnterp

nnterp.route_kernels(nnterp.families.granitemoehybrid, "torch")     # before the first trace
model = StandardizedTransformer(
    "ibm-granite/granite-4.0-h-1b-base", attn_implementation="eager", dtype=torch.float32
)

The kernels disagree in bfloat16. On H-1B, on a seven-token prompt, the routed model's bfloat16 logits differ from its float32 logits by up to 5.7, the default kernels' by up to 2.4 (with a different top token at one position), and the two bfloat16 runs differ from each other by up to 5.5. In float32 the routed and default kernels agree to 0.07 with the same top tokens, and eager and sdpa attention to the bit. Compare runs, and measure effects, in float32.

The Mamba-2 mixer

On H-1B each mixer has 48 heads of 64 channels in one group, with a state of 128: attention_queries (C) and attention_keys (B) are [batch, seq, 1, 128], attention_values (x) and attention_head_outputs are [batch, seq, 48, 64], and the state is [batch, 48, 128, 64]. A width-4 convolution runs over x, B and C before the scan, and a gated RMSNorm after it; attention_head_outputs is read before the norm. dt (betas) is clamped to time_step_limit, (0, inf), so a written zero stays zero: betas[:, t] = 0 skips token t's write exactly, and the state after t equals the state before it.

nnterp.chunk_per_token(model)
mix = model.layers[ssm[0]].linear_attn
with model.trace(prompt):
    betas = mix.betas.clone()
    betas[:, 3] = 0
    mix.betas = betas
    states = mix.states.save()                # [batch, seq, heads, 128, 64]

assert torch.equal(states[:, 3], states[:, 2])

Attention without rotary on the H models

The H models set position_embedding_type to nope: the model builds no rotary embedding, and attention_queries is q_proj's output split into heads, unchanged. Order reaches the attention only through the causal mask and the Mamba-2 blocks before it. The non-H 350M, 1B and Micro apply rotary embeddings on every block.

attn = model.layers[attn_blocks[0]].self_attn         # block 5 on H-1B
with model.trace(prompt):
    q_raw = attn.source.self_q_proj_0.output.save()
    q = attn.attention_queries.save()

b, s, _ = q_raw.shape
assert torch.equal(q_raw.view(b, s, -1, model.head_dim).transpose(1, 2), q)

The query scale is config.attention_multiplier, 1 / head_dim (1/128 on H-1B and H-Tiny). H-1B has 12 query heads over 4 key/value heads of 128, so an edit to key/value head j reaches query heads 3j to 3j + 2. The six interior values need attn_implementation="eager".

A write rounds the other positions with 0.22

A write to attention_output or mlp_output reaches the model as the whole edited copy divided by residual_multiplier, which the block multiplies again. With 0.22, for 502 of the 65280 finite bfloat16 products x * 0.22 that round trip does not give the product back, and the Mamba-2 recurrence carries the difference on. On H-1B in bfloat16, adding a unit vector to mlp_output at the last position moved the earlier positions' logits by up to 0.5 from block 1 (a KL of 0.009) and by up to 4.0 from block 5 (a KL of 0.23), which the edit cannot reach causally. The same edit on the module's output leaves them bit-identical:

m = model.config.residual_multiplier
with model.trace(prompt):
    model.layers[5].mlp.mlp_output[:, -1] += v    # every position rounds

with model.trace(prompt):
    model.layers[5].mlp.output[:, -1] += v / m    # position -1 only (dense checkpoint)

On a mixture checkpoint mlp.output is the shared MLP's: edit block_sparse_moe.output for the routed experts' term. With the multipliers of H-350M (0.246) and 350M (0.263) the round trip is exact for every bfloat16 product.

Position 0 carries a large norm

The tokenizer prepends nothing, so position 0 is the prompt's own first token. On H-1B the stream there has a norm of about 270 after block 20 against 21 to 28 elsewhere, whether the prompt starts with The, def or In, and the four attention blocks put 0.47 to 0.90 of their attention on it, averaged over heads and queries. Slice [:, 1:] before averaging activations for steering vectors, mean ablation or probes.

with model.trace(prompt):
    resid = model.layers[20].layer_output.save()

resid[0].float().norm(dim=-1)        # H-1B: about 270 at position 0, 21 to 28 elsewhere

Logits and embeddings are Granite's

model.logits is model.lm_head.output / logits_scaling and project_on_vocab divides too, so a lens at the last block equals logits. A softmax over lm_head.output is at temperature 1/logits_scaling: after "The capital of the United Kingdom is", London has probability 0.56 from logits and the top token 1.00 from lm_head.output on H-1B. token_embeddings is the plain lookup and layers[0].input is it times 12. lm_head and embed_tokens are one parameter.

with model.trace(prompt):
    emb = model.token_embeddings.save()
    first = model.layers[0].input.save()
    resid = model.layers[-1].layer_output.save()
    raw = model.lm_head.output.save()
    logits = model.logits.save()

torch.testing.assert_close(first, emb * model.config.embedding_multiplier)
torch.testing.assert_close(logits, raw / model.config.logits_scaling)
torch.testing.assert_close(model.project_on_vocab(resid), logits)

What this family module covers

Every checkpoint whose config says granitemoehybrid: all of Granite 4.0, the H models, the attention-only 350M, 1B and Micro, the Tiny previews, Granite Guardian 4.0 3B and granite-4.0-h-3b-ar. Granite 4.1 and 4.2 are granite, the Swash models granite_swa and granitemoe_swa. granitemoeshared has the same feed-forward on an attention-only block.

From the family module

04 · nnterp/families/granitemoehybrid.py

GraniteMoE-Hybrid (GraniteMoeHybridForCausalLM, IBM Granite 4.0 and 4.0-H).

Llama's containers over one block class that holds either a Mamba-2 mixer or attention, by config.layers_block_type, and then a feed-forward::

model.layers[i]          GraniteMoeHybridDecoderLayer
    .input_layernorm
    .mamba               GraniteMoeHybridMambaLayer, an SSD mixer   (or .self_attn)
    .post_attention_layernorm
    .block_sparse_moe    GraniteMoeHybridMoE (only when num_local_experts > 0)
    .shared_mlp          GraniteMoeHybridMLP (on every block)

h = x + mixer(input_layernorm(x)) * residual_multiplier
out = h + (block_sparse_moe(n) + shared_mlp(n)) * residual_multiplier     with n = post_attention_layernorm(h)

(shared_mlp(n) alone on a dense checkpoint, Granite-4.0-H-Micro's.) mamba is linear_attn (nnterp.StateSpace) and shared_mlp, the one feed-forward every block has, is mlp; each block has self_attn or linear_attn, never both. As on Granite (granite.py) the block adds each sublayer's output times residual_multiplier: both mixers' attention_output are the module's output times it, computed copies divided on assignment and carried back by a transform. No module returns the feed-forward's sum, so mlp_output is the block's binding of it (hidden_states_4 with experts, hidden_states_5 without) times the multiplier, and the family's Layer sets sourced for it; the mixer and the shared expert, whose modules keep no config, read the multiplier (and the Mlp whether routed experts run) off their parent block. A write lands on the sum, replacing both experts' contribution. mlp.output is the shared expert's output alone and layers[i].block_sparse_moe.output the routed experts'. The mixture's values live on mlp (a Moe), which reads block_sparse_moe's router and experts through its parent: shared_expert_output is mlp.output and routed_output + shared_expert_output == mlp_output / residual_multiplier. mlp.intermediate_size and the root's intermediate_size are the shared expert's width, shared_intermediate_size; the routed experts are intermediate_size wide. The attention runs the shared eager forward with attention_multiplier as its scale and no rotary under position_embedding_type nope (Granite-4.0-H); embedding_multiplier and logits_scaling are as on Granite.