Speculative Decoding & KV-Cache Compaction in Multi-Tenant LLM Serving

1import torch
2
3def speculative_step(target_model, draft_model, prefix_tokens, k=4):
4    """
5    Speculates K tokens with draft model and validates simultaneously with target model.
6    """
7    # 1. Generate K draft tokens rapidly using small draft model
8    draft_tokens = []
9    curr = prefix_tokens
10    for _ in range(k):
11        logits = draft_model(curr)
12        next_tok = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
13        draft_tokens.append(next_tok)
14        curr = torch.cat([curr, next_tok], dim=-1)
15
16    candidate_sequence = curr
17
18    # 2. Parallel single-pass verification by target model
19    target_logits = target_model(candidate_sequence)
20    target_predictions = torch.argmax(target_logits[:, prefix_tokens.shape[1]-1 : -1, :], dim=-1)
21
22    # 3. Accept matching prefix tokens
23    accepted = []
24    for i in range(k):
25        if draft_tokens[i].item() == target_predictions[:, i].item():
26            accepted.append(draft_tokens[i])
27        else:
28            # Emit target model's correction on first mismatch
29            accepted.append(target_predictions[:, i : i+1])
30            break
31
32    return torch.cat([prefix_tokens] + accepted, dim=-1)