Decide which draft tokens one speculative decoding step keeps, using the target and draft probabilities with the supplied uniform noise, and pick the correction token from the residual distribution.
The batch has B rows, and for each row a small draft model has proposed G tokens. draft_tokens[b, i] is the draft token xi at position i of row b. The vector qi = draft_probs[b, i] is the distribution the draft model sampled xi from, and pi = target_probs[b, i] is the large target model's distribution at the same position. target_probs has one extra position, pG, which is the target distribution after the last draft token. With ui = accept_uniforms[b, i], draft token i passes its test when
ui<qi(xi)pi(xi)The tests run in order and the first failure ends the step, so n = num_accepted is the number of tests passed before the first failure, or G when every test passes. One more token y is then drawn at position n from the distribution r.
r(v)=∑wmax(pn(w)−qn(w),0)max(pn(v)−qn(v),0) if n<G,r(v)=∑wpG(w)pG(v) if n=GHere v and w run over the V vocabulary entries. With s = sample_uniforms[b], the drawn token y is the smallest v whose cumulative probability r(0)+⋯+r(v) is greater than s.
Return a dict with two keys. num_accepted is an int32 JAX array of shape (B,), and tokens is an int32 JAX array of shape (B, G + 1) whose row b holds the n accepted draft tokens, then y, then -1 in every remaining position.
Input: draft_tokens = [[3,2]], draft_probs = [[[0.2,0.59,0.04,0.17],[0.12,0.43,0.42,0.03]]], target_probs = [[[0.04,0.44,0.24,0.28],[0.27,0.36,0.23,0.14],[0.19,0.34,0.07,0.4]]], accept_uniforms = [[0.27,0.79]], sample_uniforms = [0.92]
Output: {"num_accepted":[1],"tokens":[[3,3,-1]]}
Explanation: Draft token 3 has target probability 0.28 and draft probability 0.17, a ratio of about 1.647, so it passes with u = 0.27. Draft token 2 has ratio 0.23 / 0.42, about 0.548, and fails with u = 0.79, so num_accepted is 1. The residual at position 1 is max(p - q, 0) = [0.15, 0, 0, 0.11], which normalizes to about [0.577, 0, 0, 0.423], and the cumulative probability first exceeds 0.92 at token 3.
Input: draft_tokens = [[1,0,2],[2,1,1]], draft_probs = [[[0.38,0.31,0.2,0.11],[0.16,0.01,0.55,0.28],[0.26,0.33,0.06,0.35]],[[0.01,0.13,0.49,0.37],[0.26,0.12,0.37,0.25],[0.05,0.25,0.67,0.03]]], target_probs = [[[0.3,0.4,0.05,0.25],[0.63,0.1,0.2,0.07],[0.16,0.05,0.46,0.33],[0.14,0.55,0.18,0.13]],[[0.11,0.71,0.08,0.1],[0.21,0.45,0.01,0.33],[0.07,0.02,0.49,0.42],[0.22,0.25,0.28,0.25]]], accept_uniforms = [[0.83,0.99,0.77],[0.6,0.08,0.03]], sample_uniforms = [0.59,0.73]
Output: {"num_accepted":[3,0],"tokens":[[1,0,2,1],[1,-1,-1,-1]]}
Input: draft_tokens = [[2]], draft_probs = [[[0.01,0.25,0.49,0.01,0.24]]], target_probs = [[[0.24,0.31,0.06,0.05,0.34],[0.04,0.17,0.1,0.14,0.55]]], accept_uniforms = [[0.66]], sample_uniforms = [0.38]
Output: {"num_accepted":[0],"tokens":[[0,-1]]}
jnp.take_along_axis(p, draft_tokens[..., None], axis=-1) reads each draft token's probability.
jnp.sum(jnp.cumprod(passed, axis=-1), axis=-1) counts tokens before the first failure.
x[jnp.arange(B), pos] selects a different position in each row.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Decide which draft tokens one speculative decoding step keeps, using the target and draft probabilities with the supplied uniform noise, and pick the correction token from the residual distribution.
The batch has B rows, and for each row a small draft model has proposed G tokens. draft_tokens[b, i] is the draft token xi at position i of row b. The vector qi = draft_probs[b, i] is the distribution the draft model sampled xi from, and pi = target_probs[b, i] is the large target model's distribution at the same position. target_probs has one extra position, pG, which is the target distribution after the last draft token. With ui = accept_uniforms[b, i], draft token i passes its test when
ui<qi(xi)pi(xi)The tests run in order and the first failure ends the step, so n = num_accepted is the number of tests passed before the first failure, or G when every test passes. One more token y is then drawn at position n from the distribution r.
r(v)=∑wmax(pn(w)−qn(w),0)max(pn(v)−qn(v),0) if n<G,r(v)=∑wpG(w)pG(v) if n=GHere v and w run over the V vocabulary entries. With s = sample_uniforms[b], the drawn token y is the smallest v whose cumulative probability r(0)+⋯+r(v) is greater than s.
Return a dict with two keys. num_accepted is an int32 JAX array of shape (B,), and tokens is an int32 JAX array of shape (B, G + 1) whose row b holds the n accepted draft tokens, then y, then -1 in every remaining position.
Input: draft_tokens = [[3,2]], draft_probs = [[[0.2,0.59,0.04,0.17],[0.12,0.43,0.42,0.03]]], target_probs = [[[0.04,0.44,0.24,0.28],[0.27,0.36,0.23,0.14],[0.19,0.34,0.07,0.4]]], accept_uniforms = [[0.27,0.79]], sample_uniforms = [0.92]
Output: {"num_accepted":[1],"tokens":[[3,3,-1]]}
Explanation: Draft token 3 has target probability 0.28 and draft probability 0.17, a ratio of about 1.647, so it passes with u = 0.27. Draft token 2 has ratio 0.23 / 0.42, about 0.548, and fails with u = 0.79, so num_accepted is 1. The residual at position 1 is max(p - q, 0) = [0.15, 0, 0, 0.11], which normalizes to about [0.577, 0, 0, 0.423], and the cumulative probability first exceeds 0.92 at token 3.
Input: draft_tokens = [[1,0,2],[2,1,1]], draft_probs = [[[0.38,0.31,0.2,0.11],[0.16,0.01,0.55,0.28],[0.26,0.33,0.06,0.35]],[[0.01,0.13,0.49,0.37],[0.26,0.12,0.37,0.25],[0.05,0.25,0.67,0.03]]], target_probs = [[[0.3,0.4,0.05,0.25],[0.63,0.1,0.2,0.07],[0.16,0.05,0.46,0.33],[0.14,0.55,0.18,0.13]],[[0.11,0.71,0.08,0.1],[0.21,0.45,0.01,0.33],[0.07,0.02,0.49,0.42],[0.22,0.25,0.28,0.25]]], accept_uniforms = [[0.83,0.99,0.77],[0.6,0.08,0.03]], sample_uniforms = [0.59,0.73]
Output: {"num_accepted":[3,0],"tokens":[[1,0,2,1],[1,-1,-1,-1]]}
Input: draft_tokens = [[2]], draft_probs = [[[0.01,0.25,0.49,0.01,0.24]]], target_probs = [[[0.24,0.31,0.06,0.05,0.34],[0.04,0.17,0.1,0.14,0.55]]], accept_uniforms = [[0.66]], sample_uniforms = [0.38]
Output: {"num_accepted":[0],"tokens":[[0,-1]]}
jnp.take_along_axis(p, draft_tokens[..., None], axis=-1) reads each draft token's probability.
jnp.sum(jnp.cumprod(passed, axis=-1), axis=-1) counts tokens before the first failure.
x[jnp.arange(B), pos] selects a different position in each row.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: array
