GraniteMoeHybridForCausalLM

GraniteMoeHybridForCausalLM at 40 layers and hidden size 1536. 6 published checkpoints share this shape.

Layer map

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

mmmmmAmmmmmmmmmAmmmmmmmmmAmmmmmmmmmAmmmm

Implementing it

How to read this sheet: method.

From granite-4.0-h-tiny

Also covers granite-4.0-h-tiny-base.

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

for i in 0 .. 39:
    # layer_types[i] decides the mixer; both share the block shape
    h = rmsnorm(x, layers[i].input_layernorm, eps=1e-05)

    if layer_types[i] == "mamba":
        zxbcdt = h @ layers[i].mamba.in_proj.T
        z, xBC, dt = split(zxbcdt, [3072, 3328, 48])
        xBC = silu(causal_conv1d(xBC, layers[i].mamba.conv1d, width=4, bias=layers[i].mamba.conv1d.bias))
        u, B, C = split(xBC, [3072, 128, 128])
        dt = softplus(dt + layers[i].mamba.dt_bias)
        A  = -exp(layers[i].mamba.A_log)
        dA = exp(dt * A)                   # per head, in (0,1]
        S  = dA * S + dt * outer(u, B)
        y  = S @ C + layers[i].mamba.D * u   # per token, in order
        #  ^ ssd_scan, written out: one state [64 x 128] per head,
        #    a scalar decay per head, and one B/C pair shared across
        #    every head in a group the way GQA shares a KV head.
        y  = group_rmsnorm(y * silu(z), layers[i].mamba.norm,
                           groups=1, width=3072, eps=1e-5)
        #  ^ gate first, then normalise (see notes).
        #    `mamba_chunk_size: 256` is the block length transformers
        #    uses to evaluate this same recurrence in chunks.
        y  = y @ layers[i].mamba.out_proj.T
    else:
        q = h @ layers[i].self_attn.q_proj.T       # [T, 12*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]
        # no rope: position_embedding_type is "nope" (see notes)
        k, v = repeat_kv(k, v, 3)            # 12 query heads share 4 key/value heads
        a = softmax(q @ k.T * 0.0078125, mask=causal) @ v      # <- attention_multiplier
        a = a @ layers[i].self_attn.o_proj.T

    x = x + (y if layer_types[i] == "mamba" else a) * 0.22

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

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

Differs from the block above

Only the expert count differs. Everything else is as above.

From granite-4.0-h-1b

Also covers granite-4.0-h-1b-base.

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

for i in 0 .. 39:
    # layer_types[i] decides the mixer; both share the block shape
    h = rmsnorm(x, layers[i].input_layernorm, eps=1e-05)

    if layer_types[i] == "mamba":
        zxbcdt = h @ layers[i].mamba.in_proj.T
        z, xBC, dt = split(zxbcdt, [3072, 3328, 48])
        xBC = silu(causal_conv1d(xBC, layers[i].mamba.conv1d, width=4, bias=layers[i].mamba.conv1d.bias))
        u, B, C = split(xBC, [3072, 128, 128])
        dt = softplus(dt + layers[i].mamba.dt_bias)
        A  = -exp(layers[i].mamba.A_log)
        dA = exp(dt * A)                   # per head, in (0,1]
        S  = dA * S + dt * outer(u, B)
        y  = S @ C + layers[i].mamba.D * u   # per token, in order
        #  ^ ssd_scan, written out: one state [64 x 128] per head,
        #    a scalar decay per head, and one B/C pair shared across
        #    every head in a group the way GQA shares a KV head.
        y  = group_rmsnorm(y * silu(z), layers[i].mamba.norm,
                           groups=1, width=3072, eps=1e-5)
        #  ^ gate first, then normalise (see notes).
        #    `mamba_chunk_size: 256` is the block length transformers
        #    uses to evaluate this same recurrence in chunks.
        y  = y @ layers[i].mamba.out_proj.T
    else:
        q = h @ layers[i].self_attn.q_proj.T       # [T, 12*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]
        # no rope: position_embedding_type is "nope" (see notes)
        k, v = repeat_kv(k, v, 3)            # 12 query heads share 4 key/value heads
        a = softmax(q @ k.T * 0.0078125, mask=causal) @ v      # <- attention_multiplier
        a = a @ layers[i].self_attn.o_proj.T

    x = x + (y if layer_types[i] == "mamba" else a) * 0.22

    h = rmsnorm(x, layers[i].post_attention_layernorm)
    g, u = split(h @ layers[i].shared_mlp.input_linear.T, 2)   # width 4096
    y = (silu(g) * u) @ layers[i].shared_mlp.output_linear.T
    # num_local_experts: 0, so no router and no experts (see notes)
    x = x + y * 0.22

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

num_local_experts: 0: transformers builds no block_sparse_moe for such a layer (L1116–L1117) and the feed-forward is shared_mlp alone (L1202–L1207).

Notes

head_dim is not published; 1536 / 12 = 128 is used.

State-space layer. Split sizes and order follow modeling_granitemoehybrid.py at v4.57.1: the mixer width is mamba_expand × hidden_size (L409); in_proj splits into the gate z, the convolution input and one dt per head (L651–L653); after the convolution that input splits into x, B and C (L692–L696); norm multiplies by silu(z) and then normalises, with eps = rms_norm_eps (L874–L883, L416). Which layers are state-space comes from layer_types alone (L1125–L1128); a config with no layer_types makes every layer state-space (configuration_granitemoehybrid.py L252–L253).

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

position_embedding_type: "nope": transformers builds no rotary embedding for this value (L1306) and the attention layers then skip rope (L197–L199). The config still carries a rope_theta, which is unused. The recorded pass logs rope because the step is entered, as annotated in its tree.

Routing (the granite-4.0-h-tiny block and the preview row under it). The router logits are cast to float32, the top k are taken with topk, and the gates are the softmax of those k logits only, not the top k of a softmax over all experts (L1024–L1026). PyTorch does not guarantee which index topk returns for tied values (torch.topk). transformers sorts the token–expert pairs by ascending expert index (L1040) and adds the gated expert outputs back with one index_add (L1105); summation order beyond that is not part of the model's definition. The shared MLP is added to the routed sum for every token (L1205).

Geometry

Weight structure

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

The three block_sparse_moe tensors exist only on the members with routed experts; granite-4.0-h-1b and granite-4.0-h-1b-base do not carry them.

Other shapes of GraniteMoeHybridForCausalLM

Only the columns that differ are shown; a cell with several values means the checkpoints of that shape disagree.

How this was checked

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

• each of the 6 checkpoints of this shape: runs, shape-parity — a stored checkpoint of the same shape was compared against a reference

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

• a decoder implemented from the published configuration
• a state-space mixer implemented from the published configuration
• 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.