huggingface / huggingface/candle

Add a graph-replay-safe decode position seam (`Cache::set_decode_position`) for CUDA graph decode capture

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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.