Speculative decoding drops the first generated token

Speculative decoding drops the first generated token

Bug description

generate() samples the target model's first new token during prefill:

token = next_token(
    target_model,
    input_pos,
    prompt.view(1, -1),
    ...
)

That token is then passed to speculative_decoding() as context, but the output accumulator is initialized as an empty list and only records the later tokens returned by speculative decoding. The prefill token is therefore never returned.

This changes the generated text and also makes the result one token shorter than max_returned_tokens. The existing test_generate currently codifies that short output with T + max_new_tokens - 1.

The same gap means a stop token sampled during target prefill is not checked before another round of generation.

Minimal CPU reproduction

This does not require a GPU or model download:

from unittest.mock import patch

import torch
from torch import nn

import litgpt.generate.speculative_decoding as generation


model = nn.Module()
model.max_seq_length = 10
prompt = torch.tensor([1, 2])

with (
    patch.object(
        generation,
        "next_token",
        side_effect=[torch.tensor([9]), torch.tensor([7])],
    ),
    patch.object(
        generation,
        "speculative_decoding",
        side_effect=[torch.tensor([8]), torch.tensor([8])],
    ),
):
    output, _ = generation.generate(model, model, prompt, 5, speculative_k=1)

print(output.tolist())
assert output.tolist() == [1, 2, 7, 8, 8]

Current output:

[1, 2, 8, 8]

Expected behavior

The first target token sampled after the prompt should be the first returned completion token. It should count toward max_returned_tokens and should go through the same stop-token handling as tokens produced by later speculative decoding rounds.

Proposed fix

Initialize the generated-token accumulator with the target prefill token when it is not a stop token, skip further generation when it is a stop token, and handle the resulting empty completion without calling torch.stack([]). Add focused CPU regression coverage and update the existing length assertion.

Environment

• Operating system: Linux
• LitGPT: main at 7bf2960dfb26bae8e815c9a16a22732974824ac1
• PyTorch: 2.13.0