Split View: Speculative Decoding으로 LLM 추론 2~3배 빠르게: 원리부터 실전 구현까지

Token 1 generation → Token 2 generation → Token 3 generation → ...
       ↓                    ↓                    ↓
   Full model           Full model           Full model
   Forward Pass         Forward Pass         Forward Pass

70B model, FP16:
- Model size: ~140GB
- 1 token generation: 140GB memory read
- A100 80GB memory bandwidth: 2TB/s
- Theoretical maximum: 2000/140 ≈ 14 tokens/s

In practice ~10 tokens/s due to KV Cache access, etc.
→ GPU compute utilization: 1-2%

Draft Model (1B):  t1 → t2 → t3 → t4 → t5  (quickly propose 5)
                    ↓    ↓    ↓    ↓    ↓
Target Model (70B): ✅   ✅   ✅   ❌   -    (verify at once)
                                    ↓
                              t3' regeneration  (correct after rejection)

Result: [t1, t2, t3, t3'] → 4 tokens from 1 forward pass of 70B!

import torch

def speculative_decode(draft_model, target_model, input_ids, K=5):
    """Core Speculative Decoding algorithm"""

    # 1) Generate K tokens with Draft model
    draft_tokens = []
    draft_probs = []
    current = input_ids.clone()

    for _ in range(K):
        logits = draft_model(current).logits[:, -1]
        probs = torch.softmax(logits, dim=-1)
        token = torch.multinomial(probs, 1)
        draft_tokens.append(token)
        draft_probs.append(probs.gather(-1, token))
        current = torch.cat([current, token], dim=-1)

    # 2) Verify all at once with Target model
    all_tokens = torch.cat([input_ids] + draft_tokens, dim=-1)
    target_logits = target_model(all_tokens).logits

    # 3) Rejection Sampling
    accepted = []
    n = input_ids.shape[-1]

    for i in range(K):
        target_prob = torch.softmax(target_logits[:, n+i-1], dim=-1)
        p_target = target_prob.gather(-1, draft_tokens[i])
        q_draft = draft_probs[i]

        # Acceptance probability
        accept_prob = torch.min(
            torch.ones_like(p_target),
            p_target / q_draft
        )

        if torch.rand(1) < accept_prob:
            accepted.append(draft_tokens[i])
        else:
            # Rejection: sample new token from correction distribution
            residual = torch.clamp(target_prob -
                torch.softmax(draft_model(all_tokens[:, :n+i]).logits[:, -1], dim=-1),
                min=0)
            residual = residual / residual.sum(dim=-1, keepdim=True)
            new_token = torch.multinomial(residual, 1)
            accepted.append(new_token)
            break
    else:
        # Bonus token when all accepted
        bonus = torch.multinomial(
            torch.softmax(target_logits[:, n+K-1], dim=-1), 1
        )
        accepted.append(bonus)

    return torch.cat(accepted, dim=-1)

# Layer Skip approach
class SelfSpeculativeModel(nn.Module):
    def draft_forward(self, x):
        """Fast draft using only first 8 layers"""
        for layer in self.layers[:8]:
            x = layer(x)
        return self.lm_head(self.norm(x))

    def verify_forward(self, x):
        """Verification with all layers"""
        for layer in self.layers:
            x = layer(x)
        return self.lm_head(self.norm(x))

          Target LM Head → t[n+1]
Input → Hidden States →
          Medusa Head 1 → t[n+2]  (prediction)
          Medusa Head 2 → t[n+3]  (prediction)
          Medusa Head 3 → t[n+4]  (prediction)

from vllm import LLM, SamplingParams

# Specify Draft model
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    speculative_model="meta-llama/Llama-3.2-1B-Instruct",
    num_speculative_tokens=5,
    tensor_parallel_size=4,
    gpu_memory_utilization=0.9,
)

params = SamplingParams(temperature=0.7, max_tokens=512)
outputs = llm.generate(["Explain quantum computing:"], params)

# Standard decoding
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-3.1-70B-Instruct \
    --tensor-parallel-size 4

# Speculative Decoding
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-3.1-70B-Instruct \
    --speculative-model meta-llama/Llama-3.2-1B-Instruct \
    --num-speculative-tokens 5 \
    --tensor-parallel-size 4

import tensorrt_llm
from tensorrt_llm import BuildConfig

# Build Draft and Target models simultaneously
build_config = BuildConfig(
    max_batch_size=8,
    max_input_len=2048,
    max_seq_len=4096,
    speculative_decoding_mode="draft_tokens_external",
    max_draft_len=5,
)

import time

def find_optimal_k(draft_model, target_model, test_prompts, k_range=range(1, 11)):
    """Search for optimal number of speculative tokens"""
    results = {}

    for k in k_range:
        start = time.time()
        total_tokens = 0

        for prompt in test_prompts:
            output = speculative_generate(
                draft_model, target_model, prompt,
                num_speculative_tokens=k, max_tokens=256
            )
            total_tokens += len(output)

        elapsed = time.time() - start
        throughput = total_tokens / elapsed
        results[k] = throughput
        print(f"K={k}: {throughput:.1f} tokens/s")

    optimal_k = max(results, key=results.get)
    print(f"\nOptimal K = {optimal_k} ({results[optimal_k]:.1f} tokens/s)")
    return optimal_k