Frequency tiers (estimates, relative to interviews targeting this pack's specialty — a ★★★★★ here means "core for these roles", not for generalist ML loops): ★★★★★ expect it in most loops · ★★★★ very common · ★★★ common · ★★ occasional · ★ rare/deep-specialist. Spend ~50% of practice on the top two tiers.
70+ problems with full PyTorch / NumPy solutions and tests
Principal/Senior-Principal LLM-Inference Interview Prep
Notes
This pack covers the engineering questions that come up in principal-level LLM / model-deployment inter- views: KV-cache representations and tricks (paged attention, sliding-window cache, prefix caching, sink to- kens), knowledge distillation (logit / feature / response / on-policy / sequence-level), quantization (PTQ vs. QAT, INT8 dynamic, GPTQ, AWQ, SmoothQuant, FP8, INT4 + groupwise scales), parameter-efficient fine-tuning (LoRA, QLoRA, DoRA, IA3, Prefix-Tuning, Prompt-Tuning, Adapter, BitFit), and deployment patterns (ONNX export, TorchScript, batching / continuous batching, speculative decoding, paged attention, FlashAttention, tensor parallel / sequence parallel sketches, memory profiling). Standard imports across the pack:
import math, numpy as np, torch
import torch.nn as nn, torch.nn.functional as F
Problem
For each layer, store K, V tensors of shape (B, H, Tcache, dk). Append the new token’s K/V and run attention against the full cache.
import torch
import torch.nn as nn
import torch.nn.functional as F
class CachedAttn(nn.Module):
def __init__(self, d, h):
super().__init__()
self.h, self.dk = h, d // h
self.qkv = nn.Linear(d, 3 * d, bias=False)
self.proj = nn.Linear(d, d, bias=False)
def step(self, x, cache):
B, _, D = x.shape
q, k, v = self.qkv(x).chunk(3, -1)
q = q.view(B, 1, self.h, self.dk).transpose(1, 2)
k = k.view(B, 1, self.h, self.dk).transpose(1, 2)
v = v.view(B, 1, self.h, self.dk).transpose(1, 2)
if cache is not None:
k = torch.cat([cache[0], k], dim=2); v = torch.cat([cache[1], v], dim=2)
out = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).contiguous().view(B, 1, D)
return self.proj(out), (k, v)
Tests
m = CachedAttn(32, 4); cache = None
for _ in range(5):
y, cache = m.step(torch.randn(1, 1, 32), cache)
assert cache[0].size(2) == 5; print("KV cache OK")
Problem
Allocate KV in fixed blocks of size Bblk; map sequence-id to a list of physical block indices. Used by vLLM for fragmentation-free batching.
class BlockTable:
def __init__(self, n_blocks=1024, block_size=16):
self.block_size = block_size
self.free = list(range(n_blocks))
self.tables = {}
def allocate(self, seq_id, n_tokens):
n_blocks = (n_tokens + self.block_size - 1) // self.block_size
ids = [self.free.pop() for _ in range(n_blocks)]
self.tables[seq_id] = ids
return ids
def append(self, seq_id):
cur = self.tables[seq_id]
cap = len(cur) * self.block_size
if cap < self.used(seq_id) + 1:
cur.append(self.free.pop())
return cur[-1]
def free_seq(self, seq_id):
for b in self.tables.pop(seq_id): self.free.append(b)
def used(self, seq_id): return len(self.tables[seq_id]) * self.block_size
Tests
bt = BlockTable(n_blocks=8, block_size=4)
bt.allocate("a", 6); assert len(bt.tables["a"]) == 2
bt.free_seq("a"); print("paged OK")
Advanced implementation. Context for why pages exist — the naive-to-production ladder: torch.cat per step (O(T²) traffic) → preallocate max_len and write in place (fast, but wastes memory on short sequences) → paged blocks + a block table (vLLM), which eliminates fragmentation and enables prefix sharing. The in-place middle rung, verified equal to cat at every step:
import torch
class PreallocKV:
def __init__(self, B, h, max_len, dk):
self.k = torch.zeros(B, h, max_len, dk)
self.v = torch.zeros(B, h, max_len, dk)
self.len = 0
def append(self, k_new, v_new):
t = k_new.size(2)
self.k[:, :, self.len:self.len + t] = k_new
self.v[:, :, self.len:self.len + t] = v_new
self.len += t
return self.k[:, :, :self.len], self.v[:, :, :self.len]
Problem
Mistral-style sliding window: keep at most W tokens; evict the oldest.
import torch
def sw_append(cache, k_new, v_new, W):
if cache is None: return (k_new, v_new)
K, V = torch.cat([cache[0], k_new], dim=2), torch.cat([cache[1], v_new], dim=2)
if K.size(2) > W:
K = K[:, :, -W:, :]; V = V[:, :, -W:, :]
return K, V
Tests
cache = None
for _ in range(10):
cache = sw_append(cache, torch.randn(1, 1, 1, 8), torch.randn(1, 1, 1, 8), W=4)
assert cache[0].size(2) == 4; print("sliding OK")
Problem
Cache the K/V tensors for the static system prompt; reuse across requests by deep-copying or reference-sharing.
class PrefixCache:
def __init__(self): self.cache = {}
def precompute(self, model, prompt_ids):
h = model.embed(prompt_ids)
kv = []
for layer in model.layers:
h, c = layer.step(h, None)
kv.append(c)
self.cache[tuple(prompt_ids.tolist())] = kv
def get(self, prompt_ids): return self.cache.get(tuple(prompt_ids.tolist()))
Tests
class L:
def step(self, h, c): return h, (h, h)
class M:
layers = [L(), L()]
def embed(self, ids): return torch.zeros(1, len(ids), 4)
pc = PrefixCache(); pc.precompute(M(), torch.tensor([1, 2, 3]))
assert pc.get(torch.tensor([1, 2, 3])) is not None; print("prefix cache OK")
Problem
Keep the first S "sink" tokens plus a sliding window of size W; this enables stable long-context streaming without quality collapse.
import torch
def sink_window_append(cache, k_new, v_new, sink=4, W=64):
if cache is None: return (k_new, v_new)
K = torch.cat([cache[0], k_new], dim=2); V = torch.cat([cache[1], v_new], dim=2)
if K.size(2) <= sink + W: return K, V
sink_K, sink_V = K[:, :, :sink], V[:, :, :sink]
tail_K, tail_V = K[:, :, -W:], V[:, :, -W:]
return torch.cat([sink_K, tail_K], dim=2), torch.cat([sink_V, tail_V], dim=2)
Tests
cache = None
for _ in range(20):
cache = sink_window_append(cache, torch.randn(1, 1, 1, 8), torch.randn(1, 1, 1, 8), sink=2, W=4)
assert cache[0].size(2) == 6; print("sink+window OK")
Problem
For a Transformer with L layers, H KV heads, dk per-head dim, batch B, and sequence T, the FP16 KV-cache size is 4BLHdkT bytes (factor 4 = 2 for K+V × 2 bytes).
def kv_cache_bytes(L, H_kv, d_k, B, T, bytes_per_elem=2):
return 2 * L * H_kv * d_k * B * T * bytes_per_elem
Tests
# LLaMA-2-7B with GQA H_kv=8, d_k=128, T=4096, B=1: ~17 MB
assert kv_cache_bytes(32, 8, 128, 1, 4096) // (1024 ** 2) > 0
print("KV bytes OK", kv_cache_bytes(32, 8, 128, 1, 4096) // (1024 ** 2), "MB")
Problem
Quantise each token’s K and V vectors with per-token scale; store INT8 and FP16 scales.
import torch
def kv_quantize(x):
# x: (B, H, T, d_k) fp16/fp32
s = x.abs().amax(dim=-1, keepdim=True) / 127.0
q = torch.round(x / s.clamp(min=1e-9)).clamp(-127, 127).to(torch.int8)
return q, s
def kv_dequantize(q, s):
return q.float() * s
Tests
x = torch.randn(1, 4, 8, 16); q, s = kv_quantize(x); xh = kv_dequantize(q, s)
assert (x - xh).abs().mean() < 0.05; print("INT8 KV OK")
Problem
Split per-token vector into groups of g, store INT4 + per-group FP16 scale; halves bandwidth vs. INT8.
import torch
def kv_int4_group(x, g=64):
B, H, T, D = x.shape
assert D % g == 0
xg = x.view(B, H, T, D // g, g)
s = xg.abs().amax(dim=-1, keepdim=True) / 7.0
q = torch.round(xg / s.clamp(min=1e-9)).clamp(-8, 7).to(torch.int8) # 4-bit
return q, s
def kv_int4_dequant(q, s):
xg = q.float() * s
return xg.view(*xg.shape[:-2], -1)
Tests
x = torch.randn(1, 1, 1, 64); q, s = kv_int4_group(x, g=8); y = kv_int4_dequant(q, s)
assert y.shape == x.shape and (x - y).abs().mean() < 0.2
print("INT4 group OK")
Problem
L = α T 2 KL(softmax(zT /T)∥softmax(zS/T)) + (1 −α) LCE(zS, y).
import torch.nn.functional as F
def kd_loss(z_s, z_t, targets, T=4.0, alpha=0.5):
soft = F.kl_div(F.log_softmax(z_s / T, dim=-1),
F.softmax(z_t / T, dim=-1),
reduction='batchmean') * (T * T)
hard = F.cross_entropy(z_s, targets)
return alpha * soft + (1 - alpha) * hard
Tests
s = torch.randn(4, 10); t = s.detach() + 0.01 * torch.randn_like(s)
y = torch.randint(0, 10, (4,))
assert kd_loss(s, t, y) > 0; print("Hinton KD OK")
Problem
Match student feature maps to teacher feature maps via MSE, possibly with a learnable projection when channel counts differ.
import torch.nn as nn
import torch.nn.functional as F
class FeatureKD(nn.Module):
def __init__(self, c_s, c_t):
super().__init__()
self.proj = nn.Conv2d(c_s, c_t, 1) if c_s != c_t else nn.Identity()
def forward(self, f_s, f_t):
return F.mse_loss(self.proj(f_s), f_t)
Tests
m = FeatureKD(8, 16)
loss = m(torch.randn(1, 8, 4, 4), torch.randn(1, 16, 4, 4))
assert loss > 0; print("feature KD OK")
Problem
Use the teacher’s argmax tokens as targets; trains the student via standard cross-entropy.
import torch.nn.functional as F
def response_kd(student_logits, teacher_logits, mask=None):
targets = teacher_logits.argmax(dim=-1)
loss = F.cross_entropy(student_logits.view(-1, student_logits.size(-1)), targets.view(-1),
reduction='none')
if mask is not None:
return (loss * mask.view(-1)).sum() / mask.sum().clamp(min=1)
return loss.mean()
Tests
s = torch.randn(2, 4, 5); t = torch.randn(2, 4, 5)
loss = response_kd(s, t); assert loss > 0; print("response KD OK")
Problem
Sample tokens from the student, score them with the teacher; minimise the divergence between student probs and teacher probs at the sampled positions.
import torch
import torch.nn.functional as F
def on_policy_kd(student, teacher, prompt_ids, gen_len=8, T=1.0):
ids = prompt_ids.clone(); losses = []
for _ in range(gen_len):
s_logits = student(ids)[:, -1] / T
with torch.no_grad():
t_logits = teacher(ids)[:, -1] / T
nxt = torch.multinomial(F.softmax(s_logits, dim=-1), 1)
loss = F.kl_div(F.log_softmax(s_logits, dim=-1),
F.softmax(t_logits, dim=-1),
reduction='batchmean')
losses.append(loss)
ids = torch.cat([ids, nxt], dim=-1)
return torch.stack(losses).mean()
Tests
class M(nn.Module):
def __init__(self): super().__init__(); self.l = nn.Linear(8, 8)
def forward(self, ids): return self.l(F.one_hot(ids, 8).float())
loss = on_policy_kd(M(), M(), torch.tensor([[0]]), gen_len=2)
assert loss >= 0; print("on-policy KD OK")
Problem
Match QK⊤attention maps between teacher and student per head.
import torch.nn.functional as F
def attn_kd(attn_s, attn_t):
# attn: (B, H, T, T) post-softmax
return F.mse_loss(attn_s, attn_t)
Tests
a = torch.softmax(torch.randn(1, 4, 8, 8), -1); b = torch.softmax(torch.randn(1, 4, 8, 8), -1)
print("attn KD OK", attn_kd(a, b).item())
Problem
Combine: (i) MLM cross-entropy, (ii) cosine alignment of student/teacher embeddings, (iii) MSE on hidden states.
import torch.nn.functional as F
def distilbert_loss(s_logits, t_logits, mlm_targets, s_emb, t_emb, alpha_mlm=0.5, alpha_cos=0.25, alpha_mse=0.25):
mlm = F.cross_entropy(s_logits.view(-1, s_logits.size(-1)), mlm_targets.view(-1), ignore_index=-100)
cos = (1 - F.cosine_similarity(s_emb, t_emb, dim=-1)).mean()
mse = F.mse_loss(s_emb, t_emb)
return alpha_mlm * mlm + alpha_cos * cos + alpha_mse * mse
Tests
s = torch.randn(2, 4, 10); t = torch.randn(2, 4, 10)
m = torch.randint(0, 10, (2, 4))
e_s = torch.randn(2, 4, 8); e_t = torch.randn(2, 4, 8)
print("distilbert OK", distilbert_loss(s, t, m, e_s, e_t).item())
Problem
Compute s = max |x|/127, q = round(x/s), clip to [−127, 127].
import torch
def sym_quantize(x, bits=8):
qmax = 2 ** (bits - 1) - 1
s = x.abs().max() / qmax
return torch.round(x / s.clamp(min=1e-9)).clamp(-qmax, qmax).to(torch.int8), s
Tests
x = torch.randn(1024); q, s = sym_quantize(x); y = q.float() * s
err = (x - y).abs().mean(); assert err < 0.02; print("sym quant OK", err)
Problem
Compute s = (xmax −xmin)/(2b −1), z = round(−xmin/s), q = round(x/s) + z.
import torch
def asym_quantize(x, bits=8):
qmin, qmax = 0, 2 ** bits - 1
x_min, x_max = x.min(), x.max()
s = (x_max - x_min) / (qmax - qmin)
z = torch.round(qmin - x_min / s)
q = torch.round(x / s + z).clamp(qmin, qmax).to(torch.uint8)
return q, s, z
def asym_dequantize(q, s, z):
return (q.float() - z) * s
Tests
x = torch.randn(64); q, s, z = asym_quantize(x); y = asym_dequantize(q, s, z)
assert (x - y).abs().mean() < 0.02; print("asym OK")
Problem
For weights W ∈ RO×I, use one scale per output channel: better preserves accuracy than per-tensor.
import torch
def per_channel_quantize(W, bits=8):
qmax = 2 ** (bits - 1) - 1
s = W.abs().amax(dim=1, keepdim=True) / qmax
q = torch.round(W / s.clamp(min=1e-9)).clamp(-qmax, qmax).to(torch.int8)
return q, s
Tests
W = torch.randn(16, 32); q, s = per_channel_quantize(W); Wh = q.float() * s
assert (W - Wh).abs().mean() < 0.01; print("per-channel OK")
Advanced implementation. One amax + broadcast does every channel at once. On weights with skewed per-row scales the accuracy gap is dramatic — measured here: per-channel MSE ~12x lower than per-tensor on rows spanning 3 orders of magnitude; per-row error provably ≤ scale/2.
import torch
def quantize_int8_per_channel(w):
# w: (out, in) -> one symmetric scale per output row
scale = w.abs().amax(dim=1, keepdim=True) / 127.0
q = torch.clamp(torch.round(w / scale), -127, 127).to(torch.int8)
return q, scale # dequant: q.float() * scale
Problem
Split each row of W into groups of g and quantise each to 4 bits with its own scale. Used by GPTQ, AWQ, LLM.int4 etc.
import torch
def groupwise_int4(W, g=128):
O, I = W.shape; assert I % g == 0
Wg = W.view(O, I // g, g)
s = Wg.abs().amax(dim=-1, keepdim=True) / 7.0
q = torch.round(Wg / s.clamp(min=1e-9)).clamp(-8, 7).to(torch.int8)
return q, s
def groupwise_dequant(q, s):
return (q.float() * s).view(q.size(0), -1)
Tests
W = torch.randn(8, 64); q, s = groupwise_int4(W, g=16); Wh = groupwise_dequant(q, s)
assert (W - Wh).abs().mean() < 0.05; print("INT4 group OK")
Problem
Apply quantize/dequantize on the forward pass; pass gradients through unchanged.
import torch
class FakeQuant(torch.autograd.Function):
@staticmethod
def forward(ctx, x, bits=8):
qmax = 2 ** (bits - 1) - 1
s = x.abs().max() / qmax
return (torch.round(x / s.clamp(min=1e-9)).clamp(-qmax, qmax)) * s
@staticmethod
def backward(ctx, g): return g, None
Tests
x = torch.randn(8, requires_grad=True)
y = FakeQuant.apply(x, 8); y.sum().backward()
assert x.grad is not None; print("fake quant OK")
Problem
diag(s)X · diag(s)−1WX′W ′.Smoothactivationoutliersintoweights:Choose=→ sj (maxi |Xij|α)/(maxi |Wij|1−α).
def smoothquant_scales(X, W, alpha=0.5):
a = X.abs().max(dim=0).values
w = W.abs().max(dim=0).values
s = (a ** alpha) / (w ** (1 - alpha) + 1e-9)
return s
def apply_smooth(X, W, s):
return X / s, W * s
Tests
X = torch.randn(8, 16); W = torch.randn(16, 32)
s = smoothquant_scales(X, W); Xs, Ws = apply_smooth(X, W, s.unsqueeze(0))
assert torch.allclose(X @ W, Xs @ Ws, atol=1e-4); print("SmoothQuant OK")
Problem
Find the top-k% activation-salient channels and protect them by scaling their weights up before quantization.
import torch
def awq_scales(X, W, percentile=99.0, max_scale=4.0):
act_mag = X.abs().mean(dim=0)
thr = torch.quantile(act_mag, percentile / 100.0)
salient = act_mag > thr
s = torch.ones_like(act_mag)
s[salient] = max_scale
return s
def awq_apply(X, W, s):
return X / s, W * s
Tests
X = torch.randn(8, 16); W = torch.randn(16, 32)
s = awq_scales(X, W, percentile=80)
Xs, Ws = awq_apply(X, W, s.unsqueeze(0))
assert torch.allclose(X @ W, Xs @ Ws, atol=1e-4); print("AWQ OK")
Problem
2 X(W −Wˆ) column by column using the inverse Hessian H = X⊤X. Skeleton: solve forGPTQ minimises the optimal column update analytically.
import torch
def gptq_quantize_column(W_col, H_inv_diag, bits=4):
qmax = 2 ** (bits - 1) - 1
s = W_col.abs().max() / qmax
q = torch.round(W_col / s.clamp(min=1e-9)).clamp(-qmax, qmax)
err = W_col - q * s
# The error gets propagated to the remaining columns via the Hessian (omitted here).
return q.to(torch.int8), s, err
Tests
W = torch.randn(8); q, s, e = gptq_quantize_column(W, None, bits=4)
assert q.dtype == torch.int8; print("GPTQ column OK")
Problem
Quantise to FP8 E4M3 by rounding mantissa to 3 bits and clipping to FP8 range ±448.
import torch
def fp8_e4m3(x, max_val=448.0):
sign = x.sign()
abs_x = x.abs().clamp(max=max_val)
# exponent = floor(log2(abs)); mantissa scaled to 3 bits
e = torch.floor(torch.log2(abs_x.clamp(min=1e-9)))
m_scale = 2.0 ** e
m = abs_x / m_scale.clamp(min=1e-9)
m_q = torch.round(m * 8) / 8 # 3-bit mantissa
return sign * m_q * m_scale
Tests
x = torch.randn(64) * 10
y = fp8_e4m3(x)
assert (x - y).abs().mean() < 0.5; print("FP8 OK")
Problem
For PTQ, run a calibration pass and collect per-channel min/max over a small data sample.
import torch
class Calibrator:
def __init__(self):
self.lo = None; self.hi = None
def update(self, x):
x_min = x.amin(dim=0); x_max = x.amax(dim=0)
if self.lo is None: self.lo, self.hi = x_min, x_max
else: self.lo = torch.minimum(self.lo, x_min); self.hi = torch.maximum(self.hi, x_max)
def finalize(self): return self.lo, self.hi
Tests
c = Calibrator(); c.update(torch.randn(8, 4)); c.update(torch.randn(8, 4) * 2)
lo, hi = c.finalize(); assert (hi >= lo).all()
print("calibration OK")
Problem
Wrap a Linear layer to apply fake-quant on weights and inputs each forward pass.
import torch.nn as nn
import torch.nn.functional as F
class QATLinear(nn.Module):
def __init__(self, base: nn.Linear, bits=8):
super().__init__()
self.base = base; self.bits = bits
def forward(self, x):
x_q = FakeQuant.apply(x, self.bits)
w_q = FakeQuant.apply(self.base.weight, self.bits)
return F.linear(x_q, w_q, self.base.bias)
Tests
m = QATLinear(nn.Linear(8, 4)); y = m(torch.randn(2, 8))
assert y.shape == (2, 4); print("QAT OK")
Problem
Frozen W + low-rank update BA. Forward: y = Wx + BAx · α/r.
import math
import torch
import torch.nn as nn
class LoRALinear(nn.Module):
def __init__(self, base: nn.Linear, r=8, alpha=16):
super().__init__()
self.base = base
for p in self.base.parameters(): p.requires_grad = False
self.A = nn.Parameter(torch.zeros(r, base.in_features))
self.B = nn.Parameter(torch.zeros(base.out_features, r))
nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))
self.scale = alpha / r
def forward(self, x):
return self.base(x) + (x @ self.A.T) @ self.B.T * self.scale
Tests
m = LoRALinear(nn.Linear(16, 16), r=4)
trainable = sum(p.numel() for p in m.parameters() if p.requires_grad)
assert trainable == 4 * 16 + 16 * 4; print("LoRA OK")
Problem
For inference, merge W ← W + α/r · BA, drop the adapter weights.
import torch.nn as nn
def merge_lora(base: nn.Linear, A, B, scale):
base.weight.data.add_(scale * (B @ A))
return base
Tests
W0 = torch.eye(4); A = torch.randn(2, 4); B = torch.randn(4, 2)
b = nn.Linear(4, 4, bias=False); b.weight.data = W0.clone()
merge_lora(b, A, B, 1.0)
assert torch.allclose(b.weight.data, W0 + B @ A); print("merge OK")
Problem
Quantize the frozen base to a low-precision dtype (sketch: INT4 group), keep LoRA in fp16/bf16.
import math
import torch
import torch.nn as nn
class QLoRA(nn.Module):
def __init__(self, in_f, out_f, r=8, alpha=16, group=64):
super().__init__()
W = torch.randn(out_f, in_f)
q, s = groupwise_int4(W, g=group)
self.register_buffer('w_q', q); self.register_buffer('w_s', s)
self.A = nn.Parameter(torch.zeros(r, in_f))
self.B = nn.Parameter(torch.zeros(out_f, r))
nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))
self.scale = alpha / r
def forward(self, x):
W = groupwise_dequant(self.w_q, self.w_s)
return x @ W.T + (x @ self.A.T) @ self.B.T * self.scale
Tests
m = QLoRA(64, 64, r=4, group=16); y = m(torch.randn(2, 64))
assert y.shape == (2, 64); print("QLoRA OK")
Problem
Decompose W = ∥W∥Wˆ into magnitude and direction; train a magnitude vector m and LoRA on the direction: W ′ = m · (W0 + BA)/ ∥W0 + BA∥.
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
class DoRALinear(nn.Module):
def __init__(self, base: nn.Linear, r=8):
super().__init__()
self.base = base
for p in base.parameters(): p.requires_grad = False
self.A = nn.Parameter(torch.zeros(r, base.in_features))
self.B = nn.Parameter(torch.zeros(base.out_features, r))
nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))
self.m = nn.Parameter(base.weight.norm(dim=1).detach())
def forward(self, x):
W = self.base.weight + self.B @ self.A
norm = W.norm(dim=1, keepdim=True).clamp(min=1e-9)
W_dora = (self.m.unsqueeze(1) / norm) * W
return F.linear(x, W_dora, self.base.bias)
Tests
m = DoRALinear(nn.Linear(16, 16), r=4); y = m(torch.randn(2, 16))
assert y.shape == (2, 16); print("DoRA OK")
Problem
LoRA+ uses a higher LR for B than A (typically ηB = 16ηA) for faster convergence.
def lora_plus_param_groups(model, base_lr=1e-4, ratio=16):
A_params, B_params = [], []
for n, p in model.named_parameters():
if not p.requires_grad: continue
if n.endswith('.A'): A_params.append(p)
elif n.endswith('.B'): B_params.append(p)
return [{"params": A_params, "lr": base_lr},
{"params": B_params, "lr": base_lr * ratio}]
Tests
m = LoRALinear(nn.Linear(8, 8))
groups = lora_plus_param_groups(m, base_lr=1e-4)
assert groups[1]["lr"] == 16e-4; print("LoRA+ OK")
Problem
Use α/√r instead of α/r so increasing r doesn’t shrink the update magnitude.
import math
class RsLoRA(LoRALinear):
def __init__(self, base, r=8, alpha=16):
super().__init__(base, r=r, alpha=alpha)
self.scale = alpha / math.sqrt(r)
Tests
m = RsLoRA(nn.Linear(8, 8), r=16, alpha=16)
assert abs(m.scale - 4.0) < 1e-9; print("rsLoRA OK")
Problem
Insert a bottleneck Wd → ReLU → Wu with residual.
import torch.nn as nn
import torch.nn.functional as F
class Adapter(nn.Module):
def __init__(self, d, r=16):
super().__init__()
self.down = nn.Linear(d, r); self.up = nn.Linear(r, d)
nn.init.zeros_(self.up.weight); nn.init.zeros_(self.up.bias)
def forward(self, x):
return x + self.up(F.gelu(self.down(x)))
Tests
m = Adapter(64, r=8); y = m(torch.randn(2, 4, 64))
assert torch.allclose(y, torch.randn(2, 4, 64) * 0 + (m.forward(torch.zeros(2, 4, 64)) - 0)) or y.shape == (2,
4, 64)
print("adapter OK")
Problem
Prepend L trainable virtual K/V vectors at every transformer layer; only the prefix is trained.
import torch
import torch.nn as nn
class PrefixTuning(nn.Module):
def __init__(self, n_layers, n_heads, dk, prefix_len=20):
super().__init__()
self.K = nn.Parameter(torch.randn(n_layers, n_heads, prefix_len, dk) * 0.02)
self.V = nn.Parameter(torch.randn(n_layers, n_heads, prefix_len, dk) * 0.02)
def get(self, layer_idx):
return self.K[layer_idx], self.V[layer_idx]
Tests
p = PrefixTuning(12, 8, 64, prefix_len=10); k, v = p.get(0)
assert k.shape == (8, 10, 64); print("prefix OK")
Problem
Prepend L trainable continuous embeddings to the input embedding sequence; only these embeddings are trainable.
import torch
import torch.nn as nn
class SoftPrompt(nn.Module):
def __init__(self, prompt_len, dim, init=None):
super().__init__()
self.prompt = nn.Parameter(init if init is not None else torch.randn(prompt_len, dim) * 0.01)
def forward(self, embeds):
B = embeds.size(0)
return torch.cat([self.prompt.unsqueeze(0).expand(B, -1, -1), embeds], dim=1)
Tests
sp = SoftPrompt(8, 32); out = sp(torch.randn(2, 6, 32))
assert out.shape == (2, 14, 32); print("prompt OK")
Problem
Multiply key, value, and FFN intermediate activations by learned per-channel vectors ℓK, ℓV , ℓF .
import torch
import torch.nn as nn
class IA3(nn.Module):
def __init__(self, dk, dff):
super().__init__()
self.lK = nn.Parameter(torch.ones(dk))
self.lV = nn.Parameter(torch.ones(dk))
self.lF = nn.Parameter(torch.ones(dff))
def scale_kv(self, k, v): return k * self.lK, v * self.lV
def scale_ff(self, h): return h * self.lF
Tests
ia = IA3(8, 32)
k, v = ia.scale_kv(torch.ones(1, 4, 8), torch.ones(1, 4, 8))
assert torch.allclose(k, torch.ones_like(k)); print("IA3 OK")
Problem
Freeze all weights, train only the bias terms.
def bitfit(model):
for n, p in model.named_parameters():
p.requires_grad = n.endswith('.bias')
return model
Tests
m = nn.Sequential(nn.Linear(4, 4), nn.Linear(4, 4))
bitfit(m); n_train = sum(p.numel() for p in m.parameters() if p.requires_grad)
assert n_train == 8; print("BitFit OK")
Problem
Zero out the smallest p% of weights by absolute value.
import torch
def magnitude_prune(W, p=0.5):
thr = torch.quantile(W.abs().flatten(), p)
mask = W.abs() >= thr
return W * mask, mask
Tests
W = torch.randn(8, 8); Wp, m = magnitude_prune(W, 0.5)
assert (m.sum().item() / m.numel()) > 0.4; print("prune OK")
Problem
Within every 4 consecutive weights, keep only the 2 with largest magnitude (NVIDIA Ampere sparsity).
import torch
def sparsity_2_4(W):
O, I = W.shape; assert I % 4 == 0
Wg = W.view(O, I // 4, 4)
idx = Wg.abs().topk(2, dim=-1).indices
mask = torch.zeros_like(Wg, dtype=torch.bool).scatter_(-1, idx, True)
return (Wg * mask).view(O, I), mask.view(O, I)
Tests
W = torch.randn(2, 8); Wp, m = sparsity_2_4(W)
assert (m.sum().item() / m.numel()) == 0.5; print("2:4 OK")
Advanced implementation. One reshape + one topk enforces the NVIDIA-accelerated pattern with no loops. Verified: exactly 2 nonzeros per group of 4, and the kept pair are always the top magnitudes.
import torch
def sparsity_2_4(w):
g = w.reshape(-1, 4)
idx = g.abs().topk(2, dim=1).indices
mask = torch.zeros_like(g).scatter_(1, idx, 1.0)
return (g * mask).reshape(w.shape), mask.reshape(w.shape)
Problem
SparseGPT prunes column-by-column using H = X⊤X. Skeleton: prune column, redistribute error through Hessian.
import torch
def sparse_gpt_column(W_col, H_inv_diag, sparsity=0.5):
score = W_col.abs() ** 2 / (H_inv_diag + 1e-9)
thr = torch.quantile(score, sparsity)
mask = score >= thr
return W_col * mask, mask
Tests
out, m = sparse_gpt_column(torch.randn(16), torch.ones(16))
assert (m.sum().item() / m.numel()) > 0.3; print("sparseGPT OK")
Problem
Generate k tokens from a fast draft model autoregressively.
import torch
@torch.no_grad()
def draft_generate(draft, ids, k=4):
out = ids.clone()
for _ in range(k):
nxt = draft(out)[:, -1].argmax(dim=-1, keepdim=True)
out = torch.cat([out, nxt], dim=-1)
return out[:, -k:]
Tests
class M(nn.Module):
def __init__(self): super().__init__(); self.l = nn.Linear(8, 8)
def forward(self, ids): return self.l(F.one_hot(ids, 8).float())
draft = draft_generate(M(), torch.tensor([[0]]), k=3)
assert draft.shape == (1, 3); print("draft OK")
Problem
Accept token ti with probability min(1, pt(ti)/pd(ti)). On reject, sample from max(0, pt −pd) normalised.
import random
def speculative_step(draft_probs, target_probs, draft_tokens, rng):
accepted = []
for p_d, p_t, t in zip(draft_probs, target_probs, draft_tokens):
ratio = min(1.0, p_t[t].item() / max(p_d[t].item(), 1e-12))
if rng.random() < ratio:
accepted.append(int(t))
else:
adj = (p_t - p_d).clamp(min=0); adj = adj / adj.sum().clamp_min(1e-12)
return accepted + [int(torch.multinomial(adj, 1))]
return accepted + [int(torch.multinomial(target_probs[-1], 1))]
Tests
rng = random.Random(0); V = 4
dp = [F.softmax(torch.randn(V), -1) for _ in range(3)]
tp = [F.softmax(torch.randn(V), -1) for _ in range(3)]
out = speculative_step(dp, tp, [int(p.argmax()) for p in dp], rng)
assert 1 <= len(out) <= 4; print("spec accept OK")
Problem
The base model has K extra "medusa" heads predicting tokens at offsets 1, 2, . . . , K; verify all in parallel.
import torch
import torch.nn as nn
class MedusaHeads(nn.Module):
def __init__(self, d, vocab, K=3):
super().__init__()
self.heads = nn.ModuleList([nn.Linear(d, vocab) for _ in range(K)])
def forward(self, h_last):
return torch.stack([head(h_last) for head in self.heads], dim=1) # (B, K, V)
Tests
m = MedusaHeads(32, 100, K=3); out = m(torch.randn(2, 32))
assert out.shape == (2, 3, 100); print("medusa OK")
Problem
Group requests of similar length into fixed-size batches; pad to the max length per batch.
def static_batch(requests, batch_size):
requests = sorted(requests, key=len)
return [requests[i:i+batch_size] for i in range(0, len(requests), batch_size)]
Tests
batches = static_batch([[1,2], [1], [1,2,3], [1,2,3,4]], 2)
assert len(batches) == 2; print("static batch OK")
Problem
Maintain a pool of in-flight sequences; each step, decode one token for each, evict finished sequences and admit new ones.
import torch
class ContinuousScheduler:
def __init__(self, model, max_active=8):
self.model = model; self.max_active = max_active
self.active = {} # seq_id -> ids
self.results = {}
def admit(self, seq_id, prompt_ids):
if len(self.active) < self.max_active:
self.active[seq_id] = prompt_ids; return True
return False
def step(self, eos):
finished = []
for sid, ids in self.active.items():
nxt = self.model(ids)[:, -1].argmax(dim=-1, keepdim=True)
self.active[sid] = torch.cat([ids, nxt], dim=-1)
if int(nxt) == eos: finished.append(sid)
for sid in finished:
self.results[sid] = self.active.pop(sid)
Tests
class M(nn.Module):
def __init__(self): super().__init__(); self.l = nn.Linear(8, 8)
def forward(self, ids): return self.l(F.one_hot(ids, 8).float())
sched = ContinuousScheduler(M(), max_active=4)
sched.admit("a", torch.tensor([[0]])); sched.step(eos=2)
assert "a" in sched.active or "a" in sched.results
print("continuous batch OK")
Problem
Build buckets of sequence lengths so each batch has uniform length, reducing padding overhead.
import numpy as np
def bucket_by_length(seqs, n_buckets=8):
lengths = [len(s) for s in seqs]
edges = np.linspace(min(lengths), max(lengths) + 1, n_buckets + 1)
buckets = [[] for _ in range(n_buckets)]
for s, L in zip(seqs, lengths):
for b in range(n_buckets):
if edges[b] <= L < edges[b + 1]:
buckets[b].append(s); break
return buckets
Tests
b = bucket_by_length([[1]*5, [1]*7, [1]*30], 4)
assert sum(len(x) for x in b) == 3; print("bucket OK")
Problem
Implement a priority queue that selects requests by cost-aware policy score = wait + α · remaining_tokens.
import heapq
class PriorityQueue:
def __init__(self): self.h = []
def push(self, score, req):
heapq.heappush(self.h, (score, id(req), req))
def pop(self): return heapq.heappop(self.h)[-1]
Tests
q = PriorityQueue(); q.push(2.0, "a"); q.push(1.0, "b")
assert q.pop() == "b"; print("pq OK")
Problem
Trace a model with example inputs into a TorchScript module; use torch.jit.optimize_for_inference.
import torch
def trace_model(model, example):
model.eval()
traced = torch.jit.trace(model, example)
return torch.jit.optimize_for_inference(traced)
Tests
m = nn.Linear(4, 4); ts = trace_model(m, torch.randn(1, 4))
assert ts(torch.randn(1, 4)).shape == (1, 4); print("ts OK")
Problem
Export a model to ONNX with dynamic axes.
import io
def export_onnx(model, example, f, opset=17):
model.eval()
torch.onnx.export(model, example, f, opset_version=opset,
input_names=['input'], output_names=['output'],
dynamic_axes={'input': {0: 'B'}, 'output': {0: 'B'}})
Tests
buf = io.BytesIO(); export_onnx(nn.Linear(4, 4), torch.randn(1, 4), buf)
assert buf.tell() > 0; print("onnx OK")
Problem
Convert the model to FP16 except for normalisation layers, which stay FP32 to avoid overflow.
import torch.nn as nn
def to_half_safe(model):
model.half()
for m in model.modules():
if isinstance(m, (nn.LayerNorm, nn.BatchNorm1d, nn.BatchNorm2d, nn.RMSNorm if hasattr(nn, 'RMSNorm') else
nn.LayerNorm)):
m.float()
return model
Tests
m = nn.Sequential(nn.Linear(8, 16), nn.LayerNorm(16))
to_half_safe(m)
assert m[0].weight.dtype == torch.float16 and m[1].weight.dtype == torch.float32
print("half-safe OK")
Problem
Reset peak memory, run a forward pass, read torch.cuda.max_memory_allocated().
import torch
def measure_peak_memory(fn):
if torch.cuda.is_available():
torch.cuda.reset_peak_memory_stats()
fn()
return torch.cuda.max_memory_allocated() / (1024 ** 2)
return -1.0
Tests
peak = measure_peak_memory(lambda: torch.randn(1024, 1024) @ torch.randn(1024, 1024))
print("mem OK", peak)
Problem
Compute per-tensor scales using a percentile-based clipping of the activation distribution.
import torch
def percentile_scale(x, p=99.99, qmax=448.0):
s = torch.quantile(x.abs().flatten(), p / 100.0) / qmax
return s.clamp(min=1e-9)
Tests
s = percentile_scale(torch.randn(1024) * 10)
assert s > 0; print("percentile scale OK", s.item())
Problem
Split W ∈ RO×I along the output dim across P devices: each device owns O/P rows. After local matmul, an all-gather concatenates results.
import torch
import torch.nn.functional as F
def tp_column_split(W, P):
return W.chunk(P, dim=0)
def tp_column_forward(x, W_shards):
return torch.cat([F.linear(x, W) for W in W_shards], dim=-1)
Tests
W = torch.randn(8, 16); shards = tp_column_split(W, 4)
y_ref = F.linear(torch.randn(2, 16), W)
print("TP col OK")
Problem
Split W along the input dim; each device owns I/P columns; an all-reduce sums partial outputs.
import torch.nn.functional as F
def tp_row_split(W, P): return W.chunk(P, dim=1)
def tp_row_forward(x_shards, W_shards):
parts = [F.linear(x, W) for x, W in zip(x_shards, W_shards)]
return sum(parts)
Tests
W = torch.randn(8, 16); x = torch.randn(1, 16)
shards = tp_row_split(W, 4); xs = x.chunk(4, dim=-1)
y = tp_row_forward(xs, shards)
assert torch.allclose(y, F.linear(x, W), atol=1e-5); print("TP row OK")
Problem
Run stage A on device 0, transfer activation, run stage B on device 1.
def pp_forward(stage_a, stage_b, x):
h = stage_a(x); return stage_b(h)
Tests
y = pp_forward(nn.Linear(4, 8), nn.Linear(8, 16), torch.randn(2, 4))
assert y.shape == (2, 16); print("pipeline OK")
Problem
Each rank stores only 1/P of the optimizer state; gather other ranks’ params during the step. Skeleton.
def shard_optimizer_state(state_dict, rank, world):
out = {}
for k, v in state_dict.items():
flat = v.flatten()
chunk = flat.chunk(world)[rank]
out[k] = chunk.clone()
return out
Tests
state = {'w': torch.arange(8).float()}
local = shard_optimizer_state(state, rank=0, world=2)
assert local['w'].numel() == 4; print("ZeRO shard OK")
Notes
Survival tactics for live deployment / inference coding:
KV cache: shape is (B, Hkv, T, dk) per layer. With GQA, Hkv ≪ Hq. Be ready to compute its size: 2 · L · Hkv · dk · B · T · bpe.
Quantization: distinguish PTQ (calibration only, no retraining) vs. QAT (fake-quant in training loop). Per-tensor vs. per-channel vs. groupwise has a 1% accuracy delta.
INT4 with groups: typical group g = 64 or 128. AWQ adds activation-aware scaling; SmoothQuant migrates outliers from activations to weights; GPTQ uses Hessian-aware column updates.
LoRA family: LoRA (α/r scaling), rsLoRA (α/√r), DoRA (decompose into magnitude + direction),
QLoRA (4-bit base + LoRA), LoRA+ (different LR for A and B), IA3 (per-channel scaling), Adapter (bottleneck + residual), Prefix / Prompt tuning (virtual tokens), BitFit (bias-only).
Speculative decoding: 2–4× throughput typical. Medusa makes it a single forward pass; EAGLE goes further by predicting in feature space.
Continuous batching / paged attention: vLLM-style. Allocate KV in fixed blocks, schedule per token (not per request), enable mixed prefill/decode batches.
Tensor parallel: column split (Megatron) needs all-gather; row split needs all-reduce. Pipeline parallel needs micro-batching to fill bubbles.
ONNX / TorchScript: trace requires example inputs + dynamic axes;export on a model with dynamic control flow needs torch.jit.script instead.
Distillation: Hinton (logit), feature, response, sequence-level (TinyBERT, DistilBERT, MiniLM, on-policy distillation, GKD). State the loss explicitly (KL with T 2 scaling, MSE on features, soft-prompt sequence- level CE).
Always measure: profile peak memory, latency, throughput tokens/s. Sub-1% accuracy delta and 4× memory savings is a typical INT4 quant result; INT8 should be <0.3% accuracy delta.