For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Python function
build_stochastic_acceptance_sampler_graph
build_stochastic_acceptance_sampler_graph()
max.pipelines.sampling.build_stochastic_acceptance_sampler_graph(device, *, draft_proposal='argmax', vocab_size=None)
Builds a target-only stochastic rejection sampler for speculative decoding.
Accepts a draft token on coin < p_target / q_draft. How p_target is
filtered depends on the proposal mode: "argmax" applies temperature
only, while "sampled" applies temperature, top-k and top-p.
draft_proposal="argmax" (the default) means the draft proposed
deterministically, so its one-hot q needs no input and recovered
tokens are sampled from the
target distribution. "sampled" means the draft sampled its own token,
and the graph takes one more input: the distribution it drew from, which
rejection recovers from via max(p_target - q_draft, 0). That mode also
needs a concrete vocab_size, since the distribution’s trailing dim has
to be static.
The sampling RNG seed is bound as a graph input — callers refresh it per execution so RNG varies across calls.
-
Parameters:
-
Returns:
-
A graph that takes draft tokens, target logits, target logit offsets, sampling parameters, a per-execute seed, and in
"sampled"mode the draft distributions, and outputs the first rejected index, recovered tokens, and a bonus token. -
Return type: