Transformer self-attention becomes expensive as context grows because every new token can attend to every previous key and value. KV-cache eviction methods reduce memory, but many still score individual tokens frequently or discard local semantic structure. ChunkLLM proposes a different compromise: divide context into semantically coherent chunks, learn lightweight adapters that predict which chunks matter, and run full attention only over the selected chunks.
The paper frames ChunkLLM as a pluggable acceleration framework. The frozen backbone model stays intact; small QK Adapters are attached to Transformer layers to compress query/key features and estimate chunk-level attention, while a bottom-layer Chunk Adapter detects boundaries. Training uses attention distillation so the lightweight predictor recalls the important chunks that full attention would have used.
The result is a practical long-context inference pattern. Instead of asking the serving runtime to choose arbitrary individual cache tokens, it predicts a sparse set of relevant semantic blocks. The authors report that ChunkLLM preserves 98.64% of long-context benchmark performance while retaining 48.58% of the KV cache, with maximum speedup of 4.48x on 120K-token inputs.
A chunk is a contiguous region of tokens intended to carry coherent local meaning. Selecting or skipping whole chunks protects phrase, sentence, and paragraph context better than isolated token eviction. Chunk boundaries are predicted from contextual representations rather than fixed character or token counts.
QK Adapters are lightweight modules attached to each Transformer layer. They compress query and key features into a lower-dimensional representation that is cheap enough to use for chunk relevance prediction. The backbone attention weights remain frozen, so the adapter learns a side signal without rewriting the base model.
During training, full attention acts as the teacher. The adapter predicts chunk-level importance, and the loss encourages its distribution to match attention mass aggregated from the original model. This is what makes chunk selection model-aware rather than a generic text segmentation heuristic.
ChunkLLM compresses the attention decision. It keeps the expensive value path for selected regions but uses a cheap learned predictor to avoid attending over chunks that are unlikely to affect the next token.
The Chunk Adapter operates near the bottom of the model, where local semantic transitions are still visible. It predicts whether the current token begins or ends a chunk. At inference time, chunk selection is triggered when a boundary is detected, so the runtime does not repeatedly rescore every token at every step.
class ChunkAdapter:
def boundary_score(self, hidden_state):
features = layer_norm(hidden_state)
return sigmoid(linear(features))
def update_chunks(tokens, hidden_states, threshold=0.5):
chunks = []
start = 0
for i, h in enumerate(hidden_states):
if ChunkAdapter().boundary_score(h) > threshold:
chunks.append((start, i))
start = i + 1
chunks.append((start, len(tokens) - 1))
return chunksFor each layer, full attention over the original context supplies teacher weights. Those weights are aggregated by chunk, producing a target distribution over chunks. The QK Adapter learns to produce similar chunk scores from compressed query/key features.
def chunk_attention_targets(full_attention, chunks):
targets = []
for start, end in chunks:
targets.append(full_attention[:, start:end + 1].sum(axis=-1))
return normalize(stack(targets), axis=-1)At generation time, the adapter scores chunks and selects the top candidates under a cache or latency budget. The model then performs normal attention over tokens inside selected chunks. This keeps the final attention computation compatible with the frozen backbone while reducing the number of keys and values considered.
Long documents often contain locally coherent sections where only a few regions matter for a given query. Chunk-based selection can preserve the relevant section context while skipping unrelated sections that would otherwise consume attention bandwidth.
Agents accumulate tool outputs, observations, and prior decisions. Grouping these histories into semantic chunks lets the model keep complete tool-result blocks instead of retaining scattered tokens that are hard to interpret later.
Serving systems can use chunk selection to bound per-token attention cost on long prompts. Because the backbone remains frozen, the method is easier to evaluate as a serving-time adapter than a full model retrain.
Keep the original model weights fixed and train only the adapters. This limits training cost, reduces regression risk, and makes it easier to roll the acceleration path back if a workload is sensitive to skipped context.
for parameter in backbone.parameters():
parameter.requires_grad = False
trainable = list(qk_adapters.parameters()) + list(chunk_adapter.parameters())
optimizer = AdamW(trainable, lr=adapter_lr)Even with chunk selection, keep a short recent window always visible. Recent tokens carry local syntax and immediate instruction state, and skipping them can cause failures that chunk-level distillation does not catch.
The selector should be tuned for high recall of important chunks, not just maximum sparsity. Missing a crucial chunk is usually worse than including a few extra chunks, especially for retrieval, legal, or code contexts.