Problem
sample_tokens samples draft tokens through the native-dtype logits path, while speculative rejection/correction uses the float32 distribution produced by logits_to_probs.
In speculative decoding, the proposal sample must be drawn from the same draft distribution used in the acceptance ratio and correction distribution. If the sample is drawn from r but the algorithm computes acceptance as if it came from p, the final marginal distribution is generally biased.
Impact
With bfloat16 logits, the sampled draft distribution can differ from the float32 draft distribution used for rejection. That breaks the exact-distribution guarantee expected from speculative decoding.
Fix
Fixed in PR #30 by casting logits to float32 before temperature scaling and softmax in sample_tokens, so sampling and rejection use the same draft distribution.
Problem
sample_tokenssamples draft tokens through the native-dtype logits path, while speculative rejection/correction uses the float32 distribution produced bylogits_to_probs.In speculative decoding, the proposal sample must be drawn from the same draft distribution used in the acceptance ratio and correction distribution. If the sample is drawn from
rbut the algorithm computes acceptance as if it came fromp, the final marginal distribution is generally biased.Impact
With bfloat16 logits, the sampled draft distribution can differ from the float32 draft distribution used for rejection. That breaks the exact-distribution guarantee expected from speculative decoding.
Fix
Fixed in PR #30 by casting logits to float32 before temperature scaling and softmax in
sample_tokens, so sampling and rejection use the same draft distribution.