Module speculative 

Acceptance for speculative decoding, on the device: from a verification’s logits (or probabilities) and the draft, one [3] uint32 array [accepted, next, eos_hit], so a verification step costs one small read.

• accepted: the draft tokens kept (a prefix), cut after the first accepted end-of-sequence token;
• next: the token after them (the model’s own choice at that position, or a sample);
• eos_hit: 1 if an accepted draft token or next is an end-of-sequence id.

Both are compositions of array operations (CPU and GPU, any strides; vmap over a leading batch of verification steps; no gradient, as argmax).

use nuvec_core::{Array, array, speculative};

// The model's choices at three positions are 2, 0, 1: the draft [2, 0] is accepted,
// and 1 comes next.
let logits = array![[0.0, 0.0, 9.0], [9.0, 0.0, 0.0], [0.0, 9.0, 0.0]];
let no_eos = Array::from_slice(&[0u32; 0], [0]);
let out = speculative::accept_greedy(&logits, &array![2u32, 0], &no_eos);
assert_eq!(speculative::read(&out)?, (2, 1, false)); // one read per verification

Functions

accept_greedy

Greedy acceptance: logits [N + 1, vocab] (float32/16, bfloat16; any strides, such as a slice of [1, L, vocab]), draft [N] and eos [E] (uint32 or int32, maybe still on the device). The accepted prefix is the draft tokens equal to the logits’ argmax at their position; next is the argmax after it. Invalid inputs give a failed array.

accept_sampled

Lossless sampled acceptance (Leviathan et al., Chen et al.): probs [N + 1, vocab] is the target’s sampling distribution at each position (after temperature and filters), draft_probs [N, vocab] the drafter’s (None: deterministic drafts, a point mass on each draft token). Draft token i is accepted when u_i < min(1, p_i(d_i) / q_i(d_i)); after the first rejection at position a, next is sampled from norm(max(p_a − q_a, 0)) (p_a itself when that has no mass), and after a fully accepted draft from p_N. The output’s distribution is the target’s.

read

The accepted count, the next token and whether generation ended, read from an acceptance result in one host read.