huggingface / huggingface/candle

Add an additive use_flashinfer_attention seam to candle-transformers::models::llama for external decode-attention wiring

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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.