nuvec_core/
speculative.rs

1//! Acceptance for speculative decoding, on the device: from a verification's logits (or
2//! probabilities) and the draft, one `[3]` uint32 array `[accepted, next, eos_hit]`, so a
3//! verification step costs one small read.
4//!
5//! - `accepted`: the draft tokens kept (a prefix), cut after the first accepted
6//!   end-of-sequence token;
7//! - `next`: the token after them (the model's own choice at that position, or a sample);
8//! - `eos_hit`: 1 if an accepted draft token or `next` is an end-of-sequence id.
9//!
10//! Both are compositions of array operations (CPU and GPU, any strides; `vmap` over a
11//! leading batch of verification steps; no gradient, as `argmax`).
12//!
13//! ```
14//! use nuvec_core::{Array, array, speculative};
15//!
16//! # fn main() -> nuvec_core::Result<()> {
17//! // The model's choices at three positions are 2, 0, 1: the draft [2, 0] is accepted,
18//! // and 1 comes next.
19//! let logits = array![[0.0, 0.0, 9.0], [9.0, 0.0, 0.0], [0.0, 9.0, 0.0]];
20//! let no_eos = Array::from_slice(&[0u32; 0], [0]);
21//! let out = speculative::accept_greedy(&logits, &array![2u32, 0], &no_eos);
22//! assert_eq!(speculative::read(&out)?, (2, 1, false)); // one read per verification
23//! # Ok(())
24//! # }
25//! ```
26
27use std::panic::Location;
28
29use crate::{Array, DType, OpErrorKind, Result, ops};
30
31/// The rows and draft length, or the error naming the wrong input.
32fn check(op: &str, rows: &Array, draft: &Array, eos: &Array) -> std::result::Result<i32, String> {
33    if rows.ndim() != 2 || !matches!(rows.dtype(), DType::F32 | DType::F16 | DType::BF16) {
34        return Err(format!(
35            "{op}: the verification rows must be [N + 1, vocab] floats, got {:?} {}",
36            rows.shape(),
37            rows.dtype().name()
38        ));
39    }
40    for (name, a) in [("draft", draft), ("eos", eos)] {
41        if a.ndim() != 1 || !matches!(a.dtype(), DType::U32 | DType::I32) {
42            return Err(format!(
43                "{op}: `{name}` must be a [n] uint32 or int32 array, got {:?} {}",
44                a.shape(),
45                a.dtype().name()
46            ));
47        }
48    }
49    let n = draft.shape()[0];
50    if rows.shape()[0] != n + 1 {
51        return Err(format!(
52            "{op}: {} rows for a draft of {n} tokens (N + 1 needed)",
53            rows.shape()[0]
54        ));
55    }
56    Ok(n)
57}
58
59/// Whether each of `tokens` (`[n]`) is one of `eos` (`[e]`; none when empty).
60fn is_eos(tokens: &Array, eos: &Array) -> Array {
61    let t = tokens.astype(DType::U32).reshape([-1, 1]);
62    let e = eos.astype(DType::U32).reshape([1, -1]);
63    t.equal(&e).any(-1)
64}
65
66/// `[accepted, next, eos_hit]` from the draft's per-token acceptance `ok` (`[n]` bool), the
67/// next-token candidates per position (`next`, `[n + 1]`), the draft and the EOS ids.
68fn finish(ok: &Array, next: &Array, draft: &Array, eos: &Array) -> Array {
69    // The accepted prefix: positions before the first rejection.
70    let prefix = ok.astype(DType::U32).cumprod(0);
71    // Cut after the first accepted end-of-sequence token.
72    let eos_in = prefix.multiply(is_eos(draft, eos).astype(DType::U32));
73    let before = eos_in.cumsum(0).subtract(&eos_in); // EOS strictly before each position
74    let keep = prefix.multiply(before.equal(0u32).astype(DType::U32));
75    let accepted = keep.sum(..).astype(DType::U32);
76    let chosen = next
77        .take(&accepted.reshape([1]), 0)
78        .squeeze(..)
79        .astype(DType::U32);
80    let hit = eos_in
81        .multiply(&keep)
82        .sum(..)
83        .greater(0u32)
84        .add(is_eos(&chosen.reshape([1]), eos).squeeze(..)) // or, for booleans
85        .astype(DType::U32);
86    ops::stack(&[accepted, chosen, hit], 0)
87}
88
89/// Greedy acceptance: `logits` `[N + 1, vocab]` (float32/16, bfloat16; any strides, such as
90/// a slice of `[1, L, vocab]`), `draft` `[N]` and `eos` `[E]` (uint32 or int32, maybe still
91/// on the device). The accepted prefix is the draft tokens equal to the logits' argmax at
92/// their position; `next` is the argmax after it. Invalid inputs give a failed array.
93#[track_caller]
94pub fn accept_greedy(logits: &Array, draft: &Array, eos: &Array) -> Array {
95    let at = Location::caller();
96    Array::sticky("accept_greedy", at, &[logits, draft, eos], || {
97        let n = check("accept_greedy", logits, draft, eos).map_err(invalid)?;
98        let choice = logits.argmax(-1).astype(DType::U32);
99        let ok = draft.astype(DType::U32).equal(choice.i(0..n));
100        Ok(finish(&ok, &choice, draft, eos))
101    })
102}
103
104fn invalid(m: String) -> crate::Error {
105    crate::Error::Op {
106        op: "speculative",
107        at: Some(Location::caller()),
108        kind: OpErrorKind::InvalidArgument(m),
109    }
110}
111
112/// Lossless sampled acceptance (Leviathan et al., Chen et al.): `probs` `[N + 1, vocab]` is
113/// the target's sampling distribution at each position (after temperature and filters),
114/// `draft_probs` `[N, vocab]` the drafter's (`None`: deterministic drafts, a point mass on
115/// each draft token). Draft token `i` is accepted when `u_i < min(1, p_i(d_i) / q_i(d_i))`;
116/// after the first rejection at position `a`, `next` is sampled from
117/// `norm(max(p_a − q_a, 0))` (`p_a` itself when that has no mass), and after a fully
118/// accepted draft from `p_N`. The output's distribution is the target's.
119///
120/// `key` (a threefry key, [`crate::random`]) is split into two: the first draws the
121/// uniforms `u` (`[N]`), the second `next`; the same key gives the same result.
122#[track_caller]
123pub fn accept_sampled(
124    probs: &Array,
125    draft: &Array,
126    draft_probs: Option<&Array>,
127    eos: &Array,
128    key: &Array,
129) -> Array {
130    let at = Location::caller();
131    let mut inputs = vec![probs, draft, eos, key];
132    inputs.extend(draft_probs);
133    Array::sticky("accept_sampled", at, &inputs, || {
134        let n = check("accept_sampled", probs, draft, eos).map_err(invalid)?;
135        let v = probs.shape()[1];
136        let p = probs.astype(DType::F32);
137        let d = draft.astype(DType::U32).reshape([n, 1]);
138        let q = match draft_probs {
139            Some(q) => {
140                if q.shape() != [n, v] {
141                    return Err(invalid(format!(
142                        "accept_sampled: `draft_probs` must be [{n}, {v}], got {:?}",
143                        q.shape()
144                    )));
145                }
146                q.astype(DType::F32)
147            }
148            // A point mass on each draft token.
149            None => ops::arange(0, v, 1)
150                .reshape([1, v])
151                .equal(&d)
152                .astype(DType::F32),
153        };
154        let [uniform_key, next_key] = crate::random::split_n::<2>(key);
155        let p_draft = p.i(0..n).take_along_axis(&d, -1).reshape([n]);
156        let q_draft = q.take_along_axis(&d, -1).reshape([n]);
157        let u = crate::random::uniform(0.0, 1.0, [n], DType::F32, &uniform_key);
158        // u < p / q, without dividing: u · q < p (q > 0 for a draft token the drafter drew).
159        let ok = u.multiply(&q_draft).less(&p_draft);
160        // Per position: the residual (draft positions) or the target (the last row).
161        let residual = p.i(0..n).subtract(&q).maximum(0.0);
162        let mass = residual.sum_keepdims(-1);
163        let residual = mass
164            .greater(0.0)
165            .select(residual.divide(mass.maximum(f32::MIN_POSITIVE)), p.i(0..n));
166        let rows = ops::concatenate(&[residual, p.i(n..n + 1)], 0);
167        // `next` for every possible accepted count; `finish` picks the one at `accepted`.
168        let keys = crate::random::split(&next_key, (n + 1) as usize);
169        let next = ops::stack(
170            &(0..=n)
171                .map(|i| {
172                    crate::random::categorical(&rows.i(i).log(), -1, &keys.i(i)).astype(DType::U32)
173                })
174                .collect::<Vec<_>>(),
175            0,
176        );
177        Ok(finish(&ok, &next, draft, eos))
178    })
179}
180
181/// The accepted count, the next token and whether generation ended, read from an
182/// acceptance result in one host read.
183pub fn read(result: &Array) -> Result<(usize, u32, bool)> {
184    match result.to_vec::<u32>()?.as_slice() {
185        &[accepted, next, eos] => Ok((accepted as usize, next, eos != 0)),
186        v => Err(crate::Error::InvalidArgument(format!(
187            "speculative::read: an acceptance result has 3 values, got {}",
188            v.len()
189        ))),
190    }
191}
nuvec_core/speculative.rs

nuvec_core/
speculative.rs
