Kimi K3 · Read · OpenModelAtlas

93 layers › 68 KDA layers › Kimi Delta Attention (KDA) › Delta-rule memory

Read

no parameterso · 610, 629

Read each head’s memory with its query.

{\htmlData{light=read-out read}{\tilde o}}={\htmlData{light=memory write}{S^{\top}}}\hat {\htmlClass{vl-stream-q}{q}}/\sqrt{128}

in: $S$ [B batch, 96 heads, 128 key d_h, 128 value d_h] $Q$ [B batch, 96 heads, 1 positions, 128 d_h]
out: $\tilde o$ [B batch, 96 heads, 1 positions, 128 d_h]

Explanation

The read uses the memory after the current write, so a position can read what it has just stored.

Codenaive.py 12–22, 45–66

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

Sources1

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