huggingface / huggingface/candle

Add an additive Cache::Paged seam to candle-transformers::models::llama for external paged-attention wiring

Open
#3,720 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

## 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.