Fixes the error in num_accepted_tokens calculation in reject sampler

[sighingnow]

The accepted mask vector [True, False, True, True], the num_accepted_tokens should be 1 instead of 3, as the last two tokens will be rejected as the 2-nd token is not accepted.

Thus, accepted.sum() is incorrect for num_accepted_tokens..

The reject sampler was merged in #2336, so @cadedaniel could you please help to confirm this bugfix here? Thanks!

[LiuXiaoxuanPKU]

[cadedaniel]

Hi @sighingnow and @LiuXiaoxuanPKU . Thanks for investigating this behavior and your interest in speculative decoding.

This is actually by design. To run speculative decoding in production, we want an empirical measurement of the quality of our draft model. We formulate this as the number of tokens acceptable by rejection sampling assuming their prefix is also accepted. This gives us a good real-world proxy for how aligned the draft and target models are.

For the top-level metric on how many tokens are emitted by speculative decoding, one should use num_emitted_tokens. This is what should be used to model the reduction in user-perceived inter-token-latency (ITL), as ITL is approximately step_time / num_tokens_per_step.

[sighingnow]

num_accepted_tokens should be 4 here because 8 is a bonus token so it does not count.

Here, the 1 in the first sequence is a token generated by the target model, rather than drafted by the draft model. Note that in speculative decoding there is always at least 1 token to be generated even though all draft tokens are rejected.

For the top-level metric on how many tokens are emitted by speculative decoding, one should use num_emitted_tokens. This is what should be used to model the reduction in user-perceived inter-token-latency (ITL), as ITL is approximately step_time / num_tokens_per_step.

I see. Thanks for making it clear.