DEV Community

Alex Chen
Alex Chen

Posted on

Learn Linear Attention From Kimi K3's KDA Mechanism in 20 Lines of Python

Kimi K3 made headlines with 2.8 trillion parameters, but the most interesting part for learners is not the size. It is the attention mechanism.

K3 uses something called KDA-Kimi Delta Attention. It is a variant of linear attention, which is a different way of computing the relationship between tokens in a sequence. If you are learning how transformers work, understanding the difference between standard softmax attention and linear attention will help you see why architecture choices matter more than parameter counts.

Standard attention in one paragraph

In a standard transformer, every token looks at every other token. If your sequence has N tokens, the attention matrix is N x N. For a 1-million-token context window (which K3 supports), that matrix would be enormous-impossibly large to compute directly.

That is why most models use tricks like sliding windows or sparse patterns to avoid computing the full matrix.

Linear attention in one paragraph

Linear attention rewrites the formula so you do not need the full N x N matrix. Instead of computing attention between every pair of tokens, you maintain a running state that gets updated as you process each token. The cost per token becomes roughly constant instead of growing with sequence length.

This is what makes K3's 1M context window practical.

A 20-line Python illustration

This is not KDA itself-KDA adds a delta mechanism on top of linear attention that the K3 team designed to recover some of the expressiveness that linear attention loses. But this example shows the core idea of linear attention:

import numpy as np

# Simulate a small sequence of 8 tokens, each with 4-dimensional features
np.random.seed(42)
seq_len = 8
dim = 4
queries = np.random.randn(seq_len, dim)
keys = np.random.randn(seq_len, dim)
values = np.random.randn(seq_len, dim)

# Standard attention: N x N matrix
attn_scores = queries @ keys.T  # 8x8 matrix
attn_weights = np.exp(attn_scores) / np.exp(attn_scores).sum(axis=-1, keepdims=True)
standard_output = attn_weights @ values  # 8x4

# Linear attention: maintain a running state, no N x N matrix
state = np.zeros((dim, dim))
linear_output = []
for i in range(seq_len):
    state += np.outer(keys[i], values[i])  # update running state
    linear_output.append(queries[i] @ state)  # query against state

linear_output = np.array(linear_output)  # 8x4

print("Standard attention output shape:", standard_output.shape)
print("Linear attention output shape:", linear_output.shape)
print("Outputs are different (expected) but both are 8x4")
Enter fullscreen mode Exit fullscreen mode

Run this and you will see that both methods produce the same shape output, but the linear version never builds the full 8x8 matrix. For 8 tokens the difference is nothing. For 1 million tokens, it is the difference between feasible and impossible.

What KDA adds

KDA (Kimi Delta Attention) introduces a delta term that helps the model retain information that plain linear attention tends to blur over long sequences. The technical details are in Moonshot AI's published materials. For a beginner, the key takeaway is: linear attention makes long context cheap, and KDA tries to keep it good.

Why this matters for learning

If you are starting out with machine learning, do not rush to run K3. Start here:

  1. Understand standard softmax attention (build the matrix, watch it grow)
  2. Understand why N x N does not scale (try N=10000 in the example above)
  3. Understand linear attention as a workaround (the running state)
  4. Read about KDA as a specific improvement on that workaround

Parameters get the headlines. Architecture is where the actual engineering happens.

Disclosure: I'm a MonkeyCode user sharing my own experience, not affiliated with the project. MonkeyCode is an open-source AI coding platform: https://github.com/chaitin/MonkeyCode

Top comments (0)