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 class

RejectionRunner

RejectionRunner

class max.pipelines.sampling.RejectionRunner(session, device_ref)

source

Bases: Protocol

Interface for rejection sampling runners.

Parameters:

• session (InferenceSession)
• device_ref (DeviceRef)

run()

run(draft_tokens, draft_logits, target_logits, target_logit_offsets, all_draft_logits, context_batch)

source

Run the rejection sampler.

Parameters:

• draft_tokens (Buffer)
• draft_logits (Buffer | None)
• target_logits (Buffer)
• target_logit_offsets (Buffer)
• all_draft_logits (Buffer | None)
• context_batch (list[TextContext])

Return type:

tuple[Buffer, Buffer, Buffer | None]