A small issue of complexity analysis in the blogpost

[KexinFeng]

This is just to check if my understanding is right. In the blogpost https://jaykmody.com/blog/speculative-sampling complexity analysis (as well as in the code), I think the time consumption should probably be modified to t_{draft} *K + t_{target}.

In your explanation, I think the +1 * t_{draft} consumption from step 2 can be eliminated, since the forward pass p(x | x_1,...., x_n, x^\tilde_1, ..., x^\tilde_K) is not really needed, right? Only the passes up to p(x | x_1,...., x_n, x^\tilde_1, ..., x^\tilde_{K-1}) are needed, as computed in step 1.

[jaymody]

I see. Basically, p[-1] is never used if we do the forward pass with x_draft. So yeah, the following code should also work (removing the additional forward pass on the draft model):

            # Step 1: auto-regressive decode K tokens from draft model
            x_draft = x
            for _ in range(K):
                p = draft_model(x_draft)
                x_draft = np.append(x_draft, sample(p[-1]))

            # Step 2: full draft and target model forward passes on x_draft
            q = target_model(x_draft)

I'll update the blog post (I'll give you a mention) and the code.

[jaymody]

Just updated the code and the blogpost: https://jaykmody.com/blog/speculative-sampling/

Great catch, thanks for submitting the issue!