Interview Prep — File 05: Research Engineering Questions

============================================================

Paper-to-code challenges and research discussion questions.

For each: explain the paper, implement the key algorithm, discuss trade-offs.

Section A: Paper-to-Code Challenges (45 min each)

Challenge 1: GPTQ (Frantar et al., 2022)

Paper summary: GPTQ quantizes LLM weights to INT4 using Optimal Brain Quantization (OBQ). Instead of rounding each weight independently, it minimizes the output error of each layer by solving:

$$\min_{\hat{W}} |WX - \hat{W}X|_F^2$$

It processes columns left-to-right, quantizing one column at a time and updating remaining columns to compensate for the quantization error using the inverse Hessian.

Key equations: $$\hat{w}F = \arg\min{\hat{w}_F} \frac{(w_F - \hat{w}F)^2}{[H_F^{-1}]{FF}}$$

$$W_{-F} \leftarrow W_{-F} - \frac{w_F - \hat{w}F}{[H_F^{-1}]{FF}} (H_F^{-1})_{-F,F}$$

Implement:

def gptq_quantize_layer(W, H, bits=4, blocksize=128):
    """
    Quantize weight matrix W to `bits` bits using GPTQ.
    H: Hessian approximation H = 2 * X @ X.T (X = layer inputs)
    Returns: W_quantized (same shape), W_scale (per-group or per-channel)
    """
    Q = W.clone()
    Losses = torch.zeros_like(W)
    
    # Cholesky decomposition of H for numerical stability
    H = H + 1e-6 * torch.eye(H.shape[0], device=H.device)  # Tikhonov regularization
    H_inv = torch.linalg.inv(H)
    
    # Process in column blocks
    for i in range(0, W.shape[1], blocksize):
        block_end = min(i + blocksize, W.shape[1])
        W_block = Q[:, i:block_end]
        H_inv_block = H_inv[i:block_end, i:block_end]
        
        # Quantize column by column within block
        for j in range(block_end - i):
            col = j + i
            w = Q[:, col]
            d = H_inv[col, col]  # diagonal element
            
            # Quantize w
            scale = w.abs().max() / (2 ** (bits - 1) - 1)
            scale = scale.clamp(min=1e-8)
            w_q = (w / scale).round().clamp(-(2 ** (bits-1)), 2 ** (bits-1) - 1) * scale
            Q[:, col] = w_q
            
            # Error = quantization residual / Hessian diagonal
            err = (w - w_q) / d
            
            # Update remaining columns in block using off-diagonal Hessian
            Q[:, col+1:block_end] -= err.unsqueeze(1) * H_inv[col, col+1:block_end].unsqueeze(0)
    
    return Q

Key interview discussion points:

  • Why process column-by-column? Allows compensation for each quantization error before moving to next column
  • Why use Hessian? Captures input activation statistics; directions with large activation need higher precision
  • Computational cost: O(d² × n) for Hessian inversion; amortized by reuse across rows
  • vs AWQ: GPTQ minimizes reconstruction error; AWQ minimizes activation scaling separately

Challenge 2: FlashAttention-2 (Dao, 2023)

Key improvements over FA1:

  1. Sequence parallelism: parallelize over Q blocks (not just K/V blocks)
  2. Reduced non-matmul FLOPs: fewer rescaling operations
  3. Causal masking: skip entire tiles above diagonal

Tile loop structure (FA2 vs FA1):

FA1: Outer loop over Q tiles, Inner loop over K/V tiles
FA2: Outer loop over K/V tiles, Inner loop over Q tiles (GPU-friendly)
     → Each warp handles full K/V tile, threads parallelize over Q

Online softmax recurrence (FA2):

# For each K/V tile, update all Q rows in parallel
m_new = max(m_old, rowmax(Q_i @ K_j.T / sqrt(d)))
O = diag(exp(m_old - m_new)) @ O + P_ij @ V_j
l = diag(exp(m_old - m_new)) @ l + rowsum(P_ij)
# Final: O = O / l

Memory complexity analysis:

  • Standard attention: O(T²) for storing attention matrix
  • FlashAttention: O(T) — stores running statistics only
  • HBM reads/writes: Standard=4T², FA=5T (dominated by Q, K, V, O)

Interview question: "When does FlashAttention NOT improve performance?"

  • For very small T (< 32): overhead of tiling logic dominates
  • When attention matrix already fits in SRAM: no HBM savings
  • Non-square attention (e.g., cross-attention with very different T_q and T_k): tiling is suboptimal

Challenge 3: AWQ (Lin et al., 2023)

Key insight: In LLMs, a small fraction of weight channels (<1%) are "salient" — their corresponding input activations have high magnitude. Quantizing these channels more aggressively causes large errors. AWQ finds a scale factor s per input channel that makes weights easier to quantize by reducing the effective range.

Algorithm:

  1. Compute per-channel activation magnitude: a_j = mean(|x_j|) across calibration data
  2. Find optimal scale: s_j = a_j^α (grid search α ∈ [0, 1])
  3. Transform: W' = W * diag(s), x' = x * diag(s)^{-1} (absorbed into previous LayerNorm)
  4. Quantize W' (now with more uniform magnitude)

Why it works:

  • Salient channels: large a_j → large s_j → small a_j * s_j^{-1} → less quantization error
  • Non-salient channels: small a_j → small s_j → no change
  • Net: shifts the "hardness" from salient channels to non-salient channels

Code sketch:

def awq_search_scale(W, x, bits=4, grid_size=20):
    """Find optimal per-channel scale via grid search."""
    a_j = x.abs().mean(dim=0)  # [in_features]
    best_loss = float('inf')
    best_s = torch.ones_like(a_j)
    
    for alpha in torch.linspace(0, 1, grid_size):
        s = a_j.pow(alpha)
        W_scaled = W * s.unsqueeze(0)  # [out, in] * [in]
        W_q = quantize(W_scaled, bits)  # fake quant
        # Reconstruction loss: ||W_scaled - W_q||² weighted by activation
        loss = ((W_scaled - W_q) * a_j.unsqueeze(0)).pow(2).mean()
        if loss < best_loss:
            best_loss = loss
            best_s = s
    
    return best_s

Section B: Research Discussion Questions

B1: "Walk me through how torch.compile works end-to-end"

Expected answer structure:

  1. TorchDynamo (Python bytecode → FX graph):

    • Intercepts Python bytecode at the RESUME instruction
    • Converts Python operations into an FX graph via OutputGraph
    • Handles Python control flow by "breaking" the graph at dynamic branches (graph breaks)
    • Caches compiled graphs; re-compiles on guard failure

  2. AOTAutograd (forward + backward tracing):

    • Traces both forward and backward passes via make_fx and FunctionalizationInterpreter
    • Produces a single joint graph of forward + backward (for memory optimization)
    • Applies operator fusion and dead code elimination

  3. Backend dispatch (FX graph → executable):

    • Default: Inductor → generates Triton kernels for GPU, C++ for CPU
    • Alternative: eager (no optimization), custom backends
    • Inductor: applies pointwise fusion, persistent reduction, matrix tile scheduling

  4. Guard system (correctness):

    • Guards check: input shapes, dtypes, Python object identities
    • If guard fails: fall back to uncompiled execution
    • assume_static_by_default=True → recompile on new shapes unless dynamic=True

Common interview follow-up: "What causes graph breaks and how do you fix them?"

  • Dynamic Python code (loops with variable bounds, data-dependent control flow)
  • Unsupported Python builtins (print, some list ops)
  • Fix: use torch.jit.is_scripting(), rewrite with tensor ops, use torchdynamo.explain(fn, *args) to find breaks

B2: "Explain quantization-aware training (QAT) vs post-training quantization (PTQ)"

PTQQAT
Calibration data100-1000 samplesFull training set
Training requiredNoYes (fine-tuning)
Training timeMinutesHours-days
Accuracy (INT8)-0.5 to -1.5%-0.1 to -0.5%
Accuracy (INT4)-3 to -8%-1 to -3%
Use caseDeployment deadlines, no GPU timeProduction models with accuracy budget

STE (Straight-Through Estimator): The key trick in QAT. Quantization is non-differentiable (step function). STE approximates gradient as 1 in [-clip, clip]:

class FakeQuant(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, scale, zero_point, qmin, qmax):
        x_q = torch.clamp(torch.round(x / scale) + zero_point, qmin, qmax)
        x_dq = (x_q - zero_point) * scale
        ctx.save_for_backward(x, scale)
        return x_dq
    
    @staticmethod
    def backward(ctx, grad_output):
        x, scale = ctx.saved_tensors
        # STE: pass gradient through (as if no quantization happened)
        return grad_output, None, None, None, None

B3: "How would you debug a 20% accuracy regression after quantization?"

Systematic debugging protocol:

  1. Layer-by-layer sensitivity analysis:

    for i, layer in enumerate(model.layers):
    model_copy = deep_copy(model)
    quantize_only_layer(model_copy, i)
    acc = evaluate(model_copy, val_set)
    sensitivity[i] = baseline_acc - acc

    Sort by sensitivity. If top-1 sensitive layer has 15% drop → that layer is the cause.

  2. Activation range inspection:

    • Look for layers with extremely wide activation ranges (outliers)
    • Symptom: calibration scale is huge because a few outliers dominate
    • Fix: percentile calibration, SmoothQuant, or FP16 for that layer

  3. Weight distribution inspection:

    • Plot weight histogram per layer: bimodal or heavy-tailed → hard to quantize
    • Fix: channel-wise quantization, GPTQ, AWQ

  4. Cosine similarity between FP32 and INT8 outputs:

    for name, (fp32_out, int8_out) in layer_outputs.items():
    cos_sim = F.cosine_similarity(fp32_out.flatten(), int8_out.flatten(), dim=0)
    if cos_sim < 0.99:
        print(f"Layer {name}: CosSim={cos_sim:.4f} ← problem layer")

  5. Check attention layers specifically:

    • Softmax input has large variance → exp() amplifies quantization noise
    • Fix: FP16 attention with INT8 projection layers only

B4: "Explain the Pareto frontier in model accuracy vs. latency"

Definition: A set of models M* where no model is strictly better than another on ALL objectives:

  • M_i dominates M_j if: acc(M_i) ≥ acc(M_j) AND latency(M_i) ≤ latency(M_j)
  • Pareto frontier = set of non-dominated models

Practical use at Qualcomm:

  1. Quantize at INT8, INT4, FP16 → 3 points on frontier
  2. Prune to 50%, 75%, 90% sparsity → more points
  3. Distill to 0.5B, 1B, 3B, 7B → cross the efficiency boundary

Interview answer: "When picking a model for production, I'd first define the hard constraint (e.g., latency < 10ms on HTP), then choose the most accurate model satisfying that constraint from the Pareto frontier — not just the fastest or most accurate in isolation."

B5: "What is the memory bandwidth bottleneck in LLM decode, and how do you solve it?"

Analysis:

During decode (batch_size=1, generating one token at a time):

  • Every token requires reading ALL model weights (7B params × 2 bytes = 14 GB)
  • A100 HBM bandwidth: 2 TB/s → 14 GB / 2 TB/s = 7ms per token
  • A100 compute: 312 TFLOP/s; for 7B params, matmul FLOPs = 2×7B×1×1 = 14 GFLOP
  • Compute time: 14G / 312T = 0.045ms → 155× faster than memory → fully memory-bound

Solutions:

  1. Larger batch size: amortizes weight reads over more requests

    • batch=1: 7ms/token; batch=16: 7ms/16 ≈ 0.44ms/token (same weight read, 16× more tokens)

  2. Weight quantization (INT4): 7B × 0.5 bytes = 3.5 GB → 7ms/2 = 3.5ms/token

  3. Speculative decoding: draft model generates 4 tokens speculatively; target model verifies all 4 in one forward pass → effective throughput 3-5× higher

  4. Multi-query attention (GQA): KV cache reduction: 32-head → 4-head KV = 8× less KV reads

  5. Continuous batching: always keep GPU busy; never wait for slow requests

Section C: Implementation Drills (15 min each)

Drill 1: Compute MMLU accuracy from raw logits

def score_mmlu_from_logits(logits, choices_token_ids, correct_idx):
    """
    logits: [vocab_size] — model output for the position before the answer
    choices_token_ids: [4] — token IDs for " A", " B", " C", " D"
    Returns: True if predicted choice matches correct_idx
    """
    choice_logits = logits[choices_token_ids]
    predicted = choice_logits.argmax().item()
    return predicted == correct_idx

Drill 2: Compute perplexity

def compute_perplexity(model, tokenizer, text, stride=512, max_length=2048):
    """
    Compute perplexity with sliding window for long texts.
    """
    encodings = tokenizer(text, return_tensors="pt")
    seq_len = encodings.input_ids.shape[1]
    total_nll, total_tokens = 0.0, 0
    prev_end = 0
    
    for begin in range(0, seq_len, stride):
        end = min(begin + max_length, seq_len)
        tgt_len = end - prev_end
        input_ids = encodings.input_ids[:, begin:end]
        
        with torch.no_grad():
            outputs = model(input_ids, labels=input_ids)
            # outputs.loss is mean NLL; multiply back to get sum
            total_nll += outputs.loss.item() * tgt_len
        
        total_tokens += tgt_len
        prev_end = end
        if end == seq_len:
            break
    
    return math.exp(total_nll / total_tokens)

Drill 3: McNemar's test for accuracy comparison

from scipy.stats import chi2

def mcnemar_test(correct_a, correct_b):
    """
    McNemar's test for paired accuracy comparison.
    correct_a, correct_b: boolean arrays [n_examples]
    Returns: (chi2_stat, p_value)
    """
    # Discordant pairs
    b = ((correct_a) & (~correct_b)).sum()   # A correct, B wrong
    c = ((~correct_a) & (correct_b)).sum()   # A wrong, B correct
    
    # With continuity correction (Edwards)
    chi2_stat = (abs(b - c) - 1) ** 2 / (b + c) if (b + c) > 0 else 0
    p_value = 1 - chi2.cdf(chi2_stat, df=1)
    return chi2_stat, p_value