huggingface / huggingface/candle
[feature] CPU/Metal (device-agnostic) PagedAttention to complement the CUDA kernels from #3655
- Dominant language
- Rust
- Stars
- 21k
- Forks
- 1.8k
- Avg merge
- 16h 42m
- Merged PRs (30d)
- 25
Description
### Context
#3655 added paged flash-attention kernels in `candle-flash-attn`, resolving the
CUDA side of #3649. Those kernels are CUDA-only (the crate is built behind the
`cuda` feature), so there is currently no paged-attention path on CPU or Metal.
### Problem
Code that wants to use block-paged KV caches today must branch between the CUDA
flash-attn path and "something else" on CPU/Metal, where no paged primitive
exists. This hurts portability for tests, local development on Apple hardware,
and CPU fallbacks.
### Proposal
Add a device-agnostic paged-attention primitive (CPU + Metal, also runnable on
CUDA) as a complement to — not a replacement for — the fused CUDA kernels from
#3655:
- `paged_attention(q, k_cache, v_cache, block_tables, context_lens, block_size,
scale, alibi_slopes) -> Tensor`, implemented with plain candle tensor ops
(gather + dense attention), so it runs on every backend.
- `reshape_and_cache(...)` to write new K/V into a block-paged cache by slot
mapping.
- A small block-pool allocator for the KV-cache block infrastructure.
This would live in `candle-nn` (alongside the existing reference attention
variants) and serve as a portable fallback and a correctness reference for the
CUDA kernels.
### Notes
A working pure-candle reference implementation already exists (MHA/GQA/MQA +
optional ALiBi, validated on CPU against dense attention) and can be contributed
as a starting point. Happy to open a PR if this direction is welcome.
Relates to #3649, complements #3655.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start in candle-nn alongside the existing reference attention variants and review the described pure-Candle implementation for MHA, GQA, MQA, and optional ALiBi. Implement the device-agnostic paged_attention, reshape_and_cache, and KV-cache block-pool pieces, then validate the portable path against dense attention on CPU and ensure it can run on Metal and CUDA.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100