Speculative Decoding State Reversion after Rejection

[breaddaerb]

Hi, I've been working with the speculative decoding implementation in the provided codebase. I'm curious about the output sequence out when the predictions made by the smaller model (small_net) are rejected by the larger model (net). Upon a rejection, we should revert the state of out to its condition prior to the rejected predictions. However, it appears that the current implementation of speculative_decoding function does not perform this state reversion.
Say if the original seq_len is 128 and gamma is 5, if all of them are rejected then in this iteration the seq_len should be 128+1, but not 128+5+1.