afmoe¶
AFMoE (AfmoeForCausalLM, Arcee's Trinity).
Llama's tree with a sandwich block, every sublayer normed before and after::
h = x + post_attention_layernorm(self_attn(input_layernorm(x)))
out = h + post_mlp_layernorm(mlp(pre_mlp_layernorm(h)))
What the block adds is each post-norm's output, so attention_output points at
post_attention_layernorm (which follows the attention, as on Gemma-2) and
mlp_output at post_mlp_layernorm; pre_mlp_layernorm keeps its native
name and mlp.input is its output. The attention runs the shared eager forward
with q/k norms, the rotary only on its sliding-window (local) blocks, and a
sigmoid gate (gate_proj) on the interface's output before o_proj, so
attention_head_outputs are ungated. The first num_dense_layers blocks have
a dense AfmoeMLP, the rest a mixture (AfmoeSparseMoeBlock) returning its
shared expert's output plus the routed experts' sum as one tensor. The shared
expert (mlp.shared_experts) is an AfmoeMLP too, so it is an Mlp; its
mlp_output is unavailable, since the block adds the mixture's post-normed sum.