KimiLinear: A Next-Generation Linear Attention Architecture for Long-Context LLMs
KimiLinear introduces a hardware-efficient linear attention design that combines gated memory updates with chunkwise parallelism. This article explains the KDA module, its per-dimension forgetting mechanism, and the engineering practices that make linear attention viable for production-scale long-context models.
1. Background and Motivation: Why a New Attention Architecture Is Needed
As large language models continue to expand their context windows, the computational cost of traditional full attention grows quadratically with sequence length. This leads to sharply higher training and inference expenses. Sparse attention and approximate attention methods partially alleviate the problem, but they still fall short in long-sequence modeling, retrieval accuracy, and hardware utilization.
Linear attention offers a different path. By linearizing the kernel function in attention computation, it reduces complexity to a linear function of sequence length, making efficient long-text processing feasible for large models. However, early linear attention models often struggled to match full attention in quality, especially on tasks requiring precise memory and complex reasoning. The core challenge, therefore, is to preserve the efficiency benefits of linear attention while closing the gap with full attention, or even surpassing it where possible.
2. Core Innovation: Linear Attention with Gating and Forgetting Mechanisms
The central module in the KimiLinear architecture is called KDA (Kimi Delta Attention). Its design combines strengths from several linear attention variants under a unified recurrent state update framework, enabling finer-grained control over both memory writing and forgetting.
2.1 From Pure Linear Attention to DeltaNet
The state update rule for pure linear attention is:
[
St = S{t-1} + v_t k_t^\top
]
This formulation simply accumulates historical information into the state matrix. Although straightforward, it has two clear drawbacks. First, all historical information receives equal weight, with no selectivity. Second, as the sequence grows, irrelevant information continuously accumulates in the state matrix, causing memory pollution.
DeltaNet introduces a learnable gating factor (\beta_t) to control how much new information replaces old information:
[
St = (I - \beta_t k_t k_t^\top) S{t-1} + \beta_t v_t k_t^\top
]
Here, (\beta_t) determines how strongly the current input updates the state matrix. When (\beta_t) approaches 1, the model tends to overwrite old information with new information. When it approaches 0, more historical state is preserved. This mechanism gives the model the ability to write on demand.
2.2 Introducing a Forget Gate: Gated DeltaNet
However, DeltaNet still has a notable limitation. By default, fully refreshing a piece of information in the state matrix requires at least (d_k) steps, where (d_k) is the key vector dimension. This means that even when the model recognizes that certain historical information is no longer important, it cannot quickly discard it, which can hurt subsequent predictions.
Gated DeltaNet addresses this by introducing an additional forgetting coefficient (\alpha_t):
[
St = \alpha_t (I - \beta_t k_t k_t^\top) S{t-1} + \beta_t v_t k_t^\top
]
The scalar (\alpha_t) controls the overall decay rate of the historical state. By setting (\alpha_t) to a small value, the model can rapidly forget a large amount of irrelevant history in a single step, preventing outdated information from interfering with current decisions. This mechanism is conceptually similar to the forget gate in LSTM or GRU recurrent networks, but it is implemented more elegantly within the linear attention framework.
2.3 KDA’s Fine-Grained Gating Design
KDA makes a key improvement over Gated DeltaNet. Its state update rule is:
[
St = (I - \beta_t k_t k_t^\top) \text{Diag}(a_t) S{t-1} + \beta_t v_t k_t^\top
]
Unlike the scalar (\alpha_t), (a_t) in KDA is a per-dimension gating vector. This means the model can apply different forgetting rates to different rows of the state matrix, which correspond to different latent semantic features. The result is finer-grained management of information lifecycles. This design is similar in spirit to advanced gating strategies used in RWKV-7, but with further optimization for implementation details and training stability.
The value of per-dimension gating lies in the fact that not all historical information becomes outdated at the same rate. In long-document understanding, for example, document-level topic information may need to be retained over a long span, while local details need to be updated quickly. KDA’s (a_t) vector lets each feature dimension adjust its forgetting rhythm adaptively, which significantly improves the model’s ability to capture complex information structures in long sequences.
3. Hardware-Efficient Chunkwise Parallel Algorithm
The recurrent form of linear attention is theoretically sequential, processing one token at a time. This severely limits hardware utilization during large-scale training. To overcome this bottleneck, the KDA paper proposes a hardware-efficient chunkwise parallel algorithm.
3.1 Chunking Strategy
The core idea is to divide a long sequence into fixed-length chunks of size C. Computation within each chunk is parallelized, while information is propagated recursively between chunks. Specifically, the output of each chunk can be decomposed into two parts:
- Intra-chunk contribution: Attention interactions among tokens within the current chunk can be computed in parallel using matrix multiplication.
- Inter-chunk contribution: Historical state information from previous chunks is propagated recursively and applied uniformly within the current chunk.
This decomposition turns a fully sequential computation into a hybrid pattern of intra-chunk parallelism and inter-chunk sequential propagation. Parallelism increases from 1 to C, which greatly improves utilization on parallel hardware such as GPUs.
3.2 Balancing Numerical Stability and Computational Efficiency
A key challenge in implementing the chunkwise algorithm is numerical stability. Since the scale of the state matrix (S_t) can fluctuate significantly with gating values, direct matrix operations may cause floating-point overflow or precision loss. To address this, the KDA algorithm adopts proven practices from DeltaNet linear attention computation, such as state normalization and gating value constraints, and introduces reparameterization techniques to maintain numerical stability.
Chunked computation also provides another important benefit: reduced memory usage for intermediate activations. Full attention requires storing an (O(n^2)) attention score matrix. In chunkwise linear attention, only a local score matrix of size (C^2) needs to be saved for each chunk. Since C is typically much smaller than the full sequence length n, both communication overhead and memory pressure drop significantly.
4. Production Validation: Linear Attention Is Ready for Real-World Use
Another major contribution of the KDA architecture is that it demonstrates through large-scale production experiments that linear attention is not just theoretically efficient, but also practically viable in real engineering environments.
4.1 Training Efficiency Gains
During training, the chunkwise parallel algorithm allows the model to maintain linear complexity while fully utilizing modern GPU parallelism. Compared with full attention models of similar scale, KDA achieves significantly higher training throughput. When processing extremely long sequences such as contexts above 128K tokens, both memory usage and training time improve by orders of magnitude.
4.2 Inference Speed and Memory Optimization
The inference benefits are even more pronounced. The state in linear attention is a fixed-size matrix that does not grow with sequence length, enabling O(1) incremental decoding complexity. This means the cost of generating each token remains constant regardless of context length, fundamentally addressing the latency inflation problem in long-text generation.
At the same time, because there is no need to cache the full Key-Value history, memory usage drops substantially. This makes larger-batch online serving feasible and directly reduces the inference cost per request.
4.3 Model Quality Holds Up or Improves
A long-standing concern about linear attention has been whether it can match full attention in model quality. KDA’s experimental results provide a positive answer. On standard language modeling benchmarks, long-text understanding tasks, and instruction-following scenarios, KDA not only matches full attention baselines but even outperforms them on some metrics. Two factors mainly account for this:
- The per-dimension forget gate gives the model more flexible control over long-term and short-term memory, which is advantageous for long-range dependency tasks compared with fixed gating schemes.
- More refined state updates reduce the accumulation of noisy information and improve the signal-to-noise ratio of key information, positively affecting downstream task accuracy.
This does not mean linear attention can fully replace full attention. For certain specialized tasks that require precise positional awareness and complex cross-token relationship modeling, such as code execution simulation, full attention still has irreplaceable advantages. However, KDA’s results clearly show that linear attention has grown from a research prototype into a production-ready tool.
5. Summary and Outlook
Looking at the design of the KDA architecture, its success comes from deep integration of several existing approaches combined with key innovations:
- It combines DeltaNet’s write control with a gated forgetting mechanism, making state matrix updates more selective.
- It upgrades scalar gating to per-dimension vector gating, improving the model’s ability to manage complex information lifecycles.
- It uses a chunkwise parallel algorithm to overcome the training efficiency bottleneck of linear attention, opening the door to industrial deployment.
From a broader perspective, KDA represents an important direction for next-generation large model architectures: under limited computational resources, smarter state management enables models to achieve both speed and quality in long-sequence processing. In the future, as hardware support for structured sparse computation improves and gating mechanisms become more adaptive, linear attention architectures are likely to become mainstream in a wider range of applications.
For engineers looking to deploy long-context models in large-scale production environments, a solid understanding of KDA’s principles and engineering implementation is valuable not only for tracking the technology’s evolution, but also for making more informed architecture decisions in practice.
