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, read from the published configuration.

AssAsssAsssAsssAsssAsssAsssA

Implementing it

How to read this sheet: method.

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

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)             # every layer; see notes
    k, v = repeat_kv(k, v, 5)            # 20 query heads share 4 key/value heads
    # win = 128 on the 20 sliding layers (query t sees keys t-127 .. t),
    # unbounded on the other 8
    s   = q @ k.T * 0.015625 + 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.26        # <- residual_multiplier

    h = rmsnorm(x, layers[i].post_attention_layernorm)
    moe = layers[i].block_sparse_moe
    r = (h @ moe.router.weight.T).float()          # [T, 48]
    r_top, sel = topk(r, 4)                        # tie order unspecified
    w = softmax(r_top)                             # over the 4 selected, not all 48
    y = 0
    for (e, w_e) in zip(sel, w):                   # summation order: see notes
        g, u = split(h @ moe.experts.gate_up_proj[e].T, 2)     # [48, 1024, 1280]
        y += w_e * ((silu(g) * u) @ moe.experts.down_proj[e].T)  # [48, 1280, 512]
    g, u = split(h @ layers[i].shared_mlp.input_linear.T, 2)    # shared expert, width 1280
    y += (silu(g) * u) @ layers[i].shared_mlp.output_linear.T   # every token, not routed
    x = x + y * 0.26

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

Notes

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

Attention sink. self_attn.sinks ([20] per layer) is not a configuration key: GraniteMoeSWAAttention always creates it (L237) and scales each head's output by sigmoid(lse - sink), where lse is the log-sum-exp of that head's masked scores (L193–L201). 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_granitemoe_swa.py L127–L128).

Routing. The router logits are cast to float32, the top 4 are taken with topk, and the gates are the softmax of those 4 logits only, not the top 4 of a softmax over all 48 experts (L65–L67). PyTorch does not guarantee which index topk returns for tied values (torch.topk). The eager expert loop visits experts in ascending index and accumulates with index_add_ (L96–L106); the class is wrapped in use_experts_implementation (L71), so other expert backends may sum in another order. The shared MLP is added to the routed sum for every token (L396–L400).

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-3b-a600m: 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.