GraniteMoeSWAForCausalLM

GraniteMoeSWAForCausalLM at 28 layers and hidden size 1280. One published checkpoint has this shape.

Layer map

One character per layer, input on the left.

AssAsssAsssAsssAsssAsssAsssA

The pass (full forward)

The order of operations for one forward. ×N marks a position that fires once per layer.

20 steps across nested levels, recorded from one complete forward of this shape in the author's independent implementation — the structure is what that run emitted, not a reading of the configuration.

forward
├─ embed
├─ layer ×N
│  ├─ rmsnorm_attn
│  ├─ wq
│  ├─ wk
│  ├─ wv
│  ├─ rope
│  ├─ attn
│  │  ├─ scores
│  │  └─ attn_mix
│  ├─ wo
│  ├─ rmsnorm_mlp
│  ├─ moe
│  │  ├─ router
│  │  └─ expert
│  ├─ w_gate
│  ├─ w_up
│  └─ w_down
└─ head

Recorded separately from the structure: layer fired 28 times. A ×N run says only that it repeated — how long a run is is not part of the structure, so two models differing only in depth have the same pass.

The same thing in canonical form:

0(1 2*N(3 4 5 6 7 8(9 10) 11 12 13(14 15) 16 17 18) 19)

Numbers are positions in the pass, not layer indices. Two passes count as the same structure when these strings match.

Implementing it

One forward, written out, per checkpoint whose arithmetic actually differs. A shape groups checkpoints by class, depth, width and layer stack — none of which decides an activation, a rope base or a routing rule — so where the members of this shape disagree there is a block each, and each names the checkpoint it was generated from.

It is not the recorded pass above, which is captured from a real forward; it is what the published configuration says the arithmetic is. Where the configuration does not say, the line says that instead of guessing.

The names are the ones the published checkpoint uses where a name is published, and canonical otherwise. A checkpoint loader may store them differently — fusing a gate/up pair into one matrix, or splitting one published projection in two — and those are that loader's names, not the model's. Everything spelled here is what the download contains.

Generated from granite-swash-3b-a600m.

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

for i in 0 .. 27:
    h = rmsnorm(x, layers[i].input_layernorm, eps=1e-05)
    q = h @ layers[i].self_attn.q_proj.T       # [T, 20*64]
    k = h @ layers[i].self_attn.k_proj.T       # [T, 4*64]
    v = h @ layers[i].self_attn.v_proj.T       # [T, 4*64]
    q, k = rope(q, k, theta=10000)
    k, v = repeat_kv(k, v, 5)            # 20 query heads share 4 key/value heads
    # win = 128 on the 20 sliding layers, unbounded on the other 8
    p, lse = softmax_with_lse(q @ k.T * 0.015625, mask=causal, window=win)
    a = (p @ v) * sigmoid(lse - layers[i].self_attn.sinks[head])
    #  ^ one learned scalar per query head, on every layer — the
    #    config never mentions it. Equivalent to an extra softmax
    #    key with logit `sink` and value zero, so a head with
    #    nothing worth attending to emits near zero instead of
    #    averaging its window.
    a = a @ layers[i].self_attn.o_proj.T
    x = x + a * 0.26        # <- residual_multiplier

    h = rmsnorm(x, layers[i].post_attention_layernorm)
    logits_r = h @ block_sparse_moe.router.layer.T      # [48], kept f32
    sel = argtop4(logits_r)                        # ties -> lower expert index
    w   = softmax(logits_r[sel])                     # over the selected k, not all
    y = 0
    for (e, w_e) in sorted(zip(sel, w), by=e):        # ascending expert index
        g, u = split(h @ block_sparse_moe.input_linear[e].T, 2)
        y += w_e * ((silu(g) * u) @ block_sparse_moe.output_linear[e].T)
    g, u = split(h @ shared_mlp.input_linear.T, 2)       # shared expert, width 1280
    y += (silu(g) * u) @ shared_mlp.output_linear.T     # always on, not routed
    x = x + y * 0.26

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

head_dim is not published; 1280 / 20 = 64 is used.

The Granite multipliers are the part with no Llama analogue: embedding_multiplier, attention_multiplier, residual_multiplier, logits_scaling. They are ordinary floats and a port that ignores them still produces fluent text, which is what makes them worth printing where they act.

Two rules in the routing are load-bearing and neither is in the configuration. The gates are softmax of the top-k logits, not the top-k of a full softmax — the first normalises over the experts that ran, the second over all of them and then discards most of the mass, which scales the block down by whatever the losing experts held. And the selected experts are summed in ascending expert index, because floating-point addition is not associative and the reference accumulates in that order; summing in router-rank order differs by rounding at every layer.

This checkpoint carries a learned attention sink per head that the configuration never mentions. self_attn.sinks is [20] on every layer — including the unwindowed ones — and scales each head's output by sigmoid(lse - sink), where lse is that head's log-sum-exp. Read as one expression it is an extra key appended to the softmax whose logit is the sink and whose value is zero, so a head with nothing worth attending to can output near zero instead of being forced to average its window. Nothing above is derived from it because nothing in config.json says it is there; it is in the weights.

Geometry

Weight structure

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

How this was checked

implemented, evidence grade checkpoint-parity, per the assessment: this checkpoint's logits were compared against a recording made with the model's own published reference implementation. The implementation meant here and below is an unpublished independent inference implementation by the author.

What stands behind the block above, beyond the published configuration it is read from:

• 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

Where an independent implementation and the published configuration disagree, the configuration is what this page reports and the disagreement is what it says.

Checkpoints with this architecture