Efficient Attention

Kimi Linear

The first linear attention to outperform full attention at scale. Using per-channel gating and a hybrid 3:1 architecture, Kimi Linear achieves 6x faster decoding and 75% memory reduction while matching or exceeding MLA quality on million-token contexts.

Moonshot AIJanuary 2025KDA ArchitectureLong Context
6.3x
Faster Decoding
75%
Memory Reduction
1M
Token Context
3:1
KDA:MLA Ratio

Deep Dive Contents

Modern LLMs need to process increasingly long contexts—million-token documents, entire codebases, or extended conversations. But standard attention has a fatal flaw: O(n²) complexity. Double the context length and you quadruple the compute. This makes long-context inference prohibitively expensive.

Full Attention Scaling

  • 4K tokens: baseline compute
  • 64K tokens: 256x baseline compute
  • 1M tokens: 62,500x baseline compute
  • KV cache grows linearly, attention grows quadratically

What We Want

  • O(n) complexity for arbitrary sequence lengths
  • Constant memory per decoding step
  • No quality degradation vs full attention
  • GPU-efficient parallelization

The Linear Attention Dream

Linear attention promises O(n) scaling by avoiding the n×n attention matrix. Instead of computing pairwise attention scores, we maintain a compressed state that summarizes all past tokens. The challenge: previous linear attention methods lost significant quality. Kimi Linear is the first to close this gap.

Real-World Impact

Cost Reduction

6x faster decoding = 6x lower inference costs at scale

Batch Size

75% less KV cache = 4x more concurrent requests

Context Length

Same hardware now supports million-token contexts

Standard softmax attention computes: Attention(Q, K, V) = softmax(QK²/√d)V. The softmax creates an n×n matrix, forcing O(n²) compute. Linear attention replaces softmax with a kernel trick, enabling a recurrent formulation with constant memory.

Softmax Attention

y_t = ∑_i softmax(q_t · k_i) · v_i
  • Must compute all pairwise scores
  • Stores full KV cache
  • O(n²) memory for attention matrix

Linear Attention

y_t = S_t · q_t where S_t = ∑_i k_i v_i'
  • Maintains d×d state matrix
  • Constant memory regardless of sequence length
  • O(n) total compute

The State Matrix Intuition

Think of the state matrix S as a compressed memory of all past tokens. Each key-value pair gets “written” into this matrix: S = ∑ k_i ⊗ v_i. When querying, we retrieve: y = S · q. The compression is lossy—early linear attention lost too much information.

State Size

d × d matrix (~4M params for d=2048)

Information Capacity

Fixed capacity regardless of sequence

Update Cost

O(d²) per token

Why Previous Linear Attention Failed

Early methods (Linear Transformers, Performers) used simple additive updates: S_t = S_{t-1} + k_t ⊗ v_t. This causes state explosion—old information never decays, and the state becomes a noisy mess. DeltaNet added forgetting, but used a single scalar decay for all channels. Kimi's insight: different channels need different decay rates.

KDA (Kimi Delta Attention) is the core innovation. Unlike DeltaNet which uses a single scalar gate, KDA learns a separate decay rate for each channel. Some channels can have long memory (high α), others can forget quickly (low α). This dramatically improves information retention.

KDA State Update Formula

S_t = (I - βk_t k_t') · Diag(α_t) · S_{t-1} + βk_t v_t'
(I - βkk')
Forget old k direction
Diag(α_t)
Per-channel gating
S_{t-1}
Previous state
βkv'
Write new pair
Mode:

Initial State S_{t-1}

The hidden state matrix carries forward information from previous tokens

KDA Innovation

Per-channel gating allows different dimensions to have different decay rates. Some channels can have long memory (high α) while others forget quickly (low α). This lets the model learn what information to keep vs. discard per dimension.

DeltaNet (Scalar Gating)

S_t = α · S_{t-1} + ...

Single α value for all channels. Information either decays everywhere at the same rate, or is retained everywhere. No flexibility for different types of information.

KDA (Per-Channel Gating)

S_t = Diag(α_t) · S_{t-1} + ...

Each channel has its own α. Channel 1 might retain 99% of information (long memory), while Channel 2 retains only 50% (short memory). The model learns optimal patterns.

Why Per-Channel Matters

Different types of information need different retention:

  • Long-term context: Character names, facts, code structure (high α)
  • Short-term context: Recent sentence structure, local syntax (low α)
  • Dynamic needs: Same channel can adjust decay based on content

The KDA state update has four distinct operations that happen in sequence. Understanding each operation reveals why this architecture is so effective at balancing memory and efficiency.

Step 1: Directional Forgetting

(I - βk_t k_t') · S

Before writing new information, selectively forget old information in the direction of the current key. This prevents interference—the new key-value pair won't collide with old ones in the same direction. The β parameter controls forgetting strength.

Step 2: Per-Channel Decay

Diag(α_t) · S

Apply learned per-channel decay rates. Each diagonal element α_i is between 0 and 1, determining how much of that channel's information to retain. The α values are computed from the input: α_t = σ(W_α x_t), making decay content-dependent.

Step 3: Write New Information

S + βk_t v_t'

Add the new key-value pair as an outer product. The β coefficient controls write strength, typically learned or fixed. This “writes” the current token's information into the state matrix for future queries.

Step 4: Query

y_t = S_t · q_t

Retrieve information by multiplying the state with the query vector. This is a simple matrix-vector multiply—O(d²) regardless of sequence length. The output y_t contains relevant information from all past tokens, compressed through the state.

Complexity Analysis

State Update

O(d²) per token (matrix ops)

Query

O(d²) per token (mat-vec)

Total for n tokens

O(nd²) vs O(n²d) for softmax

The recurrent formulation is mathematically elegant but GPU-unfriendly—sequential updates can't be parallelized. DPLR (Diagonal-Plus-Low-Rank) reformulates the state transition to enable chunked parallel processing while maintaining exact equivalence.

DPLR State Transition Matrix

A_t = Diag(α_t) - βk_t(α_t ⊙ k_t)'

The key insight: the state transition A_t is the sum of a diagonal matrix and a rank-1 outer product. This structure enables efficient parallel scans using the Woodbury identity.

Diag(α)
Diagonal part (d operations)
-βk(α⊙k)'
Low-rank part (rank 1)

Naive Sequential

  • Process one token at a time
  • Can't utilize GPU parallelism
  • Memory bound, low throughput
  • n sequential operations

DPLR Chunked Parallel

  • Process chunks in parallel
  • Full GPU utilization
  • Compute bound, high throughput
  • log(n/chunk) sequential ops

The Woodbury Identity Trick

For DPLR matrices, the product of multiple transitions can be computed efficiently:

A_1 · A_2 · ... · A_c = Diag + Low-Rank(c)

The product of c DPLR matrices is itself DPLR with rank at most c. This enables efficient parallel scans: compute chunk-level products in parallel, then combine sequentially with only O(chunks) operations instead of O(n).

Implementation Note

Kimi Linear uses the WY representation from numerical linear algebra to maintain stability during the parallel scan. This avoids numerical issues when multiplying many near-identity matrices together.

Pure linear attention, even with per-channel gating, still loses some expressivity compared to full attention. Kimi Linear's solution: a hybrid architecture that uses KDA for most layers and MLA (Multi-head Latent Attention) for every 4th layer.

KDA Layer (9)
MLA Layer (3)
In
1
KDA
2
KDA
3
KDA
4
MLA
5
KDA
6
KDA
7
KDA
8
MLA
9
KDA
10
KDA
11
KDA
12
MLA
Out

Why 3:1 Ratio?

Experiments showed that placing MLA at every 4th layer provides the best balance between efficiency and accuracy. More MLA layers hurt speed; fewer hurt quality.

Effective Memory

With only 25% MLA layers, KV cache is reduced by ~75%. At 1M tokens, this means storing 250K token equivalents instead of 1M.

KDA Layers (75%)

  • Complexity: O(n) per layer
  • Memory: Constant state size
  • Role: Efficient local + compressed global context

MLA Layers (25%)

  • Complexity: O(n²) per layer
  • Memory: Full KV cache for that layer
  • Role: Full attention for global patterns

Why 3:1 Ratio?

The ratio was determined through ablation studies:

1:1 (50% MLA)

Best quality but only 2x speedup

3:1 (25% MLA)

Best quality/speed tradeoff

7:1 (12.5% MLA)

Faster but noticeable quality loss

Effective Memory Savings

With 3:1 ratio, only 25% of layers need KV cache. At 1M tokens:

Full MLA
100%
Kimi Linear
25%

75% memory reduction enables 4x higher batch sizes or 4x longer contexts on same hardware.

Standard transformers use position embeddings (RoPE, ALiBi, etc.) to encode sequence order. Kimi Linear takes a radical approach: no explicit position embeddings. Position information emerges naturally from the recurrent state dynamics.

Standard Transformers

  • RoPE: Rotary position encoding in attention
  • ALiBi: Position-based attention bias
  • Must extrapolate beyond training length
  • Context extension requires careful tuning

Kimi Linear Approach

  • No explicit position encoding
  • Decay rates encode relative position
  • Natural length generalization
  • No extrapolation issues

How Position Emerges

The per-channel decay rates (α) naturally encode position information:

  • Recent tokens: High contribution (less decay applied)
  • Distant tokens: Lower contribution (more decay)
  • Content-dependent: Important tokens can have higher decay to “stand out”

Benefit: Infinite Context Length

Without explicit position embeddings, there's no theoretical limit on context length. The model trained on 128K contexts naturally generalizes to 1M tokens without any position interpolation or other tricks.

Training requires processing long sequences efficiently. Kimi Linear divides sequences into chunks and uses a parallel scan algorithm to compute states across chunks, achieving training speeds comparable to standard attention.

Chunked Processing

Within Chunk

Full parallelism using CUDA cores

Across Chunks

Parallel scan with DPLR

Chunk Size

Tuned per hardware (64-256)

Each chunk computes its local attention and state contribution in parallel. The parallel scan then propagates states across chunks in O(log(n/chunk)) steps instead of O(n).

Training Performance

vs FlashAttention-20.95x
Memory efficiencyBetter at long seq
Scaling with lengthO(n) vs O(n²)

Inference Performance

Prefill (1M tokens)~1x (similar)
Decode per token6.3x faster
KV cache75% smaller

WY Representation

The parallel scan uses the WY representation from numerical linear algebra. This represents the cumulative state transitions as W·Y' where W and Y are tall-skinny matrices. This maintains numerical stability when multiplying many near-identity matrices.

Kimi Linear was trained from scratch using a careful curriculum that gradually increases context length. The training process reveals important insights about scaling linear attention.

Training Configuration

Model size~200B parameters
Training tokens~15T tokens
Initial context4K tokens
Final context128K tokens
KDA layers75% of layers
MLA layers25% of layers
State dimensiond×d per head
Chunk size64-256

Context Length Curriculum

Phase 1
4K
Phase 2
32K
Phase 3
128K

Gradual context extension helps the model learn proper attention patterns before handling long sequences.

Key Training Insights

  • Per-channel α needs careful initialization
  • MLA layers should use RoPE for position
  • Longer context training improves all lengths
  • Training speed matches FlashAttention

Data Mix

  • Web data (filtered, deduplicated)
  • Code (GitHub, multiple languages)
  • Long documents (books, papers)
  • Synthetic long-context examples

Kimi Linear matches or exceeds full attention models on standard benchmarks while dramatically outperforming on efficiency. The key finding: per-channel gating closes the quality gap that plagued previous linear attention methods.

MLA (Full Attention)
Kimi Linear
95ms71ms48ms24ms0ms
4K16K64K128K256K512K1M
Context Length (tokens)

Peak Speedup: 6.3x

At 1M tokens, Kimi Linear decodes in 15ms vs MLA's 95ms per token. The gap widens with sequence length due to linear vs quadratic scaling.

Memory Savings: 75%

KV cache is 4x smaller due to 3:1 KDA-to-MLA ratio. This enables longer contexts on the same hardware or higher batch sizes.

Language Modeling (Perplexity)

Kimi Linear
~1.02x
Full MLA
1.00x
DeltaNet
~1.10x

Lower is better. Kimi Linear is within 2% of full attention, vs 10% gap for scalar-gated methods.

Long-Context Tasks

RULER (128K)94.2%
Needle in Haystack (1M)99.1%
Multi-doc QA (256K)87.3%
Summarization (128K)91.5%

Standard Benchmarks

MMLU87.2%
HumanEval75.6%
GSM8K89.4%
MATH52.8%

Key Finding

For the first time, a linear attention model achieves parity with full attention on both short and long-context tasks. Previous methods (Linear Transformers, Performers, RWKV, Mamba) all showed measurable quality gaps. The per-channel gating is the critical innovation.

The efficiency gains compound at scale. At 1M tokens, Kimi Linear is dramatically faster and more memory-efficient than any full attention baseline.

6.3x
Decode Speed
at 1M context
75%
Memory Saved
KV cache reduction
4x
Batch Size
same hardware

Scaling Behavior

Prefill (Encoding)

Both methods are compute-bound during prefill. Kimi Linear is slightly faster due to O(n) vs O(n²) attention, but the gap is smaller than decode.

At 1M tokens: ~15% faster prefill
Decode (Generation)

Decode is memory-bound for full attention (loading full KV cache). Kimi Linear uses fixed-size state, making decode dramatically faster.

At 1M tokens: 6.3x faster decode

Cost Implications

API Cost Reduction

6x faster decode = 6x lower per-token costs for long outputs

Throughput Increase

4x batch size = 4x more concurrent users per GPU

Hardware Efficiency

The DPLR formulation maps efficiently to modern GPU architectures. Tensor cores handle the matrix operations, and the parallel scan uses efficient prefix sum patterns. No custom hardware required—pure software innovation on standard GPUs.

Kimi Linear proves that linear attention can match full attention quality. This opens exciting research directions for even more efficient architectures.

Near-Term Improvements

  • Higher ratios: Can we go to 7:1 or pure KDA?
  • Larger states: Increase d for more capacity
  • Better initialization: Optimize α init
  • Distillation: Smaller KDA models

Research Questions

  • ?What is the theoretical capacity of linear states?
  • ?Can we learn when to use KDA vs MLA dynamically?
  • ?How does per-channel gating interact with sparse attention?
  • ?Can this approach work for vision transformers?

Industry Impact

If Kimi Linear's approach becomes standard, we could see: 6x reduction in inference costs for long-context applications, enabling affordable million-token AI assistants, real-time processing of entire codebases, and practical retrieval-augmented generation at scale.

The Bigger Picture

Linear attention has been a holy grail of efficient transformers since 2020. Previous attempts (Linear Transformers, Performers, RWKV, Mamba) all showed quality gaps. Kimi Linear's per-channel gating finally closes this gap, suggesting that:

  • The softmax isn't magic—its effect can be approximated with learned gating
  • Hybrid architectures (linear + sparse attention) may be the future
  • Efficiency and quality are not fundamentally at odds

Kimi Linear: The Complete Picture

Kimi Linear proves that linear attention can match full attention quality through per-channel gating. The 3:1 hybrid architecture with DPLR parallelization delivers 6x faster decoding and 75% memory reduction at million-token scales.

KDA
Per-Channel Gating
3:1
Hybrid Ratio
6.3x
Faster Decode
1M
Token Context