Speculative Decoding · L3 Code

Implement the speculative decoding accept/reject loop in PyTorch-style pseudocode.

Draft model (small, fast, autoregressive)

Target model (large, one parallel verify pass)

Accept/reject + residual sampling logic

The three pieces we implement below.

python

import torch
import torch.nn.functional as F

@torch.no_grad()
def draft_tokens(draft_model, prefix_ids, k):
    ids = prefix_ids.clone()
    draft_probs = []
    for _ in range(k):
        logits = draft_model(ids)[:, -1, :]         # [B, vocab]
        probs = F.softmax(logits, dim=-1)
        next_id = torch.multinomial(probs, num_samples=1)
        draft_probs.append(probs.gather(-1, next_id))  # q(x_t) for the sampled token
        ids = torch.cat([ids, next_id], dim=1)
    return ids[:, prefix_ids.shape[1]:], torch.cat(draft_probs, dim=1)  # [B,k], [B,k]

Drafting k candidate tokens with the small model.