huggingface / huggingface/candle

[feature] CPU/Metal (device-agnostic) PagedAttention to complement the CUDA kernels from #3655

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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.