huggingface / huggingface/candle
Add an additive Cache::Paged seam to candle-transformers::models::llama for external paged-attention wiring
- Dominant language
- Rust
- Stars
- 21k
- Forks
- 1.8k
- Avg merge
- 16h 42m
- Merged PRs (30d)
- 25
Description
## Use case
I'm building a downstream inference-serving project on top of `candle` that needs to switch a Llama-family model from the contiguous per-request KV cache to a block-paged one (vLLM-style), so that a block allocator/eviction component can operate at block granularity instead of reallocating a whole-request cache.
`candle-flash-attn` already ships the kernel needed for the paged path — `candle_flash_attn::flash_attn_varlen_paged_windowed(q, k, v, seqlens_q, seqlens_k, block_table, mm_prefix_ranges, max_seqlen_q, max_seqlen_k, softmax_scale, window_size_left, window_size_right, page_block_size, softcap)` — with `k`/`v` laid out as `(num_blocks, page_block_size, num_heads_kv, head_size)` and `block_table` as `(batch_size, max_blocks)` physical block indices per sequence.
## The gap
`candle-transformers::models::llama` has no seam for a downstream crate to use that kernel. `CausalSelfAttention`, `Mlp`, `Block`, and their `forward` methods are private (`struct`/`fn`, not `pub`), and `Cache` hardcodes the contiguous concat-and-narrow KV path (`Cache.kvs[block_idx]`). A caller can drive `Llama::forward`/`Llama::load` end to end, but cannot substitute paged storage into an existing model without either forking this file or reimplementing the whole transformer stack (RoPE, GQA `repeat_kv`, norm placement) to get access to the attention call site.
## Proposed additive API
An additive `Cache` seam, caller-owned paged storage, no behavior change for existing callers:
```rust
/// Physical KV storage and per-sequence block table for paged attention.
/// Constructed and owned by the caller; block_table/seqlens_k are read-only
/// from the model's perspective, key_cache/value_cache are written to.
pub struct PagedKvCache {
pub key_cache: Tensor, // (num_blocks, page_block_size, num_kv_heads, head_dim)
pub value_cache: Tensor, // (num_blocks, page_block_size, num_kv_heads, head_dim)
pub block_table: Tensor, // (batch_size, max_blocks), physical block ids
pub seqlens_k: Tensor, // (batch_size + 1,), cumulative
pub page_block_size: usize,
}
```
`Cache` gains a per-layer, opt-in paged slot (e.g. `Cache::set_paged_kv(block_idx, PagedKvCache)` / `clear_paged_kv` / `paged_kv`), defaulting to `None` everywhere so `Cache::new` and every existing field/caller is unaffected. When a layer has paged storage attached, `CausalSelfAttention::forward` writes the new K/V into the caller-owned `key_cache`/`value_cache` at the slot `block_table` designates, then calls `flash_attn_varlen_paged_windowed` instead of the dense matmul/softmax or existing `flash_attn` path; otherwise it's the existing contiguous path, byte-for-byte.
Block allocation, eviction, and admission policy are intentionally **not** part of this ask — that stays a caller concern. This is scoped to the minimal seam needed to drive the existing kernel from outside the crate.
Happy to open the PR for this if the shape above looks reasonable.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start in candle-transformers::models::llama at Cache and CausalSelfAttention::forward, then inspect candle_flash_attn::flash_attn_varlen_paged_windowed and its tensor layouts. Trace the existing contiguous KV path and determine the smallest additive per-layer opt-in seam; done means paged storage uses the caller's block table while existing callers retain their current path.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100