Optimizing LLM Inference: From Knowledge Distillation to Speculative Decoding
Introduction Part 2: Knowledge Distillation, KV Caching, and Speculative Decoding In the first part, quantization and pruning were discussed as techniques to optimize Large Lang...
Introduction
Part 2: Knowledge Distillation, KV Caching, and Speculative Decoding
In the first part, quantization and pruning were discussed as techniques to optimize Large Language Models (LLMs) at the model level. This segment explores three different methods focusing on inference optimization: Knowledge Distillation, KV Caching, and Speculative Decoding. These techniques help make LLMs more efficient and ready for production use.
By the end of this article, readers will understand how these methods function at the algorithmic level, their role in an inference pipeline, and their integration with previous techniques for a fully optimized serving stack.
Key Takeaways
-
Knowledge Distillation: This process involves training a smaller model to mimic the behavior of a larger model by utilizing its probability distributions, which helps it perform beyond its apparent capabilities.
-
KV Caching: A critical optimization for runtime, KV caching avoids redundant computation by storing intermediate states, thereby reducing the need for recalculating attention over the entire context.
-
PagedAttention: This addresses memory waste issues in KV caching by using blocks of memory efficiently, significantly enhancing throughput.
-
Speculative Decoding: This method leverages parallel processing to verify multiple tokens at once, maintaining the same output distribution as the large model but at a faster rate.
Knowledge Distillation
The Core Idea
Imagine a scenario where a knowledgeable professor guides a student through an exam. Instead of merely providing answers, the professor explains the reasoning behind each answer. This is akin to knowledge distillation, where a large teacher model guides a smaller student model, not just in producing outputs but in understanding underlying reasoning patterns.
Why Probability Distributions Carry More Information Than Labels
Probability distributions from a teacher model offer more insights than simple labels. For example, a distribution might indicate that "dog" and "wolf" are semantically related, providing richer information than a mere label would.
The Loss Function
The training objective combines cross-entropy loss with KL divergence loss, where the latter is adjusted for temperature to enhance the signal of less probable but significant data points.
Distillation Flavors
-
Response-Based Distillation: The student learns from the teacher's final outputs without needing internal access.
-
Feature-Based Distillation: The student mimics the teacher's internal states, gaining deeper insight into the reasoning process.
-
Relation-Based Distillation: Focuses on preserving the similarity structure of the teacher’s representation space.
Patient Knowledge Distillation (PKD)
PKD uses multiple intermediate layers from the teacher as learning signals, leading to faster convergence and better retention of the teacher’s knowledge.
MiniLLM: Fixing the Mode-Averaging Problem
MiniLLM flips the KL divergence, encouraging the student to focus on high-confidence continuations rather than spreading probability across all plausible options.
## Python Example: Teacher-Student Distillation Training Loop
## Import necessary libraries
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
from torch.utils.data import DataLoader
## Define the DistillationTrainer class
class DistillationTrainer:
def __init__(self, teacher_model, student_model, temperature=2.0, alpha=0.5, device="cuda"):
self.teacher = teacher_model.to(device).eval()
self.student = student_model.to(device)
self.T = temperature
self.alpha = alpha
self.device = device
# Freeze teacher's parameters
for param in self.teacher.parameters():
param.requires_grad = False
def distillation_loss(self, student_logits, teacher_logits, labels):
# Hard label loss
shift_logits = student_logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
L_ce = F.cross_entropy(
shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1),
ignore_index=-100
)
# Soft label loss
soft_teacher = F.softmax(teacher_logits[..., :-1, :] / self.T, dim=-1)
soft_student = F.log_softmax(student_logits[..., :-1, :] / self.T, dim=-1)
L_kd = (self.T ** 2) * F.kl_div(
soft_student,
soft_teacher,
reduction="batchmean"
)
return self.alpha * L_ce + (1 - self.alpha) * L_kd
def train_step(self, input_ids, attention_mask, optimizer):
input_ids = input_ids.to(self.device)
attention_mask = attention_mask.to(self.device)
with torch.no_grad():
teacher_out = self.teacher(input_ids=input_ids, attention_mask=attention_mask)
teacher_logits = teacher_out.logits
student_out = self.student(input_ids=input_ids, attention_mask=attention_mask)
student_logits = student_out.logits
loss = self.distillation_loss(student_logits, teacher_logits, input_ids)
optimizer.zero_grad()
loss.backward()
optimizer.step()
return loss.item()
def run_distillation(teacher_name, student_name, dataset, epochs=3, lr=1e-4):
teacher = AutoModelForCausalLM.from_pretrained(teacher_name, torch_dtype=torch.float16)
student = AutoModelForCausalLM.from_pretrained(student_name, torch_dtype=torch.float32)
trainer = DistillationTrainer(
teacher_model=teacher,
student_model=student,
temperature=2.0,
alpha=0.3
)
optimizer = torch.optim.AdamW(student.parameters(), lr=lr)
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
for epoch in range(epochs):
total_loss = 0.0
for batch in dataloader:
loss = trainer.train_step(
input_ids=batch["input_ids"],
attention_mask=batch["attention_mask"],
optimizer=optimizer
)
total_loss += loss
avg_loss = total_loss / len(dataloader)
print(f"Epoch {epoch+1}/{epochs} | Loss: {avg_loss:.4f}")
return student
KV Cache
Why Attention Has a Memory Problem
During autoregressive text generation, the model computes Query, Key, and Value vectors for each token, attending to the entire context. This leads to redundant computations, especially for longer sequences.
The Memory Cost of Caching
KV caching trades compute for memory. For large models like LLaMA-2, the cache size can be substantial, impacting inference efficiency.
Multi-Head vs Multi-Query vs Grouped-Query Attention
These attention variants balance quality and KV cache size, with Multi-Query Attention offering significant memory savings at the cost of some quality loss.
PagedAttention: Solving the Fragmentation Problem
PagedAttention divides the KV cache into fixed-size blocks, allocating them as needed, improving memory utilization and throughput.
## Python Example: KV Cache in Attention
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass, field
from typing import Optional, Tuple
@dataclass
class KVCache:
keys: Optional[torch.Tensor] = None
values: Optional[torch.Tensor] = None
def update(self, new_keys: torch.Tensor, new_values: torch.Tensor):
if self.keys is None:
self.keys = new_keys
self.values = new_values
else:
self.keys = torch.cat([self.keys, new_keys], dim=2)
self.values = torch.cat([self.values, new_values], dim=2)
return self.keys, self.values
@property
def seq_len(self) -> int:
return self.keys.shape[2] if self.keys is not None else 0
class CachedMultiHeadAttention(nn.Module):
def __init__(self, d_model: int, n_heads: int):
super().__init__()
assert d_model % n_heads == 0
self.n_heads = n_heads
self.head_dim = d_model // n_heads
self.scale = self.head_dim ** -0.5
self.q_proj = nn.Linear(d_model, d_model, bias=False)
self.k_proj = nn.Linear(d_model, d_model, bias=False)
self.v_proj = nn.Linear(d_model, d_model, bias=False)
self.o_proj = nn.Linear(d_model, d_model, bias=False)
def _split_heads(self, x: torch.Tensor) -> torch.Tensor:
B, S, D = x.shape
return x.view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
def forward(self, x: torch.Tensor, kv_cache: Optional[KVCache] = None, use_causal_mask: bool = True) -> Tuple[torch.Tensor, KVCache]:
B, S, _ = x.shape
Q = self._split_heads(self.q_proj(x))
K = self._split_heads(self.k_proj(x))
V = self._split_heads(self.v_proj(x))
if kv_cache is not None:
K, V = kv_cache.update(K, V)
attn_scores = torch.matmul(Q, K.transpose(-2, -1)) * self.scale
if use_causal_mask:
total_len = K.shape[2]
current_len = Q.shape[2]
mask = torch.triu(torch.ones(current_len, total_len, device=x.device), diagonal=total_len - current_len + 1).bool()
attn_scores = attn_scores.masked_fill(mask, float('-inf'))
attn_weights = F.softmax(attn_scores, dim=-1)
out = torch.matmul(attn_weights, V)
out = out.transpose(1, 2).contiguous().view(B, S, -1)
return self.o_proj(out), kv_cache
def autoregressive_generate(model, tokenizer, prompt: str, max_new_tokens: int = 50):
tokens = tokenizer(prompt, return_tensors="pt")["input_ids"]
kv_cache = KVCache()
with torch.no_grad():
output, kv_cache = model(tokens, kv_cache=kv_cache)
generated = []
next_token = output[:, -1:, :]
for step in range(max_new_tokens):
new_token_id = next_token.argmax(dim=-1)
new_token_embed = model.embed(new_token_id)
with torch.no_grad():
output, kv_cache = model(new_token_embed, kv_cache=kv_cache)
token_id = output[:, -1, :].argmax(dim=-1).item()
generated.append(token_id)
next_token = output[:, -1:, :]
if token_id == tokenizer.eos_token_id:
break
return tokenizer.decode(generated)
Speculative Decoding
The Sequential Bottleneck
LLM inference is inherently sequential, relying heavily on memory bandwidth rather than compute power, which speculative decoding seeks to optimize.
The Algorithm
Speculative decoding involves a draft model generating candidate tokens, which are then verified in parallel by the target model. This maintains the same output distribution but speeds up the process significantly.
## Python Example: Speculative Decoding from Scratch
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
from typing import List, Tuple
def speculative_decode(draft_model, target_model, tokenizer, prompt: str, max_new_tokens: int = 100, K: int = 4, temperature: float = 1.0, device: str = "cuda") -> Tuple[str, dict]:
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
generated = input_ids.clone()
stats = {"accepted": 0, "rejected": 0, "rounds": 0}
draft_model.eval()
target_model.eval()
while generated.shape[1] - input_ids.shape[1] < max_new_tokens:
stats["rounds"] += 1
draft_tokens = []
draft_probs = []
draft_input = generated.clone()
with torch.no_grad():
for _ in range(K):
draft_out = draft_model(draft_input)
logits = draft_out.logits[:, -1, :] / temperature
probs = F.softmax(logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
draft_tokens.append(next_token)
draft_probs.append(probs)
draft_input = torch.cat([draft_input, next_token], dim=1)
candidate_sequence = torch.cat([generated] + draft_tokens, dim=1)
with torch.no_grad():
target_out = target_model(candidate_sequence)
target_logits = target_out.logits[:, generated.shape[1]-1:-1, :] / temperature
target_probs = F.softmax(target_logits, dim=-1)
n_accepted = 0
new_tokens = []
for i in range(K):
draft_token_id = draft_tokens[i].squeeze(-1)
p_target = target_probs[:, i, :]
p_draft = draft_probs[i]
token_idx = draft_token_id.unsqueeze(-1)
p_t = p_target.gather(1, token_idx).squeeze()
p_d = p_draft.gather(1, token_idx).squeeze()
accept_prob = torch.minimum(torch.ones_like(p_t), p_t / (p_d + 1e-10))
r = torch.rand_like(accept_prob)
accepted = r < accept_prob
if accepted.all():
new_tokens.append(draft_tokens[i])
n_accepted += 1
stats["accepted"] += 1
else:
residual = torch.clamp(p_target - p_draft, min=0.0)
residual = residual / (residual.sum(dim=-1, keepdim=True) + 1e-10)
corrected_token = torch.multinomial(residual, num_samples=1)
new_tokens.append(corrected_token)
stats["rejected"] += 1
break
if n_accepted == K:
bonus_logits = target_out.logits[:, -1, :] / temperature
bonus_probs = F.softmax(bonus_logits, dim=-1)
bonus_token = torch.multinomial(bonus_probs, num_samples=1)
new_tokens.append(bonus_token)
for t in new_tokens:
generated = torch.cat([generated, t], dim=1)
if tokenizer.eos_token_id in new_tokens[-1]:
break
total_tokens = stats["accepted"] + stats["rejected"]
alpha = stats["accepted"] / max(total_tokens, 1)
stats["acceptance_rate"] = alpha
stats["effective_tokens_per_round"] = total_tokens / max(stats["rounds"], 1)
output_ids = generated[0, input_ids.shape[1]:]
return tokenizer.decode(output_ids, skip_special_tokens=True), stats
## Usage
def demo():
draft_model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
target_model_name = "meta-llama/Llama-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(target_model_name)
draft_model = AutoModelForCausalLM.from_pretrained(draft_model_name, torch_dtype=torch.float16)
target_model = AutoModelForCausalLM.from_pretrained(target_model_name, torch_dtype=torch.float16)
prompt = "Explain the difference between a kernel and a hypervisor:"
text, stats = speculative_decode(draft_model, target_model, tokenizer, prompt, max_new_tokens=200, K=4)
print(f"Generated: {text}")
print(f"Acceptance rate: {stats['acceptance_rate']:.2%}")
print(f"Avg tokens per round: {stats['effective_tokens_per_round']:.2f}")
print(f"Theoretical speedup: ~{stats['effective_tokens_per_round']:.1f}x")
Beyond Basic Speculative Decoding
EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency)
EAGLE eliminates the need for a separate draft model by incorporating a prediction head into the target model, achieving high acceptance rates with minimal extra computation.
Medusa
Medusa incorporates multiple decoding heads into the base model, allowing for efficient parallel processing of candidate tokens.
Lookahead Decoding
This method uses Jacobi iteration to process multiple future positions in parallel without the need for a draft model, though with generally lower acceptance rates.
Speculative decoding is particularly useful when latency is a priority, and the task has a predictable structure suitable for exploitation by the draft model.
Conclusion
The techniques covered in both parts—quantization, pruning, knowledge distillation, KV caching, and speculative decoding—address different bottlenecks in LLM inference. Each method offers unique benefits, and their optimal use depends on the specific workload and hardware constraints. Understanding these techniques enables more efficient deployment of LLMs in practical applications.