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.

mmmmmAmmmmmmmmmAmmmmmmmmmAmmmmmmmmmAmmmm

The pass (full forward)

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

21 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
│  ├─ delta_mix
│  ├─ rmsnorm_mlp
│  ├─ moe
│  │  ├─ router
│  │  └─ expert
│  ├─ w_gate
│  ├─ w_up
│  ├─ w_down
│  ├─ wq
│  ├─ wk
│  ├─ wv
│  ├─ rope
│  ├─ attn
│  │  ├─ scores
│  │  └─ attn_mix
│  └─ wo
└─ head

Recorded separately from the structure: layer fired 40 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.

delta_mix is the state-space layer's token-mixing step — the recurrence written out as ssd_scan in the block below. It fires only in the state-space layers; in the attention layers the steps from wq to wo take its place.

Which checkpoint the pass comes from. The recording has a moe step with a router and an expert inside it, so it was made from one of the members with routed experts — granite-4.0-h-tiny, granite-4.0-h-tiny-base or one of the two tiny previews; the assessment does not record which. granite-4.0-h-1b and granite-4.0-h-1b-base publish num_local_experts: 0: they have no router and route nothing, so their forward has no router or expert step, only the always-on feed-forward written out in their block below.

The recorded pass fires rope and these models have no positions. Both are true. position_embedding_type: "nope" is what the configuration says, and the trace records that the rotary step was entered, not that it rotated anything — the step is logged whether or not it does any work for this configuration. A reader comparing the two sections would otherwise have to guess which one is wrong, and neither is.

The same thing in canonical form:

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

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.

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)
        S  = dA * S + dt * outer(u, B)     # dA = exp(dt * A), in (0,1]
        y  = S @ C + 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. Normalising before the gate is a different
        #    function and reads exactly as fluent.
        #    (`mamba_chunk_size: 256` is published and is not used
        #     above: the recurrence is sequential. The chunked kernel
        #     is a throughput variant with its own rounding.)
        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 position signal: position_embedding_type is "nope".
        # Not rotary_dim == 0 — a flag, and the only thing that says so.
        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)
    logits_r = h @ block_sparse_moe.router.layer.T      # [64], kept f32
    sel = argtop6(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 1024
    y += (silu(g) * u) @ shared_mlp.output_linear.T     # always on, not routed
    x = x + y * 0.22

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

From granite-4.0-tiny-base-preview

Also covers granite-4.0-tiny-preview.

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)
        S  = dA * S + dt * outer(u, B)     # dA = exp(dt * A), in (0,1]
        y  = S @ C + 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. Normalising before the gate is a different
        #    function and reads exactly as fluent.
        #    (`mamba_chunk_size: 256` is published and is not used
        #     above: the recurrence is sequential. The chunked kernel
        #     is a throughput variant with its own rounding.)
        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 position signal: position_embedding_type is "nope".
        # Not rotary_dim == 0 — a flag, and the only thing that says so.
        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)
    logits_r = h @ block_sparse_moe.router.layer.T      # [62], kept f32
    sel = argtop6(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 1024
    y += (silu(g) * u) @ shared_mlp.output_linear.T     # always on, not routed
    x = x + y * 0.22

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

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)
        S  = dA * S + dt * outer(u, B)     # dA = exp(dt * A), in (0,1]
        y  = S @ C + 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. Normalising before the gate is a different
        #    function and reads exactly as fluent.
        #    (`mamba_chunk_size: 256` is published and is not used
        #     above: the recurrence is sequential. The chunked kernel
        #     is a throughput variant with its own rounding.)
        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 position signal: position_embedding_type is "nope".
        # Not rotary_dim == 0 — a flag, and the only thing that says so.
        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 @ shared_mlp.input_linear.T, 2)       # width 4096
    y = (silu(g) * u) @ shared_mlp.output_linear.T
    # `num_local_experts: 0` — the expert names are the template's,
    # there is no router here and nothing is routed.
    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 on a class whose tensors are named for a mixture of experts. The MLP is the shared expert alone; nothing is routed.

Notes that apply to more than one block

Stated once here rather than under each block above.

head_dim is not published; 1536 / 12 = 128 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.

position_embedding_type: "nope" — the attention layers rotate nothing. It is a flag and not a rotary width of zero, and it is the only thing in the configuration that says so. Position is carried by the state-space layers instead, which is why the handful of attention layers can do without it.

For the blocks from granite-4.0-h-tiny and granite-4.0-tiny-base-preview: 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.

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

implemented, evidence grade shape-parity, per the assessment, for each checkpoint of this shape: a converted checkpoint of this shape was compared against a reference recording; whether it holds this checkpoint's weights is not established. 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 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

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