Efficiency Frontiers › Code Implementation
Language AI HandbookAll chapters Efficiency Frontiers › Code ImplementationOne section · Sections in this chapter6 / 10 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
This section implements several key efficiency techniques: quantization, KV-cache management, and a simple demonstration of speculative decoding concepts.
Setting Up
import numpy as np
# Set random seed for reproducibility
np.random.seed(42)Quantization: From Float to Integer
Let's implement a simple uniform quantization scheme to see how bit-width affects representational precision.
def quantize(x: np.ndarray, bits: int) -> tuple:
"""Uniform symmetric quantization."""
x_max = np.max(np.abs(x))
num_levels = 2 ** (bits - 1) - 1 # symmetric around 0
scale = x_max / num_levels
x_int = np.round(x / scale).astype(np.int32)
x_int = np.clip(x_int, -num_levels, num_levels)
return x_int, scale
def dequantize(x_int: np.ndarray, scale: float) -> np.ndarray:
"""Reconstruct float values from quantized representation."""
return x_int.astype(np.float32) * scale
# Generate a sample weight matrix (mimicking a small transformer weight)
W = np.random.randn(64, 64).astype(np.float32)
# Quantize at different bit widths
results = {}
for bits in [8, 4, 2]:
W_int, scale = quantize(W, bits)
W_reconstructed = dequantize(W_int, scale)
error = np.mean((W - W_reconstructed) ** 2)
# Account for packed storage at the target bit-width. The demonstration
# array remains int32 so that dequantization is straightforward.
memory_bytes = W.size * bits // 8
float_bytes = W.nbytes
results[bits] = {
"mse": error,
"memory_bytes": memory_bytes,
"float_bytes": float_bytes,
"compression": float_bytes / memory_bytes,
}Quantization Comparison (64x64 weight matrix) Bits MSE Memory (bytes) Compression ------------------------------------------------------- 8 0.000080 4096 4.0x 4 0.026122 2048 8.0x 2 0.850858 1024 16.0x Original FP32 memory: 16384 bytes
The table shows how quantization error increases as bit-width decreases. At 8 bits, the mean squared error is tiny, and the weight values are nearly indistinguishable from the original. At 4 bits, the error grows but remains small enough for most practical purposes. At 2 bits, reconstruction quality degrades substantially, which is why 2-bit quantization is rarely used in practice without special techniques. The compression ratio grows inversely with bit-width: 8-bit weights occupy 4x less memory than FP32, 4-bit weights occupy 8x less.
KV-Cache Simulation
Let's simulate the KV-cache to understand how its size grows with context length and how Grouped Query Attention reduces it.
def kv_cache_memory_mb(
n_tokens: int,
n_layers: int,
n_kv_heads: int,
head_dim: int,
bytes_per_value: int = 2, # BF16
) -> float:
"""Compute KV-cache memory in MB."""
# K and V each: n_layers * n_kv_heads * head_dim * n_tokens
kv_elements = 2 * n_layers * n_kv_heads * head_dim * n_tokens
return (kv_elements * bytes_per_value) / (1024**2)
# Model configurations
configs = {
"7B MHA (32 heads)": {"n_layers": 32, "n_kv_heads": 32, "head_dim": 128},
"7B GQA (8 KV heads)": {"n_layers": 32, "n_kv_heads": 8, "head_dim": 128},
"7B MQA (1 KV head)": {"n_layers": 32, "n_kv_heads": 1, "head_dim": 128},
}
context_lengths = [512, 1024, 2048, 4096, 8192, 16384, 32768]
# Compute memory for each config and context length
kv_memory = {}
for name, cfg in configs.items():
kv_memory[name] = [
kv_cache_memory_mb(
n, cfg["n_layers"], cfg["n_kv_heads"], cfg["head_dim"]
)
for n in context_lengths
]KV-Cache Memory (MB) by Context Length and Attention Variant Config 512 1024 2048 4096 8192 16384 32768 ------------------------------------------------------------------------------------------ 7B MHA (32 heads) 256.0 512.0 1024.0 2048.0 4096.0 8192.0 16384.0 7B GQA (8 KV heads) 64.0 128.0 256.0 512.0 1024.0 2048.0 4096.0 7B MQA (1 KV head) 8.0 16.0 32.0 64.0 128.0 256.0 512.0
The table shows how dramatically context length affects KV-cache requirements. At 32,768 tokens, a standard Multi-Head Attention (MHA) model requires several gigabytes just for the KV-cache of a single request. With Grouped Query Attention (GQA, 8 KV heads), this is reduced by 4x. Multi-Query Attention (MQA, 1 KV head) reduces it by 32x but at a higher quality cost. These tradeoffs explain why Llama 3 and Mistral models use GQA by default.
Measuring Arithmetic Intensity
Arithmetic intensity determines whether an operation is compute-bound or memory-bound. Let's compute it for key transformer operations.
def arithmetic_intensity(flops: float, bytes_accessed: float) -> float:
"""Compute arithmetic intensity in FLOPs/byte."""
return flops / bytes_accessed
# For a matrix multiplication C = A @ B with shapes (m, k) and (k, n)
# FLOPs: 2 * m * k * n (multiply and add for each output element)
# Bytes: (m*k + k*n + m*n) * bytes_per_element
def matmul_intensity(
m: int, k: int, n: int, bytes_per_element: int = 2
) -> float:
flops = 2 * m * k * n
bytes_acc = (m * k + k * n + m * n) * bytes_per_element
return arithmetic_intensity(flops, bytes_acc)
# For attention softmax on a row of n elements
# FLOPs: ~5n (max, subtract, exp, sum, divide)
# Bytes: n * bytes_per_element * 2 (read + write)
def softmax_intensity(n: int, bytes_per_element: int = 2) -> float:
flops = 5 * n
bytes_acc = n * bytes_per_element * 2
return arithmetic_intensity(flops, bytes_acc)
# Common sizes
seq_len = 2048
d_model = 4096
d_ff = 16384
batch_size = 1
ops = {
"MatMul (QK^T), seq=2048, d=128": matmul_intensity(seq_len, 128, seq_len),
"MatMul (FF up), d=4096->16384": matmul_intensity(
batch_size * seq_len, d_model, d_ff
),
"MatMul (FF down), d=16384->4096": matmul_intensity(
batch_size * seq_len, d_ff, d_model
),
"Softmax row, n=2048": softmax_intensity(seq_len),
"Softmax row, n=256": softmax_intensity(256),
}
# H100 hardware roofline (approximate)
peak_flops_per_byte = 300 # H100 ~1000 TFLOPS BF16, ~3 TB/s HBM bandwidthArithmetic Intensity Analysis Operation FLOPs/byte Bottleneck --------------------------------------------------------------------------- MatMul (QK^T), seq=2048, d=128 113.8 Memory MatMul (FF up), d=4096->16384 1260.3 Compute MatMul (FF down), d=16384->4096 1260.3 Compute Softmax row, n=2048 1.2 Memory Softmax row, n=256 1.2 Memory H100 roofline crossover: ~300 FLOPs/byte Operations above crossover are compute-bound; below are memory-bound.
This analysis reveals why FlashAttention matters so much. The matrix multiplication and softmax operations in attention are memory-bound on modern GPUs, while the feed-forward layer matrix multiplications are compute-bound. FlashAttention specifically targets the memory-bound attention operations by reducing HBM accesses, which is why it delivers such substantial speedups.
Speculative Decoding: A Simple Demonstration
Let's implement a simplified speculative decoding simulation to illustrate the acceptance-rejection mechanism.
def simulate_speculative_decoding(
n_tokens: int,
k_draft: int,
p_accept: float,
draft_latency_ms: float,
verifier_latency_ms: float,
) -> dict:
"""
Simulate speculative decoding throughput.
Parameters
----------
n_tokens : int
Total tokens to generate.
k_draft : int
Draft tokens per round.
p_accept : float
Per-token acceptance probability.
draft_latency_ms : float
Latency of draft model per token (ms).
verifier_latency_ms : float
Latency of verifier forward pass (ms), regardless of k.
"""
rng = np.random.RandomState(42)
tokens_generated = 0
total_time_ms = 0.0
rounds = 0
while tokens_generated < n_tokens:
# Draft phase: k_draft autoregressive steps
draft_time = k_draft * draft_latency_ms
# Verifier phase: one forward pass evaluating all k draft tokens
verifier_time = verifier_latency_ms
# Acceptance is sequential: after the first rejection, later draft
# tokens are discarded. The verifier supplies the replacement token,
# or one bonus token when the whole draft prefix is accepted.
accepted_prefix = 0
for draw in rng.random(k_draft):
if draw >= p_accept:
break
accepted_prefix += 1
generated_this_round = accepted_prefix + 1
tokens_generated += min(
generated_this_round, n_tokens - tokens_generated
)
total_time_ms += draft_time + verifier_time
rounds += 1
tokens_per_second = tokens_generated / (total_time_ms / 1000)
return {
"tokens_per_second": tokens_per_second,
"total_time_ms": total_time_ms,
"rounds": rounds,
"avg_tokens_per_round": tokens_generated / rounds,
}
# Baseline: verifier only (no speculation)
def baseline_generation(n_tokens: int, verifier_latency_ms: float) -> dict:
total_time_ms = n_tokens * verifier_latency_ms
return {
"tokens_per_second": n_tokens / (total_time_ms / 1000),
"total_time_ms": total_time_ms,
}
# Parameters
n_tokens = 1000
k_draft = 4
draft_latency = 5.0 # ms per token (small 7B model)
verifier_latency = 50.0 # ms per forward pass (70B model)
baseline = baseline_generation(n_tokens, verifier_latency)
scenarios = {
"p_accept=0.5": simulate_speculative_decoding(
n_tokens, k_draft, 0.5, draft_latency, verifier_latency
),
"p_accept=0.7": simulate_speculative_decoding(
n_tokens, k_draft, 0.7, draft_latency, verifier_latency
),
"p_accept=0.85": simulate_speculative_decoding(
n_tokens, k_draft, 0.85, draft_latency, verifier_latency
),
"p_accept=0.95": simulate_speculative_decoding(
n_tokens, k_draft, 0.95, draft_latency, verifier_latency
),
}Speculative Decoding Simulation (4 draft tokens, 1000 total tokens) Draft latency: 5.0 ms/token, Verifier latency: 50.0 ms/pass Scenario tok/s Speedup vs baseline Avg tokens/round ---------------------------------------------------------------------- Baseline (no spec) 20.0 1.0x 1.0 p_accept=0.5 26.9 1.34 x 1.88 p_accept=0.7 39.1 1.96 x 2.74 p_accept=0.85 52.1 2.61 x 3.65 p_accept=0.95 64.6 3.23 x 4.52
The simulation shows how speculative decoding speedup scales with acceptance rate. At 50% acceptance, the expected speedup is about 1.4x because draft-model overhead offsets part of the gain. At 85-95% acceptance (typical for a well-matched draft-verifier pair on natural text), the idealized speedup is roughly 2.6-3.2x. The key insight is that the draft model cost is small relative to the verifier, so even moderate acceptance rates produce meaningful gains.
Comparing Architecture Efficiency
Let's compute theoretical compute requirements for standard transformer versus MoE architectures.
def transformer_flops_per_token(
n_layers: int, d_model: int, d_ff: int, n_heads: int, seq_len: int
) -> dict:
"""
Estimate FLOPs per forward pass token for a dense transformer.
Counts multiply-adds as 2 FLOPs.
"""
d_head = d_model // n_heads
# Attention: QKV projections + attention + output projection
attn_qkv = 2 * 3 * d_model * d_model # QKV linear projections
attn_scores = 2 * n_heads * seq_len * d_head # Q @ K^T
attn_values = 2 * n_heads * seq_len * d_head # attn @ V
attn_out = 2 * d_model * d_model # output projection
# Feed-forward: up projection + activation + down projection
ff_up = 2 * d_model * d_ff
ff_down = 2 * d_ff * d_model
per_layer = (
attn_qkv + attn_scores + attn_values + attn_out + ff_up + ff_down
)
total = n_layers * per_layer
return {
"total_flops": total,
"attn_flops_per_layer": attn_qkv + attn_scores + attn_values + attn_out,
"ff_flops_per_layer": ff_up + ff_down,
}
def moe_flops_per_token(
n_layers: int,
d_model: int,
d_ff: int,
n_heads: int,
seq_len: int,
n_experts: int,
top_k: int,
) -> dict:
"""
Estimate active FLOPs per token for a MoE transformer.
Only top_k experts fire per token.
"""
d_head = d_model // n_heads
# Attention is the same as dense
attn_qkv = 2 * 3 * d_model * d_model
attn_scores = 2 * n_heads * seq_len * d_head
attn_values = 2 * n_heads * seq_len * d_head
attn_out = 2 * d_model * d_model
# Only top_k of n_experts are active
ff_up = 2 * top_k * d_model * d_ff
ff_down = 2 * top_k * d_ff * d_model
# Gating network
gating = 2 * d_model * n_experts
per_layer = (
attn_qkv
+ attn_scores
+ attn_values
+ attn_out
+ ff_up
+ ff_down
+ gating
)
total = n_layers * per_layer
total_params_ff = n_layers * n_experts * (d_model * d_ff + d_ff * d_model)
active_params_ff = n_layers * top_k * (d_model * d_ff + d_ff * d_model)
return {
"total_flops": total,
"total_ff_params": total_params_ff,
"active_ff_params": active_params_ff,
}
# Model configurations: 7B dense vs Mixtral-style 8x7B MoE
dense_cfg = {
"n_layers": 32,
"d_model": 4096,
"d_ff": 14336,
"n_heads": 32,
"seq_len": 512,
}
moe_cfg = {
"n_layers": 32,
"d_model": 4096,
"d_ff": 14336,
"n_heads": 32,
"seq_len": 512,
"n_experts": 8,
"top_k": 2,
}
dense_result = transformer_flops_per_token(**dense_cfg)
moe_result = moe_flops_per_token(**moe_cfg)Dense Transformer vs Mixtral-style MoE (per token) Metric Dense 7B MoE 8x7B ----------------------------------------------------------------- Active FLOPs per token (GFLOPs) 12.08 19.60 FF param count (M) 32 - Total FF params (M) - 30064M Active FF params (M) - 7516M MoE uses 1.6x the FLOPs of dense for this seq_len, but has ~4x more total FF parameters with top-2 routing. The quality improvement from larger capacity justifies the parameter overhead.
The computation confirms the MoE efficiency story. With top-2 routing across 8 experts, each MoE layer uses roughly 2x the feed-forward FLOPs of a dense layer (because 2 of 8 experts are active), but the model has 8x more total feed-forward capacity. The practical result is that Mixtral 8x7B competes with Llama 2 70B on many benchmarks while being much faster at inference.