huggingface / huggingface/candle

[feature] Opt-in block metadata for long-context KV pruning experiments

Open
#3,701 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Rust
Stars
21k
Forks
1.8k
Avg merge
16h 42m
Merged PRs (30d)
25

Description

## Motivation

Long-context decoding spends a growing fraction of time reading and attending over old K/V tokens. Candle has contiguous and rotating KV cache helpers in `candle-nn`, CPU flash attention paths, CUDA flash attention, and ongoing paged-attention work, but there is not yet an upstream way to attach cheap per-block metadata to KV caches and use it to measure candidate pruning for approximate long-context attention.

summarize groups of cached keys, use query/block bounds to choose candidate blocks, then benchmark candidate-count reduction separately from the exact attention kernels.

## Proposal

Add an opt-in, experimental block metadata helper in `candle-nn` for long-context attention experiments. The first PR would be deliberately conservative:

- build per-block metadata over cached key rows: centroid, radius, count, and simple concentration/variance statistics;
- expose a candidate-block API for a query vector and threshold/top-k budget;
- keep existing exact attention and cache behavior unchanged;
- add tests and benchmarks that report metadata build time, candidate block ratio, and attention-compute reduction potential;
- do not route production `flash_attn` or model forward paths through the pruner by default.

The API shape can be adjusted, but one possible direction is:

```rust
use candle_nn::kv_cache::BlockKvMetadata;

let metadata = BlockKvMetadata::from_keys(&k_cache, /*seq_dim=*/2, /*block_size=*/64)?;
let active_blocks = metadata.select_blocks(&query, /*max_blocks=*/128, /*threshold=*/None)?;
```

This would let downstream projects and Candle benchmarks test approximate KV pruning without forking model internals or modifying the current exact kernels.

## Evidence from prototype

In an isolated CPU prototype with real Vietoris-Rips persistent-homology tests used as a proof harness, a block/token pruning experiment measured:

- `3.25x` KV compression (`24/78` retained synthetic tokens);
- `2.76x` attention-compute speedup in the pruned synthetic attention loop;
- much lower persistence loss than LRU (`2.029` vs `17.832`).

A separate activation-routing prototype skipped `74.6%` of activations and measured `3.37x` sparse matvec compute speedup. I am not proposing to upstream those experimental PH internals as-is. They are useful as evidence that the metadata/pruning direction is worth testing, not as Candle production code.

## Why this belongs upstream

- Long-context cache pruning is easier to benchmark correctly if Candle has a standard metadata/candidate API near its KV cache and attention utilities.
- Keeping the first PR opt-in avoids changing exact model behavior.
- A small metadata layer can become a correctness/benchmarking foundation for later CPU/Metal/CUDA sparse or paged attention work.
- The implementation can be tested without model weights and without depending on a specific transformer architecture.

## Initial PR scope

A minimal first PR could include:

- `candle-nn/src/kv_cache.rs` or a new `candle-nn/src/attention/block_metadata.rs` helper;
- unit tests for centroid/radius/count, empty and partial blocks, deterministic candidate selection, reset/append invalidation behavior if tied to `KvCache`;
- a CPU benchmark over synthetic K/V/query tensors at `kv_len` 1k, 4k, 16k, and 32k, reporting metadata build time and selected block ratio;
- no default changes to `flash_attn`, model loaders, or quantized model forward paths.

## Open questions

- Should this live under `candle-nn::kv_cache`, `candle-nn::attention`, or behind a separate experimental module?
- Should the first API only expose block selection metadata, or also provide a reference approximate attention path for CPU tests?
- What benchmark shapes would be most useful for Candle maintainers to judge whether a follow-up kernel path is worthwhile?

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reading candle-nn/src/kv_cache.rs and the attention modules, then review how existing cache tests and CPU benchmarks are organized. The work is complete when an opt-in metadata and candidate-selection API has tests for block statistics and selection, benchmarks report build and candidate ratios, and exact attention and model paths remain unchanged.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
machine-learning, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.