Implement and Benchmark Speculative Decoding

Problem Statement

In latency-critical inference pipelines, autoregressive decoding on large language models is frequently memory-bandwidth bound. Speculative decoding accelerates generation by using a small draft model to generate gamma candidate tokens sequentially, which the larger target model then verifies in a single parallel forward pass.

Implement greedy speculative decoding and an accompanying benchmarking routine.

Part 1: Speculative Decoding Function

Implement the function speculative_decode using greedy verification (where a draft token is accepted if and only if it matches the target model's greedy argmax prediction):

def speculative_decode(
    target_model,
    draft_model,
    input_ids: list[int],
    max_new_tokens: int,
    gamma: int,
    eos_token_id: int
) -> tuple[list[int], float]:
    """
    Returns:
        (generated_sequence, acceptance_rate)
    """

In each speculative step:

1. The draft model generates K candidate tokens (K <= gamma), stopping early if it produces eos_token_id.
2. The target model evaluates the candidates in a single batched pass.
3. Verify candidates greedily: accept all draft tokens up to the first position i where draft_token[i] != target_token[i]. Append the target model's corrected token at position i.
4. If all K tokens are accepted, append the target model's additional next-token prediction if max_new_tokens has not been reached.
5. Roll back the draft and target models' internal KV caches to reflect only accepted tokens plus the correction.
6. Stop if eos_token_id is emitted or max_new_tokens is reached.

Part 2: Benchmarking Harness

Implement benchmark_decoding to compare speculative decoding against standard autoregressive decoding:

def benchmark_decoding(
    target_model,
    draft_model,
    prompts: list[list[int]],
    max_new_tokens: int,
    gamma: int,
    eos_token_id: int,
    warmup_runs: int = 2,
    num_runs: int = 5
) -> dict:

The benchmark must include warm-up iterations, synchronize CUDA/accelerator timers where applicable, and return a dictionary containing:

• mean_latency_speculative and mean_latency_autoregressive (in seconds)
• speedup (autoregressive / speculative)
• mean_acceptance_rate
• tokens_per_second_speculative and tokens_per_second_autoregressive

Constraints

• 1 <= len(input_ids) <= 2048
• 1 <= gamma <= 16
• 1 <= max_new_tokens <= 512
• Both models expose forward(input_ids, use_cache=True, past_key_values=...) and a KV-cache truncation method truncate_cache(past_key_values, valid_length).