nnterp · encyclopedia granite_swa · nnterp.families.granite_swa

family granite_swa

Granite Swash

Granite's block and multipliers with sliding-window attention on most blocks and a learned sink per head that scales each head's output after the softmax: the pattern's rows sum to one, and each head writes only a share of what its pattern mixes.

Scaled residual addsSliding windowSink-scaled headsEmbedding 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; this family's blocks differ by layer_types.

embed_tokenstoken_embeddings
→
layers× 24
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsscaled
i = 0

identity layers[i].input + self_attn.attention_output + mlp.mlp_output == layer_output

The API

02 · support()

Values, by host

model GraniteSWAForCausalLMlayoutwhat 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] GraniteSWADecoderLayerlayoutwhat 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 GraniteSWAAttentionlayoutwhat 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 F_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 F_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 GraniteSWAMLPlayoutwhat it isread at
mlp_activation Neurons [batch seq intermediate_size] The activation function's output

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

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

the module's own output

Printout

What print(model) shows on ibm-granite/granite-swash-2b: the native modules, standard names and every standard value with its layout.

GraniteSWAForCausalLM(
  (lm_head): Linear(in_features=2560, out_features=100352, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(100352, 2560, padding_idx=100256)
  (layers): ModuleList(
    (0-23): 24 x GraniteSWADecoderLayer(
      (self_attn): GraniteSWAAttention(
        (q_proj): Linear(in_features=2560, out_features=2560, bias=False)
        (k_proj): Linear(in_features=2560, out_features=512, bias=False)
        (v_proj): Linear(in_features=2560, out_features=512, bias=False)
        (o_proj): Linear(in_features=2560, out_features=2560, bias=False)
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (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_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): GraniteSWAMLP(
        (gate_proj): Linear(in_features=2560, out_features=8192, bias=False)
        (up_proj): Linear(in_features=2560, out_features=8192, bias=False)
        (down_proj): Linear(in_features=8192, out_features=2560, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: its output 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
      )
      (input_layernorm): GraniteSWARMSNorm((2560,), eps=1e-05)
      (post_attention_layernorm): GraniteSWARMSNorm((2560,), eps=1e-05)
      (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): GraniteSWARMSNorm((2560,), 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_layers24
model.hidden_size2560
model.num_heads20
model.num_kv_heads4
model.head_dim128
model.qk_head_dim128
model.vocab_size100352
model.intermediate_size8192
architecturesGraniteSWAForCausalLM
num_hidden_layers24
hidden_size2560
intermediate_size8192
num_attention_heads20
num_key_value_heads4
vocab_size100352
max_position_embeddings8192
sliding_window128
hidden_actsilu
rms_norm_eps1e-05
tie_word_embeddingsTrue
embedding_multiplier12
residual_multiplier0.28
attention_multiplier0.0078125
logits_scaling10
attention_biasFalse

The Notes

03 · what to know before you run

The block, in order

x      = embed_tokens(ids) * embedding_multiplier
h      = x + self_attn(input_layernorm(x)) * residual_multiplier
out    = h + mlp(post_attention_layernorm(h)) * residual_multiplier
logits = lm_head(norm(out)) / logits_scaling

Granite's block, names and multipliers, with its own attention. On granite-swash-2b embedding_multiplier is 12, residual_multiplier 0.28, attention_multiplier 1/128 and logits_scaling 10. Of its 24 blocks, 7 attend over the whole prefix (0, 3, 7, 11, 15, 19 and 23) and 17 over the last 128 tokens; config.layer_types says which, and layers[i].self_attn._module.sliding_window is 128 on a sliding block and None on a full one. Every block applies rotary embeddings with base 10000.

The sink scales the head outputs, not the pattern

Each head has a learned sink, self_attn._module.sinks. The attention takes the softmax over the real keys, mixes the values with it, and then multiplies each query's head output by sigmoid(logsumexp(scores) - sink): the share of the mass a softmax with the sink as one more key would leave on the real keys. So attention_probabilities rows sum to one, softmax(attention_scores) reproduces them, and attention_head_outputs are read after the scale.

attn = model.layers[1].self_attn
with model.trace(prompt):
    v = attn.attention_values.save()
    scores = attn.attention_scores.save()
    probs = attn.attention_probabilities.save()
    heads = attn.attention_head_outputs.save()     # [batch, seq, num_heads, head_dim]

share = torch.sigmoid(scores.logsumexp(-1) - attn._module.sinks.view(1, -1, 1))
v = v.repeat_interleave(model.num_heads // model.num_kv_heads, dim=1)
torch.testing.assert_close(scores.softmax(-1), probs)
torch.testing.assert_close(((probs @ v) * share[..., None]).transpose(1, 2), heads)

The share is small on most blocks. On a code prompt its median over heads and queries is 0.68 in block 0 and between 0.05 and 0.37 in every other block, and its largest value in a block is 0.67 to 1.00: most queries of most heads write a fraction of what the pattern alone would give. A pattern read as "how much this head moves from token j" is probs * share[..., None]. An edit to the pattern does not change the share, which is computed from the scores; to set what a head writes, edit attention_head_outputs.

Eager is the default

transformers' Granite SWA has no sdpa path, so a default load runs the eager attention and the six interior values are available without attn_implementation="eager".

The query scale is config.attention_multiplier, 1/128, which is 1 / head_dim (Llama's would be 128 ** -0.5). The 2B has 20 query heads over 4 key/value heads of 128: an edit to key/value head j reaches query heads 5j to 5j + 4.

The contributions are scaled copies of the modules' outputs

attention_output is self_attn.output[0] * residual_multiplier and mlp_output is mlp.output * residual_multiplier, so the plain identity holds, bit for bit in bfloat16 on 2B:

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn_raw = model.layers[1].self_attn.output[0].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()

m = model.config.residual_multiplier                  # 0.28
torch.testing.assert_close(attn, attn_raw * m)
torch.testing.assert_close(x + attn + mlp, out)

A write to a contribution reaches the model as the whole edited copy divided by residual_multiplier, which the block multiplies again. With 0.28 that round trip is exact in bfloat16: every product x * 0.28 of a bfloat16 x comes back from / 0.28 * 0.28 unchanged, so the positions you did not edit come back bit-identical. Adding a unit vector to mlp_output or attention_output at the last position of block 1, 12 or 23 left every earlier position's logits unchanged. Editing mlp.output[:, -1] += v / m has the same effect.

Logits are divided by 10, and the embeddings multiplied by 12

model.logits is model.lm_head.output / logits_scaling exactly, and project_on_vocab divides too, so a logit lens at the last block equals logits. A softmax over lm_head.output is at temperature 1/10: after "The capital of the United Kingdom is", London has probability 0.79 from logits and 1.00 from lm_head.output. token_embeddings is the plain lookup, and layers[0].input is it times 12; to set what block 0 reads, edit layers[0].input.

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)

lm_head and embed_tokens are one parameter, so an edit to embed_tokens.weight is an edit to the unembedding.

The tokenizer

The vocabulary is Granite 4's, 100352 tokens. No beginning-of-sequence token is prepended. " Paris" (12366), "Paris" (60704) and " London" (7295) are single tokens.

What this family module covers

Every checkpoint whose config says granite_swa: granite-swash-2b. The module also serves a block without rotary (layer_rope_theta 0 on that block), which no released checkpoint uses. granitemoe_swa is the Swash mixture, the same attention in a block whose MLP is a mixture of experts.

From the family module

04 · nnterp/families/granite_swa.py

Granite SWA (GraniteSWAForCausalLM).

Granite's tree, block and multipliers (granite.py: scaled residual adds, embedding_multiplier, logits_scaling), with layer_types mixing sliding-window and full-attention blocks, a rotary embedding per block (or none, NoPE, where the model passes no position embeddings), and a learned attention sink per head. The sink does not join the softmax as a column (as on GPT-OSS): the eager forward takes the softmax over the real keys, mixes the values with it, and then scales each query's head output by sigmoid(logsumexp(scores) - sink), the share of the mass a softmax with the sink as one more logit would leave on the real keys. So attention_probabilities is that softmax (F_dropout_0), whose rows sum to one, attention_scores its input, and attention_head_outputs the interface's output, after the sink scale. The eager forward spells its softmax and dropout F.softmax / F.dropout, so both are redefined on those operations.