I just finished overhauling our speculative decoding implementation for Mixlayer, so maybe I can help. I think the piece of information that might make this cl

I just finished overhauling our speculative decoding implementation for Mixlayer, so maybe I can help. I think the piece of information that might make this click for you is the model outputs the probability distribution for all intermediate tokens even during prefill. So for example, let's say you prefill the prompt "The quick brown fox" (and for the sake of simplicity, let's say each word is a single token). The model outputs a tensor that is [4, $vocabulary_size]. The first dimension is a token index into the input and the 2nd dimension assigns a probability to each token in the vocabulary. So even during prefill, we can look at the prediction logits for all of the intermediate tokens. That is, we can look at what the model would have predicted after "quick" and "brown", not just the tail token "fox". In the single token autoregressive case, we just look at the next token prediction for "fox". But in the speculative decoding case we can use this information to compare the distribution of the draft model against the target model. In the greedy decoding case (ie, no sampling) we just make sure the highest probability token matches in draft and target. If we have sampling params like temperature and top-P, we have to apply something called Leviathan rejection sampling to the distribution. This basically allows us make sure the distribution is the same even if the exact probabilities are not and accept or reject draft tokens on that.