Gated MLA

Checkpoint configuration and implementation moonshotai/Kimi-K3 @ f831ab6, fla-org/flash-linear-attention @ 07ca1e4 · parameter count checked against the checkpoint’s tensor headers · text model only · How to read

93 layers › 24 MLA layers › Gated MLA

Output gate

×

no parametersattn_output * g · 471–472

Scale every merged feature by a sigmoid gate computed from the layer’s input.

Why

The gate lets each token modulate the channels it reads from global attention. Its projection is full rank, matching the output gate Kimi K3 gives its delta attention. Kimi K3 report §2.1.2

o=σ(g)⊙o~{\htmlData{light=gate}{o}}=\sigma({\htmlData{light=gate-projection}{g}})\odot{\htmlData{light=merge}{\tilde o}}
in
o~\tilde o [B batch, T positions, 12288 heads × value d_h]gg [B batch, T positions, 12288 heads × value d_h]
out
oo [B batch, T positions, 12288 heads × value d_h]
Codemodeling_kimi_linear.py 468–474
modeling_kimi_linear.py 468–474
        attn_output = attn_output.reshape(
            batch_size, seq_length, -1).contiguous()
        if self.use_output_gate:
            g = self.g_proj(hidden_states).sigmoid()
            attn_output = attn_output * g
        attn_output = self.o_proj(attn_output)
        return attn_output
Sources1

Direct evidence

Head outputs are flattened, multiplied by the sigmoid of a full-rank projection of the layer’s input, and projected back to the model width.

mla-output · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 468–474

Read code excerpt
        attn_output = attn_output.reshape(
            batch_size, seq_length, -1).contiguous()
        if self.use_output_gate:
            g = self.g_proj(hidden_states).sigmoid()
            attn_output = attn_output * g
        attn_output = self.o_proj(attn_output)
        return attn_output