For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/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)

source

Builds a stochastic rejection sampler for speculative decoding.

Accepts draft tokens based on coin < p_target(draft_token) where p_target is computed after applying temperature, top-k, and top-p filtering, and on rejection recovers from the positive residual (p_target with the draft token’s mass removed). The draft proposes deterministically (argmax), so its one-hot proposal needs no draft probabilities as input.

The sampling RNG seed is bound as a graph input — callers refresh it per execution so RNG varies across calls.

Parameters:

• device (DeviceRef) – Device for the graph.
• draft_proposal (Literal['argmax']) – Proposal distribution used to produce draft_tokens. Only "argmax" is supported by this point-mass residual graph.

Returns:

A graph that takes draft tokens, target logits, target logit offsets, sampling parameters, and a per-execute seed, and outputs the first rejected index, recovered tokens, and a bonus token.

Return type:

Graph

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit