Kimi K3

All models, with previews →

Moonshot AI · Release date not pinnedNew parts in the block · after Kimi K2 in the atlas

KDA layer

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

Full sequencePrefillDecode

93 layers › 68 KDA layers

Kimi Delta Attention (KDA)

Open inside →×

learnedself.self_attn · KimiDeltaAttention

Each of 96 heads keeps a 128 × 128 memory that every position updates by a delta rule and reads with its query.

Why

Linear. Along the sequence, the authors want an efficient mechanism for mixing tokens over a long context, and keep global attention only in some layers. Kimi K3 report §2

This is the linear part of that hybrid: a delta-rule recurrence with a forget gate for every key channel. Kimi K3 report §2.1.1

Fixed state. It replaces the growing key–value cache of softmax attention with a recurrent state of fixed size, which is cheap to transfer and reuse. Kimi K3 report §5.1

Hybrid. Three of these layers are followed by one MLA layer, a 3:1 ratio repeated through the stack, and one more MLA layer closes it so that the last layer always attends globally. Kimi K3 report §2.1

Δattn=oWO\Delta_{\mathrm{attn}}=oW_{O}
All 25 steps
Query heads:\text{Query heads:}
qp=xattnWQ{\htmlClass{vl-stream-q}{q^{\mathrm p}}}=x_{\mathrm{attn}}W_{{\htmlClass{vl-stream-q}{Q}}}
qtc=SiLU⁡ ⁣(∑j=03wj⊙qt−jp){\htmlClass{vl-stream-q}{q^{\mathrm c}_t}}=\operatorname{SiLU}\!\Big(\sum_{j=0}^{3}w_j\odot {\htmlClass{vl-stream-q}{q^{\mathrm p}_{t-j}}}\Big)
q^=qh/∥qh∥2\hat {\htmlClass{vl-stream-q}{q}}={\htmlClass{vl-stream-q}{q^{\mathrm h}}}/\lVert {\htmlClass{vl-stream-q}{q^{\mathrm h}}}\rVert_2
Key heads:\text{Key heads:}
kp=xattnWK{\htmlClass{vl-stream-k}{k^{\mathrm p}}}=x_{\mathrm{attn}}W_{{\htmlClass{vl-stream-k}{K}}}
ktc=SiLU⁡ ⁣(∑j=03wj⊙kt−jp){\htmlClass{vl-stream-k}{k^{\mathrm c}_t}}=\operatorname{SiLU}\!\Big(\sum_{j=0}^{3}w_j\odot {\htmlClass{vl-stream-k}{k^{\mathrm p}_{t-j}}}\Big)
k^=kh/∥kh∥2\hat {\htmlClass{vl-stream-k}{k}}={\htmlClass{vl-stream-k}{k^{\mathrm h}}}/\lVert {\htmlClass{vl-stream-k}{k^{\mathrm h}}}\rVert_2
Value heads:\text{Value heads:}
vp=xattnWV{\htmlClass{vl-stream-v}{v^{\mathrm p}}}=x_{\mathrm{attn}}W_{{\htmlClass{vl-stream-v}{V}}}
vtc=SiLU⁡ ⁣(∑j=03wj⊙vt−jp){\htmlClass{vl-stream-v}{v^{\mathrm c}_t}}=\operatorname{SiLU}\!\Big(\sum_{j=0}^{3}w_j\odot {\htmlClass{vl-stream-v}{v^{\mathrm p}_{t-j}}}\Big)
Retention gate:\text{Retention gate:}
z↓=xattnWα↓z^{\downarrow}=x_{\mathrm{attn}}W^{\downarrow}_{\alpha}
z=z↓Wα↑z=z^{\downarrow}W^{\uparrow}_{\alpha}
α=exp⁡ ⁣(−5 σ(eA(z+bα)))\alpha=\exp\!\big(-5\,\sigma\big(e^{A}(z+b_{\alpha})\big)\big)
Write gate:\text{Write gate:}
w=xattnWβw=x_{\mathrm{attn}}W_{\beta}
β=σ(w)\beta=\sigma(w)
Delta-rule memory:\text{Delta-rule memory:}
S←Diag⁡(α) SS\leftarrow\operatorname{Diag}(\alpha)\,S
S←S+β k^ (v−S⊤k^)⊤S\leftarrow S+\beta\,\hat {\htmlClass{vl-stream-k}{k}}\,({\htmlClass{vl-stream-v}{v}}-S^{\top}\hat {\htmlClass{vl-stream-k}{k}})^{\top}
o~=S⊤q^/128\tilde o=S^{\top}\hat {\htmlClass{vl-stream-q}{q}}/\sqrt{128}
g=xattnWgg=x_{\mathrm{attn}}W_{g}
Gate and project:\text{Gate and project:}
og=σ(g)⊙RMSNorm⁡(o~)o^{g}=\sigma(g)\odot\operatorname{RMSNorm}(\tilde o)
Δattn=oWO\Delta_{\mathrm{attn}}=oW_{O}
in
xattnx_{\mathrm{attn}} [B batch, 1 positions, 7168 d]
out
Δattn\Delta_{\mathrm{attn}} [B batch, 1 positions, 7168 d]
parameters
443,740,416 · × 68 KDA layers = 30,174,348,288

The checkpoint stores the per-head log-scale AA with 128 elements per layer; the pinned code creates it with 96, one per head (modeling_kimi_linear.py 520–521). The atlas counts the checkpoint’s.

Parameters

443,740,416 parameters per instance, across its learned components

Exact per-instance accounting
ComponentCalculationParameters
Query heads88,129,536 across its components88,129,536
Key heads88,129,536 across its components88,129,536
Value heads88,129,536 across its components88,129,536
Retention gate2,502,784 across its components2,502,784
Write gate688,128 across its components688,128
Output gate projection7,168×12,2887{,}168 \times 12{,}28888,080,384
Gate and project88,080,512 across its components88,080,512
Per-instance total443,740,416

Across 68 independent KDA layers: 443,740,416×68=30,174,348,288443{,}740{,}416 \times 68 = 30{,}174{,}348{,}288 parameters.

Derived from the pinned configuration and implementation, with the per-head log-scale at the checkpoint’s 128 elements where the code creates 96; equal to the checkpoint’s own tensors. The quantisation scales are not counted.

Explanation

The memory has a fixed size, so the cost of a position does not grow with the length of the sequence; the 24 MLA layers provide attention over all earlier positions.

Codemodeling_kimi_linear.py 478–659 · naive.py 12–22, 45–66
modeling_kimi_linear.py 478–541
    def __init__(self, config: KimiLinearConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.mode = "chunk"

        self.hidden_size = config.hidden_size
        self.conv_size = config.linear_attn_config["short_conv_kernel_size"]
        self.head_dim = config.linear_attn_config["head_dim"]
        self.num_heads = config.linear_attn_config["num_heads"]
        self.head_k_dim = self.head_dim
        self.num_k_heads = self.num_heads

        self.layer_idx = layer_idx

        assert self.mode in [
            'chunk', 'fused_recurrent'], f"Not supported mode `{self.mode}`."

        projection_k_size = self.head_k_dim * self.num_k_heads
        projection_size = self.head_dim * self.num_heads

        self.q_proj = nn.Linear(
            self.hidden_size, projection_k_size, bias=False)
        self.k_proj = nn.Linear(
            self.hidden_size, projection_k_size, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False)

        self.q_conv1d = ShortConvolution(
            hidden_size=projection_k_size,
            kernel_size=self.conv_size,
            activation='silu',
        )
        self.k_conv1d = ShortConvolution(
            hidden_size=projection_k_size,
            kernel_size=self.conv_size,
            activation='silu',
        )
        self.v_conv1d = ShortConvolution(
            hidden_size=projection_size,
            kernel_size=self.conv_size,
            activation='silu',
        )

        self.A_log = torch.nn.Parameter(torch.log(torch.empty(
            self.num_heads, dtype=torch.float32).uniform_(1, 16)))

        self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
        self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)

        self.dt_bias = nn.Parameter(
            torch.empty(projection_size, dtype=torch.float32))

        self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False)

        self.use_full_rank_gate = config.linear_attn_config.get("use_full_rank_gate", False)
        self.gate_lower_bound = config.linear_attn_config.get("gate_lower_bound", None)
        if self.use_full_rank_gate:
            self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
        else:
            self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
            self.g_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)

        self.o_norm = FusedRMSNormGated(
            self.head_dim, eps=config.rms_norm_eps, activation='sigmoid')
        self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False)
modeling_kimi_linear.py 559–561
        use_cache = cache_params is not None
        batch_size, q_len, _ = hidden_states.shape
        mode = 'fused_recurrent' if use_cache and q_len == 1 else self.mode
modeling_kimi_linear.py 572–649
        conv_state_q, conv_state_k, conv_state_v = None, None, None
        recurrent_state = None
        if cache_params is not None:
            if cache_params.conv_states[self.layer_idx] is not None:
                conv_state_q, conv_state_k, conv_state_v = cache_params.conv_states[
                    self.layer_idx]
            recurrent_state = cache_params.recurrent_states[self.layer_idx]

        q_proj_states = self.q_proj(hidden_states)
        k_proj_states = self.k_proj(hidden_states)
        v_proj_states = self.v_proj(hidden_states)
        q, conv_state_q = self.q_conv1d(
            x=q_proj_states,
            cache=conv_state_q,
            output_final_state=use_cache,
            cu_seqlens=cu_seqlens,
        )
        k, conv_state_k = self.k_conv1d(
            x=k_proj_states,
            cache=conv_state_k,
            output_final_state=use_cache,
            cu_seqlens=cu_seqlens,
        )
        v, conv_state_v = self.v_conv1d(
            x=v_proj_states,
            cache=conv_state_v,
            output_final_state=use_cache,
            cu_seqlens=cu_seqlens,
        )
        g = self.f_b_proj(self.f_a_proj(hidden_states))
        g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim)
        beta = self.b_proj(hidden_states).float()

        q, k = map(lambda x: rearrange(
            x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))
        v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)

        if mode == 'chunk':
            o, recurrent_state = chunk_kda(
                q=q,
                k=k,
                v=v,
                g=g,
                beta=beta,
                A_log=self.A_log,
                dt_bias=self.dt_bias,
                initial_state=recurrent_state,
                output_final_state=True,
                use_qk_l2norm_in_kernel=True,
                use_gate_in_kernel=True,
                use_beta_sigmoid_in_kernel=True,
                safe_gate=self.gate_lower_bound is not None,
                lower_bound=self.gate_lower_bound,
                transpose_state_layout=True,
                cu_seqlens=cu_seqlens,
            )
        else:
            o, recurrent_state = fused_recurrent_kda(
                q=q,
                k=k,
                v=v,
                g=g,
                beta=beta,
                A_log=self.A_log,
                dt_bias=self.dt_bias,
                initial_state=recurrent_state,
                output_final_state=True,
                use_qk_l2norm_in_kernel=True,
                use_gate_in_kernel=True,
                use_beta_sigmoid_in_kernel=True,
                lower_bound=self.gate_lower_bound,
                transpose_state_layout=True,
                cu_seqlens=cu_seqlens,
            )
        if cache_params is not None:
            cache_params.recurrent_states[self.layer_idx] = recurrent_state
            cache_params.conv_states[self.layer_idx] = (
                conv_state_q, conv_state_k, conv_state_v)
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)
naive.py 12–22
def naive_recurrent_kda(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    scale: float | None = None,
    initial_state: torch.Tensor | None = None,
    output_final_state: bool = False,
):
    r"""
naive.py 45–66
    dtype = v.dtype
    B, T, H, K, HV, V = *q.shape, v.shape[2], v.shape[-1]
    G = HV // H
    if scale is None:
        scale = K ** -0.5

    q, k, v, g, beta = map(lambda x: x.to(torch.float), [q, k, v, g, beta])
    q = q.repeat_interleave(G, dim=2) * scale   # [B, T, HV, K]
    k = k.repeat_interleave(G, dim=2)           # [B, T, HV, K]

    S = k.new_zeros(B, HV, K, V).to(q)
    if initial_state is not None:
        S += initial_state
    o = torch.zeros_like(v)
    for i in range(0, T):
        q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i]
        S = S * g_i[..., None].exp()
        S = S + torch.einsum('b h k, b h v -> b h k v', b_i[..., None] * k_i, v_i - (k_i[..., None] * S).sum(-2))
        o[:, i] = torch.einsum('b h k, b h k v -> b h v', q_i, S)
    if not output_final_state:
        S = None
    return o.to(dtype), S
Sources4

Direct evidence

Kimi Delta Attention owns query, key and value projections to 96 heads of 128 features, a causal short convolution on each, a per-head log-scale A_log, a rank-128 pair of projections and a bias for the decay logits, a per-head write-strength projection, a full-rank output gate projection, a gated RMSNorm over each head and an output projection.

kda-init · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 478–541

Read code excerpt
    def __init__(self, config: KimiLinearConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.mode = "chunk"

        self.hidden_size = config.hidden_size
        self.conv_size = config.linear_attn_config["short_conv_kernel_size"]
        self.head_dim = config.linear_attn_config["head_dim"]
        self.num_heads = config.linear_attn_config["num_heads"]
        self.head_k_dim = self.head_dim
        self.num_k_heads = self.num_heads

        self.layer_idx = layer_idx

        assert self.mode in [
            'chunk', 'fused_recurrent'], f"Not supported mode `{self.mode}`."

        projection_k_size = self.head_k_dim * self.num_k_heads
        projection_size = self.head_dim * self.num_heads

        self.q_proj = nn.Linear(
            self.hidden_size, projection_k_size, bias=False)
        self.k_proj = nn.Linear(
            self.hidden_size, projection_k_size, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False)

        self.q_conv1d = ShortConvolution(
            hidden_size=projection_k_size,
            kernel_size=self.conv_size,
            activation='silu',
        )
        self.k_conv1d = ShortConvolution(
            hidden_size=projection_k_size,
            kernel_size=self.conv_size,
            activation='silu',
        )
        self.v_conv1d = ShortConvolution(
            hidden_size=projection_size,
            kernel_size=self.conv_size,
            activation='silu',
        )

        self.A_log = torch.nn.Parameter(torch.log(torch.empty(
            self.num_heads, dtype=torch.float32).uniform_(1, 16)))

        self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
        self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)

        self.dt_bias = nn.Parameter(
            torch.empty(projection_size, dtype=torch.float32))

        self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False)

        self.use_full_rank_gate = config.linear_attn_config.get("use_full_rank_gate", False)
        self.gate_lower_bound = config.linear_attn_config.get("gate_lower_bound", None)
        if self.use_full_rank_gate:
            self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
        else:
            self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
            self.g_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)

        self.o_norm = FusedRMSNormGated(
            self.head_dim, eps=config.rms_norm_eps, activation='sigmoid')
        self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False)

Direct evidence

Queries, keys and values are projected and convolved; the decay logits and the write strength are projected from the layer’s input; the kernel returns the output and the recurrent state, in the recurrent form for one new position with a cache and chunked otherwise; the state and the convolution states are kept per layer.

kda-forward · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 559–649

Read code excerpt
        use_cache = cache_params is not None
        batch_size, q_len, _ = hidden_states.shape
        mode = 'fused_recurrent' if use_cache and q_len == 1 else self.mode

… lines 562–571 omitted …

        conv_state_q, conv_state_k, conv_state_v = None, None, None
        recurrent_state = None
        if cache_params is not None:
            if cache_params.conv_states[self.layer_idx] is not None:
                conv_state_q, conv_state_k, conv_state_v = cache_params.conv_states[
                    self.layer_idx]
            recurrent_state = cache_params.recurrent_states[self.layer_idx]

        q_proj_states = self.q_proj(hidden_states)
        k_proj_states = self.k_proj(hidden_states)
        v_proj_states = self.v_proj(hidden_states)
        q, conv_state_q = self.q_conv1d(
            x=q_proj_states,
            cache=conv_state_q,
            output_final_state=use_cache,
            cu_seqlens=cu_seqlens,
        )
        k, conv_state_k = self.k_conv1d(
            x=k_proj_states,
            cache=conv_state_k,
            output_final_state=use_cache,
            cu_seqlens=cu_seqlens,
        )
        v, conv_state_v = self.v_conv1d(
            x=v_proj_states,
            cache=conv_state_v,
            output_final_state=use_cache,
            cu_seqlens=cu_seqlens,
        )
        g = self.f_b_proj(self.f_a_proj(hidden_states))
        g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim)
        beta = self.b_proj(hidden_states).float()

        q, k = map(lambda x: rearrange(
            x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))
        v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)

        if mode == 'chunk':
            o, recurrent_state = chunk_kda(
                q=q,
                k=k,
                v=v,
                g=g,
                beta=beta,
                A_log=self.A_log,
                dt_bias=self.dt_bias,
                initial_state=recurrent_state,
                output_final_state=True,
                use_qk_l2norm_in_kernel=True,
                use_gate_in_kernel=True,
                use_beta_sigmoid_in_kernel=True,
                safe_gate=self.gate_lower_bound is not None,
                lower_bound=self.gate_lower_bound,
                transpose_state_layout=True,
                cu_seqlens=cu_seqlens,
            )
        else:
            o, recurrent_state = fused_recurrent_kda(
                q=q,
                k=k,
                v=v,
                g=g,
                beta=beta,
                A_log=self.A_log,
                dt_bias=self.dt_bias,
                initial_state=recurrent_state,
                output_final_state=True,
                use_qk_l2norm_in_kernel=True,
                use_gate_in_kernel=True,
                use_beta_sigmoid_in_kernel=True,
                lower_bound=self.gate_lower_bound,
                transpose_state_layout=True,
                cu_seqlens=cu_seqlens,
            )
        if cache_params is not None:
            cache_params.recurrent_states[self.layer_idx] = recurrent_state
            cache_params.conv_states[self.layer_idx] = (
                conv_state_q, conv_state_k, conv_state_v)

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)

Direct evidence

The delta-rule recurrence per head: the state decays channel by channel, takes a write of the key times the difference between the value and what the state already returns for that key, scaled by the write strength, and the output reads the state with the scaled query.

kda-kernel · fla-org/flash-linear-attention / fla/ops/kda/naive.py
Pinned revision 07ca1e49ef6c ↗ · lines 12–66

Read code excerpt
def naive_recurrent_kda(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    scale: float | None = None,
    initial_state: torch.Tensor | None = None,
    output_final_state: bool = False,
):
    r"""

… lines 23–44 omitted …

    dtype = v.dtype
    B, T, H, K, HV, V = *q.shape, v.shape[2], v.shape[-1]
    G = HV // H
    if scale is None:
        scale = K ** -0.5

    q, k, v, g, beta = map(lambda x: x.to(torch.float), [q, k, v, g, beta])
    q = q.repeat_interleave(G, dim=2) * scale   # [B, T, HV, K]
    k = k.repeat_interleave(G, dim=2)           # [B, T, HV, K]

    S = k.new_zeros(B, HV, K, V).to(q)
    if initial_state is not None:
        S += initial_state
    o = torch.zeros_like(v)
    for i in range(0, T):
        q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i]
        S = S * g_i[..., None].exp()
        S = S + torch.einsum('b h k, b h v -> b h k v', b_i[..., None] * k_i, v_i - (k_i[..., None] * S).sum(-2))
        o[:, i] = torch.einsum('b h k, b h k v -> b h v', q_i, S)
    if not output_final_state:
        S = None
    return o.to(dtype), S