Study the paper

Kimi Linear: An Expressive, Efficient Attention Architecture

Lessons, visuals, quizzes, flashcards, and resources—organized in teaching order.

Efficiency Analysis and DPLR Comparison

Comparing KDA with General DPLR Formulations

Sources

block

Kimi Linear: An Expressive, Efficient Attention Architecture T ECHNICAL R EPORT       =             + ⊙ ! During the output stage, we adopt an inter-block recurrent and intra-block parallel strategy to maximize matrix multiplication throughput, thereby fully utilizing the computational potential of Tensor Cores. O [ t ] = Γ 1 → C [ t ] ⊙ Q [ t ] S [ t ] | {z } inter chunk + Tril   Γ 1 → C [ t ] ⊙ Q [ t ] K [ t ] Γ 1 → C [ t ] ! ⊤   | {z } intra chunk U [ t ] − W [ t ] S [ t ] | {z } “pseudo”-value term ∈ R C × d v (9) = ⊙ !       + " # 3.2 Efficiency Analysis 2K 4K 8K 16K 32K 64K 0 16 32 48 64 Input length Execution Time (ms) DPLR KDA (ours) Figure 2: Execution time of kernels for vary- ing input lengths, with a uniform batch size of 1 and 16 heads. In terms of representational capacity, KDA aligns with the generalized DPLR formulation, i.e., S t = ( D − a t b ⊤ t ) S t − 1 + k t v ⊤ t , both exhibiting fine-grained decay behavior. However, such fine-grained decay introduces numerical precision issues during division operations (e.g., the intra-chunk computation in Eq. 9). To address this, prior work such as GLA [114] per- forms computations i…

The state transition of general Diagonal-Plus-Low-Rank (DPLR) models is defined as $D - a_t b_t^\top$ . While highly expressive, general DPLR formulations suffer from high computational costs and poor parallelizability because they require complex diagonalization or secondary chunking steps to maintain numerical stability during division operations.

Sources

block

Kimi Linear: An Expressive, Efficient Attention Architecture T ECHNICAL R EPORT       =             + ⊙ ! During the output stage, we adopt an inter-block recurrent and intra-block parallel strategy to maximize matrix multiplication throughput, thereby fully utilizing the computational potential of Tensor Cores. O [ t ] = Γ 1 → C [ t ] ⊙ Q [ t ] S [ t ] | {z } inter chunk + Tril   Γ 1 → C [ t ] ⊙ Q [ t ] K [ t ] Γ 1 → C [ t ] ! ⊤   | {z } intra chunk U [ t ] − W [ t ] S [ t ] | {z } “pseudo”-value term ∈ R C × d v (9) = ⊙ !       + " # 3.2 Efficiency Analysis 2K 4K 8K 16K 32K 64K 0 16 32 48 64 Input length Execution Time (ms) DPLR KDA (ours) Figure 2: Execution time of kernels for vary- ing input lengths, with a uniform batch size of 1 and 16 heads. In terms of representational capacity, KDA aligns with the generalized DPLR formulation, i.e., S t = ( D − a t b ⊤ t ) S t − 1 + k t v ⊤ t , both exhibiting fine-grained decay behavior. However, such fine-grained decay introduces numerical precision issues during division operations (e.g., the intra-chunk computation in Eq. 9). To address this, prior work such as GLA [114] per- forms computations i…

KDA addresses this bottleneck by introducing a constrained variant of DPLR where the low-rank update vectors $a_t$ and $b_t$ are bound directly to the key vector $k_t$ :

Sources

block

Kimi Linear: An Expressive, Efficient Attention Architecture T ECHNICAL R EPORT       =             + ⊙ ! During the output stage, we adopt an inter-block recurrent and intra-block parallel strategy to maximize matrix multiplication throughput, thereby fully utilizing the computational potential of Tensor Cores. O [ t ] = Γ 1 → C [ t ] ⊙ Q [ t ] S [ t ] | {z } inter chunk + Tril   Γ 1 → C [ t ] ⊙ Q [ t ] K [ t ] Γ 1 → C [ t ] ! ⊤   | {z } intra chunk U [ t ] − W [ t ] S [ t ] | {z } “pseudo”-value term ∈ R C × d v (9) = ⊙ !       + " # 3.2 Efficiency Analysis 2K 4K 8K 16K 32K 64K 0 16 32 48 64 Input length Execution Time (ms) DPLR KDA (ours) Figure 2: Execution time of kernels for vary- ing input lengths, with a uniform batch size of 1 and 16 heads. In terms of representational capacity, KDA aligns with the generalized DPLR formulation, i.e., S t = ( D − a t b ⊤ t ) S t − 1 + k t v ⊤ t , both exhibiting fine-grained decay behavior. However, such fine-grained decay introduces numerical precision issues during division operations (e.g., the intra-chunk computation in Eq. 9). To address this, prior work such as GLA [114] per- forms computations i…

$S_t = (D - a_t b_t^\top) S_{t-1} + k_t v_t^\top \quad \text{s.t.} \quad D = \text{Diag}(\alpha_t), \; a_t = \beta_t k_t, \; b_t = k_t \odot \alpha_t$

Sources

block

Kimi Linear: An Expressive, Efficient Attention Architecture T ECHNICAL R EPORT       =             + ⊙ ! During the output stage, we adopt an inter-block recurrent and intra-block parallel strategy to maximize matrix multiplication throughput, thereby fully utilizing the computational potential of Tensor Cores. O [ t ] = Γ 1 → C [ t ] ⊙ Q [ t ] S [ t ] | {z } inter chunk + Tril   Γ 1 → C [ t ] ⊙ Q [ t ] K [ t ] Γ 1 → C [ t ] ! ⊤   | {z } intra chunk U [ t ] − W [ t ] S [ t ] | {z } “pseudo”-value term ∈ R C × d v (9) = ⊙ !       + " # 3.2 Efficiency Analysis 2K 4K 8K 16K 32K 64K 0 16 32 48 64 Input length Execution Time (ms) DPLR KDA (ours) Figure 2: Execution time of kernels for vary- ing input lengths, with a uniform batch size of 1 and 16 heads. In terms of representational capacity, KDA aligns with the generalized DPLR formulation, i.e., S t = ( D − a t b ⊤ t ) S t − 1 + k t v ⊤ t , both exhibiting fine-grained decay behavior. However, such fine-grained decay introduces numerical precision issues during division operations (e.g., the intra-chunk computation in Eq. 9). To address this, prior work such as GLA [114] per- forms computations i…

Implementation detail

By binding $a_t = b_t = k_t$ (up to scaling factors), KDA achieves major computational savings over general DPLR:

Elimination of Secondary Chunking: General DPLR requires four second-level chunk matrix computations to avoid numerical instability under fine-grained decay. KDA reduces this to just two, completely eliminating the need for secondary chunking in full precision.
Fewer Matrix Multiplications: KDA eliminates three additional matrix multiplications during the inter-chunk and output computation stages.
Speedup: As shown in the efficiency benchmarks, these optimizations allow KDA to achieve approximately a $2\times$ speedup (100% improvement in operator efficiency) compared to the standard DPLR formulation.

Sources

block

Kimi Linear: An Expressive, Efficient Attention Architecture T ECHNICAL R EPORT 1 def chunk_dplr(q, k, v, a, b, g, chunk_size): 2 B, H, T, K, V, BT = *q.shape, v.shape[-1], chunk_size 3 NT, S = T // BT, k.new_zeros(B, H, K, V) 4 q, k, v, a, b, g = map(lambda x: rearrange(x, ' b h (n c) d -> b h n c d ' , c=BT), [q, k, v, a, b, g]) , → 5 gc = g.cumsum(-2) 6 - Aab, Aak, Aqb, Aqk = (torch.zeros(B, H, NT, BT, BT) for _ in range(4)) , → 7 8 for i in range(BT): 9 a_i, q_i, g_i = (x[:,:,:,i,None] for x in (a, q, gc)) , → 10 mask = (torch.arange(BT) <= i)[..., None] 11 s1_i = (g_i - gc).exp().where(mask, 0) 12 s2_i= (g_i - g[:,:,:,i,None] - gc).where(mask, 0) 13 - Aqk[..., i, :] = (q_i * k * s1_i).sum(-1) 14 - Aqb[..., i, :] = (q_i * b * s1_i).sum(-1) 15 - Aab[..., i, :] = (a_i * b * s2_i).sum(-1) 16 - Aak[..., i, :] = (a_i * k * s2_i).sum(-1) 17 for i in range(1, BT): 18 Aab[..., i, :i] = Aab[..., i, :i] + (Aab[..., i, :, None] * Aab[..., :, :i]).sum(-2) , → 19 Aab = Aab + torch.eye(BT) 20 u, w = Aab @ (Aak @ v), Aab @ ((gc-g).exp() * a) 21 o = torch.zeros_like(v) 22 mask = torch.triu(torch.ones(BT, BT), diagonal=1) 23 for i in range(0, NT): 24 q_i, k_i, v_i, u_i, w_i, b_i = (x[:, :, i]…