spec-verify-xpu

Speculative decoding verification for Intel GPUs, loadable through kernels: the standard rejection-sampling acceptance rule applied to a batch of draft tokens, returning both the accepted prefix length and the next token. The reference baseline is a scalar Python implementation of the same rule, matched bit-exactly on both outputs. CPU and CUDA siblings: spec-verify, spec-decode-ops.

Speculative decoding is how a small GPU runs a large model at a usable rate, and the verification step has to be both exact and cheap or the technique gives back what it gains. Written eagerly the residual distribution costs a full renormalization and a multinomial per sequence; here it is one work-group per sequence with a cooperative draw, and the target model's output distribution is preserved exactly.

Measured live on an Intel Iris Xe: 0.39 ms against the eager residual-and-multinomial route's 0.90, with both outputs bit-exact against a scalar reference.

Usage

from kernels import get_kernel

sv = get_kernel("phanerozoic/spec-verify-xpu", version=1, trust_remote_code=True)

accepted, next_token = sv.verify(
    draft_tokens,   # [B, K]        int64
    draft_probs,    # [B, K, V]     float32, rows normalized
    target_probs,   # [B, K + 1, V] float32, rows normalized
)
# accepted [B] int64 in [0, K]; next_token [B] int64

Pass uniform_accept [B, K] and uniform_sample [B, K + 1] for deterministic draws. version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

Method

Draft token t_k is accepted with probability min(1, p_target(t_k) / p_draft(t_k)). On the first rejection the next token is drawn from the normalized residual max(0, p_target - p_draft) at that position. If every draft token is accepted, the bonus token is drawn from the target's trailing distribution. This preserves the target model's output distribution exactly.

One work-group per sequence. The acceptance walk is evaluated identically by every work-item so control flow stays uniform across the group's collectives. The residual distribution is then drawn cooperatively: the vocabulary is cut into one contiguous chunk per work-item, an exclusive scan over the chunk sums locates the chunk holding the target mass, and only that chunk is walked serially, reducing the serial span to V / 256. Randomness is supplied by the caller, so a run is reproducible and can be compared against a reference draw for draw.

Measured

On Intel Iris Xe (Gen12LP, 96 EUs), torch 2.13.0+xpu, 30 iterations after 5 warmup, batch 8, 5 draft tokens, vocabulary 32000:

The cooperative draw is what makes this viable: walking the cumulative distribution on a single work-item leaves the rest of the group idle across the whole vocabulary. The table is reproduced by benchmarks/bench.py.

Validation

Nine checks against a scalar Python reference implementing the same rule, on Intel Iris Xe, torch 2.13.0+xpu. Both outputs are bit-exact in every case. Coverage includes mixed accept/reject with random draws, forced accept-all (which must return K and draw the bonus token), forced reject-at-zero (which must return 0 and draw from the residual), and identical draft and target distributions, where the acceptance ratio is 1 and nothing may be rejected. A vocabulary of 4096 exercises the path where the reduction strides past the work-group size.

Requirements and limits

• float32 probabilities with rows normalized, int64 tokens.
• Intel GPU with a working Level Zero driver and a torch XPU build.
• Distribution preservation assumes the caller warps both models' distributions identically.
• The ops ship fake implementations, so they trace under torch.compile without breaking the graph.

References

Leviathan, Kalman, Matias, "Fast Inference from Transformers via Speculative Decoding" (2023); Chen et al., "Accelerating Large Language Model Decoding with Speculative Sampling" (2023).

License

Apache-2.0.