Gate and project

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 › First layer › Kimi Delta Attention (KDA) › Gate and project

Merge heads

×

no parametersrearrange(o, 'b t h d -> b t (h d)') · 658

Concatenate the 96 heads into 12,288 features.

[B, 96, T, 128][B,\,96,\,T,\,128]
→  [B, T, 96, 128]\to\;[B,\,T,\,96,\,128]
→  [B, T, 12288]\to\;[B,\,T,\,12288]
in
ogo^{g} [B batch, 96 heads, T positions, 128 d_h]
out
oo [B batch, T positions, 12288 heads × d_h]
Explanation

Heads are concatenated in order: head ii fills features 128i,…,128i+127{128i},\dots,{128i+127}.

Codemodeling_kimi_linear.py 651–659
modeling_kimi_linear.py 651–659
        if self.use_full_rank_gate:
            g = self.g_proj(hidden_states)
        else:
            g = self.g_b_proj(self.g_a_proj(hidden_states))
        g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim)
        o = self.o_norm(o, g)

        o = rearrange(o, 'b t h d -> b t (h d)')
        o = self.o_proj(o)
Sources1

Direct evidence

The kernel’s output passes a gated RMSNorm, whose gate is the sigmoid of a full-rank projection of the layer’s input, and an output projection.

kda-output · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 651–659

Read code excerpt
        if self.use_full_rank_gate:
            g = self.g_proj(hidden_states)
        else:
            g = self.g_b_proj(self.g_a_proj(hidden_states))
        g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim)
        o = self.o_norm(o, g)

        o = rearrange(o, 'b t h d -> b t (h d)')
        o = self.o_proj(o)