LeanK: Learnable K-Cache Channel Pruning

• LeanK is a learnable K-cache channel pruning method that statically prunes key channels to reduce memory bandwidth in long-context decoding.
• It employs a two-stage training process with L2 distillation and L1 regularization to learn hardware-aligned binary masks that preserve accuracy.
• Empirical results show up to 70% reduction in K-cache memory and a 1.3× speedup in attention computation on Llama-3.1-8B-Instruct and Qwen2.5-7B-Instruct.

LeanK denotes a learnable K-cache channel pruning method for efficient long-context decoding in LLMs. It targets the channel dimension of the key cache, learns a static binary mask offline, and couples that mask with hardware-aligned deployment and a custom decoding kernel so that memory reduction translates into decoding speedup. In the literature considered here, similar strings also appear in distinct settings, but in LLM-serving research the term refers specifically to “LeanK: Learnable K Cache Channel Pruning for Efficient Decoding” (Zhang et al., 4 Aug 2025).

1. Scope and referential ambiguity

LeanK is unambiguously the title of a KV-cache efficiency method for long-context inference (Zhang et al., 4 Aug 2025). At the same time, nearby literature uses closely related names for different systems. This makes disambiguation useful, especially because “LeanK” can surface in searches that also retrieve work on Lean theorem proving or on other “Lean”-prefixed compression methods.

Term	Domain	Relation
LeanK	LLM decoding	K-cache channel pruning (Zhang et al., 4 Aug 2025)
LeanKV	LLM serving	KV-cache compression framework (Zhang et al., 2024)
LeanContext	Domain QA	Adaptive top-$k$ sentence selection (Arefeen et al., 2023)
Kimina Lean Server	Lean 4 infrastructure	Not LeanK; batch verification backend (Santos et al., 29 Apr 2025)

This suggests that the meaning of the string depends strongly on context. In efficient decoding, LeanK is about static sparsity in the key-cache channel dimension; in other settings, the same or a similar string may refer to top-$k$ context reduction, Lean 4 tooling, or other unrelated methods (Arefeen et al., 2023).

2. Problem setting: KV-cache growth and channel redundancy

LeanK is motivated by the cost of autoregressive long-context inference. For an autoregressive transformer, each generated token requires attention over all past tokens stored in the KV cache, and as context length grows the cache grows linearly. The paper emphasizes that decoding becomes increasingly expensive because the GPU must repeatedly read large K/V tensors from memory, so the bottleneck is increasingly tied to KV-cache growth rather than model weights (Zhang et al., 4 Aug 2025).

The method is positioned against three established classes of KV-cache optimization: eviction, selection, and quantization. Those approaches typically operate at the token, head, or precision level, and they typically assume that all channels within a key vector are equally necessary when computing attention. LeanK challenges that assumption. Its central claim is that the K-cache itself exhibits structured redundancy in the channel dimension, an axis that had been comparatively underexplored relative to token eviction, token selection, or cache quantization (Zhang et al., 4 Aug 2025).

A defining empirical premise is that channel importance is largely static rather than highly input-specific. The paper supports this with three observations: Pearson correlation of channel-importance distributions across tasks and lengths is close to 1; a static norm-based mask performs comparably to a dynamic method like ThinK; and some channels can have large norm but low downstream impact. On that basis, LeanK learns a global channel-wise pruning pattern offline rather than recomputing a prompt-specific pattern at serving time (Zhang et al., 4 Aug 2025).

3. Two-stage learning of a static channel mask

LeanK learns pruning masks in two stages. The first stage introduces a learnable scaling tensor

$$\boldsymbol{\alpha} \in \mathbb{R}^{L \times n \times d},$$

where $L$ is the number of layers, $n$ is the number of heads or groups, and $d$ is the channel dimension per head. Each element of $\alpha$ represents the global importance of a K channel. During training, the model computes full attention for context tokens, but for answer tokens it uses a channel-wise scaled K cache of the form

$$\mathbf{K}\,\mathrm{diag}(\boldsymbol{\alpha}_{i,j}).$$

This scaling is applied only to the middle attention region. The attention sink and local sliding window are left intact because they are considered more critical and their memory cost is constant with respect to sequence length. The first-stage objective combines distillation and sparsity: $$\mathcal{L}_{\text{1st}} = \left\| \mathbf{H}_{\text{full}} - \mathbf{H}_{\text{scaled}} \right\|_2^2 + \lambda \left\| \boldsymbol{\alpha} \right\|_1.$$ The paper describes this as an L2 distillation loss plus an L1 regularizer. Stage 1 is trained on retrieval-style tasks, specifically dense retrieval and multi-value retrieval, to preserve long-context retrieval ability (Zhang et al., 4 Aug 2025).

The second stage converts the continuous scores into a binary deployment mask

$$\boldsymbol{\beta} \in \{0,1\}^{L \times n \times d},$$

with $k$0 indicating that the $k$1-th K channel in layer $k$2, head $k$3, is pruned. The construction is written as

$k$4

where the operator first performs global top-$k$5 selection across heads to satisfy a target pruning ratio, and then enforces per-head hardware alignment by rounding the retained channel count to a multiple of an alignment factor $k$6, such as 16 or 32. If a head initially keeps $k$7 channels, the rounded count becomes

$k$8

The second-stage objective removes the sparsity penalty and retains only the distillation term: $k$9

The rationale for the two-stage design is explicit: direct Top-K on the first-stage scores does not reliably satisfy the desired sparsity ratio or hardware alignment requirements, and omitting the second stage hurts performance. This makes the deployed mask simultaneously learned, sparse, and hardware-aware (Zhang et al., 4 Aug 2025).

4. Inference-time cache organization and kernel design

At inference time, LeanK partitions the cache as

$$\boldsymbol{\alpha} \in \mathbb{R}^{L \times n \times d},$$0

where $$\boldsymbol{\alpha} \in \mathbb{R}^{L \times n \times d},$$1 contains the full sink and local-window keys, and $$\boldsymbol{\alpha} \in \mathbb{R}^{L \times n \times d},$$2 contains the pruned keys for the middle region. Decoding attention is then computed as

$$\boldsymbol{\alpha} \in \mathbb{R}^{L \times n \times d},$$3

If an attention head has all of its K channels pruned, the corresponding V cache can also be removed for that head, yielding additional V-cache savings: $$\boldsymbol{\alpha} \in \mathbb{R}^{L \times n \times d},$$4 This is the mechanism behind the reported V-cache reduction (Zhang et al., 4 Aug 2025).

A major implementation component is a custom fused decoding kernel implemented with TileLang. The kernel groups attention heads by remaining channel count after pruning, reorders Q/K/V/O projection weights to match those groups, maintains separate pruned K caches per group, and fuses the decoding step so that it can directly read grouped caches and perform attention efficiently. The paper’s explanation is explicitly bandwidth-oriented: the main bottleneck in decoding is often memory bandwidth rather than arithmetic, and reading fewer K channels reduces memory traffic (Zhang et al., 4 Aug 2025).

The deployment path is therefore not merely a mask application. It is a coordinated procedure that combines static pruning, per-head alignment, grouped execution, and a decoding kernel specialized to the learned sparsity pattern. The method is also described as compatible with existing cache optimizations such as DuoAttention, Quest, and KIVI, indicating that channel pruning is treated as an additional compression dimension rather than a replacement for all other KV-cache methods (Zhang et al., 4 Aug 2025).

5. Empirical results and interpretability findings

LeanK is evaluated on Llama-3.1-8B-Instruct and Qwen2.5-7B-Instruct across LongBench, RULER, and GSM-Infinite. The headline memory results are up to 70% K-cache memory reduction and about 16%–18% V-cache memory reduction. The abstract summarizes the V-cache figure as “16%-18%,” and the paper attributes it to heads that become fully pruned and can drop their V-cache entirely (Zhang et al., 4 Aug 2025).

At a 70% pruning ratio, the method is reported to preserve accuracy closely. On RULER, LeanK stays close to original accuracy across 4K–128K lengths, while ThinK degrades badly at 70%. On LongBench, it remains close to full-model performance. On GSM-Infinite, it substantially outperforms ThinK and in some settings preserves or even improves performance relative to baseline. The article’s empirical emphasis is robustness under aggressive pruning, especially relative to prior dynamic approaches (Zhang et al., 4 Aug 2025).

The custom kernel provides up to 1.3× speedup for attention computation on Llama-3.1-8B-Instruct and about 1.6× on Qwen2.5-7B-Instruct. The abstract highlights the 1.3× figure as the representative result. The end-to-end serving gain is smaller but still material: the reported setup yields about a 1.2× end-to-end throughput improvement, driven by a combination of memory reduction, better batching, and faster attention (Zhang et al., 4 Aug 2025).

The paper also offers mechanistic observations about channel and head behavior. Because RoPE assigns frequencies to channel pairs, LeanK analyzes channel importance by channel-pair index and finds that low-frequency channels are generally more important, whereas high-frequency channels are more prunable. The authors note exceptions: some high-frequency channel pairs still matter substantially, so frequency alone does not fully determine importance. At the head level, the paper defines a high-frequency ratio

$$\boldsymbol{\alpha} \in \mathbb{R}^{L \times n \times d},$$5

Heads with low $$\boldsymbol{\alpha} \in \mathbb{R}^{L \times n \times d},$$6 are described as more important for retrieval and long-context understanding, whereas heads with high $$\boldsymbol{\alpha} \in \mathbb{R}^{L \times n \times d},$$7 can often be pruned or converted to streaming heads with little accuracy loss. These analyses are used to support the broader thesis that K-channel importance is a static property of the pretrained model rather than a highly prompt-specific one (Zhang et al., 4 Aug 2025).

6. Relation to adjacent methods and common misconceptions

A common confusion is to conflate LeanK with LeanKV because both address KV-cache efficiency. The two systems are related only at a high level. LeanKV is a unified KV-cache compression framework that differentiates keys vs. values, token importance, and attention-head sparsity patterns, and it relies on heterogeneous precision, dynamic token-level compression, and an on-GPU memory manager (Zhang et al., 2024). LeanK, by contrast, focuses specifically on learned static pruning of K-cache channels and couples that pruning with a hardware-aligned binary mask and a custom decoding kernel (Zhang et al., 4 Aug 2025).

A second source of ambiguity comes from non-serving uses of similar strings. In the QA literature, the “$$\boldsymbol{\alpha} \in \mathbb{R}^{L \times n \times d},$$8” in LeanContext refers to adaptive top-$$\boldsymbol{\alpha} \in \mathbb{R}^{L \times n \times d},$$9 sentence selection rather than to K-cache pruning, and the paper explicitly states that this is “not a different algorithm” called LeanK (Arefeen et al., 2023). In the Lean theorem-proving ecosystem, search results may also surface infrastructure papers such as Kimina Lean Server, which provides a unified REST API, server-side parallelization, LRU caching of imports, and infotree-based extraction for Lean 4 workflows, but that system addresses scalable theorem-proving interaction rather than decoding efficiency (Santos et al., 29 Apr 2025).

These contrasts clarify the place of LeanK within the broader literature. It belongs to the systems line of work on efficient LLM inference, where its distinctive contribution is to treat channel pruning inside the K-cache as a deployable, learned, and hardware-aware compression mechanism. A plausible implication is that LeanK is most useful when long-context efficiency is limited by memory bandwidth and cache growth, while other “Lean”-prefixed systems belong to unrelated lines of work in retrieval, theorem proving, or general KV-cache compression (Zhang et al., 4 Aug 2025).

7. Significance and limitations

LeanK’s significance lies in adding a new compression axis for long-context decoding. Instead of asking only which tokens to keep, which heads to preserve, or how many bits to use, it asks which key channels are globally worth storing and reading. The method therefore reduces both memory footprint and memory bandwidth by permanently removing low-value K channels while preserving accuracy through learned static masks (Zhang et al., 4 Aug 2025).

Its practical profile is also unusually explicit. The paper ties the learned mask to hardware alignment requirements, a grouped execution strategy, and a custom TileLang kernel. The result is a method whose claimed gains are not limited to abstract sparsity or compression ratios; they are expressed in K-cache reduction, V-cache reduction, attention speedup, and end-to-end throughput. The code is reported as available at https://aka.ms/LeanK (Zhang et al., 4 Aug 2025).

The limitations stated or implied in the paper are equally important. LeanK depends on an offline training procedure, assumes that K-channel importance is largely static across inputs and tasks, and preserves the attention sink and local sliding window rather than pruning them. Its strongest empirical claims are reported for Llama-3.1-8B-Instruct and Qwen2.5-7B-Instruct on LongBench, RULER, and GSM-Infinite, so broader generalization would need to be assessed within those boundaries. Even so, the evidence presented positions LeanK as a specialized and technically distinctive method for efficient long-context decoding in which K-cache channel sparsity is treated as a first-class systems primitive (Zhang et al., 4 Aug 2025).