huggingface / huggingface/candle
Add an additive use_flashinfer_attention seam to candle-transformers::models::llama for external decode-attention wiring
- Dominant language
- Rust
- Stars
- 21k
- Forks
- 1.8k
- Avg merge
- 16h 42m
- Merged PRs (30d)
- 25
Description
## Use case
A downstream project needs to switch a Llama-family model's **decode-step** attention (one query token per sequence) from the dense matmul+softmax (or the existing `flash_attn` path) to a FlashInfer-style single-token decode attention kernel — a reference, numerically-stable implementation with a genuine CPU fallback and F32/F16/BF16 support (no dtype restriction, unlike some paged-attention kernels).
## The gap
`candle_transformers::models::llama::CausalSelfAttention::forward` is private, and its decode-step attention dispatch (`use_flash_attn ? flash_attn(...) : dense_matmul_softmax(...)`) has no seam for a downstream crate to substitute a decode-optimized kernel.
This doesn't need any new external state: the existing contiguous `Cache.kvs[block_idx]: Option<(Tensor, Tensor)>` already stores K/V as `(b_sz, num_kv_heads, seqlen_so_far, head_dim)` — exactly the shape a decode-attention kernel's `k`/`v` arguments expect, and `q` at the decode step (`seq_len == 1`) is already `(b_sz, num_heads_q, 1, head_dim)`, one squeeze away from the expected `(batch, num_heads_q, head_dim)`. Grouped-query attention (`num_heads_kv < num_heads_q`) can be handled internally by the kernel, so this path should use the **pre**-`repeat_kv` `k`/`v`, not the repeated ones the dense/flash_attn branches use.
## Proposed additive change
A new boolean, analogous to the existing `use_flash_attn`, threaded the same way through `Config`/`CausalSelfAttention`:
```rust
// Config
pub use_flashinfer_attention: bool,
```
In `CausalSelfAttention::forward`, after the existing contiguous-cache update (`cache.kvs[block_idx] = Some((k.clone(), v.clone()))`) and *before* `repeat_kv`, add a branch that only applies at the decode step (one query token per sequence — this kernel isn't meant for multi-token prefill):
```rust
if self.use_flashinfer_attention && seq_len == 1 {
let q_dec = q.squeeze(2)?.contiguous()?; // (b_sz, num_heads_q, head_dim)
let softmax_scale = 1f32 / (self.head_dim as f32).sqrt();
let y = flashinfer_decode_attention(&q_dec, &k, &v, softmax_scale)?;
let y = y.reshape(&[b_sz, 1, hidden_size])?;
return self.o_proj.forward(&y);
}
```
Existing `use_flash_attn`/dense paths stay byte-for-byte unchanged when this flag is unset (default `false`), and prefill (`seq_len > 1`) always uses the existing path regardless of the flag, since the decode kernel is decode-only by design.
## Scope
Additive only — no existing behavior changes when the new field/flag is left at its default. Feature-gated behind a new `flashinfer-kernels` cargo feature and an optional `candle-flashinfer-kernels` crate (CPU/CUDA/Metal backends), so nothing changes for existing callers who don't opt in.
I have a working, additive implementation (including the new crate, the `Config`/`CausalSelfAttention` seam, a CPU regression test comparing decode-step output against the existing dense path, and CI validation on real CUDA hardware) and am happy to open the PR.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start in candle-transformers::models::llama, reading Config and CausalSelfAttention::forward around the contiguous cache update, repeat_kv, and existing flash-attention/dense dispatch. Review the proposed feature-gated candle-flashinfer-kernels crate and run the mentioned CPU regression comparison; done means additive default-off behavior, decode-only routing, prefill preservation, and F32/F16/BF16 backend coverage.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- backend, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 30/100