TensorTonic
Study Plans
Study Plan
Perform Exact Speculative Sampling
Language ModelsInference
Hard

Verify a sequence of proposed tokens using supplied draft and target probabilities and random values. For zero-based draft position iii, let xix_ixi​ be the proposed token, qi(v)q_i(v)qi​(v) its draft probability for token vvv, and pi(v)p_i(v)pi​(v) its target probability. The acceptance probability is:

ai=min⁡(1,pi(xi)qi(xi)).a_i=\min\left(1,\frac{p_i(x_i)}{q_i(x_i)}\right).ai​=min(1,qi​(xi​)pi​(xi​)​).

Accept when the supplied acceptance value at position iii is strictly less than aia_iai​. When the draft probability is zero, use acceptance one if the target probability is positive and zero otherwise. At the first rejection, the replacement distribution is:

ri(v)=max⁡(pi(v)−qi(v),0)∑umax⁡(pi(u)−qi(u),0).r_i(v)=\frac{\max(p_i(v)-q_i(v),0)}{\sum_u\max(p_i(u)-q_i(u),0)}.ri​(v)=∑u​max(pi​(u)−qi​(u),0)max(pi​(v)−qi​(v),0)​.

Here, ri(v)r_i(v)ri​(v) is the replacement probability and uuu ranges over the vocabulary. Choose the smallest token index whose cumulative replacement probability is strictly greater than residual_random_value; if rounding leaves the cumulative sum too small, use the last index with positive replacement probability. Emit the accepted prefix followed by this replacement and stop. If all proposals are accepted, return the draft unchanged without an extra token.

Return a Python dictionary with tokens, an int64 PyTorch tensor on the draft-token device; accepted_count, a Python integer counting accepted proposals only; and rejected_at, the zero-based rejected position as a Python integer, or None when all proposals are accepted.

Loading visualization...

Examples

Input: draft_tokens = [0], draft_probabilities = [[0.75,0.25]], target_probabilities = [[0.25,0.75]], acceptance_random_values = [0.9], residual_random_value = 0.5

Output: {"tokens":[1],"accepted_count":0,"rejected_at":0}

Explanation: The proposal is rejected, and all positive replacement probability belongs to token 1.

Input: draft_tokens = [0,1], draft_probabilities = [[0.8,0.1,0.1],[0.2,0.7,0.1]], target_probabilities = [[0.2,0.5,0.3],[0.2,0.7,0.1]], acceptance_random_values = [0.9,0.1], residual_random_value = 0.4

Output: {"tokens":[1],"accepted_count":0,"rejected_at":0}

Input: draft_tokens = [0,1], draft_probabilities = [[0.5,0.3,0.2],[0.1,0.8,0.1]], target_probabilities = [[0.6,0.2,0.2],[0.5,0.2,0.3]], acceptance_random_values = [0.2,0.8], residual_random_value = 0.75

Output: {"tokens":[0,2],"accepted_count":1,"rejected_at":1}

Hint 1

torch.clamp can remove negative entries from a tensor.

Hint 2

torch.cumsum builds cumulative probabilities along an axis.

Hint 3

torch.searchsorted locates a value in a sorted tensor.

Requirements

  • Use PyTorch and the supplied random values; do not draw new random samples.
  • Use float64 arithmetic for the probability ratio and residual distribution.
  • Preserve token order and do not modify input tensors.

Constraints

  • draft_tokens is a one-dimensional int64 PyTorch tensor containing valid vocabulary indices.
  • draft_probabilities and target_probabilities have shape (K, V), where K is the proposal count and V is the vocabulary size.
  • Probability rows are normalized, nonnegative and finite.
  • acceptance_random_values is a floating-point PyTorch tensor of shape (K,); residual_random_value is a Python number.
  • All random values lie in [0, 1), and every rejected row has positive residual mass.
  • Floating tensors use float32 or float64 and may be noncontiguous.
Try Similar Problems
Decode Attention with a KV CacheHardPlan KV-Cache CapacityMediumRoute Cache-Aware Prefill RequestsHardSimulate Continuous Decode BatchingHardFind Reusable Prefixes in a KV TrieMediumRecover an IsoFLOP Compute-Optimal FrontierHardFit a Data Power Law with a Noise FloorHardTop-k MoE Router with Load StatisticsHardImplement Speculative Decoding and Token VerificationHardBeam Search DecodingHardImplement Prefix Cache Matching and ReuseMediumText Generation SamplingMedium

Accepts: array

Accepts: array

Accepts: array

Accepts: array

Accepts: number

You must run your code first.

Perform Exact Speculative Sampling

Language ModelsInference
Hard

Verify a sequence of proposed tokens using supplied draft and target probabilities and random values. For zero-based draft position iii, let xix_ixi​ be the proposed token, qi(v)q_i(v)qi​(v) its draft probability for token vvv, and pi(v)p_i(v)pi​(v) its target probability. The acceptance probability is:

ai=min⁡(1,pi(xi)qi(xi)).a_i=\min\left(1,\frac{p_i(x_i)}{q_i(x_i)}\right).ai​=min(1,qi​(xi​)pi​(xi​)​).

Accept when the supplied acceptance value at position iii is strictly less than aia_iai​. When the draft probability is zero, use acceptance one if the target probability is positive and zero otherwise. At the first rejection, the replacement distribution is:

ri(v)=max⁡(pi(v)−qi(v),0)∑umax⁡(pi(u)−qi(u),0).r_i(v)=\frac{\max(p_i(v)-q_i(v),0)}{\sum_u\max(p_i(u)-q_i(u),0)}.ri​(v)=∑u​max(pi​(u)−qi​(u),0)max(pi​(v)−qi​(v),0)​.

Here, ri(v)r_i(v)ri​(v) is the replacement probability and uuu ranges over the vocabulary. Choose the smallest token index whose cumulative replacement probability is strictly greater than residual_random_value; if rounding leaves the cumulative sum too small, use the last index with positive replacement probability. Emit the accepted prefix followed by this replacement and stop. If all proposals are accepted, return the draft unchanged without an extra token.

Return a Python dictionary with tokens, an int64 PyTorch tensor on the draft-token device; accepted_count, a Python integer counting accepted proposals only; and rejected_at, the zero-based rejected position as a Python integer, or None when all proposals are accepted.

Loading visualization...

Examples

Input: draft_tokens = [0], draft_probabilities = [[0.75,0.25]], target_probabilities = [[0.25,0.75]], acceptance_random_values = [0.9], residual_random_value = 0.5

Output: {"tokens":[1],"accepted_count":0,"rejected_at":0}

Explanation: The proposal is rejected, and all positive replacement probability belongs to token 1.

Input: draft_tokens = [0,1], draft_probabilities = [[0.8,0.1,0.1],[0.2,0.7,0.1]], target_probabilities = [[0.2,0.5,0.3],[0.2,0.7,0.1]], acceptance_random_values = [0.9,0.1], residual_random_value = 0.4

Output: {"tokens":[1],"accepted_count":0,"rejected_at":0}

Input: draft_tokens = [0,1], draft_probabilities = [[0.5,0.3,0.2],[0.1,0.8,0.1]], target_probabilities = [[0.6,0.2,0.2],[0.5,0.2,0.3]], acceptance_random_values = [0.2,0.8], residual_random_value = 0.75

Output: {"tokens":[0,2],"accepted_count":1,"rejected_at":1}

Hint 1

torch.clamp can remove negative entries from a tensor.

Hint 2

torch.cumsum builds cumulative probabilities along an axis.

Hint 3

torch.searchsorted locates a value in a sorted tensor.

Requirements

  • Use PyTorch and the supplied random values; do not draw new random samples.
  • Use float64 arithmetic for the probability ratio and residual distribution.
  • Preserve token order and do not modify input tensors.

Constraints

  • draft_tokens is a one-dimensional int64 PyTorch tensor containing valid vocabulary indices.
  • draft_probabilities and target_probabilities have shape (K, V), where K is the proposal count and V is the vocabulary size.
  • Probability rows are normalized, nonnegative and finite.
  • acceptance_random_values is a floating-point PyTorch tensor of shape (K,); residual_random_value is a Python number.
  • All random values lie in [0, 1), and every rejected row has positive residual mass.
  • Floating tensors use float32 or float64 and may be noncontiguous.
Try Similar Problems
Decode Attention with a KV CacheHardPlan KV-Cache CapacityMediumRoute Cache-Aware Prefill RequestsHardSimulate Continuous Decode BatchingHardFind Reusable Prefixes in a KV TrieMediumRecover an IsoFLOP Compute-Optimal FrontierHardFit a Data Power Law with a Noise FloorHardTop-k MoE Router with Load StatisticsHardImplement Speculative Decoding and Token VerificationHardBeam Search DecodingHardImplement Prefix Cache Matching and ReuseMediumText Generation SamplingMedium

Accepts: array

Accepts: array

Accepts: array

Accepts: array

Accepts: number

You must run your code first.