GraniteSWAForCausalLM

GraniteSWAForCausalLM at 24 layers and hidden size 2560. One published checkpoint has this shape.

Layer map

One character per layer, input on the left, read from the published configuration.

AssAsssAsssAsssAsssAsssA

Implementing it

How to read this sheet: method.

Generated from granite-swash-2b. Library behaviour below is sourced to transformers v5.15.0, the first release with this class.

x = embed[ids] * 12                    # [T, 2560]   <- embedding_multiplier

for i in 0 .. 23:
    h = rmsnorm(x, layers[i].input_layernorm, eps=1e-05)
    q = h @ layers[i].self_attn.q_proj.T       # [T, 20*128]
    k = h @ layers[i].self_attn.k_proj.T       # [T, 4*128]
    v = h @ layers[i].self_attn.v_proj.T       # [T, 4*128]
    q, k = rope(q, k, theta=10000)             # every layer; see notes
    k, v = repeat_kv(k, v, 5)            # 20 query heads share 4 key/value heads
    # win = 128 on the 17 sliding layers (query t sees keys t-127 .. t),
    # unbounded on the other 7
    s   = q @ k.T * 0.0078125 + mask(causal, window=win)  # <- attention_multiplier
    lse = logsumexp(s)                                    # per query, per head
    a = (softmax(s) @ v) * sigmoid(lse - layers[i].self_attn.sinks[head])
    #  ^ one learned scalar per query head, on every layer, including
    #    the unwindowed ones. See notes.
    a = a @ layers[i].self_attn.o_proj.T
    x = x + a * 0.28        # <- residual_multiplier

    h = rmsnorm(x, layers[i].post_attention_layernorm)
    g = silu(h @ layers[i].mlp.gate_proj.T)           # [T, 8192]
    u =       h @ layers[i].mlp.up_proj.T
    y = (g * u) @ layers[i].mlp.down_proj.T
    x = x + y * 0.28

x = rmsnorm(x, model.norm)
logits = (x @ embed.T) / 10        # tied to the input embedding; <- logits_scaling

Notes

head_dim is not published; 2560 / 20 = 128 is used.

Attention sink. self_attn.sinks ([20] per layer) is not a configuration key: GraniteSWAAttention always creates it (L153) and scales each head's output by sigmoid(lse - sink), where lse is the log-sum-exp of that head's masked scores (L109–L117). Since exp(lse) / (exp(lse) + exp(sink)) = sigmoid(lse - sink), this equals a softmax with one extra logit sink whose value vector is zero, so the weights on real keys sum to less than 1.

Window and positions. A sliding layer keeps key j for query t when j > t - 128 (masking_utils.py L98–L99). The class supports a per-layer rope base, with 0 meaning no positions, but this config sets no layer_rope_theta, so every layer uses the global base of 10,000 (configuration_granite_swa.py L115–L116).

Multiplier placement: embedding, attention scale, residual, logits.

Geometry

Weight structure

The tensors one element of the repeating stack holds, by the names the published checkpoint uses.

How this was checked

Status and evidence grade in the assessment (see what the grades mean):

• granite-swash-2b: runs, checkpoint-parity — logits compared against this model's own reference implementation

Parts of the author's unpublished implementation that exist for this shape:

• a decoder implemented from the published configuration
• a batched execution path over the same arithmetic
• a checkpoint loader, which is where published tensor names are read
• an independent reference implementation of this architecture in Python, driving the published modelling code — an executable statement of what the model should compute, written against the publication rather than against any one implementation of it

Checkpoints with this architecture

Recorded pass: passes.