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)

source

Builds a target-only 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. No draft probabilities are needed.

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.

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