src/sd_core.py

#!/usr/bin/env python3
"""
Shared speculative-decoding core for the SPEED-Bench reproduction (Rl2uQlCoQX).

Faithful lossless greedy speculative decoding with a genuine draft/target pair:
  target = gpt2 (124M), draft = distilgpt2 (82M).

Acceptance rule (lossless, greedy target): a drafted token is ACCEPTED iff it
equals the target model's greedy argmax at that position; on the first mismatch
we take the target's token and stop the round (standard SD). A round that
accepts all k draft tokens appends one bonus token from the target.

"Acceptance length" (AL) per round = number of draft tokens accepted (0..k).
Mean AL drives throughput: target calls per generated token ~ 1/(mean_AL+1).
"""
import numpy as np
import torch

def load_pair(device):
    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained("gpt2")
    tgt = AutoModelForCausalLM.from_pretrained("gpt2").to(device).eval()
    drf = AutoModelForCausalLM.from_pretrained("distilgpt2").to(device).eval()
    return tok, tgt, drf

@torch.no_grad()
def spec_decode(tgt, drf, tok, prompt_ids, k=4, max_new=24, device="cpu"):
    """Run greedy lossless SD from prompt_ids (1,L) tensor. Returns list of per-round AL."""
    ids = prompt_ids.to(device)
    L0 = ids.shape[1]
    als = []
    while ids.shape[1] < L0 + max_new:
        cur = ids.shape[1]
        # draft k tokens greedily from the draft model (deterministic)
        d = ids.clone()
        dtoks = []
        for _ in range(k):
            dl = drf(d).logits[:, -1, :]
            nt = dl.argmax(dim=-1)
            dtoks.append(int(nt.item()))
            d = torch.cat([d, nt.unsqueeze(0)], dim=1)
        # one target forward over prompt+draft; verify greedily
        verify = torch.cat([ids, torch.tensor([dtoks], device=device)], dim=1)
        tl = tgt(verify).logits
        accepted = 0
        for i, dt in enumerate(dtoks):
            pos = cur + i - 1
            ttok = int(tl[:, pos, :].argmax(dim=-1).item())
            if ttok == dt:
                accepted += 1
                ids = torch.cat([ids, torch.tensor([[dt]], device=device)], dim=1)
            else:
                ids = torch.cat([ids, torch.tensor([[ttok]], device=device)], dim=1)
                break
        else:
            # all accepted -> bonus token from target's last position
            btok = int(tl[:, -1, :].argmax(dim=-1).item())
            ids = torch.cat([ids, torch.tensor([[btok]], device=device)], dim=1)
        als.append(accepted)
    return als

def prompt_to_ids(tok, text, max_len=48):
    ids = tok.encode(text, return_tensors="pt")
    if ids.shape[1] > max_len:
        ids = ids[:, :max_len]
    if ids.shape[1] == 0:
        ids = tok.encode(tok.eos_token or ".", return_tensors="pt")
    return ids

def random_prompt_ids(tok, length=24, rng=None):
    """Synthetic prompt: uniform random token ids (the 'synthetic input' baseline)."""
    if rng is None:
        rng = np.random.default_rng(0)
    vocab = tok.vocab_size
    ids = rng.integers(0, vocab, size=(1, length)).astype(np.int64)
    return torch.tensor(ids)