huggingface / huggingface/candle
Add a graph-replay-safe decode position seam (`Cache::set_decode_position`) for CUDA graph decode capture
- Dominant language
- Rust
- Stars
- 21k
- Forks
- 1.8k
- Avg merge
- 16h 42m
- Merged PRs (30d)
- 25
Description
## Use case
During autoregressive decoding, the steady-state single-token decode step is identical from one step to the next except for the KV position, which makes it a natural candidate for capture-once/replay-many via a CUDA graph (`cuStreamBeginCapture`/`cuGraphLaunch`): capture the sequence of kernel launches once, then replay it every subsequent step instead of relaunching each kernel individually. This can meaningfully cut per-token launch overhead once the model is small relative to the GPU.
A CUDA graph records the exact kernel launches and memory addresses used during capture — it does not re-derive them. Any buffer whose *contents* need to change between replays must be pre-allocated and updated in place; a value that only exists as a host-side scalar baked into a kernel launch at capture time will not change on replay.
## The gap: rotary position handling is not graph-replay-safe
`candle_transformers::models::llama`'s rotary embedding application computes, per decode step:
```rust
let cos = cache.cos.narrow(0, index_pos, seq_len)?; // index_pos: usize, host-side
let sin = cache.sin.narrow(0, index_pos, seq_len)?;
```
`Tensor::narrow` is a zero-copy view: it clones the same underlying storage and adjusts a host-computed offset into the existing `Layout`. That offset is derived from `index_pos` *at the Rust/host level, before the kernel launch*, and becomes part of the actual device pointer passed to whatever kernel consumes the narrowed tensor. A CUDA graph capturing this sequence bakes in that specific pointer. Since `index_pos` increments every decode step, replaying a graph captured at step *N* at step *N+1* would silently read the rotary embeddings for step *N* again — a correctness bug, not just a missed optimization, if graph capture is layered naively on top of the existing API.
(The existing contiguous KV cache, `Cache.kvs[block_idx]`, has the same class of problem for a different reason: it grows via `Tensor::cat` — a fresh, larger allocation every step — which is likewise disallowed for a captured buffer, since a tensor allocated inside the captured closure becomes a graph-owned allocation whose device memory is only valid while the graph is executing. A pre-allocated, in-place-updated KV cache design solves this half separately; this issue only concerns the rotary position lookup, and is expected to compose with such a cache rather than solve the contiguous-cache case itself.)
## Proposed additive change
A way to drive the rotary embedding lookup from a persistent, in-place-updatable device tensor instead of a host `usize`, e.g.:
```rust
impl Cache {
/// Attaches a persistent, caller-owned position tensor (shape `(1,)`,
/// dtype matching `cos`/`sin`'s index requirements) that graph-captured
/// decode steps read from via gather instead of `narrow`. The caller
/// updates this tensor's *contents* in place (e.g. via `Tensor::slice_set`)
/// before each graph replay; the tensor's identity/address must not
/// change across replays.
pub fn set_decode_position(&mut self, block_idx: usize, position: Tensor) -> Result<()>;
}
```
and a corresponding branch in the rotary-embedding application: when a decode position tensor is attached for a layer, use `index_select`/gather against the precomputed `cos`/`sin` table with that device tensor instead of `narrow(0, index_pos, seq_len)`. Existing callers that never attach a decode-position tensor are unaffected — this is purely additive.
This is scoped to the **decode step only** (`seq_len == 1`, one query token). Prefill keeps using the existing `narrow`-based path unconditionally.
## Scope
Additive only — default behavior (no decode-position tensor attached) is unchanged. The Rust API shape above is illustrative, not literal; the actual implementation may need adjustment once real constraints are worked through (e.g. exact dtype requirements for the index tensor, or whether the branch belongs in `apply_rotary_emb` vs. a higher call site).
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by reading the llama rotary-embedding path, especially Cache, apply_rotary_emb, and the existing narrow-based lookup. Trace the Tensor index_select/gather and slice_set constraints before settling the position-tensor API. Done means decode-only graph replay can use a persistent position tensor while prefill and callers without one retain the existing behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100