« All Roles

Model Accuracy & AI Performance — Senior Staff Engineer Curriculum

Target Roles:

• Qualcomm — Senior Staff Machine Learning Engineer, AI Frameworks & Performance
• Apple — Senior Staff Engineer, Core ML Accuracy & Inference
• NVIDIA — Principal Engineer, Model Optimization & Accuracy
• Google DeepMind — Senior Staff Engineer, ML Systems Performance
• Meta — Staff Engineer, AI Infrastructure & Model Quality
• Microsoft — Principal Engineer, AI Model Optimization
• Arm — Senior Staff ML Compiler & Accuracy Engineer
• AMD — Principal Engineer, ROCm Model Optimization
• Cerebras / Groq / Tenstorrent — Staff ML Systems & Accuracy Engineer

Duration: 24 weeks core (6 months) — extendable to 12 months with advanced research projects
Seniority Target: Senior Staff / Principal IC (L6–L7 equivalent)
Goal: Equip you to lead end-to-end model accuracy and AI performance work on NPU/GPU/CPU hardware — from ML framework internals through hardware-aware graph optimization, quantization with accuracy recovery, profiling, eval systems, and production deployment on edge silicon.

Why This Curriculum Exists

The "Model Accuracy & AI Performance" discipline sits at the hardest intersection in AI engineering: you must simultaneously understand model internals deeply enough to reason about why accuracy changes after optimization, and hardware/compiler internals deeply enough to know where performance is left on the table. Most engineers are strong in one dimension; Senior Staff engineers are expected to own both.

The Qualcomm JD (and every equivalent role at other silicon companies) makes this explicit: they want someone who can take a research-grade PyTorch model, understand its architecture from literature, apply quantization and graph transformations, run it through a compiler toolchain targeting an NPU, profile the result, recover lost accuracy, and repeat — all while influencing the ML framework design itself. This requires 12+ years of compounding depth.

This curriculum builds that depth from first principles. Every phase ends with quantified artifacts you can demo in interviews and publish in your GitHub portfolio. The final capstone projects are designed to be genuine open-source contributions, not toy exercises.

Reference Job Targets

What You Will Build

By the end of this curriculum you will have shipped and can demo:

1. Custom autograd engine — implement reverse-mode autodiff from scratch in Python+C++, matching PyTorch semantics on standard test cases within 1e-8 numerical tolerance
2. Custom C++ PyTorch operator — registered via TORCH_LIBRARY, with CPU+CUDA kernels, torch.compile-compatible, passing all OpCheck tests
3. Transformer with full architectural variants — RoPE, ALiBi, GQA, MoE routing, all swappable — benchmarked on Wikitext-103
4. Full PTQ → QAT → GPTQ → AWQ pipeline — automated, with accuracy-latency Pareto curves and regression gating
5. torch.compile custom backend — FX graph lowering to a mock NPU ISA with >95% op coverage on LLaMA-style models
6. TVM compilation pipeline — Relay → TIR → ARM/x86 compiled module with AutoTVM tuning, 2x+ speedup over baseline
7. SNPE/QNN model deployment pipeline — ONNX → DLC → quantized NPU inference with full accuracy benchmarking vs FP32 baseline
8. Roofline analysis toolkit — measure model arithmetic intensity, plot compute/memory ceilings, automatically identify bottleneck ops
9. FlashAttention v2 from scratch in Triton — tiled SRAM-efficient implementation matching official FA2 output within 1e-5
10. Speculative decoding engine — draft + verifier with acceptance rate profiling, 2–4x throughput improvement demonstrated
11. Eval harness from scratch — MMLU, HellaSwag, GSM8K, perplexity; pluggable task registry, 3 sampling strategies, statistical significance reporting
12. Accuracy regression CI — PR-gate pipeline that fails builds if accuracy drops >0.5% on key tasks after any optimization change
13. On-device LLM end-to-end — LLaMA-3.2-1B → INT4 quantized → edge-benchmarked with accuracy delta vs FP16 documented
14. Multi-model evaluation dashboard — Pareto frontier visualization across quantization levels, latency budgets, and accuracy metrics

Folder Structure

model-accuracy-ai-performance/
├── README.md                                          ← this file
├── phase-01-pytorch-tf-framework-internals/
│   ├── README.md
│   ├── WARMUP.md                                      ← zero-to-expert primer
│   ├── HITCHHIKERS-GUIDE.md
│   ├── lab-01-custom-autograd-engine/
│   ├── lab-02-custom-cpp-operator/
│   └── lab-03-tf-function-tracing/
├── phase-02-model-architecture-mastery/
│   ├── README.md
│   ├── HITCHHIKERS-GUIDE.md
│   ├── DEEP-DIVE-MOE-SSM.md                           ← MoE + Mamba from first principles
│   ├── lab-01-transformer-architectural-variants/
│   ├── lab-02-mixture-of-experts/
│   └── lab-03-state-space-models/
├── phase-03-quantization-model-compression/
│   ├── README.md
│   ├── HITCHHIKERS-GUIDE.md
│   ├── DEEP-DIVE-QAT-GPTQ.md                          ← PTQ/QAT/GPTQ from first principles
│   ├── lab-01-ptq-pipeline/
│   ├── lab-02-qat/
│   └── lab-03-gptq/
├── phase-04-torch-compile-fx-graph/
│   ├── README.md
│   ├── HITCHHIKERS-GUIDE.md
│   ├── lab-01-fx-graph-passes/
│   ├── lab-02-custom-compile-backend/
│   └── lab-03-triton-fused-kernel/
├── phase-05-ml-compilers-ir/
│   ├── README.md
│   ├── HITCHHIKERS-GUIDE.md
│   ├── lab-01-onnx-export-optimization/
│   ├── lab-02-tvm-compilation/
│   └── lab-03-mlir-python-bindings/
├── phase-06-qualcomm-npu-hardware/
│   ├── README.md
│   ├── HITCHHIKERS-GUIDE.md
│   ├── lab-01-ai-hub-deployment/
│   ├── lab-02-hardware-aware-quantization/
│   └── lab-03-on-device-profiling/
├── phase-07-profiling-performance-engineering/
│   ├── README.md
│   ├── HITCHHIKERS-GUIDE.md
│   ├── lab-01-roofline-analyzer/
│   ├── lab-02-profiling-context/
│   └── lab-03-transformer-perf-bugs/
├── phase-08-inference-optimization-depth/
│   ├── README.md
│   ├── HITCHHIKERS-GUIDE.md
│   ├── lab-01-flash-attention/
│   ├── lab-02-speculative-decoding/
│   └── lab-03-continuous-batching/
├── phase-09-model-accuracy-evaluation/
│   ├── README.md
│   ├── HITCHHIKERS-GUIDE.md
│   ├── DEEP-DIVE-EVALUATION-AT-SCALE.md               ← Pareto, McNemar, CI gates
│   ├── lab-01-eval-harness/
│   ├── lab-02-pareto-analyzer/
│   └── lab-03-regression-ci/
├── phase-10-genai-on-edge/
│   ├── README.md
│   ├── HITCHHIKERS-GUIDE.md
│   ├── lab-01-awq-quantization/
│   ├── lab-02-qlora-adapter/
│   └── lab-03-layer-streaming/
├── phase-11-capstone/
│   ├── README.md
│   ├── HITCHHIKERS-GUIDE.md
│   ├── lab-01-regression-ci-pipeline/
│   ├── lab-02-npu-graph-optimization/
│   ├── lab-03-custom-compile-backend/
│   ├── lab-04-production-quant-suite/
│   └── lab-05-multi-model-dashboard/
├── interview-prep/
│   ├── README.md
│   ├── 01-concepts-cheatsheet.md
│   ├── 02-framework-internals-coding.py
│   ├── 03-optimization-coding.py
│   ├── 04-systems-design-questions.md
│   ├── 05-research-engineering-questions.md
│   └── 06-behavioral-star-stories.md
└── system-design/
    ├── README.md
    ├── 01-model-accuracy-regression-platform.md
    ├── 02-npu-graph-optimization-pipeline.md
    ├── 03-quantization-accuracy-recovery-system.md
    ├── 04-distributed-eval-platform.md
    └── 05-model-perf-benchmarking-infra.md

24-Week Schedule

Each Lab Structure

Prerequisites

• Python 3.10+ with pip
• PyTorch 2.3+ (pip install torch torchvision)
• CUDA 12.x (or MPS on Apple Silicon; labs work on CPU with reduced scale)
• CMake 3.20+ and g++ / clang for C++ labs
• Triton 2.3+ (comes with PyTorch nightly for some labs)
• Basic familiarity with: calculus (chain rule), linear algebra, Python, C++ basics
• Git — all labs tracked in version control
• Optional: Qualcomm AI Hub account (free) for Phase 06 simulator runs

Hardware Recommendations

Glossary — Model Accuracy & AI Performance

Every term the track (and the Qualcomm Senior Staff JD) assumes, defined from first principles. Grouped by phase area. Cross-references in [[brackets]] point to related entries.

Framework Internals (P01)

• Dispatcher (c10::Dispatcher) — PyTorch's C++ router that, for every operator call, selects the kernel based on dispatch keys. The layer between the Python API and the backend kernels.
• Dispatch key — a tag (CPU, CUDA, Autograd, Functionalize, Meta, …) that determines which registered kernel handles an op. A custom op must behave correctly under all relevant keys.
• Autograd engine — builds a DAG during the forward pass (each op records its backward) and walks it in reverse-topological order on .backward(), accumulating gradients at the leaves.
• AccumulateGrad — the autograd node that sums gradients into a leaf tensor's .grad.
• retain_graph / create_graph — keep the graph after backward (for multiple backwards) / build a differentiable backward graph (for higher-order gradients / double backward).
• torch.autograd.Function — a custom op with hand-written forward/backward; save_for_backward stashes tensors for the backward pass.
• TORCH_LIBRARY / TORCH_LIBRARY_IMPL — C++ macros to declare an op schema and register a kernel for a given dispatch key.
• Meta kernel — a kernel that computes only shapes/dtypes (no data), used for tracing, torch.compile, and shape inference. Missing it is a common custom-op bug.
• Functionalization — a pass that rewrites in-place/mutating ops into functional equivalents so the graph is pure (needed by compile/export).
• tf.function — TensorFlow's tracing decorator: runs the Python once to build a graph (FuncGraph), then executes the graph. Python side-effects happen at trace time only.
• Retracing — TF re-runs tracing when it sees a new input signature; excessive retracing is a performance bug.
• BF16 / FP16 / FP32 — brain-float16 (8-bit exponent, 7-bit mantissa — FP32's range, less precision), half (5-bit exp, 10-bit mantissa), single. [[accumulation-dtype]].
• Catastrophic cancellation — large precision loss when subtracting nearly-equal floats; a driver of low-precision instability.

Model Architecture (P02)

• Self-attention — each token attends to all others; cost is O(N²) in sequence length N (the motivation for [[FlashAttention]] and sparse/linear variants).
• FFN / MLP block — the per-token feed-forward in a transformer; the bulk of the params.
• Mixture-of-Experts (MoE) — replace the dense FFN with N experts; a router sends each token to the top-k experts. Cuts FLOPs/token but adds a dynamic gather/scatter the NPU dislikes.
• Load-balancing loss — auxiliary loss that prevents MoE router collapse (a few experts hogging all tokens). Without it, "8 experts" becomes effectively 2.
• State-Space Model (SSM) / Mamba — sequence model using a selective scan recurrence with near-linear cost instead of quadratic attention; the scan is sequential (hardware tradeoff).
• ViT (Vision Transformer) — images as patch sequences fed to a transformer; trades CNN inductive bias for scale.
• DINO — self-supervised distillation: an EMA teacher trains a student with no labels; centering + sharpening prevent representation collapse.
• CLIP bridge — projecting vision encoder features into an LLM's token embedding space to make a multimodal model.

Quantization & Compression (P03, P10)

• Quantization — mapping FP values to a small set of integers via a scale and zero-point.
• Scale (s) — the step size: s = (x_max − x_min) / (q_max − q_min). Smaller scale → smaller error.
• Zero-point (z) — the integer that maps to real 0 (asymmetric quant). Symmetric quant has z = 0.
• Symmetric vs asymmetric — range centered at 0 (z=0, simpler hardware) vs arbitrary [min,max] (uses z). Watch the dtype: 8-bit unsigned is [0,255], signed is [−128,127].
• Per-tensor vs per-channel — one scale for the whole tensor vs one per output channel (per-channel is more accurate; the hardware must support the layout).
• PTQ (Post-Training Quantization) — quantize a trained model using a little calibration data. Fast; first resort.
• QAT (Quantization-Aware Training) — simulate quantization during training (fake-quant forward, [[straight-through-estimator]] backward) so the model learns robustness. For when PTQ drops too much.
• Straight-through estimator (STE) — treat the non-differentiable round() as identity in the backward pass so gradients flow through fake-quant.
• Calibration — choosing the clipping range from observed activations (min/max, percentile, MSE-optimal, entropy/KL).
• GPTQ — second-order weight quantization: quantize column-by-column, compensating remaining weights for the error using the (approx) Hessian. Holds accuracy at 4-bit.
• AWQ (Activation-aware Weight Quantization) — protect the small set of salient weight channels (found via activation magnitudes) and crush the rest to 4-bit.
• NF4 — a 4-bit "normal float" datatype tuned to the normal distribution of NN weights (used by [[QLoRA]]).
• QLoRA — fine-tune a frozen 4-bit base model via small high-precision low-rank adapters.
• Mixed precision — keep sensitive layers high-precision, crush the rest; the real size/accuracy win.
• Sensitivity analysis — quantize one layer at a time, measure the accuracy hit, rank layers.

Compilers & IR (P04, P05)

• torch.compile — PyTorch's JIT: [[Dynamo]] captures the graph, a backend lowers it.
• TorchDynamo — captures Python bytecode into an [[FX graph]]; graph-breaks on code it can't trace (data-dependent control flow, .item(), unsupported ops).
• Graph break — Dynamo splitting a function into multiple graphs with eager glue; the #1 reason torch.compile underdelivers.
• TorchInductor — the default backend; lowers FX to Triton (GPU) / C++ (CPU) kernels.
• FX graph — PyTorch's inspectable IR of nodes (placeholder, call_function, call_module, call_method, output). Dialect matters: symbolic_trace of modules gives call_module; aot/make_fx gives aten ops.
• Operator fusion — combining ops (e.g. Linear→GELU) into one to cut memory traffic / kernel launches; only safe when the producer's output has a single consumer.
• Custom backend — a function registered with Dynamo that receives an FX GraphModule, transforms it (classify NPU-eligible nodes, offload, fall back to CPU), and returns a callable.
• ONNX — Open Neural Network Exchange; a cross-framework model interchange format. Export can drop/decompose ops — verify numerically after export.
• opset — the ONNX operator-set version; mismatches cause export/runtime failures.
• TVM — an ML compiler with auto-tuning (searches loop schedules for the fastest kernel on a given target). IRs: Relay (high-level), TIR (low-level).
• MLIR — Multi-Level IR: a compiler infrastructure of progressively-lowered dialects; the backbone of many modern hardware compilers.
• Lowering — translating a higher-level IR to a lower-level one (coarse ops → fine loops); each step optimizes and risks bugs.

Hardware / NPU (P06)

• NPU (Neural Processing Unit) — a fixed-function accelerator for tensor math; on Qualcomm, the Hexagon Tensor Processor (HTP) in the AI Engine. Loves static, dense, fixed-point work.
• Coverage / offload rate — fraction of model compute that runs on the NPU vs falls back to CPU. The key health metric; low coverage silently wrecks latency/power.
• CPU fallback — running an unsupported op on the CPU mid-graph; the round-trip is expensive.
• Fixed-point accumulator — the NPU's integer accumulator; finite width → can overflow if scales are mis-set.
• AI Hub — Qualcomm's service for compiling/profiling/deploying models to target devices.
• Operator coverage / workload synthesis — ensuring (or making) every op in a model runnable on the hardware; co-design feedback to the HW team.
• Energy per inference (joules) — the real edge budget; power and latency couple via thermal throttling.

Profiling & Inference (P07, P08)

• Roofline model — plots achievable performance vs arithmetic intensity (FLOPs/byte); classifies a kernel as compute-bound or memory-bound, which dictates the fix.
• Arithmetic intensity — FLOPs per byte of memory traffic; low intensity → memory-bound.
• Compute-bound — limited by FLOPs/s; helped by faster math, not by moving fewer bytes.
• Memory-bound — limited by bytes/s; helped by fusion, quantization, layout — not by more FLOPs. Most LLM inference is here.
• CPU↔device sync — a .item()/.cpu() in a hot loop that serializes execution; the most common "accelerator at 30% utilization" cause.
• Warmup — the first iterations (compile/alloc/cache) you must discard before timing steady-state; report the median of many runs.
• FlashAttention — computes attention in tiles using online softmax, never materializing the N×N matrix; cuts memory traffic O(N²)→O(N). Exact, not approximate.
• Online softmax — running-max + running-sum recurrence that lets softmax be computed incrementally over tiles.
• KV cache — cached keys/values for past tokens so autoregressive decode is O(1)/token in attention; grows with sequence length (a memory killer; can be quantized/paged).
• Speculative decoding — a small draft model proposes several tokens; the big model verifies them in one parallel pass. Exact (identical output), pure latency win.
• Continuous batching — evicting finished sequences and admitting new ones every step to keep the accelerator saturated (server-side throughput).
• Paged KV cache (vLLM) — managing KV memory in fixed pages to avoid fragmentation.

Accuracy & Evaluation (P09)

• Pareto frontier — the set of models where you can't gain accuracy without adding latency; deployment picks a point on it.
• Knee point — the frontier's point of diminishing returns; usually the deployment sweet spot.
• Per-example delta / churn — which specific inputs changed verdicts between two models; hides regressions a matching aggregate score conceals.
• Regression CI — automated accuracy gates that block a merge if an "optimization" drops accuracy below threshold.
• Bootstrap CI / McNemar's test — statistical tools to decide whether an accuracy difference is real vs noise (paired model comparison).
• Reproducibility — same seed/data/config → same number; the precondition for any comparison.

The numbers behind these terms (bit-widths, quant formulas, roofline math, memory sizing) live in CHEATSHEET.md.

Cheat Sheet — Model Accuracy & AI Performance

The numbers, formulas, and decision rules to have at your fingertips — the interview-whiteboard quick reference. Definitions are in GLOSSARY.md.

Quantization formulas

Affine (asymmetric):   Q(x) = clamp(round(x/s) + z, q_min, q_max)
                       x̂   = s · (Q(x) − z)
Scale:                 s = (x_max − x_min) / (q_max − q_min)
Zero-point:            z = q_min − round(x_min / s)
Symmetric (z=0):       s = max(|x|) / q_max
Quant error bound:     |Q(x) − x| ≤ s/2     → smaller s = smaller error

PTQ vs QAT: try PTQ first (fast, little data). Reach for QAT when PTQ's accuracy drop is too big (often INT4, or sensitive layers). Mixed precision (sensitive layers high, rest low) is usually the best size/accuracy point.

Memory math: params × bytes/param. 7B model → FP16 ≈ 14 GB, INT8 ≈ 7 GB, INT4 ≈ 3.5 GB. (Plus the KV cache, which grows with sequence length — quantize/bound it on edge.)

Floating-point formats

Rule: accumulate in higher precision than you store. Reductions (softmax, sum) in FP32 even when activations are FP16/BF16.

Roofline (the master diagram)

Arithmetic Intensity (AI) = FLOPs / bytes moved
Attainable perf = min( peak_FLOPs/s , peak_BW × AI )

AI low  → MEMORY-bound  → fix: fusion, quantization, layout, FlashAttention
AI high → COMPUTE-bound → fix: better matmul, more cores, higher clocks

Most LLM inference is memory-bound (weight loading at batch 1; the N×N attention matrix) — which is why FlashAttention and quantization win. Don't optimize FLOPs on a memory-bound kernel.

Profiling discipline

1. WARM UP (discard first iters — compile/alloc/cache)
2. measure MANY iters, take the MEDIAN (not mean — outliers)
3. know your clock: CPU time ≠ device time ≠ wall time; self ≠ total
4. find the regime (roofline) BEFORE choosing a fix

Top recurring perf bugs: CPU↔device sync (.item()/.cpu() in a hot loop), hidden dtype casts, non-contiguous copies from reshapes, kernel-launch overhead (many tiny ops), unfused elementwise chains.

Inference optimizations (when each wins)

Edge (batch 1, power-limited): FlashAttention ✅, quantization ✅✅, KV-cache quant ✅, spec-decode (if draft fits memory) ☑, continuous batching ➖ (server-side).

FX graph dialects (the silent bug)

A pass that matches the wrong dialect silently does nothing (no fusion, 0 NPU nodes). Match by function identity AND aten string to be robust. Fuse only when the producer has 1 user.

NPU mental model (Qualcomm)

LOVES : static shapes · dense regular dataflow · fixed-point (INT8/INT4) · supported ops · on-chip
HATES : dynamic shapes · data-dependent control flow · gather/scatter (MoE) · unsupported ops · DRAM round-trips

• Coverage = % compute on NPU. Treat as a vital sign; one unsupported hot op → CPU fallback → latency death.
• Fixed-point accumulator can overflow — size scales carefully.
• Budget is joules, not just ms (power ↔ thermal throttle ↔ latency).
• Sim validates accuracy; trust the device for performance; reconcile when they diverge.

Consistency: every optimization must be verified

quantize  → eval accuracy delta (P09)
compile/fuse/export → numerical-equivalence check vs eager (atol/rtol)
spec-decode → bit-exact to base model
deploy to NPU → eval ON DEVICE (numerics differ from sim)

Fast-but-wrong is a shipped bug. No optimization is "done" until its accuracy is proven.

Evaluation rigor

report: score + variance/CI + baseline delta + per-example churn (not a lone number)
reproducible: fixed seed/data/config
is-it-real?: bootstrap CI / McNemar (paired); effect must exceed the noise band
Pareto: accuracy vs latency is a FRONTIER; ship at the KNEE for the device's budget
regression CI: block merges that drop accuracy below threshold

The lab verification convention (this track)

# tests import `lab` (your code) if present, else `solution`:
pytest test_*.py            # against your lab.py (NotImplementedError until you fill it)
# to verify the reference solution, shadow lab.py with solution.py, then pytest.

Heavy/proprietary deps (C++ toolchain, TensorFlow, TVM, Qualcomm AI Hub SDK) make a few labs skip gracefully rather than fail — read the code and run them where the toolchain exists.

Phase 01 — PyTorch & TensorFlow Framework Internals

Difficulty: ⭐⭐⭐⭐☆
Estimated Time: 4 weeks (80–100 hours)
Roles Supported: All — this is the bedrock phase. Every downstream topic (quantization, torch.compile, NPU deployment) requires internalizing what happens below the Python API.

Why This Phase Exists

Every Senior Staff ML engineer who influences Qualcomm's framework stack will be asked — in an interview or on the first week of the job — to trace execution from loss.backward() down to the CUDA kernel. If you can't do that, you cannot meaningfully review PRs touching the dispatcher, write a custom operator that behaves correctly under torch.compile, or diagnose why a quantized model produces NaN after a specific graph transformation.

PyTorch's execution model has three distinct layers that every practitioner conflates: the Python eager frontend, the dispatcher (C++ c10), and the backend kernels. TensorFlow's execution model has evolved through three paradigms (eager, tf.function + AutoGraph, tf.experimental.mlir_bridge) that coexist in production code. Understanding both frameworks at this depth is what separates framework users from framework engineers.

This phase will make you dangerous enough to write test cases that expose correctness bugs in framework patches, write custom ops that are safe under all dispatch keys (CPU, CUDA, autograd, torch.compile, functional), and reason about TF's tracing semantics when a colleague reports that "the model gives different results in graph mode."

Concepts

• PyTorch dispatch stack: Python → C++ dispatcher (c10::Dispatcher) → dispatch keys (CPU, CUDA, Autograd, Functionalize, Meta) → kernel registration via TORCH_LIBRARY / TORCH_LIBRARY_IMPL
• Autograd engine: DAG construction during forward pass, AccumulateGrad, Engine::execute(), gradient accumulation, retain_graph, create_graph for higher-order gradients
• torch.autograd.Function: custom forward/backward, save_for_backward, composite functions, double backward
• Eager vs compiled mode: what torch.compile changes at the boundary — TorchDynamo captures bytecode into FX IR, TorchInductor lowers to Triton/C++
• torch.fx symbolic tracing: Interpreter, Transformer, Graph, Node, proxy mechanism, limitations (control flow, data-dependent shapes)
• Custom C++ operators: TORCH_LIBRARY schema, TORCH_LIBRARY_IMPL, pybind11 vs torch::Library, OpCheck, structured kernels
• TF graph execution: Session.run() (TF1) vs tf.function + tf.Graph (TF2), ConcreteFunction, FuncGraph, AutoGraph transformation rules
• tf.function tracing semantics: Python-side effects, tf.Variable vs Python variables, retracing triggers (new input signatures, Python closures)
• TF SavedModel format: tf.saved_model.save, SignatureDef, asset handling, interop with TFLite/ONNX
• Memory management: PyTorch caching allocator, CUDA unified memory, TF tensor lifetimes in sessions vs eager
• Numerical precision: FP32 vs FP16 vs BF16 accumulation, catastrophic cancellation, Kahan summation, determinism flags
• Operator decomposition: how PyTorch's _decomp registry breaks complex ops into primitive ops for portability

Labs

Lab 01 — Custom Autograd Engine from Scratch

Lab 02 — Custom C++ PyTorch Operator

Lab 03 — TF tf.function Tracing & Graph Analysis

Deliverables Checklist

• minigrad package with all 15 gradient checks passing
• warp_softmax.so with OpCheck passing and benchmarks documented
• TF graph inspector tool with FuncGraph visualization output
• All test_lab.py suites pass: pytest phase-01-*/lab-0*/test_lab.py
• Personal notes in each HITCHHIKERS-GUIDE.md section annotated with your insights

Interview Relevance

• "Explain how PyTorch's backward() traverses the computation graph." (you can draw the DAG, explain topological sort, explain AccumulateGrad nodes)
• "What is a dispatch key in PyTorch? How does torch.compile use it?" (you can explain Meta kernels, Functionalize key, and why OpCheck tests all of them)
• "When does tf.function retrace?" (you can enumerate all 5 retracing scenarios from memory after Lab 03)
• "How would you debug a NaN appearing after quantization in only graph mode?" (you have the mental model to distinguish Python-side bugs from kernel-side bugs)
• "Write a custom backward for a fused attention operation." (Lab 01 gives you the foundations to extend on a whiteboard)

Warmup Guide — PyTorch & TensorFlow Framework Internals

Who this is for: Someone who knows Python basics and wants to reach principal ML engineer level — able to contribute to PyTorch, write CUDA kernels, debug framework internals, and reason about TensorFlow's compilation model. No prior ML, C++, or GPU experience assumed. Nothing is skipped. Every concept is built from the ground up.

Table of Contents

Part I — The Mathematical Substrate

• Chapter 1: Tensors — From Linear Algebra to Memory
• Chapter 2: Derivatives and the Chain Rule — The Mathematics of Learning
• Chapter 3: Jacobians — Why Backpropagation Works at Scale
• Chapter 4: Automatic Differentiation — Two Modes, One Winner

Part II — How Computers Execute Code

• Chapter 5: CPython — How Python Actually Runs
• Chapter 6: NumPy — Arrays, Memory Layout, and Vectorization
• Chapter 7: C++ and pybind11 — The Bridge to Performance

Part III — PyTorch Architecture

• Chapter 8: ATen, c10, and libtorch — The C++ Foundation
• Chapter 9: Tensors in PyTorch — Storage, Strides, and Memory Management
• Chapter 10: The Dispatcher — PyTorch's Universal Routing System
• Chapter 11: The Autograd Engine — Building and Traversing the Backward Graph
• Chapter 12: torch.autograd.Function — Custom Differentiable Operations

Part IV — Graph Capture and Compilation

• Chapter 13: torch.fx — Intermediate Representation and Graph Transformation
• Chapter 14: torch.compile — The Full Compilation Stack
• Chapter 15: GPU Architecture — From Silicon to CUDA Kernels
• Chapter 16: Custom C++ Operators — TORCH_LIBRARY

Part V — TensorFlow

• Chapter 17: TF1 to TF2 — Two Execution Philosophies
• Chapter 18: tf.function — AutoGraph, Tracing, and ConcreteFunctions
• Chapter 19: GradientTape, XLA, and MLIR

Part VI — Numerical Precision

• Chapter 20: IEEE 754 — How Computers Store Numbers
• Chapter 21: FP16, BF16, and Mixed Precision Training

Part VII — Synthesis and Practice

• Chapter 22: How Everything Connects — The Principal Engineer's Map
• Chapter 23: Lab Guidance — Building minigrad and warp_softmax
• Chapter 24: The Interview Gauntlet — Complete Answers

References

Part I — The Mathematical Substrate

Chapter 1: Tensors — From Linear Algebra to Memory

1.1 What Is a Tensor, Mathematically?

The word "tensor" comes from Latin tendere — to stretch. In mathematics, a tensor is a generalization of scalars, vectors, and matrices to any number of dimensions.

Start with the simplest case:

• A scalar is a single number: 5.0. It has no dimensions — rank 0.
• A vector is an ordered list of numbers: [1.0, 2.0, 3.0]. It has one dimension — rank 1. A vector of length 3 lives in 3D space. Each number describes magnitude along one axis.
• A matrix is a 2D grid of numbers, indexed by row and column. Shape (3, 4) means 3 rows and 4 columns — rank 2.
• A rank-3 tensor is a 3D array. Shape (batch, height, width) for a batch of grayscale images. You can picture it as a stack of matrices.
• A rank-4 tensor extends to (batch, channels, height, width) for RGB images.

The generalization: a rank-\(n\) tensor is indexed by \(n\) integers and holds a number at each index position.

Scalar:  5.0                                 shape ()
Vector:  [1, 2, 3]                           shape (3,)
Matrix:  [[1, 2], [3, 4], [5, 6]]            shape (3, 2)
Rank-3:  [[[1,2],[3,4]], [[5,6],[7,8]]]      shape (2, 2, 2)

Why tensors matter for ML: A neural network's weight matrix is a rank-2 tensor. A batch of images is rank-4. The activations inside a transformer attention layer are rank-3 (batch × sequence_length × d_model). Everything in deep learning lives in tensor space. The framework's entire job is to manipulate tensors efficiently.

1.2 Tensors as Multi-Linear Maps (The Deep Math)

In formal mathematics, a tensor is a multi-linear map — a function that takes multiple vectors as inputs and produces a number, and is linear in each of its inputs separately.

A matrix M of shape (m, n) is a rank-2 tensor. It maps a vector v of length n to a vector Mv of length m. The mapping is linear: M(αv + βw) = αMv + βMw.

Why does this matter for engineers? Because the transpose rule for backpropagation — the fact that the gradient of y = Ax with respect to x is A^T grad_y — comes directly from this multi-linear structure. If you know what kind of object a tensor is mathematically, you can derive its gradient mechanically instead of looking it up.

1.3 Tensors in Memory: The Physical Reality

When you write torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), what exists in your computer's RAM?

Physical memory (a flat array of bytes):
  [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]   ← 6 float32s = 24 bytes, contiguous

The tensor's 2D shape is an illusion maintained by metadata:

The strides are the key. To access element at row r, column c, the index into the flat memory is:

flat_index = storage_offset + r × strides[0] + c × strides[1]
           = 0 + r × 3 + c × 1

This formula works for any shape and any strides. Operations like .T (transpose), .reshape(), and slicing x[1:, ::2] only change the metadata — they do not copy memory. Two tensors can point to the same physical bytes with different shapes and strides.

x = torch.tensor([[1, 2, 3], [4, 5, 6]])   # shape (2, 3), strides (3, 1)
y = x.T                                       # shape (3, 2), strides (1, 3) — no copy
z = x[::2, :]                                # shape (1, 3), strides (6, 1) — no copy

Consequence you must internalize: when you call x.contiguous(), PyTorch checks if the strides match the canonical row-major (C-contiguous) layout where each row is stored consecutively. If they don't — because you transposed or sliced — PyTorch allocates a new buffer and copies data into the canonical layout. This is a hidden allocation cost that appears in profiles as unexpected memory spikes.

1.4 The Storage Object

The physical memory block is owned by a Storage object (in C++: at::StorageImpl). Multiple tensors can share the same Storage with different offsets/strides.

x = torch.arange(12).reshape(3, 4)
y = x[1, :]        # row 1 — shares storage, offset = 4 * sizeof(int) = 16 bytes

# Same underlying memory:
print(x.storage().data_ptr() == y.storage().data_ptr())  # True
print(y.storage_offset())  # 4  (4 elements offset into shared storage)

This sharing is why in-place operations are dangerous: modifying y modifies x because they share a Storage.

Chapter 2: Derivatives and the Chain Rule — The Mathematics of Learning

2.1 What Is a Derivative?

A derivative measures rate of change. For a function \(f(x)\), the derivative at a point \(x_0\) is:

\[ f'(x_0) = \lim_{h \to 0} \frac{f(x_0 + h) - f(x_0)}{h} \]

In plain English: if you nudge the input by a tiny amount \(h\), the derivative tells you how much the output changes per unit of nudge. If \(f'(x_0) = 6\), then nudging \(x\) by 0.001 changes \(f\) by approximately 0.006.

Geometrically: the derivative is the slope of the tangent line to the curve \(f(x)\) at the point \(x_0\).

For ML: the derivative of the loss with respect to a weight tells you — if I increase this weight slightly, does the loss go up or down, and by how much? This is exactly the information gradient descent needs.

2.2 Partial Derivatives — Functions of Many Variables

A neural network loss depends on millions of weights simultaneously. For a function \(f(x, y, z, \ldots)\), the partial derivative with respect to \(x\) is written \(\frac{\partial f}{\partial x}\) and measures the rate of change of \(f\) when only \(x\) changes, holding all other inputs fixed.

Example: \(f(x, y) = x^2 y + 3y\)

\[ \frac{\partial f}{\partial x} = 2xy \qquad \frac{\partial f}{\partial y} = x^2 + 3 \]

The gradient of \(f\) is the vector of all partial derivatives:

\[ \nabla f = \left[\frac{\partial f}{\partial x},\ \frac{\partial f}{\partial y}\right] \]

For a neural network with weights \(W_1, W_2, b_1, b_2, \ldots\), the gradient of the loss is the vector:

\[ \nabla_\theta \mathcal{L} = \left[\frac{\partial \mathcal{L}}{\partial W_1},\ \frac{\partial \mathcal{L}}{\partial W_2},\ \frac{\partial \mathcal{L}}{\partial b_1},\ \ldots\right] \]

Each entry tells you how sensitive the loss is to that weight. Gradient descent subtracts a small multiple of this gradient from the weights:

\[ \theta \leftarrow \theta - \eta \nabla_\theta \mathcal{L} \]

where \(\eta\) (eta) is the learning rate. This moves the weights in the direction that reduces loss.

2.3 The Chain Rule — The Engine of Everything

The chain rule is a rule for differentiating composed functions — functions of functions. This is exactly what a neural network is: a composition of layers.

Single-variable form: if \(z = f(y)\) and \(y = g(x)\), then:

\[ \frac{dz}{dx} = \frac{dz}{dy} \cdot \frac{dy}{dx} \]

Read this as: "how much does \(x\) affect \(z\)?" equals "how much does \(y\) affect \(z\)?" times "how much does \(x\) affect \(y\)?".

Derivation from first principles: Using the limit definition:

\[ \frac{dz}{dx} = \lim_{h \to 0} \frac{f(g(x+h)) - f(g(x))}{h} \]

Let \(\Delta y = g(x+h) - g(x)\). When \(h \to 0\), \(\Delta y \to 0\) as well (assuming \(g\) is continuous). So:

\[ = \lim_{\Delta y \to 0} \frac{f(g(x) + \Delta y) - f(g(x))}{\Delta y} \cdot \lim_{h \to 0} \frac{g(x+h) - g(x)}{h} = f'(g(x)) \cdot g'(x) \]

Multi-variable form: if \(z = f(x, y)\) where \(x = g(t)\) and \(y = h(t)\), then:

\[ \frac{dz}{dt} = \frac{\partial f}{\partial x} \cdot \frac{dx}{dt} + \frac{\partial f}{\partial y} \cdot \frac{dy}{dt} \]

This "sum over all paths" form is exactly how autograd handles a variable like a that is used in multiple downstream operations — it sums up the gradient contributions from each usage.

Concrete example — trace d = (a × b) + a:

a = 2.0, b = 3.0
c = a × b = 6.0        (∂c/∂a = b = 3,  ∂c/∂b = a = 2)
d = c + a  = 8.0        (∂d/∂c = 1,      ∂d/∂a = 1)

Chain rule to find ∂d/∂a (two paths: direct through add, and indirect through mul):

∂d/∂a = ∂d/∂c × ∂c/∂a   +   ∂d/∂a (direct)
       =    1  ×   3     +       1
       = 4

Verify: d = (a × b) + a = ab + a, so ∂d/∂a = b + 1 = 3 + 1 = 4. ✓

This sum-over-paths structure is not a coincidence — it is exactly what the backward pass computes when it accumulates gradients into leaf tensors.

Chapter 3: Jacobians — Why Backpropagation Works at Scale

3.1 The Jacobian Matrix

When a function maps a vector to a vector — \(\mathbf{y} = f(\mathbf{x})\) where \(\mathbf{x} \in \mathbb{R}^n\) and \(\mathbf{y} \in \mathbb{R}^m\) — the derivative is no longer a single number. It is a matrix called the Jacobian:

\[ J = \frac{\partial \mathbf{y}}{\partial \mathbf{x}} = \begin{bmatrix} \frac{\partial y_1}{\partial x_1} & \cdots & \frac{\partial y_1}{\partial x_n} \ \vdots & \ddots & \vdots \ \frac{\partial y_m}{\partial x_1} & \cdots & \frac{\partial y_m}{\partial x_n} \end{bmatrix} \]

Entry \(J_{ij} = \frac{\partial y_i}{\partial x_j}\) — how much output \(i\) changes when input \(j\) changes.

Example: linear layer \(\mathbf{y} = W\mathbf{x}\) where \(W\) is \((m \times n)\). Then \(y_i = \sum_j W_{ij} x_j\), so \(\frac{\partial y_i}{\partial x_j} = W_{ij}\). The Jacobian is exactly \(W\). The gradient of \(\mathbf{y}\) with respect to \(\mathbf{x}\) IS the weight matrix.

3.2 Why Backpropagation Is a Sequence of Jacobian-Vector Products

The chain rule in matrix form: given compostion \(f \circ g\) mapping \(\mathbf{x} \to \mathbf{y} \to \mathbf{z}\):

\[ \frac{\partial \mathbf{z}}{\partial \mathbf{x}} = \frac{\partial \mathbf{z}}{\partial \mathbf{y}} \cdot \frac{\partial \mathbf{y}}{\partial \mathbf{x}} = J_f \cdot J_g \]

For a neural network with loss \(\mathcal{L}\) (a scalar), backprop computes:

\[ \frac{\partial \mathcal{L}}{\partial \mathbf{x}} = \left(\frac{\partial \mathcal{L}}{\partial \mathbf{y}}\right)^T J_g \]

This is a vector-Jacobian product (VJP). The vector is the upstream gradient (already computed); the Jacobian is the local derivative of this operation. We never form the full Jacobian matrix — we compute the product directly.

For a Linear layer y = Wx:

• The Jacobian with respect to x is W
• The VJP (upstream gradient g times Jacobian) is g @ W or equivalently W.T @ g

This is why grad_x = grad_y @ W.T in every linear layer's backward pass. It's not magic — it's the VJP of the Jacobian W.

3.3 Why This Justifies Reverse Mode

Forward mode computes Jacobian-vector products — how outputs change per input direction. To get gradients of a scalar loss w.r.t. all inputs, you'd need one forward pass per input dimension.

Reverse mode computes vector-Jacobian products — how the scalar loss flows back through each operation. One reverse pass gives you gradients w.r.t. all inputs simultaneously.

For neural networks: millions of inputs (weights), one output (loss). Reverse mode is \(O(\text{forward pass cost})\) regardless of how many parameters you have. Forward mode would be \(O(\text{parameters} \times \text{forward pass cost})\) — completely infeasible for 7-billion-parameter models.

Chapter 4: Automatic Differentiation — Two Modes, One Winner

4.1 What Is Automatic Differentiation?

Automatic differentiation (autodiff) is not symbolic differentiation (like a CAS doing algebra) and not numerical differentiation (finite differences). It is the mechanical application of the chain rule to the actual program that computes the function, tracking derivatives alongside values.

Three approaches exist:

PyTorch's autograd IS automatic differentiation. It achieves exact derivatives (up to floating-point rounding) in O(forward pass cost) time.

4.2 The Execution Tape

Reverse-mode autodiff works by recording a tape (also called the Wengert list) of all operations executed during the forward pass. The tape stores:

• The operation performed (add, mul, matmul, relu, …)
• The inputs that were consumed
• Any saved intermediate values needed for the backward

After the forward pass, the tape is traversed in reverse to compute gradients.

Forward (recording):
  a = 2.0 (leaf)
  b = 3.0 (leaf)
  c = a * b    → tape: [MUL, inputs=(a,b), saved=(a=2.0, b=3.0)]
  d = c + a    → tape: [ADD, inputs=(c,a)]
  loss = d

Backward (replay in reverse):
  d.backward(grad=1.0):
    ADD.backward(grad=1.0)  → sends 1.0 to c, 1.0 to a
    MUL.backward(grad=1.0)  → grad_a += 1.0*b=3.0, grad_b = 1.0*a=2.0
  
  a.grad = 1.0 (from ADD) + 3.0 (from MUL) = 4.0
  b.grad = 2.0

PyTorch does not literally have a "tape list" — it uses a DAG of Function objects attached to tensors. But the concept is identical.

4.3 Forward Mode vs Reverse Mode — The Full Story

Forward mode propagates derivatives from inputs to outputs alongside the forward computation. Each variable carries a "tangent" value representing its derivative w.r.t. a chosen input direction.

For f: R^n → R^m with n inputs and m outputs:

• Forward mode costs O(n) × forward pass — one pass per input
• Efficient when n < m (few inputs, many outputs)

Reverse mode propagates gradients from the output back to inputs.

• Reverse mode costs O(m) × forward pass — one pass per output
• Efficient when m < n (many inputs, one output) — exactly the ML case

Neural network training: n = millions of parameters, m = 1 (the loss scalar). Reverse mode is literally a million times more efficient than forward mode for this setting. This is the entire reason deep learning is tractable.

JAX supports both modes (jax.jvp for forward, jax.vjp for reverse) as first-class citizens. PyTorch only exposes reverse mode through autograd but supports forward mode via torch.autograd.functional.jvp.

Part II — How Computers Execute Code

Chapter 5: CPython — How Python Actually Runs

5.1 Python Is Not Executed Directly

When you run python script.py, Python does NOT interpret the text of your file line by line. It compiles it first.

Step 1 — Compilation to bytecode: Python's parser turns your source code into an AST (Abstract Syntax Tree), then the compiler turns the AST into bytecode — a sequence of simple, stack-based instructions. This bytecode is stored in .pyc files in __pycache__/.

def add(a, b):
    return a + b

Compiled bytecode (shown via dis module):

  LOAD_FAST    0 (a)    # push local variable 'a' onto the stack
  LOAD_FAST    1 (b)    # push local variable 'b' onto the stack
  BINARY_ADD           # pop two values, push their sum
  RETURN_VALUE         # pop and return the top of the stack

Step 2 — The CPython interpreter loop: CPython (the standard Python implementation, written in C) executes bytecode in a giant C while loop called ceval.c. For each bytecode instruction, it dispatches to a C function that performs the operation.

The interpreter loop is why Python is "slow" — for every a + b in your code, CPython:

1. Fetches the BINARY_ADD instruction
2. Pops two PyObject* pointers from the stack
3. Looks up the __add__ method on the left object
4. Calls through multiple C function layers
5. Allocates a new Python object for the result
6. Pushes the pointer back onto the stack

This overhead is ~100-1000× more expensive than the equivalent C code doing the same addition.

5.2 Why This Matters for ML Frameworks

ML training loops repeat billions of simple arithmetic operations. If each operation paid the CPython interpreter overhead, training would be impossibly slow. The solution: batch operations into kernels.

When you write y = torch.matmul(a, b):

1. CPython executes ONE bytecode instruction (CALL_FUNCTION)
2. This calls a C++ function registered via pybind11
3. The C++ function dispatches to a highly optimized BLAS routine (Intel MKL or cuBLAS)
4. The kernel runs billions of multiplications and additions without returning to CPython
5. A single Python object (the output tensor) is returned

The Python layer is thin glue. The heavy lifting is in C++/CUDA.

5.3 The GIL — Global Interpreter Lock

CPython has a fundamental limitation called the GIL (Global Interpreter Lock). Only one Python thread can execute Python bytecode at a time. The GIL is a mutex that prevents concurrent Python execution.

Why does this exist? CPython's memory management (reference counting) is not thread-safe. Making it thread-safe would require fine-grained locks on every object, which would be slow. The GIL is a coarse lock that serializes all Python execution.

Impact on ML: Python-level parallelism doesn't help. But ML frameworks work around this by releasing the GIL during C++ computation:

// PyTorch's pattern for long-running C++ calls:
{
  pybind11::gil_scoped_release release;   // release the GIL
  // This C++ code runs without holding Python's lock.
  // Other Python threads can run while we're computing.
  result = do_expensive_computation();
}
// GIL is re-acquired here automatically

During a CUDA kernel launch or a matmul call, Python threads are free to run (e.g., data loading threads) because the GIL is released. This is the mechanism behind PyTorch's DataLoader with num_workers > 0.

5.4 Why TorchDynamo Needs to Understand Bytecode

torch.compile uses TorchDynamo to capture computation graphs. Dynamo works by installing a frame evaluation callback in CPython — a hook that intercepts bytecode execution before it runs.

When Python is about to execute a @torch.compile-decorated function, Dynamo's callback fires. Dynamo reads the raw bytecode instructions (LOAD_FAST, BINARY_MULTIPLY, CALL_FUNCTION, etc.) and converts them into a graph of tensor operations. To do this, it must understand:

• Which bytecode instructions produce tensor values
• Which instructions are Python-only (don't involve tensors)
• Where control flow (jumps, conditionals) happens

This is why torch.compile is more powerful than torch.fx.symbolic_trace — it operates at the bytecode level where it can see everything, not at the Python function level where it only sees what you call explicitly.

Chapter 6: NumPy — Arrays, Memory Layout, and Vectorization

6.1 Why NumPy Exists

Python lists are arrays of Python objects. A list of 1000 floats stores 1000 PyObject* pointers, each pointing to a heap-allocated Python float. The floats are scattered in memory. Iterating over them requires following 1000 pointer indirections.

NumPy arrays store raw values contiguously in C-allocated memory — no Python objects, no pointers. A NumPy array of 1000 float64s is exactly 8000 bytes in one contiguous block. Accessing element 500 is a single memory read at offset 500 × 8 = 4000 bytes. No pointer indirection.

This contiguity is what enables vectorization: modern CPUs have SIMD (Single Instruction, Multiple Data) instructions (SSE, AVX, AVX-512) that operate on 4, 8, or 16 floats simultaneously. A well-optimized loop over a contiguous float array uses these instructions automatically.

6.2 Broadcasting — The Rules You Must Internalize

Broadcasting is NumPy's mechanism for applying operations between arrays of different shapes without copying. The rules:

1. If arrays have different numbers of dimensions, the smaller is padded with 1s on the left
2. Two dimensions are compatible if they are equal OR one of them is 1
3. The output shape is the element-wise maximum of the input shapes

a shape: (     3, 1)
b shape: (  4, 3, 5)
───────────────────
         (  4, 3, 5)  ← output: a is repeated along dim 0 and 1, b along dim 2

Why this matters for autograd: when you compute a + b and a was broadcast, the gradient of the loss w.r.t. a must be summed over the broadcast dimensions — because the same a values contributed to multiple output positions. This is _unbroadcast in the Lab 01 solution:

def _unbroadcast(grad, target_shape):
    # Sum over dimensions that were broadcast during forward
    ...

Forgetting to unbroadcast is one of the most common bugs in custom autograd implementations.

6.3 Vectorization vs Python Loops

# Slow: Python loop, one operation per iteration
result = []
for i in range(len(a)):
    result.append(a[i] * 2.0)

# Fast: vectorized, operates on whole array at once
result = a * 2.0  # single NumPy/PyTorch call → one C kernel

The fast version:

1. Calls into C code
2. Loops in C (100-200× faster than Python loops)
3. May use SIMD to process 8 floats per CPU clock cycle
4. Has no Python object allocation per element

For 1 million elements: the Python loop might take 500ms; the vectorized version might take 0.5ms — 1000× faster.

Chapter 7: C++ and pybind11 — The Bridge to Performance

7.1 Why PyTorch Is Written in C++

C++ gives two things Python cannot: predictable memory management and zero-overhead abstractions.

In C++, you control when memory is allocated and freed. You can allocate a tensor buffer once and reuse it without the garbage collector deciding to reclaim it at an inopportune moment. In deep learning, where a single forward pass may allocate and free gigabytes of GPU memory, this control is critical.

C++ also compiles to native machine code. The compiler can inline function calls, unroll loops, and use SIMD instructions automatically. Python cannot do any of this because it must interpret bytecode at runtime.

PyTorch's core (libtorch) is entirely C++. The Python interface is generated via pybind11.

7.2 What pybind11 Does

pybind11 is a header-only C++ library that exposes C++ classes and functions to Python. It bridges the two type systems:

// C++ function
at::Tensor my_relu(const at::Tensor& input) {
    return input.clamp(0);
}

// pybind11 module definition
PYBIND11_MODULE(my_module, m) {
    m.def("relu", &my_relu,
          py::arg("input"),
          "Apply ReLU activation");
}

After building this as a .so file, Python can import and call it:

import my_module
y = my_module.relu(x)  # calls C++ function with Python tensor

pybind11 handles the conversion between Python objects (PyObject*) and C++ objects (at::Tensor) transparently. When you pass a PyTorch tensor from Python to the C++ function, pybind11 unwraps the Python object to get the underlying at::Tensor C++ object.

7.3 RAII — How C++ Manages Memory

C++ does not have garbage collection. Instead, it uses RAII (Resource Acquisition Is Initialization): resources (memory, file handles, GPU buffers) are tied to C++ object lifetimes. When an object goes out of scope, its destructor automatically releases resources.

{
    at::Tensor t = torch::randn({1024, 1024});  // GPU memory allocated
    // ... use t ...
}  // t goes out of scope → destructor runs → GPU memory freed (returned to pool)

This deterministic cleanup is why PyTorch can manage GPU memory so precisely. Python's garbage collector is non-deterministic — you don't know when del tensor will actually free memory. C++ destructors fire immediately when the object leaves scope.

Part III — PyTorch Architecture

Chapter 8: ATen, c10, and libtorch — The C++ Foundation

8.1 The Library Layers

PyTorch's C++ codebase is structured as three layers:

libtorch (the full C++ distribution)
    ├── ATen (A Tensor Library) — tensor operations
    └── c10 (Core 10 / Caffe2 + Torch)  — foundational types
            ├── Dispatcher
            ├── OperatorRegistry
            ├── IValue (universal value type)
            └── Storage, Device, Dtype

c10 is the foundation. It provides:

• c10::Dispatcher — the universal routing system (Chapter 10)
• c10::IValue — a type-erased container that holds any value (tensor, int, string, list, dict) — used to pass arguments through the dispatch system
• c10::Device — cpu:0, cuda:0, cuda:1, meta, etc.
• c10::ScalarType — float32, float16, int8, etc.
• c10::Storage / c10::StorageImpl — the physical memory block (Chapter 9)

ATen (A Tensor Library) sits on top of c10 and provides:

• at::Tensor — the C++ tensor class (the actual object Python's torch.Tensor wraps)
• ~1600 tensor operations defined in the aten:: namespace (aten::add, aten::mm, aten::relu, etc.)
• The schema system — each operation has a typed schema string used by the dispatcher

libtorch is the distributable package containing compiled ATen + c10 + autograd, exposing a C++ API (torch::Tensor, torch::nn::Module, etc.) that mirrors PyTorch's Python API.

8.2 The aten:: Namespace — Where Operations Live

Every PyTorch operation ultimately corresponds to a function in the aten:: namespace. When you call torch.relu(x) in Python:

1. Python's torch.relu is a Python wrapper generated from the ATen operator registry
2. It calls at::relu(x) in C++
3. at::relu is registered with the dispatcher under the key aten::relu
4. The dispatcher selects the correct implementation based on the tensor's dispatch keys

You can see the full schema of any operation:

print(torch.ops.aten.relu)
# aten::relu(Tensor self) -> Tensor

The schema string is not just documentation — it is parsed by the dispatcher at registration time to understand argument types, mutability, and aliasing. This is how PyTorch's compiler can reason about operations without running them.

Chapter 9: Tensors in PyTorch — Storage, Strides, and Memory Management

9.1 The C++ Object Behind torch.Tensor

Python's torch.Tensor is a thin Python wrapper around at::Tensor (C++). The C++ object contains:

struct TensorImpl {
    c10::Storage   storage_;      // the physical memory block
    int64_t        storage_offset_; // offset into storage in elements
    IntArrayRef    sizes_;         // shape, e.g., [3, 4]
    IntArrayRef    strides_;       // strides, e.g., [4, 1]
    c10::ScalarType dtype_;        // float32, int8, etc.
    c10::Device    device_;        // cpu, cuda:0, meta
    AutogradMeta*  autograd_meta_; // grad_fn, requires_grad, etc. (null if not needed)
    DispatchKeySet key_set_;       // bit-field of active dispatch keys
};

The DispatchKeySet is a 64-bit integer where each bit represents an active dispatch key. When the dispatcher receives an operation, it inspects the key_set_ of all input tensors (bitwise OR) to find the highest-priority active key.

9.2 The Caching Allocator — Why GPU Memory Isn't Really Freed

Calling cudaMalloc (NVIDIA's API to allocate GPU memory) is expensive: it may require the GPU driver to find a free region, update page tables, and sometimes synchronize with the GPU. A single cudaMalloc can take hundreds of microseconds — orders of magnitude longer than a fast CUDA kernel.

PyTorch's caching allocator (c10/cuda/CUDACachingAllocator.cpp) avoids this by maintaining a pool of already-allocated GPU memory:

When you create a CUDA tensor:
  1. Request size S from the allocator
  2. Allocator checks its free list for a block ≥ S
  3. If found: return that block (no cudaMalloc call — ~1μs)
  4. If not found: call cudaMalloc to expand the pool (~200μs)

When the tensor is deallocated:
  1. Reference count drops to zero
  2. Destructor runs: marks block as "free" in the allocator's pool
  3. Does NOT call cudaFree
  4. The memory stays in the pool for future allocations

Why torch.cuda.empty_cache() exists: it flushes the pool back to the OS via cudaFree. You'd call this before allocating a very large model to reclaim fragmented small blocks. Under normal operation, never call it — the caching is intentional.

The consequence you'll see in profiling: CUDA memory usage grows until it stabilizes. The allocator is holding memory it has already allocated to avoid future cudaMalloc overhead. This is not a leak — it's intentional pooling.

Chapter 10: The Dispatcher — PyTorch's Universal Routing System

10.1 Why the Dispatcher Exists

Before the dispatcher, PyTorch had a simpler design: each tensor operation had a set of if/else branches choosing the implementation based on tensor type:

// Old approach (simplified pseudocode):
Tensor relu(const Tensor& self) {
    if (self.is_cuda()) {
        return relu_cuda(self);
    } else if (self.is_cpu()) {
        return relu_cpu(self);
    } else {
        throw std::runtime_error("Unknown device");
    }
}

This broke every time someone wanted to add:

• A new device (ROCm GPU, XLA, NPU)
• A new feature (autograd, quantization, AMP)
• A new transformation (functorialize mutations, meta tensors)

Each addition required modifying every operation. With 1600+ operations, this was unmaintainable.

The dispatcher solves this with a table-driven routing system where:

• Operations are registered once with a schema
• Implementations are registered independently for each (operation, dispatch_key) pair
• The dispatcher selects the implementation at runtime without any if/else in the operation itself

10.2 The DispatchKey Enum

A DispatchKey is an integer in an enum. PyTorch has ~50 dispatch keys. The dispatcher maintains a table of size (num_operators × num_keys) mapping each pair to a function pointer.

The most important keys in priority order:

Highest priority:
  Functionalize   — converts in-place/aliasing ops to pure functional form.
                    TorchScript and torch.compile require functional semantics.
                    Without this, tracking aliasing through graph transformations
                    is intractable.

  InferenceMode   — completely disables gradient tracking.
                    Faster than no_grad for inference-only code.

  AutogradMeta    — handles FakeTensors (tensors with only shape/dtype, no data).
                    Used during torch.compile's shape analysis phase.

  Autograd        — THE gradient-recording layer. Before calling any compute
                    kernel, wraps the call to record it in the autograd graph.
                    Sets grad_fn on the output tensor.

  AutocastCPU/    — AMP (Automatic Mixed Precision). Intercepts calls and
  AutocastCUDA      promotes/demotes dtypes (e.g., matmul in float16, norms
                    in float32) without user code changes.

  CUDA            — NVIDIA GPU kernels. cuBLAS, cuDNN, custom CUDA code.
  CPU             — CPU kernels. Intel MKL, Eigen, vectorized C++.

  Meta            — Shape-only kernels. No data, no computation. Returns
                    tensors with correct shape/dtype but no storage.

Lowest priority:
  BackendSelect   — Selects which device backend to use (CPU vs CUDA vs ...)
                    based on input tensor devices.

10.3 A Complete Dispatch Trace

Call: torch.relu(x) where x is a CUDA tensor with requires_grad=True.

x.key_set_ has bits set for: {CUDA, Autograd}.

Dispatcher checks keys from highest to lowest:
  Functionalize? Bit not set. Skip.
  InferenceMode? Bit not set. Skip.
  AutogradMeta? Bit not set. Skip.
  Autograd? Bit IS set → call Autograd implementation.

Autograd implementation of relu:
  1. Calls torch::autograd::generated::ReluBackward::apply_with_saved(...)
  2. This wraps the call: records that relu ran, creates ReluBackward0 grad_fn
  3. Re-dispatches DOWN, excluding the Autograd key:
     new_key_set = x.key_set_ & ~DispatchKeySet{Autograd}
     → {CUDA}
  4. Calls CUDA implementation of relu

CUDA implementation of relu:
  1. Launches CUDA kernel: relu_kernel<<<grid, block>>>(x.data_ptr<float>(), ...)
  2. Returns output tensor

Back in Autograd:
  output.grad_fn = ReluBackward0  (records: "relu was applied to x")
  return output

Total: one Python call → one Autograd interception (sets grad_fn) → one CUDA kernel launch.

10.4 Registering a New Device Backend

This is the extension point for NPU vendors like Qualcomm:

// Register a new PrivateUse1 key for a hypothetical NPU
TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
    m.impl("relu", &npu_relu_impl);
    m.impl("mm", &npu_mm_impl);
    // ... etc. for all supported ops
}

A tensor on the NPU device has PrivateUse1 in its key_set_. The dispatcher routes to npu_relu_impl automatically. The Autograd key still sits above it, so gradient recording works automatically unless explicitly disabled.

The ordering matters: PrivateUse1 is below Autograd. If a vendor incorrectly registered above Autograd, gradient recording would be silently bypassed — a catastrophic bug that would only surface during training.

Chapter 11: The Autograd Engine — Building and Traversing the Backward Graph

11.1 The Data Structures

PyTorch's autograd engine (in torch/csrc/autograd/) is built around these C++ classes:

torch::autograd::Node (also called Function): represents one differentiable operation in the graph. Stores:

• next_edges_: a list of Edge objects pointing to upstream nodes
• input_metadata_: shape/dtype info for gradient validation
• Virtual method apply(variable_list&& inputs) — the backward function

torch::autograd::Edge: a directed edge in the backward graph. Contains:

• function: pointer to the upstream Node
• input_nr: which input of that Node this edge corresponds to

torch::autograd::AccumulateGrad: a leaf Node that accumulates gradients into a leaf tensor's .grad field. Every requires_grad=True leaf tensor has one.

AutogradMeta: attached to every tensor that participates in autograd. Contains:

• grad_fn_: the Node that produced this tensor (null for leaves)
• grad_: accumulated gradient (null until backward is called)
• requires_grad_: whether this tensor needs gradients
• output_nr_: which output of grad_fn_ this tensor is

11.2 DAG Construction During Forward Pass

When the Autograd dispatch key intercepts an operation, it:

1. Calls the compute kernel (CPU/CUDA) to get the output values
2. Creates a new Node subclass specific to this operation (e.g., MmBackward0 for matmul)
3. Saves any tensors the backward will need (via save_for_backward)
4. Sets next_edges_ to point to the grad_fn_ of each input (or AccumulateGrad for leaves)
5. Attaches this new Node as the grad_fn_ of the output tensor

After forward pass of:
  c = a @ b    (where a, b require_grad)

Memory layout:
  a → AutogradMeta { grad_fn=null, requires_grad=true }
  b → AutogradMeta { grad_fn=null, requires_grad=true }
  c → AutogradMeta { grad_fn=MmBackward0 }
         MmBackward0 {
           next_edges: [AccumulateGrad(a), AccumulateGrad(b)]
           saved_tensors: [a, b]   ← for computing grad_a = grad_c @ b.T
         }

11.3 The Backward Pass: Engine::execute()

c.backward() calls torch::autograd::Engine::execute(). The C++ engine:

1. Creates a priority queue (or thread pool) of nodes to process
2. Initializes the root node (c.grad_fn = MmBackward0) with grad = ones_like(c)
3. Processes nodes in reverse topological order:
   • For each node, calls node.apply(incoming_gradients) → returns gradient tuples
   • For each output gradient, sends it to the corresponding next_edge.function
   • When a function receives gradients from ALL its outputs (tracked by a counter), it's ready to run
4. When AccumulateGrad runs, adds incoming gradient to leaf.grad

The readiness counter is what implements topological ordering. Each node tracks how many downstream outputs still need to send it a gradient. A node runs only when this counter reaches zero — i.e., all downstream contributions have arrived.

Initial state:
  MmBackward0.ready_count = 1  (one consumer: the starting loss gradient)

After loss.backward(grad=1.0):
  Engine sends 1.0 to MmBackward0
  MmBackward0.ready_count decrements to 0 → runs now

MmBackward0.apply([grad_c]):
  grad_a = grad_c @ b.T    (saved b from forward)
  grad_b = a.T @ grad_c    (saved a from forward)
  Returns [grad_a, grad_b]

Engine sends grad_a to AccumulateGrad(a) → a.grad += grad_a
Engine sends grad_b to AccumulateGrad(b) → b.grad += grad_b

11.4 Why torch.no_grad() Is Not Just "Disabling Gradients"

When you enter torch.no_grad(), PyTorch sets GradMode::is_enabled() to false. The Autograd dispatch key's code checks this flag before creating Function objects:

// In the autograd dispatch implementation:
if (!GradMode::is_enabled()) {
    // Skip all autograd recording, call compute kernel directly
    return at::AutoDispatchBelowADInplaceOrView::redispatch(...);
}

What is NOT allocated when no_grad() is active:

• No Node objects (no MmBackward0, etc.)
• No AutogradMeta on output tensors (or stripped-down version)
• No saved tensors (the inputs to the backward function)
• No next_edges_ (the DAG edges)

For a large transformer forward pass, this can be gigabytes of saved activations that are simply never allocated. The speedup comes from: (1) no memory allocation for the graph, (2) no reference counting overhead for saved tensors, (3) no function dispatch overhead for the autograd wrapper.

Chapter 12: torch.autograd.Function — Custom Differentiable Operations

12.1 When and Why to Write One

PyTorch knows how to differentiate through all built-in operations. But sometimes you need to:

• Implement an operation that doesn't exist in PyTorch (a custom NPU intrinsic)
• Override the default backward with a more numerically stable or efficient version
• Implement a non-standard gradient (e.g., straight-through estimator for quantization)

torch.autograd.Function lets you inject a custom (forward, backward) pair into the autograd graph as a single opaque node.

12.2 The API, Fully Explained

class MySoftmax(torch.autograd.Function):

    @staticmethod
    def forward(ctx, x, dim):
        # ctx is a FunctionContext object. It is the storage for the backward.
        # It is NOT a Python object you create — it's provided by the engine.
        #
        # RULE: Never store tensors directly as ctx attributes.
        #       Use ctx.save_for_backward() for tensors.
        #       Use ctx.dim = dim for non-tensor metadata.
        
        x_max = x.max(dim=dim, keepdim=True).values
        x_shift = x - x_max                  # numerically stable: all values ≤ 0
        exp_x = x_shift.exp()
        sum_exp = exp_x.sum(dim=dim, keepdim=True)
        out = exp_x / sum_exp
        
        ctx.save_for_backward(out)            # saves output (not input) for backward
        ctx.dim = dim                         # saves non-tensor metadata as attribute
        return out

    @staticmethod
    def backward(ctx, grad_output):
        # grad_output: the gradient of the loss w.r.t. the output of forward()
        #              shape matches the output of forward()
        
        out, = ctx.saved_tensors              # retrieves saved tensors
        dim = ctx.dim
        
        # Softmax Jacobian-vector product:
        # ∂L/∂x_i = out_i × (∂L/∂y_i - Σ_j (∂L/∂y_j × out_j))
        dot = (grad_output * out).sum(dim=dim, keepdim=True)
        grad_x = out * (grad_output - dot)
        
        # RULE: return one gradient per argument to forward().
        # forward() took (ctx, x, dim). dim is a Python int (non-differentiable).
        # Return grad_x for x, and None for dim.
        return grad_x, None

# Usage:
result = MySoftmax.apply(x, dim=-1)   # Note: .apply(), not direct call

12.3 The Memory Management Rule: Why save_for_backward

If you write ctx.my_tensor = x instead of ctx.save_for_backward(x):

1. PyTorch's memory-efficient gradient checkpointing cannot track this tensor
2. The version counter check (detecting in-place mutations) is bypassed — you may compute wrong gradients silently if x is modified in-place after the forward
3. The tensor is not freed when the backward graph is freed (memory leak)

ctx.save_for_backward registers the tensor with the engine. The engine:

• Tracks the version counter at save time
• Checks the version counter at backward time (raises error if tensor was modified in-place)
• Frees the saved tensors as soon as backward() processes this node (unless retain_graph=True)

12.4 The Straight-Through Estimator — A Real Production Use Case

Quantization maps a float value to a small integer (e.g., float32 → int8). The quantize-dequantize operation:

def quantize(x, scale, zero_point):
    return torch.clamp(torch.round(x / scale) + zero_point, 0, 255)

The gradient of round() is zero almost everywhere (it's a staircase). This means no gradient flows back through quantized layers — training with quantization would not work.

The Straight-Through Estimator (STE): pretend in the backward pass that quantization didn't happen — pass gradients straight through as if the operation were an identity. This is physically inaccurate but empirically works for training quantized models.

class QuantizeSTE(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, scale, zero_point):
        # Real quantization in forward
        x_int = torch.clamp(torch.round(x / scale) + zero_point, 0, 255)
        return (x_int - zero_point) * scale   # dequantize

    @staticmethod
    def backward(ctx, grad_output):
        # Straight-through: gradient passes as if forward were identity
        return grad_output, None, None

This pattern is used in every QAT (Quantization-Aware Training) framework — AIMET, PyTorch native QAT, and custom NPU training toolchains. You will implement exactly this in Phase 03.

Part IV — Graph Capture and Compilation

Chapter 13: torch.fx — Intermediate Representation and Graph Transformation

13.1 What Is an Intermediate Representation?

In compiler design, an intermediate representation (IR) is a data structure that sits between the source language (Python) and the target (native code). Good IRs are:

• Analyzable: easy to traverse and inspect programmatically
• Transformable: easy to modify (insert, delete, reorder nodes)
• Lowerable: can be compiled/interpreted efficiently

Famous IRs: LLVM IR (used by Clang/Rust), JVM bytecode (used by Java/Kotlin), WebAssembly. torch.fx is PyTorch's ML-specific IR.

13.2 The Problem FX Solves

Before FX, if you wanted to fuse conv2d → batch_norm → relu into one kernel:

• You'd have to modify PyTorch's C++ source code
• Or hook into Python's metaclass system in fragile ways
• There was no clean programmatic way to inspect and modify a model's computation

FX provides a first-class Python object representing the computation — a GraphModule containing a Graph of Node objects. You can write Python code that traverses the graph, finds patterns, and rewrites them.

13.3 How Symbolic Tracing Works — The Proxy Mechanism

torch.fx.symbolic_trace(model) does something clever: it runs the model's forward() function, but instead of passing real tensors, it passes Proxy objects.

A Proxy is a Python object that:

• Behaves like a tensor (has .shape, .dtype, .device)
• Records every operation performed on it instead of executing anything
• Returns a new Proxy representing the output of each operation

class Proxy:
    def __add__(self, other):
        # Don't compute! Record that addition happened.
        node = self.graph.create_node(
            op='call_function',
            target=operator.add,
            args=(self, other)
        )
        return Proxy(node, self.graph)
    
    def __matmul__(self, other):
        # Same pattern for all operations
        node = self.graph.create_node(...)
        return Proxy(node, self.graph)

When model.forward(proxy_input) executes:

• self.linear(proxy_input) → records call_module[linear]
• torch.relu(intermediate_proxy) → records call_function[relu]
• return proxy.sum() → records call_method[sum]

After the function returns, the graph object contains all recorded nodes. FX wraps it in a GraphModule that re-executes the graph using real tensors.

13.4 The Four Node Types and Their Fields

Every node has:

• node.op: one of placeholder, call_function, call_module, call_method, output
• node.name: a unique string identifier
• node.target: what is being called (a function, module name, or method name)
• node.args: positional arguments (may reference other nodes)
• node.kwargs: keyword arguments
• node.users: set of downstream nodes that consume this node's output

# After tracing a model:
for node in graph.nodes:
    if node.op == 'call_function' and node.target == torch.relu:
        # Access predecessors:
        for arg in node.args:
            if isinstance(arg, Node):
                print(f"{arg.name} → relu → {node.name}")
        # Access successors:
        for user in node.users:
            print(f"relu output used by: {user.name}")

13.5 Transforming Graphs: The Pattern-Replace Loop

The core graph transformation pattern:

def insert_quantization_observers(model):
    traced = torch.fx.symbolic_trace(model)
    
    # Collect nodes to transform (can't modify graph while iterating)
    nodes_to_wrap = []
    for node in traced.graph.nodes:
        if node.op == 'call_module':
            module = traced.get_submodule(node.target)
            if isinstance(module, (torch.nn.Conv2d, torch.nn.Linear)):
                nodes_to_wrap.append(node)
    
    # Apply transformations
    for node in nodes_to_wrap:
        with traced.graph.inserting_after(node):
            # Insert an observer node after each conv/linear
            observer_name = node.name + '_observer'
            traced.add_module(observer_name, MinMaxObserver())
            new_node = traced.graph.call_module(observer_name, (node,))
        
        # Replace all uses of node with new_node
        node.replace_all_uses_with(new_node)
        # Move new_node to immediately after node in the graph
        new_node.args = (node,)
    
    traced.recompile()  # regenerate Python code from modified graph
    return traced

This pattern is exactly how AIMET, PyTorch's quantize_fx, and torch.compile's optimization passes work.

13.6 FX's Fundamental Limitations

Limitation 1: Data-dependent control flow is opaque.

def forward(x):
    n = x.shape[0]          # Python int — OK
    if x.sum() > 0:          # x.sum() returns a Proxy, > 0 raises TraceError
        return x * 2
    return x + 1

The Proxy for x.sum() is a symbolic object with no concrete value. Python's if statement calls bool() on it, which raises TraceError: Tried to use FX tracing to trace a dynamic value.

Limitation 2: Shapes are concrete at trace time.

If you trace with x.shape = (4, 128), shape information is baked into the graph. Running on x.shape = (8, 128) may work (if no shapes are used as constants) or silently produce wrong results (if shapes were used in computation).

Why torch.compile avoids these limitations: TorchDynamo operates at the bytecode level. It sees the raw COMPARE_OP bytecode instruction for x.sum() > 0 and handles it differently — it evaluates the condition with the actual tensor (breaking the graph if needed) rather than trying to symbolically trace through it.

Chapter 14: torch.compile — The Full Compilation Stack

14.1 What Problem torch.compile Solves

A pure eager PyTorch model executes one operation at a time:

1. Python calls relu(x) → dispatch → CUDA kernel launch
2. Python calls linear(x) → dispatch → cuBLAS call
3. Python calls layer_norm(x) → dispatch → another CUDA kernel

Each kernel launch has overhead: Python→C++ dispatch (~1μs), CUDA launch latency (~5μs), GPU memory round-trip for intermediate results. For a transformer layer with 100+ operations, these overheads accumulate.

torch.compile captures a sequence of operations as a static graph and compiles it to optimized native code:

• Fuses relu + linear into one kernel that reads input once
• Eliminates intermediate tensor allocations between fused ops
• Schedules computation to maximize GPU utilization

14.2 TorchDynamo — Bytecode Interception

TorchDynamo installs a Python frame evaluation callback via CPython's _PyEval_EvalFrameDefault override API (PEP 558). This is a CPython hook that lets C code intercept bytecode execution.

When Python is about to execute a decorated function:

Normal Python execution:
  Python → CPython eval loop → execute bytecodes

With TorchDynamo:
  Python → Dynamo's callback fires
         → Dynamo reads bytecodes using dis/bytecode analysis
         → Classifies each instruction: tensor op? Python op? control flow?
         → Extracts tensor operations into an FX graph
         → For unsupported ops: breaks graph, falls back to Python
         → Compiles captured graph via AOTAutograd + Inductor
         → Installs guards (fast runtime checks)
         → On future calls: checks guards, dispatches to compiled artifact

Guards are the key innovation. When Dynamo captures a graph, it records every assumption it made:

• "x.shape[0] == 4"
• "x.dtype == torch.float32"
• "model.training == False"

These become fast Python checks on each call:

# Generated guard code (simplified):
def check_guards(x, model):
    return (x.shape[0] == 4 and
            x.dtype == torch.float32 and
            not model.training)

If any guard fails, Dynamo recompiles for the new configuration. The compiled artifacts are cached by guard configuration, so "x.shape[0] == 4 and float32" and "x.shape[0] == 8 and float32" are two separate compiled graphs, both cached.

14.3 AOTAutograd — Capturing the Backward Graph

By default, PyTorch builds the backward graph lazily during the forward pass. For compilation, we need the backward graph ahead of time so it can be optimized along with the forward.

AOTAutograd (Ahead-Of-Time Autograd) traces both forward and backward simultaneously:

1. Runs the forward with FakeTensors (no data, just shapes/dtypes)
2. Runs autograd on the traced forward to symbolically derive the backward
3. Produces a joint FX graph containing both forward and backward as one graph
4. Applies functionalization — converts all in-place ops and aliasing ops into pure functional equivalents (required for correct compilation)

The joint graph is then split into a forward graph (returned to the caller) and a backward graph (stored for when .backward() is called).

14.4 TorchInductor — Code Generation

Inductor takes the FX graph from AOTAutograd and generates code:

For NVIDIA GPUs: generates Triton code. Triton is a Python-like DSL for writing GPU kernels. It handles thread indexing, memory coalescing, and warp management automatically — you write scalar code; Triton vectorizes it.

# What Inductor might generate for relu + add:
@triton.jit
def relu_add_kernel(x_ptr, y_ptr, out_ptr, n: tl.constexpr):
    pid = tl.program_id(0)
    block = tl.arange(0, n)
    x = tl.load(x_ptr + pid * n + block)
    y = tl.load(y_ptr + pid * n + block)
    out = tl.maximum(x, 0) + y          # fused relu + add in one pass
    tl.store(out_ptr + pid * n + block, out)

Fusion: instead of launching two kernels (one for relu, one for add), Inductor fuses them into one. The intermediate result (after relu) stays in GPU registers — no round-trip to global memory.

For CPUs: generates C++ with OpenMP parallelism and AVX vectorization hints.

14.5 The Meta Device — Shape Analysis Without Data

Before generating any kernel, Inductor must know the output shape of every operation. It uses the Meta device for this:

A Meta tensor has shape, dtype, device = "meta", but no storage — no memory allocated, no bytes of data. Operations on Meta tensors execute "Meta kernels" that compute output shapes without performing any arithmetic.

x = torch.empty(4, 128, device='meta')   # 0 bytes allocated
y = torch.nn.Linear(128, 64)(x)          # Meta kernel: (4,128) @ (128,64) → (4,64)
print(y.shape)   # torch.Size([4, 64])
print(y.numel()) # 256 — but no data exists

Every custom operator MUST register a Meta kernel. Without one, torch.compile cannot propagate shapes and will raise:

NotImplementedError: Could not run 'myops::my_op' with Meta key.

The Meta kernel need only return at::empty_like(input) for element-wise ops, or compute the correct output shape for reshape/reduction ops.

Chapter 15: GPU Architecture — From Silicon to CUDA Kernels

15.1 Why GPUs Exist: The Throughput vs Latency Tradeoff

A modern CPU (e.g., Intel Xeon with 32 cores) is optimized for latency — executing a single complex instruction stream as fast as possible. It has:

• Large caches (L1: 32KB, L2: 256KB, L3: 16MB per core)
• Branch prediction hardware (avoids stalls from if/else)
• Out-of-order execution (reorders instructions for efficiency)
• Speculative execution (pre-executes likely branches)

This complexity lets a CPU execute a 10,000-instruction sequential program in microseconds. But it scales poorly to parallel workloads.

A GPU (e.g., NVIDIA A100) is optimized for throughput — executing thousands of identical operations simultaneously. It has:

• 6912 CUDA cores (simple integer ALUs)
• 432 Tensor Cores (4×4 matrix multiply units)
• Minimal caching (L1: 192KB/SM, L2: 40MB shared)
• No branch prediction, no out-of-order execution per core
• Massive memory bandwidth: 2 TB/s (vs CPU's ~100 GB/s)

The tradeoff: a GPU cannot run a sequential 10,000-instruction program faster than a CPU. But running the same operation on 1 million data points is 100× faster on GPU because all 1 million operations run simultaneously.

15.2 The SM — Streaming Multiprocessor

The A100 has 108 Streaming Multiprocessors (SMs). Each SM is an independent processing unit containing:

One SM:
  ├── 4 Warp Schedulers     (each schedules 32-thread warps)
  ├── 64 INT32 units        (integer arithmetic)
  ├── 64 FP32 units         (floating-point arithmetic)
  ├── 32 FP64 units         (double precision)
  ├── 4 Tensor Cores        (4×4 matrix multiply in one cycle)
  ├── 192 KB L1 cache / shared memory (configurable split)
  └── Register file: 256 KB (65,536 × 32-bit registers)

An SM can hold up to 2048 threads in flight simultaneously ("resident threads"). It switches between warps each clock cycle — when one warp stalls waiting for memory, the SM immediately executes another warp. This latency hiding is how GPUs tolerate their slower memory system.

15.3 The CUDA Execution Model: SIMT

CUDA uses a programming model called SIMT (Single Instruction, Multiple Threads). You write code for one thread; the hardware runs that same code on thousands of threads simultaneously.

Thread:  The smallest unit. Runs one instance of your kernel.
Warp:    32 threads that execute the SAME instruction simultaneously.
         This is the fundamental dispatch unit of the GPU.
Block:   1–1024 threads grouped together. Threads in a block:
           - Run on the same SM
           - Share the SM's shared memory
           - Can synchronize with __syncthreads()
Grid:    Many blocks. Blocks run independently across all SMs.

When you launch a kernel:

my_kernel<<<grid_dim, block_dim, shared_mem, stream>>>(args...);

grid_dim × block_dim total threads are created. The GPU scheduler distributes blocks across SMs. Within each SM, threads are grouped into warps of 32.

15.4 Warp Execution — Why 32 Threads Must Stay Synchronized

A warp's 32 threads share one program counter. At each clock cycle, the warp executes one instruction, and all 32 threads execute it in parallel (on different data).

Warp divergence: if threads in a warp take different branches:

if (thread_id % 2 == 0) {
    // Even threads do this
    x = x * 2.0f;
} else {
    // Odd threads do this
    x = x + 1.0f;
}

The GPU CANNOT execute both branches simultaneously. It:

1. Executes x = x * 2.0f with even threads active, odd threads masked off
2. Executes x = x + 1.0f with odd threads active, even threads masked off

Result: 2× longer execution. Each half of the warp idles while the other executes. Minimizing warp divergence is a primary GPU optimization concern.

15.5 The Memory Hierarchy — The Critical Performance Factor

The golden rule: every read from global memory costs ~800 cycles. If your kernel reads each piece of data once and computes many operations on it, global memory cost is amortized. If your kernel reads data once, computes one operation, and writes back — you're memory-bound: the bottleneck is memory, not compute.

FlashAttention's breakthrough was recognizing that the standard attention algorithm is memory-bound (reads Q, K, V matrices multiple times) and restructuring the computation to keep everything in L1/shared memory using tiling.

15.6 Memory Coalescing — Why Access Pattern Matters

When 32 threads in a warp each read one float from global memory, the GPU issues a memory transaction. If the 32 addresses are contiguous (consecutive elements of an array), the transaction fetches one 128-byte cache line — one transaction total.

If the 32 addresses are strided (e.g., every other element) or random (scattered), the GPU must issue multiple transactions — 2× or 32× slower respectively.

Coalesced access (good):
  Thread 0: address 0×4 = float at byte 0
  Thread 1: address 1×4 = float at byte 4
  Thread 2: address 2×4 = float at byte 8
  ...
  Thread 31: address 31×4 = float at byte 124
  → One 128-byte transaction covers all 32 reads.

Strided access (bad):
  Thread 0: address 0×8 = float at byte 0
  Thread 1: address 1×8 = float at byte 8
  ...every other float, stride=2
  → Two transactions, 50% bandwidth wasted

This is why matrix transposition before matmul is sometimes applied — to ensure the memory access pattern in the CUDA kernel is coalesced.

15.7 Warp Shuffle Instructions — Register-Level Communication

Normally, threads share data via shared memory (write → barrier → read, ~5 cycles). Warp shuffle instructions (__shfl_*) let threads exchange register values directly within a warp without touching memory — ~1 cycle.

Available shuffles:

__shfl_sync(mask, val, src_lane)     // all threads get val from lane src_lane
__shfl_up_sync(mask, val, delta)     // get val from lane (current - delta)
__shfl_down_sync(mask, val, delta)   // get val from lane (current + delta)
__shfl_xor_sync(mask, val, laneMask) // get val from lane (current XOR laneMask)

mask = 0xffffffff means all 32 lanes participate.

The butterfly reduction using __shfl_xor_sync:

Goal: find the sum of values across all 32 threads.

Round 1 (laneMask=16):
  Lane 0 gets lane 16's value. Lane 0 now holds sum of elements 0 and 16.
  Lane 1 gets lane 17's value. Lane 1 now holds sum of elements 1 and 17.
  ... (simultaneously for all 32 pairs)

Round 2 (laneMask=8):
  Lane 0 gets lane 8's value. Lane 0 now holds sum of elements 0, 8, 16, 24.
  ...

Round 3 (laneMask=4): Each thread holds sum of 8 elements.
Round 4 (laneMask=2): Each thread holds sum of 16 elements.
Round 5 (laneMask=1): Each thread holds sum of all 32 elements.

5 rounds = log₂(32) = 5 communication steps. Every thread simultaneously holds the global sum. No shared memory, no barriers.

15.8 The warp_softmax Kernel Step by Step

Lab 02's kernel computes softmax over a 2D tensor [rows, vocab_size] where each warp handles one row:

Tensor: [batch × seq_len, vocab_size] = [rows, V]
  Lane k handles elements: k, k+32, k+64, ..., k+32*(V/32-1)
  Each lane processes V/32 elements (1000 for vocab=32000)

Step 1 — Find row max (prevents exp() overflow):
  Each lane: local_max = max(input[k], input[k+32], ..., input[k+32*(V/32-1)])
  Butterfly reduction: all lanes now hold row_max

Step 2 — Compute exp(x - row_max) and accumulate sum:
  Each lane: for i in range(k, V, 32):
               val = exp(input[i] - row_max)
               output[i] = val
               local_sum += val
  Butterfly reduction: all lanes now hold row_sum

Step 3 — Normalize:
  Each lane: inv_sum = 1.0 / row_sum   (multiply cheaper than divide)
             for i in range(k, V, 32): output[i] *= inv_sum

Three passes over the data. All intermediate values (row_max, partial sums) live in registers — no shared memory needed. This is why the kernel is fast: minimal memory traffic, maximum parallelism.

Chapter 16: Custom C++ Operators — TORCH_LIBRARY

16.1 The Two-Level Registration System

PyTorch separates what an operator is from how it's implemented on each device. This separation is the dispatcher's core design.

Level 1: Schema registration — defines the operator's name and type signature, visible to all backends:

TORCH_LIBRARY(myops, m) {
    // Schema language:
    //   "namespace::name(arg_type arg_name, ...) -> return_type"
    // Arg types: Tensor, Tensor?, Tensor[], int, int[], float, bool,
    //            ScalarType, Device, str, Scalar
    m.def("warp_softmax(Tensor self, int dim) -> Tensor");
}

This creates torch.ops.myops.warp_softmax in Python. The operator exists in the registry. Calling it without an implementation registered for the current dispatch key will error.

Level 2: Implementation registration — binds implementations to specific dispatch keys:

// CPU implementation
TORCH_LIBRARY_IMPL(myops, CPU, m) {
    m.impl("warp_softmax", &warp_softmax_cpu);   // function pointer
}

// CUDA implementation  
TORCH_LIBRARY_IMPL(myops, CUDA, m) {
    m.impl("warp_softmax", &warp_softmax_cuda);
}

// Meta implementation — REQUIRED for torch.compile
TORCH_LIBRARY_IMPL(myops, Meta, m) {
    m.impl("warp_softmax", &warp_softmax_meta);   // returns empty_like(input)
}

16.2 How Registration Actually Works at Runtime

TORCH_LIBRARY and TORCH_LIBRARY_IMPL expand to static initializers — C++ code that runs when the .so file is loaded. They call into c10::Dispatcher::singleton() to register the operator:

// What TORCH_LIBRARY expands to (simplified):
static auto registration = c10::RegisterOperators()
    .op("myops::warp_softmax(Tensor self, int dim) -> Tensor",
        c10::DispatchKey::CPU, &warp_softmax_cpu)
    .op("myops::warp_softmax(Tensor self, int dim) -> Tensor",
        c10::DispatchKey::CUDA, &warp_softmax_cuda);

When your .so is loaded via torch.ops.load_library("warp_softmax_ext.so"), these static initializers run, and the operator is available immediately.

16.3 The Two Autograd Paths

Path A — Composite (automatic differentiation):

If your forward kernel is built from differentiable ATen ops (like the CPU kernel in Lab 02 that uses .amax(), .exp(), .sum(), /), PyTorch differentiates through them automatically. The Autograd key uses CompositeImplicitAutograd — it traces through your implementation to find the gradient.

Path B — Explicit autograd (manual backward):

If your CUDA kernel is not built from differentiable ATen ops, you must register an explicit backward:

torch::autograd::Function<SoftmaxBackward>::apply = [](
    const at::Tensor& grad_output,
    const at::Tensor& output) -> at::Tensor {
    // Softmax backward: grad_input = output * (grad - (grad * output).sum(-1, true))
    auto dot = (grad_output * output).sum(-1, true);
    return output * (grad_output - dot);
};

TORCH_LIBRARY_IMPL(myops, Autograd, m) {
    m.impl("warp_softmax", torch::autograd::autogradNotImplementedFallback());
    // OR: m.impl("warp_softmax", &warp_softmax_autograd);
}

16.4 OpCheck — The Complete Validation Suite

torch.library.opcheck tests your operator against all dispatch keys automatically:

import torch

torch.ops.load_library("warp_softmax_ext.so")

# Test with a simple input
x = torch.randn(4, 128, dtype=torch.float32, requires_grad=True)
if torch.cuda.is_available():
    x = x.cuda()

results = torch.library.opcheck(
    torch.ops.myops.warp_softmax,
    args=(x, -1),
    kwargs={},
)
# Raises if any dispatch key implementation is wrong

What OpCheck verifies:

1. CPU/CUDA correctness: output matches reference (F.softmax) within 1e-5
2. Autograd correctness: torch.autograd.gradcheck on differentiable inputs
3. Meta kernel: shape propagation with device='meta' tensors
4. FakeTensor mode: compatibility with torch.compile's shape tracking
5. Out-of-place invariant: operator does not modify its inputs in-place unexpectedly

If OpCheck passes, your operator is safe in all contexts: eager mode, training, inference, torch.compile, quantization pipelines.

Part V — TensorFlow

Chapter 17: TF1 to TF2 — Two Execution Philosophies

17.1 The TF1 Philosophy: Define-Then-Run

TensorFlow 1.x was designed by Jeff Dean and the Google Brain team in 2015, influenced by Theano and the needs of large-scale production ML.

The core insight: if you separate graph definition from graph execution, you can:

• Ship the graph (a serialized Protocol Buffer file) to mobile devices, servers, or TPUs
• Optimize the graph globally before running it (constant folding, CSE, operation fusion)
• Distribute it across machines without re-serializing Python code
• Cache the compiled form and reuse it for millions of inference requests

The penalty: you cannot inspect intermediate values during execution. The graph runs entirely in C++ without Python touching it.

# TF1 workflow
import tensorflow as tf

# PHASE 1: Build the graph (nothing executes)
x = tf.placeholder(tf.float32, shape=[None, 784], name='input')
W = tf.Variable(tf.random.normal([784, 256]), name='weights')
b = tf.Variable(tf.zeros([256]), name='bias')
h = tf.nn.relu(tf.matmul(x, W) + b, name='hidden')
loss = tf.reduce_mean(tf.square(h - target), name='loss')
optimizer = tf.train.AdamOptimizer(lr=0.001)
train_op = optimizer.minimize(loss)

# PHASE 2: Execute (everything runs in C++)
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for step in range(1000):
        batch_x, batch_y = get_batch()
        loss_val, _ = sess.run([loss, train_op], {x: batch_x, target: batch_y})

The graph is a Protocol Buffer — Google's binary serialization format. The graph can be saved with tf.saved_model.save() and loaded on any device with the TF runtime, including mobile phones (TFLite) and Google TPUs.

17.2 TF2's Eager Execution: Define-By-Run

Python programmers found TF1's two-phase model confusing and hard to debug. TF2 (released 2019) switched to eager execution by default:

# TF2 eager — operations execute immediately
import tensorflow as tf

W = tf.Variable(tf.random.normal([784, 256]))
b = tf.Variable(tf.zeros([256]))
x = tf.constant(batch_x)               # a real tensor, not a placeholder

h = tf.nn.relu(x @ W + b)              # runs NOW, returns a real tensor
print(h.shape)                          # (batch_size, 256) — immediate
print(h.numpy()[:2])                    # can inspect values directly

Like PyTorch, operations execute immediately and return concrete tensors. Debugging is trivial — add print() statements anywhere.

The cost: no static optimization, no graph export. Each Python call has interpreter overhead.

The bridge: @tf.function re-introduces graph compilation as an opt-in feature, giving you the best of both worlds.

Chapter 18: tf.function — AutoGraph, Tracing, and ConcreteFunctions

18.1 What Happens When Python Hits @tf.function

The first time you call a @tf.function-decorated function, TF runs a multi-stage pipeline:

Stage 1 — AutoGraph transformation: TF rewrites your Python source code. Python control flow structures are converted to TF graph ops:

This transformation happens at the Python source level — TF uses Python's inspect module to get the source code, then rewrites it using the gast library (Generic AST).

Stage 2 — Tracing: The AutoGraph-transformed function runs with SymbolicTensor objects (analogous to PyTorch's Proxy objects). Every TF op call records a node in a FuncGraph — TF's internal graph representation.

Stage 3 — Shape inference: TF propagates shapes through all nodes in the FuncGraph. Operations that don't know their output shape at trace time are marked with unknown dimensions.

Stage 4 — Compilation to ConcreteFunction: The FuncGraph is compiled into a ConcreteFunction — a C++ execution artifact bound to a specific input signature. XLA may further compile this to hardware-specific machine code.

Stage 5 — Caching: The ConcreteFunction is cached. Subsequent calls with matching signatures (same shapes and dtypes) route to the cached artifact — no Python execution, no tracing overhead.

18.2 The Complete List of Retrace Triggers

A retrace is expensive (re-runs AutoGraph + tracing + compilation). You must understand every trigger:

Trigger 1: New input shape (without input_signature)

@tf.function
def f(x): return x * 2

f(tf.ones([3]))           # trace 1: shape (3,)
f(tf.ones([4]))           # RETRACE: shape (4,) is new

Prevention: @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.float32)])

Trigger 2: New input dtype

f(tf.ones([3], dtype=tf.float32))  # trace 1
f(tf.ones([3], dtype=tf.float64))  # RETRACE: different dtype

Trigger 3: Change in Python non-tensor arguments

@tf.function
def f(x, n_heads):          # n_heads is Python int, not tf.Tensor
    return tf.split(x, n_heads, axis=-1)

f(x, n_heads=4)   # trace 1: n_heads=4 baked in as compile-time constant
f(x, n_heads=8)   # RETRACE: n_heads=8 requires new graph structure

AutoGraph needs Python values at trace time to build the graph (how many splits to create). Changing a Python value creates a different graph.

Trigger 4: Changed Python closure

threshold = 0.5
@tf.function
def f(x): return x > threshold  # captures Python variable 'threshold'

threshold = 0.9
f(x)  # RETRACE: closure value changed

Trigger 5: tf.Variable created inside the function (error, not retrace)

@tf.function
def bad():
    v = tf.Variable(0.0)  # FORBIDDEN after first trace
    return v
# Second call raises: ValueError: tf.function only supports tf.Variables
# created on the first call.

Prevention pattern: always create Variables outside @tf.function:

v = tf.Variable(0.0)        # created once, outside
@tf.function
def good(x):
    v.assign_add(x)          # use the existing Variable
    return v

18.3 AutoGraph's Transformation in Detail

Let's trace exactly what AutoGraph does to a loop:

# Original Python:
@tf.function
def sum_to_n(n):
    total = tf.constant(0)
    for i in tf.range(n):      # tf.range, not Python range!
        total = total + i
    return total

AutoGraph rewrites this to (approximately):

def sum_to_n(n):
    total = tf.constant(0)
    # Loop body as a separate function
    def loop_body(i, total):
        total = total + i
        return total
    # TF while_loop handles the actual iteration in C++
    total = tf.while_loop(
        cond=lambda i, _: i < n,
        body=loop_body,
        loop_vars=[tf.constant(0), total]
    )[1]
    return total

tf.while_loop is a TF graph op that encodes the loop entirely within the graph. The loop runs in C++ without returning to Python. This is how TF achieves Python-free execution of loops.

Note: if you use Python range(n) instead of tf.range(n), AutoGraph unrolls the loop at trace time — creating n nodes in the graph. For variable n, this fails. Always use tf.range for dynamic loops inside @tf.function.

Chapter 19: GradientTape, XLA, and MLIR

19.1 GradientTape — TF's Explicit Backward Recording

PyTorch records gradients implicitly (every operation on a requires_grad tensor is automatically recorded). TF2 uses explicit recording via tf.GradientTape:

x = tf.Variable([1.0, 2.0, 3.0])   # Variables are automatically watched

with tf.GradientTape() as tape:
    # Everything inside this block is recorded
    y = x ** 2                        # element-wise square
    loss = tf.reduce_sum(y)           # sum: loss = x[0]² + x[1]² + x[2]²

# Compute gradients AFTER the tape context closes
grads = tape.gradient(loss, x)
print(grads)   # [2.0, 4.0, 6.0]  — ∂loss/∂xᵢ = 2xᵢ

tf.Variable objects are automatically watched. For tf.Tensor constants, you must explicitly watch:

with tf.GradientTape() as tape:
    tape.watch(constant_tensor)
    loss = f(constant_tensor)
grad = tape.gradient(loss, constant_tensor)

By default, a tape can only be used once — the recorded operations are freed after gradient() is called. Use tape = tf.GradientTape(persistent=True) to call gradient() multiple times (like PyTorch's retain_graph=True).

19.2 XLA — Accelerated Linear Algebra

XLA is Google's compiler for linear algebra operations. When you annotate with @tf.function(jit_compile=True) or use tf.device('/GPU:0') with XLA enabled, the FuncGraph is lowered to HLO (High-Level Operations) and compiled by XLA.

HLO is XLA's IR — a set of ~100 primitive operations (dot product, element-wise, reduce, reshape, select, etc.). Every TF op eventually lowers to one or more HLO ops.

XLA performs aggressive optimizations on HLO:

• Operation fusion: merges adjacent element-wise ops into one kernel
• Buffer allocation: assigns memory to each HLO op's output, reusing buffers when possible
• Loop tiling: tiles computations to fit in cache, improving memory locality
• Layout optimization: chooses the best tensor layout (row-major vs column-major) for each op to minimize memory transpose costs

XLA then lowers HLO to machine code via LLVM for CPU or NVPTX for NVIDIA GPUs (or the TPU backend for Google's hardware). The generated code is highly optimized, often outperforming manually written CUDA kernels.

19.3 MLIR — Multi-Level Intermediate Representation

MLIR (Multi-Level Intermediate Representation) is Google's framework for building compilers, introduced in 2020. TF now uses MLIR as a bridge between TF's high-level graph and XLA's HLO.

Why multiple levels? Different optimization passes work at different abstraction levels:

• At the TF graph level: common subexpression elimination, shape inference, dead code elimination
• At the XLA HLO level: fusion, buffer assignment, layout optimization
• At the machine code level: register allocation, instruction scheduling

MLIR allows each optimization to operate at the right abstraction level using dialects — pluggable extensions that define new operations and types. TF defines:

• tf dialect: TF-level operations
• mhlo dialect: the HLO equivalent in MLIR
• linalg dialect: linear algebra operations with explicit loop structure
• memref dialect: explicit memory operations

The compilation pipeline is a sequence of lowering passes:

tf dialect → mhlo dialect → linalg dialect → affine dialect → LLVM IR → machine code

Each lowering step progressively descends toward machine code while applying level-appropriate optimizations. This multi-level approach is why TF on TPUs and GPUs often outperforms manually tuned libraries — the compiler has complete knowledge of the computation at every level.

Part VI — Numerical Precision

Chapter 20: IEEE 754 — How Computers Store Numbers

20.1 Why Computers Can't Represent All Real Numbers

Computers store everything in binary (base 2). A 32-bit integer can represent exactly 2³² values. But real numbers are infinite — there are infinitely many values between 0 and 1 alone.

Floating-point is the compromise: represent a wide range of values approximately, with the precision concentrated near 1.0 and thinning out at extreme values (very large or very small).

20.2 The IEEE 754 Standard — Bit by Bit

IEEE 754 defines how floating-point numbers are stored. Every float is represented as:

\[ \text{value} = (-1)^{\text{sign}} \times 1.\text{mantissa}_2 \times 2^{\text{exponent} - \text{bias}} \]

For float32 (32 bits total):

Bit 31      | Bits 30-23   | Bits 22-0
sign (1 bit)| exponent (8) | mantissa (23)

• Sign: 0 = positive, 1 = negative
• Exponent: stored with a bias of 127. Actual exponent = stored_value − 127. Range: 2^(0-127) = 2^-127 to 2^(255-127) = 2^128
• Mantissa: the fractional part after the leading 1 (which is implicit — not stored). 23 bits gives about 7 decimal digits of precision.

Example: how is −6.5 stored?

1. Sign: negative → bit 31 = 1
2. 6.5 in binary: 110.1 = 1.101 × 2^2
3. Exponent: 2 + 127 = 129 = 10000001 in binary
4. Mantissa: 101 followed by 20 zeros: 10100000000000000000000

Full bit pattern: 1 10000001 10100000000000000000000

20.3 Special Values — NaN, Infinity, Denormals

IEEE 754 reserves special exponent values for non-numeric cases:

Infinity: 1.0 / 0.0 = +Inf. Operations on Infinity propagate: Inf + 1 = Inf, Inf - Inf = NaN.

NaN: result of undefined operations: 0.0 / 0.0, sqrt(-1), Inf - Inf. NaN is "poisonous" — any arithmetic involving NaN produces NaN. This is why NaN in one layer propagates through the entire network.

Denormals (subnormal numbers): when the exponent is all zeros, the leading 1 is absent, allowing extremely small values near zero. Operations on denormals are 10-100× slower on some hardware because the FPU falls back to software emulation.

20.4 Precision Limits — When 1 + ε = 1

float32 has 23 mantissa bits. The smallest value that, when added to 1.0, gives a result different from 1.0 is called machine epsilon:

\[ \varepsilon_{\text{fp32}} = 2^{-23} \approx 1.19 \times 10^{-7} \]

This means: 1.0 + 1e-8 == 1.0 in float32. The 1e-8 is below the precision threshold and is simply lost.

Catastrophic cancellation: subtracting two nearly-equal large numbers loses significant digits:

x = 1.000001  (float32 representation: 1.0 + 1e-6)
y = 1.000000  (float32 representation: 1.0)
x - y = 1e-6  (correct)

But if x = 1000000.001 and y = 1000000.000:
float32 precision at scale 1e6 is ε × 1e6 = 1.19e-7 × 1e6 = 0.119
The difference 0.001 is below the representable precision!
x - y = 0.0  in float32  (wrong — should be 0.001)

This happens in:

• Softmax with large logits: sum(exp(large_values)) is huge; individual exp(x_i) may be small relative to the sum
• Layer normalization: subtracting the mean from large-valued activations
• Log probabilities: log(p) where p is very close to 0 or 1

20.5 The ULP — Measuring Floating-Point Error

"Units in the Last Place" (ULP) is the standard measure of floating-point error. 1 ULP is the distance between a float and the next representable float.

Near 1.0: 1 ULP ≈ 1.19e-7 (float32 epsilon) Near 1e8: 1 ULP ≈ 8.0 (much larger!)

When your gradient check requires tolerance 1e-5, you're saying "allow up to about 100 ULP error at values near 1.0". For values near 100, the absolute tolerance should be 100 × 1.19e-7 × 100 = 1.19e-3. This is why gradient checks use relative tolerance, not absolute.

Chapter 21: FP16, BF16, and Mixed Precision Training

21.1 Why Not Always Use float32?

float32 has excellent precision and range. But it uses 4 bytes per value. For a 7B parameter model:

• float32: 7e9 × 4 bytes = 28 GB just for weights
• float16: 7e9 × 2 bytes = 14 GB
• int8: 7e9 × 1 byte = 7 GB

Beyond storage: float32 arithmetic is 2× slower than float16 on modern GPUs with Tensor Cores. The A100's Tensor Core throughput is 312 TFLOPs (float16) vs 77 TFLOPs (float32) — 4× faster for matrix multiplications.

21.2 float16 — High Performance, Dangerous Range

float16 (half precision):

Bit 15    | Bits 14-10 | Bits 9-0
sign (1)  | exponent(5)| mantissa(10)

• Exponent: 5 bits, bias=15. Actual range: 2^-14 to 2^15 = 0.00006 to 65504
• Mantissa: 10 bits ≈ 3.3 decimal digits of precision
• Maximum value: 65504

The danger: during training, gradient norms, loss values, and large weight matrices can easily exceed 65504 → overflow → Infinity → NaN → corrupted training.

Additionally, very small gradients (common in early training or deep networks) can underflow to 0 → parameters stop updating → training stalls.

21.3 Loss Scaling — The Solution to FP16 Overflow and Underflow

The trick: multiply the loss by a large scalar (e.g., 2^15 = 32768) before backpropagation. This shifts all gradients upward by the same factor, preventing underflow. Before applying gradients to weights, divide by the same scalar.

scaler = torch.cuda.amp.GradScaler()

# Training loop
optimizer.zero_grad()
with torch.autocast(device_type='cuda', dtype=torch.float16):
    output = model(input)
    loss = criterion(output, target)

scaler.scale(loss).backward()         # multiplies loss by scaler factor (e.g., 32768)
scaler.step(optimizer)                # unscales gradients, applies them
scaler.update()                       # adjusts scale factor based on inf/nan detection

The GradScaler dynamically adjusts the scale factor: if gradients overflow (contain Inf/NaN), it halves the scale and skips this optimizer step; if gradients are clean for a while, it doubles the scale to catch any underflowing gradients.

21.4 BF16 — The Training-Safe Compromise

BF16 (Brain Float 16, developed by Google for TPUs):

Bit 15    | Bits 14-7  | Bits 6-0
sign (1)  | exponent(8)| mantissa(7)

• Exponent: 8 bits — same as float32! Range: 2^-126 to 2^127 ≈ 10^-38 to 10^38
• Mantissa: 7 bits ≈ 2.3 decimal digits of precision
• No overflow or underflow risk during training (same range as float32)
• No loss scaling needed

BF16 vs FP16 for training:

• BF16: safe range, lower precision. Training is stable.
• FP16: dangerous range, higher precision. Needs loss scaling.

Modern NVIDIA hardware (Ampere A100, Hopper H100) supports BF16 natively with the same Tensor Core throughput as FP16. BF16 is now the default training format for most LLM training runs.

21.5 AMP — Automatic Mixed Precision

torch.autocast handles the dtype selection automatically:

with torch.autocast(device_type='cuda', dtype=torch.float16):
    # These ops run in float16 (faster, sufficient precision):
    output = model.layers(input)     # matmul, conv2d
    
    # These ops are forced to float32 (numerically sensitive):
    loss = F.cross_entropy(logits, targets)   # softmax + log inside
    # layer norm, softmax, batch norm also run in float32

The autocast context uses a table of operations and their recommended dtypes. Matmul and conv can lose precision without impacting final model quality. Softmax, layer norm, and loss functions are sensitive to precision — they stay in float32.

Part VII — Synthesis and Practice

Chapter 22: How Everything Connects — The Principal Engineer's Map

                     MATH LAYER
     Chain Rule → Jacobians → VJPs → Reverse Mode AD
           ↓
     PYTHON LAYER  
     CPython bytecode → GIL → pybind11 → C++
           ↓
     PYTORCH C++ LAYER
     c10 types → ATen ops → Dispatcher → DispatchKeySet
           ↓              ↓
     Autograd Key    CUDA/CPU Keys
     (records DAG)   (compute kernels)
           ↓
     BACKWARD DAG
     Node→Edge→AccumulateGrad → leaf.grad
           ↓
     GRAPH CAPTURE
     FX IR (symbolic trace) → torch.compile (Dynamo bytecodes + guards)
           ↓
     COMPILATION
     AOTAutograd → joint forward+backward graph
     TorchInductor → Triton (GPU) / C++ (CPU)
           ↓
     CUSTOM EXTENSIONS
     TORCH_LIBRARY schema → TORCH_LIBRARY_IMPL per key
     (CPU, CUDA, Meta, PrivateUse1/NPU)

How each lab maps to production work:

Lab 01 (minigrad) → Production skill: Debugging gradient issues

• When a gradient is wrong: you know to check each Function.backward for the VJP formula
• When a gradient is None: you know to check requires_grad, no_grad() scope, or detach()
• When a model doesn't converge: you know to run gradcheck to find the broken backward
• When writing a custom backward for a fused attention kernel: you apply Lab 01 patterns directly

Lab 02 (warp_softmax) → Production skill: Writing NPU kernels

• The TORCH_LIBRARY registration pattern is identical for CPU, CUDA, and NPU backends
• The Meta kernel requirement is universal — every operator in the ecosystem needs it
• The warp reduction technique applies to any per-row reduction: max-pooling, layer norm, etc.
• opcheck is used to validate every operator in PyTorch's own test suite

Chapter 23: Lab Guidance — Building minigrad and warp_softmax

23.1 Lab 01: minigrad Implementation Order

Start from the simplest operations and build up. Each step gives you a working test before the next.

Step 1: Implement _build_topo on Tensor. This is the topological sort. Without it, backward() cannot run. Start here:

def _build_topo(self):
    topo, visited = [], set()
    def dfs(t):
        if id(t) in visited: return
        visited.add(id(t))
        if t.grad_fn:
            for inp in t.grad_fn._input_tensors:
                if isinstance(inp, Tensor): dfs(inp)
        topo.append(t)
    dfs(self)
    return topo

Step 2: Implement backward() on Tensor. Initialize gradient, call _build_topo, traverse in reverse, call each grad_fn's backward, accumulate into leaves. Verify with the smoke test at bottom of lab.py.

Step 3: Implement Add (forward + backward). Backward must handle broadcasting via _unbroadcast. Test: a + b where a.shape=(3,), b.shape=(4,3).

Step 4: Implement Mul, MatMul, Neg, Div, Pow. These teach: saving inputs for backward, the VJP formulas for each op.

Step 5: Implement ReLU, Sigmoid, Tanh, Exp, Log. Key insight: Sigmoid.forward should save the output (not input) because d(sigmoid)/dx = sigmoid(x)(1 - sigmoid(x)) is expressed in terms of output.

Step 6: Implement softmax using Max, Exp, Sum, Div. Apply the max-subtraction trick for numerical stability. All ops you implement in steps 3-5 compose naturally.

Step 7: Implement cross_entropy_loss using log_softmax. The full MLP training loop should then converge.

Numerical gradient check — what it is: approximate the derivative numerically using finite differences, compare with your analytic gradient:

numerical_grad[i] ≈ (f(x + ε·eᵢ) - f(x - ε·eᵢ)) / 2ε

If abs(analytic - numerical) / max(abs(analytic), abs(numerical)) < 1e-5, your backward is correct.

23.2 Lab 02: warp_softmax — Reading Before Building

Read the source files in this order:

1. csrc/warp_softmax.cpp: understand TORCH_LIBRARY (schema), TORCH_LIBRARY_IMPL (CPU + Meta), warp_softmax_cpu (ATen ops for CPU), warp_softmax_meta (returns empty_like).

2. csrc/warp_softmax.cu: read warp_reduce_max and warp_reduce_sum first — these are the butterfly reduction utilities. Then read warp_softmax_kernel — trace through the three steps (find max, compute exp+sum, normalize). Map each line to the hardware concepts from Chapter 15.

3. setup.py: understand how torch.utils.cpp_extension.CUDAExtension compiles your C++/CUDA and links against libtorch.

4. solution.py: understand how opcheck is called and what it verifies.

23.3 Lab 03: tf.function Tracing — Implementation Order

Chapter 18 is the theory; the lab makes it mechanical. Order matters here too:

1. trace_and_inspect first — until you can see a FuncGraph's ops, captures, and signature, the rest is blind. Key API: fn.get_concrete_function(*inputs).graph.
2. RetraceProbe — the cleverness is that a Python counter inside the traced body is the retrace counter, because Python side effects fire only at trace time (Chapter 18's central fact, now weaponized as instrumentation).
3. demonstrate_retracing — predict the three trace counts on paper before running; if your prediction misses, re-read §18's retracing rules.
4. stabilize — tf.TensorSpec([None, ...]) makes shape symbolic; verify one trace serves all batch sizes.
5. autograph_diff, then the SavedModel/TFLite round-trips — SavedModel must be bit-exact (same graph, same weights); TFLite dynamic-range must be close-but-not-exact (weights were perturbed to INT8). Both assertions are in the tests; understand why each bound holds.

23.4 Success Criteria

You are ready for Phase 02 when you can answer all of these from memory:

1. Trace a 5-op PyTorch expression's full backward pass with specific gradient values at each node.
2. Name every dispatch key in priority order and explain what each intercepts.
3. Give 3 distinct triggers for tf.function retracing and the prevention for each.
4. Explain what breaks in torch.compile if your custom op has no Meta kernel.
5. Draw the butterfly reduction for warp_reduce_max with 8 threads (3 rounds).
6. Explain why BF16 doesn't need loss scaling but FP16 does.
7. Write the VJP (backward formula) for matmul from the Jacobian definition.

If any answer is fuzzy, return to that chapter and re-read the internal mechanism — not the API, the mechanism.

Chapter 24: The Interview Gauntlet — Complete Answers

These are actual questions asked at senior/staff ML framework engineering interviews. The expected answer length and depth is shown.

Q: Walk me through what happens, step by step, when I call loss.backward().

A complete answer (2-3 minutes verbal):

loss.backward() calls into torch::autograd::Engine::execute(). The engine starts from loss.grad_fn with an initial gradient of torch.ones_like(loss). It builds a ready-count for each Function node — how many downstream nodes still need to send gradients to it. Processing order is determined by decrementing these counts: a node runs when its count reaches zero (all downstream contributions received).

For each node, the engine calls node.apply(input_gradients) — the virtual backward function. For built-in ops this is a generated class like MmBackward0; for custom ops it's the static backward method you defined. The function returns a tuple of gradients — one per input. Each gradient is sent to the corresponding upstream node (via next_edges_). Leaf nodes have AccumulateGrad as their node, which adds the incoming gradient to leaf.grad (or initializes it).

After backward completes, the Function nodes and their saved tensors are freed (unless retain_graph=True). This is why you can't call backward() twice on the same graph — the intermediate data is gone.

Q: What is a dispatch key and why is its priority ordering load-bearing?

The DispatchKeySet is a 64-bit bitfield attached to every tensor. Each bit position corresponds to one DispatchKey. The dispatcher takes the bitwise OR of all input tensors' key sets and calls the highest-priority key that has a registered implementation for the operation.

The priority ordering is load-bearing because the keys are semantic wrappers. Autograd sits above CUDA so that gradient recording always wraps the compute kernel — you can't record gradients without the Autograd key seeing the operation first. Functionalize sits above Autograd so that aliasing analysis (needed for torch.compile) can see the original mutation intent before autograd creates any graph nodes.

If an NPU vendor registered their key above Autograd, their device would silently bypass gradient recording. Training would appear to work (forward pass correct) but gradients would never be computed.

Q: When does tf.function retrace, and what is the performance cost?

Retrace triggers: (1) new input shapes when input_signature is not set, (2) new input dtypes, (3) change in non-tensor Python argument values (they are compile-time constants baked into the graph), (4) changed Python closures captured by the function.

The cost is: re-running AutoGraph (Python source transformation), re-tracing the function with SymbolicTensors (building a new FuncGraph), optionally re-running XLA compilation (extremely expensive — can take minutes for large models). In production, unexpected retracing is a severe latency spike. You diagnose with tf.autograph.set_verbosity(10) to print retrace events.

Prevention: always use input_signature with None for variable dimensions; convert Python-controlled iteration counts to tf.Tensor arguments; avoid Python closures that change.

Q: A model works in eager mode but gives NaN in torch.compile. How do you debug it?

Systematic approach:

Step 1: Isolate forward vs backward. Add torch.autograd.set_detect_anomaly(True) before the training loop — it adds expensive runtime checks but identifies exactly which backward operation produces NaN.

Step 2: Bisect the graph. Use torch.compile(fullgraph=False) — this allows graph breaks. Introduce a deliberate graph break at the suspected location (e.g., torch._dynamo.graph_break()) and check if NaN appears before or after the break.

Step 3: Check custom ops. If you have a custom operator, verify it has a Meta kernel and run opcheck. An incorrect Meta kernel can cause shape mismatches that produce NaN in compiled mode.

Step 4: Check numerical precision. Compiled mode may fuse operations that change floating-point evaluation order. Add torch.use_deterministic_algorithms(True) to force bit-identical behavior across modes.

Step 5: Check in-place ops. The Functionalize dispatch key converts in-place ops to functional form. If your custom backward incorrectly modifies saved tensors in-place, Functionalize may not handle it correctly, producing NaN.

Q: What is the Meta device, and why does every custom operator need a Meta kernel?

The Meta device is a special device type where tensors have shape, dtype, device=meta, but zero bytes of storage. Operations on Meta tensors execute "Meta kernels" that propagate shapes and dtypes without doing any computation.

torch.compile's TorchDynamo captures the computation as an FX graph. AOTAutograd traces through this graph using FakeTensors — which are Meta tensors augmented with symbolic size information. Every operation in the graph must have a Meta kernel so the compiler can determine output shapes without running compute kernels.

A custom operator without a Meta kernel causes torch.compile to raise NotImplementedError: Could not run 'myops::op' with the 'Meta' dispatch key. The fix is always:

TORCH_LIBRARY_IMPL(myops, Meta, m) {
    m.impl("my_op", [](const at::Tensor& input, int64_t dim) {
        return at::empty_like(input);  // same shape, same dtype, no data
    });
}

For ops that change shape (e.g., a custom reshape), the Meta kernel must compute the correct output shape from the input shape and arguments.

Q: Explain catastrophic cancellation with a concrete example, and how you'd prevent it in a softmax implementation.

Catastrophic cancellation: subtracting two nearly-equal large numbers produces a result with far fewer significant digits than either operand.

Concrete example: computing softmax([1000.0, 999.0]) in float32.

exp(1000.0) → inf (overflows float32 max ≈ 3.4e38)
exp(999.0) → inf

softmax = inf / (inf + inf) = inf / inf = NaN

Even if the values don't overflow: computing exp(100.0) / (exp(100.0) + exp(99.0)) in float16:

• exp(100) ≈ 2.69e43 → overflows float16 (max 65504) → Inf

Prevention — the max-subtraction trick:

softmax(x_i) = exp(x_i) / Σ exp(x_j)
             = exp(x_i - M) / Σ exp(x_j - M)    where M = max(x)

Subtracting M ensures all exponents are ≤ 0, so exp(x_i - M) ∈ (0, 1] — no overflow, no underflow, no cancellation. The maximum value is always exp(0) = 1.

For the normalization sum itself: compute in float32 even when the inputs are float16, to maintain precision. PyTorch's F.softmax does this automatically.

References

Foundational Papers

1. Reverse-mode AD: Baydin, A.G., Pearlmutter, B.A., Radul, A.A., Siskind, J.M. (2018). Automatic Differentiation in Machine Learning: a Survey. JMLR 18(153). — The definitive survey of automatic differentiation; explains forward/reverse modes rigorously.

2. PyTorch: Paszke, A. et al. (2019). PyTorch: An Imperative Style, High-Performance Deep Learning Library. NeurIPS 2019. — The paper describing PyTorch's design philosophy and dispatcher.

3. TorchDynamo: Ansel, J. et al. (2023). PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. PLDI 2023. — The paper behind torch.compile; explains Dynamo's bytecode approach and guard system.

4. Triton: Tillet, P., Kung, H.T., Cox, D. (2019). Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations. MAPL 2019. — The paper behind TorchInductor's code generation backend.

5. FlashAttention: Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. — Shows how GPU memory hierarchy analysis leads to 2-4× speedups; the same analysis applies to the warp_softmax kernel.

6. JAX: Bradbury, J. et al. (2018). JAX: Composable Transformations of Python+NumPy Programs. — JAX's jit, grad, vmap show an alternative design to PyTorch's; understanding both gives perspective on framework design tradeoffs.

7. MLIR: Lattner, C. et al. (2021). MLIR: Scaling Compiler Infrastructure for Domain Specific Computation. CGO 2021. — The paper behind TF's MLIR compilation infrastructure.

8. XLA: Leary, C. et al. (2017). XLA: TensorFlow, Compiled. TF Dev Summit 2017. — XLA's design and its HLO IR.

Books

9. Deep Learning: Goodfellow, I., Bengio, Y., Courville, A. (2016). Deep Learning. MIT Press. — Chapters 6 (feedforward networks) and 6.5 (backpropagation) are the mathematical foundation for everything in Part I.

10. Programming Massively Parallel Processors: Kirk, D.B., Hwu, W.W. (4th ed., 2022). — The definitive textbook on CUDA programming; Chapters 4 (memory), 10 (reduction), and 15 (performance optimization) are directly relevant.

11. Matrix Computations: Golub, G.H., Van Loan, C.F. (4th ed., 2013). — The mathematical foundation for understanding why matrix operations (matmul, SVD, QR) are designed as they are.

12. What Every Computer Scientist Should Know About Floating-Point Arithmetic: Goldberg, D. (1991). ACM Computing Surveys. — The definitive reference for IEEE 754, rounding modes, and catastrophic cancellation.

Technical Documentation and Blog Posts

13. PyTorch Internals — Edward Yang's blog: http://blog.ezyang.com/2019/05/pytorch-internals/ — The most accessible deep-dive into PyTorch's C++ architecture written by one of its core developers.

14. TorchDynamo Deep Dive — PyTorch Dev Discuss: https://dev-discuss.pytorch.org/t/torchdynamo-an-experiment-in-dynamic-python-bytecode-transformation/361 — The original announcement and explanation of Dynamo's bytecode approach.

15. AOTAutograd Walkthrough — https://pytorch.org/functorch/stable/notebooks/aot_autograd_optimizations.html — Official documentation with worked examples.

16. CUDA Programming Guide — NVIDIA developer docs: https://docs.nvidia.com/cuda/cuda-c-programming-guide/ — The authoritative reference for warp execution, memory hierarchy, and shuffle instructions.

17. torch.fx Documentation — https://pytorch.org/docs/stable/fx.html — Official API docs with examples of graph transformation patterns.

18. tf.function Guide — https://www.tensorflow.org/guide/function — TF's official guide covering tracing, retracing, and AutoGraph in depth.

19. Mixed Precision Training — Micikevicius, P. et al. (2018). Mixed Precision Training. ICLR 2018. — The paper that introduced loss scaling for FP16 training.

20. The TORCH_LIBRARY Tutorial — https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html — PyTorch's official guide to custom operator registration, with C++ examples.

Hitchhiker's Guide — PyTorch & TensorFlow Framework Internals

"You don't understand a system until you've broken it." — every production on-call engineer

This guide goes beneath the public API. It is the answer to "why does this work?" when the official docs say "just call .backward()."

1. PyTorch Execution: The Full Stack

1.1 From Python to Kernel

When you write y = torch.matmul(a, b) in Python, here is the exact call chain:

Python: torch.matmul(a, b)
   → C++: at::matmul(self, other)                          # aten namespace
   → c10::Dispatcher::call(op_handle, stack)               # dispatch table lookup
   → Dispatch key selection: max(a.key_set(), b.key_set()) # e.g., CUDA
   → at::native::mm_cuda(self, other)                      # CUDA kernel dispatch
   → cublasSgemm / cublasGemmEx                            # vendor BLAS

The dispatcher (c10::Dispatcher) is PyTorch's universal function registry. Every operator is registered as an entry in the dispatcher table, keyed by a DispatchKey. The dispatcher picks the highest-priority key in the tensor's DispatchKeySet.

1.2 Dispatch Key Priority (highest to lowest)

Functionalize          ← wraps mutations as functional ops (for torch.compile)
InferenceMode          ← disables autograd
AutogradMeta           ← for meta/fake tensors in torch.compile
Autograd               ← wraps kernel in grad recording logic
AutocastCPU/CUDA       ← AMP dtype promotion
CUDA                   ← GPU kernels
CPU                    ← CPU kernels
Meta                   ← shape-only "meta" device (used in compilation)

Why this matters for NPU deployment: When you register a custom NPU dispatch key (e.g., Qualcomm's QnnBackend), it must be inserted at the right priority. Below Autograd means the kernel won't record gradients — which is correct for inference-only NPU deployments.

1.3 Autograd Engine Deep Dive

Forward pass — DAG construction:

a = torch.tensor([2.0], requires_grad=True)
b = torch.tensor([3.0], requires_grad=True)
c = a * b   # grad_fn = MulBackward0
d = c + a   # grad_fn = AddBackward0

After these four lines, the autograd engine has built:

d  (AddBackward0)
├── c  (MulBackward0)
│   ├── a  (leaf)
│   └── b  (leaf)
└── a  (leaf)

Each non-leaf tensor has a grad_fn that stores:

1. The saved tensors needed for the backward pass
2. References to the next_functions (upstream nodes in the DAG)

Backward pass — Engine::execute():

The backward pass is triggered by d.backward(). The engine:

1. Starts from d.grad_fn with initial gradient torch.ones(1)
2. Calls AccumulateGrad for leaf tensors
3. Calls each grad_fn.__call__(gradient_from_downstream) in reverse topological order
4. Accumulates gradients via add_ into each leaf's .grad

Why topological order? A node can only compute its upstream gradient after it receives all gradients from downstream nodes that depend on it. This is Dijkstra's algorithm in reverse — we process nodes when their contribution count reaches zero.

Critical invariant: grad_fn is attached during the forward pass. The backward graph is fully determined by what ops ran in forward. This is why torch.no_grad() is a massive speedup — no Function objects are allocated.

1.4 torch.autograd.Function — The Correct Mental Model

class ClampedSoftmax(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, min_val, max_val):
        output = torch.softmax(input, dim=-1).clamp(min_val, max_val)
        ctx.save_for_backward(output)           # saves tensor refs (not copies)
        ctx.min_val = min_val                   # non-tensor metadata OK
        return output

    @staticmethod
    def backward(ctx, grad_output):
        output, = ctx.saved_tensors
        # Jacobian of softmax: J_ij = o_i * (delta_ij - o_j)
        # With clamping: gradient is zero where output was clamped
        mask = (output > ctx.min_val) & (output < ctx.max_val)
        # Softmax backward: grad_input = output * (grad_output - (grad_output * output).sum(-1, keepdim=True))
        s = (grad_output * output).sum(-1, keepdim=True)
        grad_input = output * (grad_output - s)
        grad_input = grad_input * mask.float()
        return grad_input, None, None  # None for non-tensor args

Common pitfalls:

1. ctx.save_for_backward(x) → storing x directly in ctx (leaks memory, misses memory-efficient inplace handling)
2. Returning wrong number of gradients — must match number of forward inputs exactly (use None for non-differentiable)
3. create_graph=True — if your backward uses non-differentiable ops, higher-order gradients will fail silently with wrong values

2. torch.fx — Symbolic Tracing Mechanics

2.1 What FX Actually Does

torch.fx.symbolic_trace(model) runs the model's forward() with Proxy objects instead of real tensors. Every operation on a Proxy records a Node in a Graph:

import torch.fx as fx

class MyModel(torch.nn.Module):
    def forward(self, x):
        return torch.relu(x + 1)

model = MyModel()
graph_module = fx.symbolic_trace(model)
print(graph_module.graph)
# graph():
#   %x : [#users=1] = placeholder[target=x]
#   %add : [#users=1] = call_function[target=operator.add](args = (%x, 1))
#   %relu : [#users=1] = call_function[target=torch.relu](args = (%add,))
#   return relu

The result is an FX IR — a data structure you can inspect, modify, and re-compile.

2.2 FX Graph Transformation — The Correct Pattern

class FuseConvBN(fx.Transformer):
    def call_module(self, target, args, kwargs):
        # Check if this is a Conv2d followed by BatchNorm2d
        module = self.fetch_attr(target)
        if isinstance(module, torch.nn.Conv2d):
            # Look ahead in graph — check next user
            node = self.current_node
            if len(node.users) == 1:
                next_node = list(node.users)[0]
                if (next_node.op == 'call_module' and
                    isinstance(self.fetch_attr(next_node.target), torch.nn.BatchNorm2d)):
                    # Fuse: fold BN params into Conv weights
                    fused = torch.nn.utils.fusion.fuse_conv_bn_eval(
                        module, self.fetch_attr(next_node.target)
                    )
                    # Register fused module, return its output
                    self.root.add_module(target + '_fused', fused)
                    return self.tracer.call_module(target + '_fused', args, kwargs)
        return super().call_module(target, args, kwargs)

2.3 FX Limitations You MUST Know for Interviews

1. No data-dependent control flow: if x.sum() > 0: cannot be traced — the Proxy doesn't know the value
2. No dynamic shapes by default: shapes are captured at trace time and become constants
3. Closures and globals: Python closures are NOT captured in the graph — they become constants at trace time, which is a semantic footgun
4. torch.compile vs FX: torch.compile uses TorchDynamo (bytecode analysis) instead of symbolic tracing — it CAN handle control flow by creating multiple graphs with guards

3. torch.compile Stack

3.1 Three-Layer Architecture

User Python code
     ↓
TorchDynamo           ← bytecode analysis, guards, graph breaks
     ↓  (FX Graph)
AOTAutograd           ← joint forward+backward graph
     ↓  (FX Graph)
TorchInductor         ← code generation (Triton for GPU, C++ for CPU)
     ↓
Compiled artifact     ← .so or PTX or Metal shader

3.2 TorchDynamo — How Bytecode Analysis Works

TorchDynamo installs a frame evaluation callback that intercepts Python bytecode execution. When it encounters a function decorated with torch.compile, it:

1. Disassembles the bytecode with dis.get_instructions()
2. Tracks every operation that produces a tensor
3. When it hits an unsupported operation (Python control flow on tensor values), it breaks the graph and falls back to Python
4. For captured graph segments, it installs guards — runtime checks that verify the assumptions made during tracing are still valid

@torch.compile
def f(x):
    if x.shape[0] > 10:  # <- GUARD: shape[0] > 10
        return x * 2
    return x + 1

First call with x.shape = (15,) → compiles graph 1 (x * 2) with guard shape[0] > 10
Second call with x.shape = (5,) → guard fails → compiles graph 2 (x + 1) with guard shape[0] <= 10
Third call with x.shape = (20,) → guard 1 passes → reuses compiled graph 1

3.3 Custom Backend Registration

from torch._dynamo import register_backend

@register_backend
def my_npu_backend(gm: torch.fx.GraphModule, example_inputs):
    # gm is the FX graph module
    # example_inputs is a list of example tensors for shape/dtype inference
    
    # Step 1: Optimize FX graph
    gm = optimize_for_npu(gm)
    
    # Step 2: Lower to NPU IR
    npu_program = lower_to_npu_ir(gm, example_inputs)
    
    # Step 3: Return a callable
    def run(*args):
        return npu_program.execute(args)
    
    return run

model_compiled = torch.compile(model, backend="my_npu_backend")

4. TensorFlow Graph Execution Mechanics

4.1 tf.function Tracing Protocol

When Python hits a @tf.function-decorated function for the first time:

1. Tracing: Python executes with SymbolicTensor objects (like PyTorch Proxies)
2. AutoGraph transformation: Python if/for/while are converted to tf.cond/tf.while_loop
3. Graph construction: a FuncGraph is built representing the TF ops
4. Compilation: the graph is compiled to a ConcreteFunction with a specific input signature
5. Execution: subsequent calls with matching signatures execute the cached ConcreteFunction

4.2 Retracing Triggers — The Complete List

@tf.function
def f(x, training=False):
    return x * 2 if training else x + 1

f(tf.ones([3]))          # trace 1: input_signature=(TensorSpec(shape=[3], dtype=float32),), training=False
f(tf.ones([3]))          # reuse trace 1 ✓
f(tf.ones([4]))          # trace 2: new shape [4] ← RETRACE (no input_signature specified)
f(tf.ones([3]), training=True)  # trace 3: new Python value for training ← RETRACE
f(tf.constant(1.0))     # trace 4: scalar shape ← RETRACE

Prevention: use @tf.function(input_signature=[tf.TensorSpec(shape=[None, 128], dtype=tf.float32)]) to fix dtype/rank and allow variable batch/sequence dimensions.

4.3 tf.Variable Inside tf.function

# WRONG — creates a new Variable on every trace:
@tf.function
def bad(x):
    v = tf.Variable(0.0)  # ← Python-side effect, runs at trace time only!
    v.assign_add(x.sum())
    return v

# CORRECT — Variable must be created outside tf.function:
v = tf.Variable(0.0)
@tf.function
def good(x):
    v.assign_add(x.sum())
    return v

5. Numerical Precision Deep Dive

5.1 FP32 vs BF16 vs FP16

Key insight for NPU work: BF16 has the same dynamic range as FP32 (same exponent bits) but lower precision. This makes BF16 safer for training (no overflow/underflow) but potentially less accurate for narrow distribution outputs. Qualcomm's HTP natively supports INT8/INT16/FP16/BF16 — knowing which precision per-layer gives the best accuracy-performance tradeoff is core to this role.

5.2 Catastrophic Cancellation and Kahan Summation

Summing [1e8, 1.0, -1e8] in FP32:

• Naive: (1e8 + 1.0) = 1e8 (1.0 is below precision of 1e8 in FP32), then 1e8 - 1e8 = 0.0 ← WRONG
• Kahan: maintains a running compensation term, gets 1.0 ← CORRECT

def kahan_sum(values):
    total = 0.0
    compensation = 0.0
    for v in values:
        y = v - compensation
        t = total + y
        compensation = (t - total) - y
        total = t
    return total

Interview answer: "Why does your softmax accumulate error in FP16?" → explain that the normalization sum in softmax has large magnitude (sum of exps), and the numerator single-exp is small — catastrophic cancellation. Fix: compute in FP32, cast output to FP16.

6. Common Interview Pitfalls

7. Resources & Further Reading

• PyTorch Internals Blog — Edward Yang
• TorchDynamo — How it Works (PyTorch Dev Blog)
• Autograd: Automatic Differentiation in ML (JMLR 2017) — the paper behind modern autodiff engines
• TF tf.function Guide
• PyTorch FX API Docs
• AOTAutograd internals
• Paper to read: "JAX: Composable Transformations of Python+NumPy Programs" (Bradbury et al., 2018) — JAX solves the same problems as torch.compile but differently; understanding both is a superpower

👨🏻 Brother Talk — Phase 01, Off the Record

Framework internals — the phase that decides whether you're a framework user or a framework engineer. At Qualcomm Senior Staff level, only the second one gets hired.

Brother, let me tell you the single thing that separates the people who pass this interview loop from the people who don't, because it took me too long to see it: almost everyone can use PyTorch, and almost nobody can explain it. They can write loss.backward() a thousand times and have no idea what happens after they hit enter. And the Qualcomm Senior Staff role is specifically about what happens after you hit enter — the dispatcher, the autograd graph, the kernel that actually runs. You are being hired to work below the API that everyone else lives on top of. So this phase isn't "review PyTorch." It's "stop being a tourist in the framework you've used for years."

Here's the mental reframe that unlocks it. PyTorch is not one thing — it's three layers pretending to be one: the Python eager frontend (what you type), the C++ dispatcher (the c10::Dispatcher that routes every op by dispatch key — CPU, CUDA, Autograd, Functionalize, Meta), and the backend kernels (the code that does the math). Every bug you'll ever debug at this job lives at a boundary between these layers: a custom op that works eager but breaks under torch.compile (boundary: dispatcher ↔ Dynamo), a quantized model that NaNs after a graph pass (boundary: Autograd ↔ Functionalize), a TF model that gives different answers in graph mode (boundary: eager ↔ tf.function tracing). If you can name the layer a bug lives in, you've already done 80% of the debugging. Build the autograd engine by hand in Lab 01 and this stops being abstract — you'll feel the graph being constructed on the forward pass and walked on the backward.

The thing nobody tells you about the autograd engine: it's simpler than you fear and more beautiful than you expect. It's just a DAG built during forward (each op records how to compute its gradient), walked in reverse-topological order during backward, accumulating gradients at the leaves. That's it. Once you've written Engine::execute() in miniature — even in pure Python — you'll never again be intimidated by "how does PyTorch know the gradient?" And you'll be the person in the room who can explain retain_graph, create_graph for double-backward, and why AccumulateGrad exists, while everyone else mumbles. That clarity is worth the week it takes.

Now the custom-operator stuff (Labs 02), because this is where the Qualcomm job actually lives. When you write a custom op — say, an NPU-accelerated kernel — it has to behave correctly under every dispatch key: CPU, CUDA, autograd (does it have a gradient?), functionalization (is it functional or does it mutate?), torch.compile (can Dynamo trace through it?). The classic disaster: someone registers a custom op that works in eager, ships it, and three months later it silently produces garbage under torch.compile because they never registered a Meta kernel or a functionalization rule. That is the bug you're being hired to prevent. The C++ operator lab feels like plumbing; it's actually the core skill — writing ops that are safe under all the ways PyTorch can call them. (And yes, the lab may skip on a machine without a C++ toolchain — that's expected; read the code, understand the TORCH_LIBRARY schema, build it for real when you have a compiler.)

On TensorFlow: don't skip it because PyTorch won. Qualcomm ships both, and the tf.function tracing model is a gold mine of interview questions precisely because it's subtle. The trap everyone falls into: Python side-effects inside a tf.function run once (at trace time), not every call, and retracing triggers on new input signatures. The "model gives different results in graph mode" bug is almost always a Python int where a tf.Variable belonged, or a closure capturing a stale value. Understand tracing-vs-execution once and you'll debug a class of bugs that mystifies whole teams.

A word on numerical precision, because it's the quiet thread through this entire role: FP32 vs FP16 vs BF16 is not just "smaller is faster." It's where you accumulate, what cancels, and when you get NaN. BF16 has FP32's exponent range with fewer mantissa bits — great for training stability, lossy for precision. The day you debug a model that's fine in FP32 and NaNs in FP16, you'll be grateful you understand catastrophic cancellation and why accumulation dtype matters. This thread runs straight into quantization (P03) and accuracy (P09) — it's the same fight at lower and lower precision.

Career truth, brother: this phase is the foundation that makes every later phase possible, and it's also the most legible signal of seniority in an interview. When you can trace loss.backward() to the kernel, explain dispatch keys, and write a custom op that's safe under compile — the interviewer knows you've operated below the API. That credibility carries the whole loop. Spend the time here. Build the autograd engine by hand. It's the difference between "I know PyTorch" and "I know how PyTorch works," and only the second sentence gets the Senior Staff offer.

Build the labs. Trace one op end to end. Then come to P02, where we go up a level to the architectures — and you'll understand them more deeply because you know what they compile down to.

— your brother 👨🏻

Lab 01 — Custom Autograd Engine

Phase: 01 — PyTorch & TF Framework Internals | Difficulty: ⭐⭐⭐☆☆ | Time: 2–3 hours

Read ../HITCHHIKERS-GUIDE.md before starting. This lab builds intuition for how PyTorch computes gradients under the hood.

What you build

A minimal scalar-valued autograd engine that mirrors PyTorch's torch.Tensor autograd mechanism:

• Value class wrapping scalars with +, *, tanh, relu, exp, pow operations
• Reverse-mode automatic differentiation via .backward()
• Dynamic computation graph construction (define-by-run)
• Gradient accumulation across shared nodes (topological sort)

Key concepts

Files

Run

pip install -r requirements.txt
python lab.py      # should print results after you fill in TODOs
pytest test_lab.py -v

Lab 02 — Custom C++ Operator

Phase: 01 — PyTorch & TF Framework Internals | Difficulty: ⭐⭐⭐⭐☆ | Time: 2–4 hours

Read ../HITCHHIKERS-GUIDE.md before starting.

What you build

A custom PyTorch operator written in C++ that registers with the dispatcher:

• C++ extension using pybind11 and torch/extension.h
• Forward pass: fused linear + ReLU in a single C++ kernel
• Registration via TORCH_LIBRARY / PYBIND11_MODULE
• Python test verifying numeric equality with the PyTorch reference

Key concepts

Files

Run

pip install -r requirements.txt
python setup.py build_ext --inplace
python solution.py

Lab 03 — TF tf.function Tracing & Graph Analysis

Phase: 01 — PyTorch & TF Framework Internals | Difficulty: ⭐⭐⭐⭐☆ | Time: 3–4 hours

tf.function is the boundary between eager Python and TensorFlow's graph runtime. Every "the model behaves differently in graph mode" bug lives at this boundary. This lab makes the tracing machinery visible.

What you build

• trace_and_inspect — capture the ConcreteFunction for given inputs and produce a GraphReport: node counts by op type, captured external tensors, structured input signature
• RetraceProbe — a wrapper that counts how many times Python tracing actually runs, used to demonstrate (and then prevent) the classic retracing triggers
• stabilize — apply an explicit input_signature so shape/value changes reuse one trace
• autograph_diff — show the AutoGraph-transformed source of a Python function with if/for control flow (what your Python actually becomes)
• roundtrip_savedmodel — save → reload → verify outputs match exactly
• convert_tflite_dynamic — TFLite dynamic-range quantization with max-abs output delta vs the TF model on random probe inputs

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v          # validate your lab.py
python solution.py             # run the full demo: retrace report + TFLite delta

Suggested TODO order

1. trace_and_inspect — everything else builds on being able to see a FuncGraph
2. RetraceProbe — prove to yourself exactly which calls retrace and why
3. stabilize — make the retraces stop with an input_signature
4. autograph_diff — look at what AutoGraph did to your if statement
5. roundtrip_savedmodel, then convert_tflite_dynamic

Success criteria

• You can enumerate the retracing triggers from memory and predict the retrace count of a call sequence before running it
• You can explain why a print() inside tf.function fires once but tf.print fires every call
• Your TFLite dynamic-range conversion shows max output delta < 1e-2 on the small MLP (weights-only quantization on a well-conditioned model is nearly lossless)

Interview Q&A

Q: When does tf.function retrace? New input dtype; new rank or (without a signature) new concrete shape; a new Python-native argument value (every distinct Python int/str is a new signature!); a new closure/object identity for weakly-referenced captures; calling with kwargs vs args inconsistently.

Q: Why did my counter only increment once even though I call the function in a loop? The increment is a Python side effect — it runs at trace time. The traced graph contains only TF ops. Use tf.Variable.assign_add to make mutation part of the graph.

Q: Why is passing a Python int as batch_size a performance bug? Each distinct int is a different trace signature → unbounded retracing. Pass it as a tf.Tensor or pin an input_signature with tf.TensorSpec([None, ...]).

Q: When is TFLite conversion not lossless? Dynamic-range quantization perturbs weights (INT8); ops without TFLite kernels get rewritten (e.g., erf-based GELU → tanh approximation); fused kernels may use different accumulation order; resource variables get frozen to constants.

Extensions

• Repeat the analysis on BERT-tiny and evaluate the MNLI accuracy delta after dynamic-range quantization (the full-scale version of this lab, see phase README)
• Inspect the XLA HLO with tf.function(jit_compile=True) and fn.experimental_get_compiler_ir(...)('hlo')
• Profile trace time vs execution time with the TF Profiler

References

• tf.function guide — the retracing rules, authoritative
• Introduction to graphs and tf.function
• AutoGraph reference
• TFLite post-training quantization

Phase 02 — Model Architecture Mastery Across Modalities

Difficulty: ⭐⭐⭐⭐☆
Estimated Time: 3 weeks (60–80 hours)
Roles Supported: All Senior Staff ML roles — architecture expertise underpins every optimization decision

Why This Phase Exists

You cannot optimize what you don't understand. Before a Senior Staff engineer at Qualcomm can ask "why does this model's accuracy degrade more under INT4 quantization than that one?", they need to know: GQA uses fewer KV heads → fewer unique weight rows → friendlier to per-group quantization. MoE models have sparse activation → different quantization strategies per expert. SSMs (Mamba) have recurrent structure → completely different memory access patterns than attention.

The JD requires "expertise in model architecture, backed by literature rationale and intuition" and "understanding of broad-spectrum model architectures across modalities." This phase builds that expertise through implementation — you will implement every major architectural variant that appears in production models (LLaMA, Mistral, Mixtral, ViT, DINO, LLaVA, Mamba) from scratch or near-scratch.

After this phase, when asked in an interview "how would you adapt your NPU tiling strategy for a Mixture-of-Experts model?", you can answer specifically because you understand the routing mechanism and its sparsity implications.

Concepts

• Scaled Dot-Product Attention (SDPA): $\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$ — computational complexity $O(n^2 d)$, memory $O(n^2)$
• Multi-Head Attention (MHA): split $d_{model}$ into $h$ heads, concatenate outputs — fully parallelizable across heads
• Grouped Query Attention (GQA): $g$ query groups share one KV head pair — reduces KV cache size by factor $h/g$; used in LLaMA-3, Mistral
• Multi-Query Attention (MQA): extreme GQA with $g=1$ — single KV head, all query heads; MQA = GQA with groups=1
• Rotary Position Embeddings (RoPE): encode position via rotation of Q,K vectors in complex plane: $f_q(x_m) = R^m_d x_m$ — length extrapolation via RoPE scaling (LLaMA)
• ALiBi (Attention with Linear Biases): position bias added to attention logits as $-m \cdot |i-j|$ — train short, generalize long
• SwiGLU / GeGLU: gated feedforward: $\text{FFN}(x) = (\sigma(W_1 x) \odot W_2 x) W_3$ — used in PaLM, LLaMA; ~25% more params but better performance per FLOP
• Pre-norm vs Post-norm: pre-norm (RMSNorm before attention) = more stable training; used in all modern LLMs
• RMSNorm: $\hat{x} = \frac{x}{\text{RMS}(x)} \cdot \gamma$ where $\text{RMS}(x) = \sqrt{\frac{1}{n}\sum x_i^2}$ — no mean subtraction, faster than LayerNorm
• Mixture of Experts (MoE): sparse activation — each token routes to top-K of E experts; active params per token = dense_params + (K/E × expert_params)
• Vision Transformer (ViT): divide image into non-overlapping patches → linear projection → standard transformer; patch size controls trade-off between sequence length and feature resolution
• DINO / DINOv2: self-supervised ViT — student-teacher distillation with no labels; produces strong dense features for segmentation/detection
• CLIP: contrastive vision-language pretraining — align image and text embeddings via InfoNCE loss; zero-shot classification and image-text retrieval
• LLaVA-style multimodal: visual tokens from CLIP vision encoder → linear projection → prepend to LLM token sequence → standard LLM forward
• Mamba / SSMs (Structured State Space Models): recurrent models with selective state transitions; linear time in sequence length; competitive with transformers on long sequences

Labs

Lab 01 — Transformer with Full Architectural Variants

Lab 02 — Mixture of Experts (MoE)

→ Lab folder: lab-02-mixture-of-experts/

Lab 03 — State Space Models (S4 / Mamba)

→ Lab folder: lab-03-state-space-models/

Lab 04 — ViT + DINO Self-Distillation

Implemented at testable scale in lab-04-vit-dino/ — ViT from scratch, DINO loss with centering+sharpening, momentum teacher, multi-crop step, 13-test suite that runs on synthetic data in seconds. The full CIFAR-10 training run described below is the lab's extension target.

Lab 05 — Multimodal Bridge (CLIP → LLM Projection)

Implemented at testable scale in lab-05-multimodal-bridge/ — trainable MLP projector between a frozen vision tower and a frozen causal LM, with the loss masking and causal-shift logic identical to LLaVA's, 12-test suite, no downloads. The full CLIP + TinyLlama + LLaVA-Instruct version below is the extension target.

Deliverables Checklist

• Configurable transformer with 5 tested architecture configs and perplexity table
• MoE block with balanced expert utilization under aux loss
• Mamba block with ZOH discretization and selective scan, memory comparison vs attention
• ViT + DINO lab passing (13 tests); (stretch) CIFAR-10 run with attention map visualizations
• Multimodal bridge lab passing (12 tests); (stretch) CLIP+TinyLlama with VQAv2 score
• All test_lab.py suites pass

Guides in This Phase

• HITCHHIKERS-GUIDE.md — attention math, ViT, MoE, Mamba overview, interview numbers
• DEEP-DIVE-MOE-SSM.md — Mixture of Experts and State Space Models from first principles (supports Labs 02–03)

Interview Relevance

• "What is the memory advantage of GQA over MHA at 32K context?" — you can compute: (h×2×n×d) vs (h/g × 2×n×d), show 8x reduction for g=8
• "How does DINO work without labels?" — explain centering + sharpening, why teacher EMA is key, why augmentation multi-crop works
• "What is the bottleneck in deploying LLaVA on an NPU?" — you understand the KV cache growth with visual tokens, the projection cost, and the LLM decode bottleneck
• "Why does RoPE generalize to longer sequences than learned position embeddings?" — explain the frequency decomposition and length-independent rotation property

Warmup Guide — Model Architecture Mastery

Zero-to-expert primer for Phase 02. Read this before the labs: it builds every concept from first principles — attention, position encodings, normalization, Mixture of Experts, State Space Models, Vision Transformers, and multimodal bridges — assuming only basic Python and linear algebra.

Table of Contents

• Chapter 1: Why Architecture Knowledge Is an Optimization Skill
• Chapter 2: Attention from Zero
• Chapter 3: The KV Cache — Where Inference Cost Lives
• Chapter 4: MHA, MQA, GQA — The KV-Head Spectrum
• Chapter 5: Position Encodings — Learned, RoPE, ALiBi
• Chapter 6: Normalization and FFN Variants
• Chapter 7: Mixture of Experts
• Chapter 8: State Space Models and Mamba
• Chapter 9: Vision Transformers and DINO
• Chapter 10: Multimodal Bridges
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: Why Architecture Knowledge Is an Optimization Skill

Zero background: a neural network architecture is the fixed computational skeleton — which matrices multiply which activations in what order. Weights change during training; the architecture doesn't.

Why it matters for this curriculum: every optimization decision downstream — which layers to quantize to INT4, how to tile a matmul for an NPU, why accuracy collapses after a graph transformation — is a function of architectural structure. Three concrete examples:

• GQA shares KV heads → fewer unique weight rows → per-group quantization has fewer groups to calibrate → quantizes more gracefully than MHA.
• MoE activates 2 of 8 experts per token → 75% of expert weights are cold per token → weight streaming and per-expert quantization become viable; but the router is tiny and precision-sensitive — quantize it last.
• Mamba carries a recurrent state → no KV cache, O(1) memory per token — but the state is numerically delicate; INT8-quantizing the recurrence accumulates error step by step in a way attention never does.

Misconception to kill now: "architectures are interchangeable backends for the same function." The function may be similar; the hardware cost surface is wildly different, and that surface is your job.

Chapter 2: Attention from Zero

The problem attention solves: process a sequence so each position can use information from any other position, with content-dependent (not fixed) routing.

Mechanism, from nothing: each token's embedding $x_i \in \mathbb{R}^d$ produces three vectors via learned projections: a query $q_i = W_Q x_i$, a key $k_i = W_K x_i$, a value $v_i = W_V x_i$. Token $i$ scores every token $j$ by dot product $q_i \cdot k_j$ — "how relevant is $j$ to me" — then turns scores into a probability distribution and takes the weighted average of values:

$$\text{Attention}(Q,K,V) = \text{softmax}!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$

Every term, justified:

• $QK^\top$ is an $n \times n$ matrix of all pairwise scores — this is where the $O(n^2)$ cost in both compute and memory comes from.
• $\sqrt{d_k}$: dot products of random $d_k$-dim vectors have variance $\propto d_k$; unscaled, softmax inputs grow with dimension, saturating it (gradients vanish). Dividing by $\sqrt{d_k}$ keeps variance ≈ 1 regardless of head size.
• softmax row-wise: each query gets a distribution over keys, so the output is a convex combination of values — bounded, differentiable routing.
• Causal masking (decoders): set scores $j > i$ to $-\infty$ before softmax, so position $i$ cannot see the future. After softmax those entries are exactly 0.

Multi-head: instead of one attention in $\mathbb{R}^d$, split into $h$ heads of dimension $d/h$, attend independently, concatenate, project. Heads learn different relation types (syntax, coreference, position). Cost is the same FLOPs as single-head at full width; the win is representational.

Production significance: the $n^2$ attention matrix is the reason long-context is hard, the reason FlashAttention exists (Phase 08), and the reason KV caching dominates inference memory (next chapter).

Chapter 3: The KV Cache — Where Inference Cost Lives

The setup: autoregressive generation produces one token at a time. Naively, generating token $t+1$ recomputes attention over all $t$ previous tokens from scratch — $O(t^2)$ work per token, $O(T^3)$ total.

The fix: $K$ and $V$ for past tokens never change (causality). Cache them. Each new token computes only its own $q, k, v$, appends $k, v$ to the cache, and attends: $O(t)$ per token.

The cost you bought: memory. Per token, per layer: $2 \times h_{kv} \times d_{head}$ values. For LLaMA-2-7B (32 layers, 32 KV heads, head dim 128, FP16):

$$32 \times 2 \times 32 \times 128 \times 2,\text{bytes} = 512,\text{KB per token}$$

A 4K-token context costs 2 GB — for batch size 1. This single number explains: why GQA exists (Chapter 4), why batch sizes are memory-bound, why PagedAttention manages KV like virtual memory (Phase 08), and why edge deployment quantizes the KV cache itself.

Misconception: "inference cost = model weights." For long contexts and large batches, KV cache rivals or exceeds weight memory. Always compute both.

Chapter 4: MHA, MQA, GQA — The KV-Head Spectrum

One design axis: how many KV head pairs do $h$ query heads share?

Mechanism: in GQA, query heads within a group attend using the same $K, V$ projections. Implementation is just repeat_interleave of KV heads to match query heads (or better: a broadcasted attention kernel that never materializes the repetition).

Why this is the consensus design (LLaMA-3, Mistral, Qwen): KV cache (Chapter 3) is the inference bottleneck; GQA cuts it 4–8× for ~zero quality loss. It also helps quantization: fewer distinct KV projections → fewer calibration distributions.

Lab 01 implements all three behind one config flag — the test checks GQA output matches the naive repeat-expansion within 1e-5, which forces you to get the grouping exactly right.

Chapter 5: Position Encodings — Learned, RoPE, ALiBi

The problem: attention is permutation-invariant — shuffle the tokens and the scores are identical. Order must be injected.

Learned absolute embeddings (GPT-2): a trainable vector per position, added to token embeddings. Simple; cannot extrapolate past training length (position 2049 was never trained); positions don't compose ("distance 5" isn't represented).

RoPE (Rotary Position Embeddings) — used by LLaMA, Qwen, most modern LLMs. Idea: encode position as rotation. Pair up dimensions of $q$ and $k$; rotate each pair $(x_1, x_2)$ at position $m$ by angle $m\theta_i$, with per-pair frequencies $\theta_i = 10000^{-2i/d}$:

$$\begin{pmatrix} \cos m\theta_i & -\sin m\theta_i \ \sin m\theta_i & \cos m\theta_i \end{pmatrix}\begin{pmatrix} x_{2i} \ x_{2i+1} \end{pmatrix}$$

The magic: a rotation by $m\theta$ on $q$ and $n\theta$ on $k$ leaves their dot product depending only on $m - n$ — relative position falls out of absolute rotations. Multi-frequency = a "clock" with many hands: fast hands resolve nearby distances, slow hands resolve far ones. Extrapolation beyond training length works partially and is extended by position interpolation / NTK-aware scaling (scaling $\theta$ at inference).

ALiBi: skip embeddings entirely; subtract a per-head linear penalty $m_h \cdot |i - j|$ from attention logits. Recency bias, baked in. Trains short, generalizes long remarkably well; but no notion of position beyond distance, and weaker at tasks needing precise absolute positions.

Choosing: RoPE is the default; ALiBi when train-short-infer-long matters most. Lab 01's test checks RoPE's relative-position property explicitly.

Chapter 6: Normalization and FFN Variants

LayerNorm from zero: deep nets suffer when activation scales drift across layers. LayerNorm re-standardizes each token's vector: subtract its mean, divide by its std, then apply learned scale/shift. RMSNorm drops the mean subtraction: $\hat{x} = x / \sqrt{\tfrac{1}{d}\sum x_i^2} \cdot \gamma$ — one fewer reduction, ~10-15% faster, and empirically equivalent in transformers. All modern LLMs use RMSNorm.

Pre-norm vs post-norm: normalize before the sublayer (pre-norm, modern) and the residual path stays an identity highway — gradients flow unimpeded, training is stable without warmup tricks. Post-norm (original Transformer) normalizes the sum, which can amplify or attenuate the residual signal — deeper stacks destabilize.

SwiGLU: standard FFN is $W_2,\sigma(W_1 x)$. Gated variants compute $W_3(\sigma(W_1x) \odot W_2x)$ — one path gates the other multiplicatively, letting the FFN model input-dependent feature selection. ~25% more parameters per FFN at the same hidden dim (three matrices, not two); in practice hidden dim is shrunk by ⅔ to compensate. PaLM/LLaMA measurements: better loss per FLOP. SiLU ($x\sigma(x)$) is the usual $\sigma$.

Quantization note: the elementwise multiply in SwiGLU creates wider activation dynamic range than ReLU FFNs — one reason LLM activations are harder to quantize than CNN activations (foreshadowing Phase 03's SmoothQuant).

Chapter 7: Mixture of Experts

Zero background: a dense FFN applies the same weights to every token. MoE replaces one FFN with $E$ parallel FFNs ("experts") plus a tiny router that picks $K$ experts per token (K=1 Switch, K=2 Mixtral). Total parameters scale with $E$; compute per token scales with $K$ — parameters and FLOPs are decoupled.

Router mechanics: logits $= W_r x$ (an $E$-way linear classifier per token); take top-K; softmax over the selected logits gives combination weights; output = weighted sum of the chosen experts' outputs.

The failure mode — expert collapse: routing is self-reinforcing. An expert that gets more tokens trains more, improves, attracts more tokens; others atrophy. End state: one expert does everything and you paid for $E$.

The fix — auxiliary load-balancing loss (Switch Transformer): with $f_i$ = fraction of tokens routed to expert $i$ and $p_i$ = mean router probability for expert $i$,

$$\mathcal{L}{aux} = E \cdot \sum{i=1}^{E} f_i , p_i$$

This is minimized when both distributions are uniform ($f_i = p_i = 1/E$ gives $\mathcal{L}_{aux}=1$). $f_i$ is non-differentiable (a count); $p_i$ carries the gradient — the product trick lets gradient flow push probability away from overloaded experts.

Production significance: capacity factors (token budgets per expert with overflow dropping), all-to-all communication in expert parallelism, and per-expert quantization. On edge: MoE means most weights are cold per token — pair with layer/expert streaming (Phase 10).

Chapter 8: State Space Models and Mamba

The motivation: attention pays $O(n^2)$ compute and $O(n)$ KV memory. An RNN pays $O(n)$ compute and $O(1)$ memory but can't be parallelized over time during training and forgets. SSMs are the modern reconciliation.

Continuous origin: a linear time-invariant system $h'(t) = Ah(t) + Bx(t),; y(t) = Ch(t) + Dx(t)$. To run on discrete tokens, discretize with step $\Delta$ using Zero-Order Hold (input held constant within a step):

$$\bar{A} = e^{\Delta A}, \qquad \bar{B} = (\bar{A} - I)A^{-1}B$$

giving the recurrence $h_t = \bar{A}h_{t-1} + \bar{B}x_t$, $y_t = Ch_t + Dx_t$. Intuition: $\bar{A}$ entries (< 1) are memory decay coefficients; $\bar{B}$ is input gain. Large $\Delta$ = step fast, forget fast; small $\Delta$ = retain.

S4's contribution: with fixed $A$ (HiPPO-initialized, diagonal), the whole sequence output is a convolution with a precomputable kernel — train in parallel like a CNN, run as an RNN at inference.

Mamba's contribution — selectivity: make $B, C, \Delta$ functions of the input $x_t$. Now the state update is content-dependent: the model can choose, per token, to store or ignore. This breaks the convolution trick (kernel is input-dependent), so Mamba ships a hardware-aware parallel associative scan computed in SRAM with recomputation in backward — an architecture and a kernel co-designed.

Why NPU/edge people care: O(1) inference state (a $d_{state} \times d$ matrix, not a growing KV cache) is the difference between bounded and unbounded memory for streaming workloads. The risk: recurrent state quantization error compounds over time — evaluate long-sequence accuracy, not just per-op SQNR.

Chapter 9: Vision Transformers and DINO

ViT from zero: chop the image into $p \times p$ patches (16×16 typical), flatten each to a vector, project linearly — now you have a "sentence" of patch tokens; run a standard transformer. The projection is implemented as Conv2d with kernel = stride = $p$ (identical math, one fused op — Lab 04 tests that non-overlap property). A learnable CLS token is prepended as the aggregation point; learned position embeddings restore spatial order.

Inductive bias tradeoff: CNNs hard-code locality and translation equivariance; ViTs learn relations from data — worse on small datasets, better at scale, and their global receptive field from layer 1 is why their attention maps are interesting.

DINO — labels-free training: two networks, student and teacher, same architecture. Teacher weights are an EMA of the student (never gradient-trained). Both see different augmented crops of the same image; the student is trained to predict the teacher's output distribution (cross-entropy). Collapse is prevented by two opposing mechanisms: centering (subtract a running mean from teacher logits — kills the one-hot collapse) and sharpening (teacher temperature τ_t ≈ 0.04 < student τ_s ≈ 0.1 — kills the uniform collapse). Multi-crop makes the task "predict global view from local crop", which forces object-level features — and produces the famous segmentation-quality attention maps with zero labels. Lab 04 implements every piece and tests each property in isolation.

Chapter 10: Multimodal Bridges

The design question: you have a pretrained vision encoder (CLIP) and a pretrained LLM. How do images get into the LLM?

LLaVA's answer (Lab 05): project CLIP's patch tokens through a small MLP into the LLM's embedding space and prepend them as if they were text tokens. Train only the projector first (the bridge must learn to "speak embedding-ese" before touching either tower); optionally unfreeze the LLM later for instruction tuning. The loss is computed on answer text positions only — visual positions and the prompt are masked with ignore_index — and there's a one-position causal shift at the modality boundary that is the canonical off-by-one bug (Lab 05 has a test that catches precisely it).

Alternatives: BLIP-2's Q-Former compresses patches into a fixed small set of query tokens (cheaper sequences); Flamingo interleaves gated cross-attention layers (LM untouched, text throughput preserved, but new layers to train). The engineering axis: prepended tokens inflate the KV cache by $N_{patches}$ at every layer — 576 CLIP tokens often dominate a short VQA prompt — which is why token compression matters on edge.

Lab Walkthrough Guidance

Order: Lab 01 → 02 → 03 → 04 → 05 (attention foundations first; MoE/SSM build on it; vision and multimodal reuse your attention implementation mentally).

• Lab 01 (variants): implement MHA first and get the causal mask + KV cache correct before adding GQA/MQA (the GQA test compares against naive expansion — write that naive expansion yourself to understand it). RoPE before ALiBi.
• Lab 02 (MoE): router → expert → block → utilization. Verify the aux-loss formula against Chapter 7's math by hand on a 2-expert toy.
• Lab 03 (SSM): discretize_zoh is pure math — check it against scalar examples ($a=-1, \Delta=0.1 \Rightarrow \bar{a} \approx 0.905$). Then the sequential scan, then selectivity.
• Lab 04 (DINO): follow the README's order; the no-grad-to-teacher and centering tests are the conceptual heart.
• Lab 05 (bridge): draw the [v_1..v_N | t_1..t_T] position layout on paper before writing loss().

Success Criteria

You are ready for Phase 03 when you can, from memory:

1. Compute the KV cache size of any (layers, kv_heads, head_dim, context, dtype) config and state the GQA reduction factor.
2. Derive why $\sqrt{d_k}$ scaling exists from the variance of random dot products.
3. Explain RoPE's relative-position property (rotations compose; dot product depends on $m-n$) without notes.
4. State both MoE collapse modes and how $\sum f_i p_i$ fixes them.
5. Write the ZOH discretization from the continuous SSM and explain what selectivity changes computationally (convolution trick breaks → scan).
6. Explain DINO's two anti-collapse mechanisms and what happens if either is removed.
7. Name the multimodal causal-shift bug and the KV-cache cost of prepended visual tokens.

Interview Q&A

Q: 32K context, LLaMA-3-8B (32 layers, 8 KV heads, head dim 128, FP16). KV cache size? $32 \times 2 \times 8 \times 128 \times 2,\text{B} = 131,\text{KB/token}$; × 32768 tokens ≈ 4.3 GB per sequence. With MHA (32 KV heads) it would be 17 GB — that 4× is GQA's entire reason for existing.

Q: Why does GQA quantize more gracefully than MHA? Fewer distinct KV projection matrices → fewer weight distributions to calibrate, and the shared KV heads see more diverse activation statistics (averaging over query groups), so per-channel ranges are better conditioned. Also the KV cache itself, if quantized, has fewer head-specific distributions.

Q: Your Mamba model passes per-op quantization checks but fails on long sequences. Why? Recurrent state error compounds: each step multiplies state by $\bar{A}$ and adds quantized input contributions; unlike attention (which reads exact cached K/V), small per-step errors accumulate over thousands of steps. Evaluate sequence-level metrics vs length; consider keeping the state and $\Delta$-path in higher precision.

Q: When would you pick ALiBi over RoPE? When the dominant requirement is train-short/infer-long generalization with no inference- time scaling tricks, and the tasks are recency-dominated (streaming ASR post-processing, log modeling). RoPE wins when precise relative structure matters or ecosystem tooling (RoPE-scaling for context extension) is needed.

Q: How would you adapt NPU tiling for an MoE model? Router runs dense and tiny (keep high precision); expert FFNs are the bulk — tile them as independent matmuls with token gathering by expert (sort tokens by expert id to batch); since only K of E experts are hot per token, weights can be streamed/paged per expert, and per-expert quantization scales are natural tile boundaries.

References

• Vaswani et al., Attention Is All You Need (2017) — arXiv:1706.03762
• Ainslie et al., GQA: Training Generalized Multi-Query Transformer (2023) — arXiv:2305.13245
• Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding (2021) — arXiv:2104.09864
• Press et al., Train Short, Test Long: ALiBi (2021) — arXiv:2108.12409
• Shazeer, GLU Variants Improve Transformer (2020) — arXiv:2002.05202
• Fedus et al., Switch Transformers (2021) — arXiv:2101.03961
• Jiang et al., Mixtral of Experts (2024) — arXiv:2401.04088
• Gu & Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces (2023) — arXiv:2312.00752
• Gu et al., Efficiently Modeling Long Sequences with S4 (2021) — arXiv:2111.00396
• Dosovitskiy et al., An Image is Worth 16x16 Words (ViT) (2020) — arXiv:2010.11929
• Caron et al., Emerging Properties in Self-Supervised ViTs (DINO) (2021) — arXiv:2104.14294
• Liu et al., Visual Instruction Tuning (LLaVA) (2023) — arXiv:2304.08485
• The Illustrated Transformer — Jay Alammar
• DEEP-DIVE-MOE-SSM.md — this phase's companion deep dive

Hitchhiker's Guide — Model Architecture Mastery

"Architecture is not what you add; it is what you understand well enough to subtract."

1. The Attention Equation and Its Variants — Precise Math

1.1 Standard Multi-Head Attention

$$\text{MHA}(X) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W^O$$

$$\text{head}_i = \text{Attention}(XW^Q_i, XW^K_i, XW^V_i)$$

$$\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V$$

where $M$ is the causal mask ($-\infty$ for future positions).

Memory during inference: For sequence length $n$, each layer stores KV cache of shape $(2, \text{batch}, n, h, d_k)$ = $2 \cdot b \cdot n \cdot h \cdot d_k$ floats. For LLaMA-7B: $b=1, n=4096, h=32, d_k=128$ → 32 layers × $2 × 4096 × 32 × 128 × 2$ bytes = 8 GB in FP16.

1.2 Grouped Query Attention

Reduce KV heads from $h$ to $g$ heads:

$$\text{GQA}(X) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W^O$$

$$\text{head}i = \text{Attention}(XW^Q_i, XW^K{\lfloor i/r \rfloor}, XW^V_{\lfloor i/r \rfloor})$$

where $r = h/g$ is the repetition factor. When $g=h$: standard MHA. When $g=1$: MQA.

KV cache reduction: $g/h$ fraction of MHA's KV cache. LLaMA-3 uses $g=8, h=32$ → 4× reduction → 2 GB for same scenario.

Implementation trick: no need to actually repeat the KV tensors — use torch.repeat_interleave(k, h//g, dim=1) only during attention computation, not in cache.

1.3 Rotary Position Embeddings (RoPE)

RoPE encodes position $m$ by rotating query/key vectors in 2D subspaces:

$$f_q(x_m, m) = \begin{pmatrix} \cos(m\theta_1) & -\sin(m\theta_1) \ \sin(m\theta_1) & \cos(m\theta_1) \end{pmatrix} \ldots \begin{pmatrix} x_1 \ x_2 \end{pmatrix}$$

where $\theta_i = 10000^{-2(i-1)/d}$ (same base as sinusoidal embeddings).

Key property: $\langle f_q(x_m), f_k(x_n) \rangle$ depends only on $x_m$, $x_n$, and $(m-n)$ — position enters only as relative offset. This is why RoPE can generalize to unseen sequence lengths.

Implementation (the complex number trick):

def apply_rotary_emb(xq, xk, freqs_cis):
    # xq: (batch, seq, n_heads, head_dim)
    # freqs_cis: (seq, head_dim//2) complex
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)

1.4 ALiBi (Attention with Linear Biases)

Instead of position embeddings, add bias $-m \cdot |i - j|$ to attention logits:

$$A_{ij} = \frac{Q_i K_j^T}{\sqrt{d}} - m \cdot |i - j|$$

where $m$ is a per-head slope: $m_h = 2^{-8h/H}$ for head $h$ of $H$ total heads.

Key advantage: completely eliminates position embedding parameters. Generalizes to longer sequences at inference time (slopes penalize large distances). Disadvantage: slightly lower peak performance than RoPE on standard benchmarks.

1.5 SwiGLU vs Standard FFN

Standard FFN: $\text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2$
SwiGLU: $\text{FFN}(x) = (\text{Swish}(xW_1) \odot xW_2)W_3$

where $\text{Swish}(x) = x \cdot \sigma(x)$.

Parameter count: SwiGLU adds $W_2$ (same size as $W_1$) but reduces $d_{ff}$ from $4d$ to $\frac{2}{3} \times 4d \approx 2.67d$ to maintain FLOPs parity. LLaMA uses $d_{ff} = \frac{8}{3} d_{model}$ rounded to multiple of 256.

2. Vision Transformer Internals

2.1 Patch Embedding

An image of size $(H, W, C)$ is divided into $N = \frac{H}{P} \times \frac{W}{P}$ patches of size $(P, P, C)$.

Each patch is linearly projected to $d_{model}$: nn.Conv2d(C, d_model, kernel_size=P, stride=P) — this is equivalent to dividing into patches and projecting each independently, but implemented efficiently as a single strided convolution.

Sequence length: $N + 1$ (N patches + 1 CLS token). For ViT-L with $P=16$, $H=W=224$: $N = 196$.

2.2 Why ViT Attention Heads Learn to Segment

Each patch is a $(P \times P)$ region. The class token attends to all patches to aggregate a global representation. In DINO training, the self-supervised objective forces the model to be invariant to local augmentations while sensitive to semantic content.

Mechanistically: early layers learn local edge/texture features (similar to CNN filters), while deep layers learn to cluster semantically related patches (foreground vs background). The CLS token's attention weights in the last layer correspond to object foreground because:

1. The DINO loss forces the CLS token to represent the "semantic essence" of the full image
2. Background patches carry less semantic information and receive lower attention weight
3. Multiple views (global + local crops in DINO) force the model to find consistent semantic regions

This is not hand-designed — it emerges purely from the training objective.

3. Mixture of Experts Architecture

3.1 Routing Mechanism

Each token independently selects the top-K experts out of E total:

def route(x, W_gate):
    # x: (batch, seq, d_model)
    logits = x @ W_gate.T  # (batch, seq, num_experts)
    gates, indices = logits.topk(k=2, dim=-1)  # select top-2 experts
    gates = F.softmax(gates, dim=-1)  # normalize selected expert weights
    return gates, indices

Auxiliary load balancing loss: without it, all tokens collapse to 1-2 experts ("router collapse"). The auxiliary loss penalizes unequal expert utilization:

$$L_{\text{aux}} = \alpha \cdot E \cdot \sum_{i=1}^{E} f_i \cdot P_i$$

where $f_i$ = fraction of tokens assigned to expert $i$, $P_i$ = average routing probability to expert $i$.

3.2 MoE Impact on NPU Deployment

MoE models (Mixtral 8x7B, GPT-4-style) are problematic for NPUs because:

1. Sparse activation: different tokens activate different experts → no batching across tokens for expert computation
2. Expert parameter loading: each expert's weights must be available — 8 experts × 7B params = 56B total (even though only 14B active per token)
3. Load imbalance: if multiple tokens route to the same expert simultaneously, that expert becomes a bottleneck

NPU mitigation strategies:

• Expert pruning: drop low-frequency experts for a specific deployment domain
• Expert merging: average 8 experts into 4 (reduces params at slight accuracy cost)
• MoE → dense distillation: distill the sparse model into a single dense FFN

4. Mamba / SSM Architecture

4.1 The Core SSM Equation

Mamba models are discrete-time state space models:

$$h_t = \bar{A} h_{t-1} + \bar{B} x_t$$ $$y_t = C h_t$$

where $h_t \in \mathbb{R}^N$ is the hidden state, $x_t, y_t \in \mathbb{R}$ are input/output, and $\bar{A}, \bar{B}$ are discretized from continuous-time parameters.

Selective SSM (what makes Mamba work): $B$, $C$, and $\Delta$ (discretization timescale) are input-dependent:

$$\bar{B}t = f_B(x_t), \quad C_t = f_C(x_t), \quad \Delta_t = \text{softplus}(f\Delta(x_t))$$

This gives the model the ability to selectively forget or remember based on content — equivalent to attention's dynamic weighting but in $O(n)$ time.

4.2 Mamba vs Attention for NPU Deployment

5. Key Interview Numbers

6. Resources

• Attention is All You Need (Vaswani et al., 2017) — the foundation
• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., 2023)
• RoFormer: Enhanced Transformer with Rotary Position Embedding (Su et al., 2021)
• Train Short, Test Long: Attention with Linear Biases (ALiBi) (Press et al., 2021)
• GLU Variants Improve Transformer (Noam Shazeer, 2020) — SwiGLU origin
• An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale (Dosovitskiy et al., 2020)
• Emerging Properties in Self-Supervised Vision Transformers (DINO) (Caron et al., 2021)
• Mamba: Linear-Time Sequence Modeling with Selective State Spaces (Gu & Dao, 2023)
• Mixtral of Experts (Mistral AI, 2024)

👨🏻 Brother Talk — Phase 02, Off the Record

Architecture mastery — where the JD says "expertise in model architecture, backed by literature rationale and intuition." That word intuition is the whole game.

Brother, read the JD line again: "Expertise in model architecture, backed by literature rationale and intuition." Most candidates hear "know the architectures" and go memorize that a transformer has attention and an FFN. That is not what this asks. It asks for rationale — why this design and not another — and intuition — what happens to accuracy and latency when you change it. At Senior Staff level you're not implementing someone else's paper; you're the person who looks at a new architecture and says "this MoE will be a nightmare on our NPU because of the gather/scatter, but here's the variant that won't." That judgment is what this phase builds.

Here's the reframe: architectures are not a zoo of unrelated animals; they're a sequence of answers to a small set of recurring problems. Attention is quadratic → so we get sparse attention, linear attention, and eventually SSMs (Mamba) that ditch attention for a recurrence with near-linear cost. Dense FFNs are expensive → so we get Mixture-of-Experts that activates only k of N experts. Vision needs inductive bias → ViT throws it away and recovers it with scale and self-distillation (DINO). Once you see architectures as moves in response to a constraint, you can reason about ones you've never seen — which is exactly the interview test.

Now the part that matters for this job specifically, and that pure-ML people miss: the best architecture on a GPU benchmark is often the worst architecture on an NPU. This is the Qualcomm lens, and it's your edge. MoE looks great on FLOPs-per-token, but the expert routing is a dynamic gather/scatter that NPUs (which love static, dense, predictable dataflow) hate — you get memory-bound stalls and CPU fallbacks. SSMs have a beautiful linear recurrence, but the scan is sequential and maps awkwardly to hardware built for parallel matmuls. A candidate who can say "yes the paper shows X, but on a mobile NPU the real bottleneck is the irregular memory access, so here's what I'd actually ship" — that person gets hired, because that's the literal job: adapt GenAI techniques for Qualcomm SOCs.

The MoE lab will teach you the thing that trips everyone: load balancing. Naive top-k routing collapses — a few experts get all the tokens, the rest die, and your "8 experts" is really 2. The auxiliary load-balancing loss isn't a footnote; it's what makes MoE work at all. And on hardware, an unbalanced MoE is a latency disaster because the busiest expert gates the whole layer. Build it, watch the collapse happen, watch the aux loss fix it. That's intuition you can't get from reading.

The SSM/Mamba lab is the one that'll make you feel smart, because the selective-scan recurrence is genuinely elegant — and it's the future-facing skill (the JD wants "latest GenAI techniques adapted for NPUs"). But hold the Qualcomm skepticism: the parallel scan that makes Mamba fast on a GPU is exactly the part that's hard to map to an NPU's dataflow. Knowing both the elegance and the hardware friction is the Senior Staff take.

On ViT + DINO (the lab I had to fix the test for — the loss genuinely decreases, the old test just used an unfair baseline): the lesson beyond the architecture is self-distillation without labels. A teacher that's an EMA of the student, centering to prevent collapse, sharpening to make confident targets — it's a deceptively deep training dynamic. The collapse modes (everything maps to one vector; everything maps to uniform) are real failure modes you'll see in production self-supervised training, and centering vs sharpening are the two knobs that fight them. Understanding why it doesn't collapse is more valuable than the architecture itself.

Multimodal (CLIP→LLM bridge): this is where the field is going and where edge GenAI gets interesting. The bridge — projecting vision features into the LLM's token space — is a small module with outsized importance, and on-device it's where a lot of the accuracy/latency tradeoff lives. Know how the modalities get glued, because "run a multimodal model on a phone" is a sentence you'll hear at this job.

Career truth, brother: architecture knowledge ages, but architecture intuition compounds. The specific models will be obsolete in two years; the ability to look at a new one and reason about its accuracy, its cost, and its hardware-friendliness is permanent. In interviews, the tell of a senior person is that they don't just describe an architecture — they critique it, they know its failure modes, and (for this role) they know how it'll behave on silicon. Read the papers for the rationale, build the labs for the intuition, and always ask "how does this map to the NPU?" — that question is your whole differentiator.

Build the variants. Make the MoE collapse and then fix it. Then come to P03, where we make these architectures small enough to ship — quantization, where accuracy goes to die if you're not careful.

— your brother 👨🏻

Deep Dive — Mixture of Experts & State Space Models (Phase 02)

"MoE and SSMs are not incremental improvements — they are fundamental rethinks of how sequence models should scale."

Section 1: Mixture of Experts — From Dense to Sparse

1.1 The Scaling Problem with Dense Models

Every parameter in a standard transformer is activated for every input token. A 70B dense model runs 70B multiplications per token. This is expensive, and evidence from scaling laws (Hoffmann et al., Chinchilla) shows that most parameters could be replaced by conditional computation — only activate what you need.

Key question: Can we get the capacity of a large model while spending compute equivalent to a smaller one?

Yes: Mixture of Experts.

1.2 The MoE Layer Architecture

Replace the Feed-Forward Network (FFN) in each Transformer block with $N$ independent FFN "experts" and a router that chooses which experts process each token.

Standard Transformer Block:
  Token x → MHA → LayerNorm → FFN → LayerNorm → output

MoE Transformer Block:
  Token x → MHA → LayerNorm → [Router] → [Expert 0, Expert 1, ..., Expert N-1]
                                              ↑ only K experts activated per token

The architecture for a single transformer layer:

class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads, n_experts=8, top_k=2):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, n_heads)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.moe   = MoEBlock(d_model, d_ff=4*d_model, n_experts=n_experts, top_k=top_k)
    
    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        moe_out, aux_loss = self.moe(self.norm2(x))
        x = x + moe_out
        return x, aux_loss

Parameter count:

• Dense FFN: $2 \times d_{model} \times d_{ff}$ (e.g., 2×4096×16384 = 134M params)
• MoE with 8 experts: $8 \times 2 \times d_{model} \times d_{ff/8}$ = same total parameters but only top-2 experts activate → 2/8 = 25% of compute

This is how Mixtral 8×7B has 47B parameters but computes like a ~13B model.

1.3 The Router: Soft Routing

The router is a learned linear projection from $d_{model}$ to $N$ expert logits, normalized with softmax:

Input x ∈ ℝ^{d_model}
Logits:  g = W_router @ x ∈ ℝ^N
Probs:   p = softmax(g) ∈ ℝ^N
Top-K:   (indices, weights) = topk(p, K)
Renorm:  weights = weights / sum(weights)

In code:

class TopKRouter(nn.Module):
    def __init__(self, d_model, n_experts, top_k=2):
        super().__init__()
        self.n_experts = n_experts
        self.top_k = top_k
        # NOTE: shape is (n_experts, d_model) so output is logit per expert
        self.router_weight = nn.Linear(d_model, n_experts, bias=False)

    def forward(self, x):
        # x: [batch*seq, d_model]
        logits = self.router_weight(x)           # [N, n_experts]
        probs  = F.softmax(logits, dim=-1)       # [N, n_experts]
        topk_w, topk_idx = probs.topk(self.top_k, dim=-1)  # [N, K] each
        
        # Renormalize weights to sum to 1
        topk_w = topk_w / topk_w.sum(dim=-1, keepdim=True)
        return topk_idx, topk_w

Walkthrough with numbers: Suppose N=6 tokens, 4 experts, top-K=2.

x = [6, d_model]
logits = router(x) → [6, 4]

Example for token 0:  logits = [2.1, 0.3, -0.5, 1.8]
                      probs  = [0.50, 0.09, 0.04, 0.37]
                      topk   = indices=[0,3], weights=[0.50, 0.37]
                      renorm = weights=[0.575, 0.425]
→ Token 0 goes to experts 0 and 3 with weights 0.575, 0.425

The MoE block output for each token: $$y_i = \sum_{k \in \text{TopK}(i)} w_{ik} \cdot E_k(x_i)$$

where $E_k$ is the $k$-th expert FFN.

1.4 Load Balancing: The Critical Problem

Expert collapse: without any penalty, the router quickly learns to always route to the same 1-2 experts (because they get more training signal and become better). Other experts starve and become useless. Your MoE degenerates to a dense FFN.

Measurement: define load imbalance = (tokens to most popular expert) / (ideal even split).

For 8 experts with 100 tokens: ideal = 100/8 = 12.5 tokens each. If expert 0 gets 80 tokens: imbalance = 80/12.5 = 6.4× — near-total collapse.

Solution: Switch Transformer auxiliary loss (Fedus et al., 2021):

For expert $i$:

• $f_i$ = fraction of tokens assigned to expert $i$ (counting top-K assignments)
• $P_i$ = mean router probability for expert $i$ across all tokens

$$\mathcal{L}{aux} = N{experts} \cdot \sum_{i=1}^{N_{experts}} f_i \cdot P_i$$

The minimum of $\mathcal{L}_{aux}$ is $1.0$ (achieved when routing is perfectly uniform). Multiplying $f_i \cdot P_i$ creates a smooth gradient: even though the top-K selection is non-differentiable, $P_i$ is differentiable and carries gradient back to the router weights.

In code:

def compute_aux_loss(probs, topk_indices, n_experts):
    """
    probs:       [N, n_experts]  — soft router probs (differentiable)
    topk_indices:[N, top_k]      — hard selections (non-differentiable)
    """
    N = probs.shape[0]
    
    # f_i: fraction of tokens dispatched to expert i (from hard selections)
    # Count each top-k assignment equally
    one_hot = torch.zeros(N, n_experts, device=probs.device)
    for k in range(topk_indices.shape[1]):
        one_hot.scatter_add_(1, topk_indices[:, k:k+1],
                             torch.full((N,1), 1.0/topk_indices.shape[1], device=probs.device))
    f = one_hot.mean(dim=0)   # [n_experts] — fraction of tokens to each expert
    
    # P_i: mean router probability (differentiable!)
    P = probs.mean(dim=0)     # [n_experts]
    
    aux_loss = n_experts * (f * P).sum()
    return aux_loss

Walkthrough: 4 experts, 4 tokens, top-1:

Routing: [E0, E1, E0, E3]
f = [0.5, 0.25, 0.0, 0.25]    (E0 gets 2/4=0.5)
P = [0.35, 0.20, 0.15, 0.30]  (average softmax probs)
aux = 4 * (0.5*0.35 + 0.25*0.20 + 0.0*0.15 + 0.25*0.30)
    = 4 * (0.175 + 0.05 + 0 + 0.075)
    = 4 * 0.30 = 1.20  (>1.0 → imbalanced, gradient pushes toward 1.0)

In training: add α * L_aux to the main loss. Typical α = 0.01.

1.5 Expert Execution: Dispatching Tokens

The naive implementation loops over every token × every K expert — O(N×K) sequential calls. Production implementations batch tokens per expert:

# Production MoE dispatch (simplified)
def moe_forward_batched(x, router, experts):
    B, T, d = x.shape
    x_flat = x.reshape(B*T, d)        # [N, d]
    N = x_flat.shape[0]
    K = router.top_k
    
    indices, weights = router(x_flat)  # [N, K], [N, K]
    output = torch.zeros_like(x_flat)
    
    # Group tokens by expert for batched execution
    for e in range(router.n_experts):
        # Find all (token, k) pairs that route to expert e
        token_mask = (indices == e)  # [N, K] bool
        if not token_mask.any():
            continue
        
        # Gather tokens going to this expert
        tok_ids, k_ids = token_mask.nonzero(as_tuple=True)
        expert_input = x_flat[tok_ids]    # [M, d] where M = tokens to this expert
        expert_out   = experts[e](expert_input)   # [M, d]
        
        # Scatter back with weights
        output.scatter_add_(0,
            tok_ids.unsqueeze(1).expand_as(expert_out),
            expert_out * weights[tok_ids, k_ids].unsqueeze(1))
    
    return output.reshape(B, T, d)

1.6 MoE in Real Models

Interview question: "Why does Mixtral use top-2, not top-1?"
Answer: Top-1 is simpler but creates extreme routing variance — small perturbations change which expert is used entirely. Top-2 provides smoother gradients and more stable training while still maintaining sparsity.

Section 2: State Space Models — The Linear Alternative to Attention

2.1 Why Attention Has a Quadratic Problem

For a sequence of length $T$:

• Attention matrix: $T \times T$ — quadratic memory, quadratic compute
• At T=32768 and FP16: 32768² × 2 bytes = 2 GB just for attention weights
• This makes long-context training extremely expensive

Can we do sequence modeling in O(T) memory and O(T) compute? State Space Models say yes.

2.2 Classical SSMs — The Continuous Foundation

State Space Models from control theory describe a dynamical system:

$$\dot{h}(t) = Ah(t) + Bx(t)$$ $$y(t) = Ch(t) + Dx(t)$$

Where:

• $x(t) \in \mathbb{R}$ — scalar input signal at time $t$
• $h(t) \in \mathbb{R}^N$ — hidden state vector (the "memory" of the system)
• $y(t) \in \mathbb{R}$ — output
• $A \in \mathbb{R}^{N \times N}$ — state transition matrix
• $B \in \mathbb{R}^{N \times 1}$ — input projection
• $C \in \mathbb{R}^{1 \times N}$ — output projection
• $D$ — skip connection scalar (often D=1 or D=0)

Concrete example: $N=2$ state, like a damped oscillator:

A = [[-0.5, 1.0],    B = [[0.0],    C = [[1.0, 0.0]]
     [-1.0, -0.5]]        [1.0]]

This system: state h = [position, velocity]
             x = external force
             y = position

The negative eigenvalues of A (-0.5 ± i) ensure the system is stable (oscillations decay).

2.3 Discretization via Zero-Order Hold (ZOH)

Computers process discrete sequences, not continuous signals. We need to convert $A, B$ to their discrete counterparts $\bar{A}, \bar{B}$ for time step $\Delta$:

$$\bar{A} = e^{\Delta A}$$ $$\bar{B} = (e^{\Delta A} - I) A^{-1} B$$

For diagonal A (which is what S4/Mamba use for efficiency), this simplifies to element-wise:

$$\bar{A}_i = e^{\Delta \cdot A_i}$$ $$\bar{B}_i = \frac{e^{\Delta \cdot A_i} - 1}{A_i} B_i$$

Why diagonal A? Computing $e^{A}$ for a dense $N \times N$ matrix is $O(N^3)$. For diagonal A, it's $O(N)$.

Stability requirement: $A_i < 0$ ensures $|\bar{A}_i| = |e^{\Delta \cdot A_i}| < 1$ — the state decays over time (like a leaky memory). S4 initializes A using the HIPPO matrix theory (approximating Legendre polynomials for long-range memory), then parameterizes $A_i = -\text{softplus}(A_\log_i)$ to maintain negativity.

Code walkthrough:

def discretize_zoh(A, B, delta):
    """
    A:     [d_state]  diagonal of continuous-time state matrix (negative)
    B:     [d_state]  input projection
    delta: scalar     time step size (learned or fixed)
    
    Returns:
    A_bar: [d_state]  discrete state transition
    B_bar: [d_state]  discrete input projection
    """
    # A_bar = exp(delta * A)    — element-wise since A is diagonal
    A_bar = torch.exp(delta * A)   
    
    # B_bar = (exp(delta*A) - 1) / A * B
    # Note: A < 0, so (A_bar - 1) < 0, and A < 0, giving B_bar > 0
    B_bar = (A_bar - 1.0) / A * B
    
    return A_bar, B_bar

# Example:
A = torch.tensor([-1.0, -2.0, -0.5])  # 3-state SSM, all negative
B = torch.tensor([1.0, 1.0, 1.0])
delta = 0.1  # time step

A_bar, B_bar = discretize_zoh(A, B, delta)
# A_bar ≈ [0.905, 0.819, 0.951]  — memory coefficients, all < 1
# B_bar ≈ [0.095, 0.181, 0.049]  — input gain

2.4 The Discrete Recurrence

Once discretized, the SSM runs as a linear recurrence:

$$h_t = \bar{A} \cdot h_{t-1} + \bar{B} \cdot x_t$$ $$y_t = C \cdot h_t + D \cdot x_t$$

For each time step $t$ we need $h_{t-1}$, so this must run sequentially in RNN mode.

But: for training, we can also run it as a convolution! Unrolling the recurrence:

$$h_1 = \bar{B} x_1$$ $$h_2 = \bar{A}\bar{B}x_1 + \bar{B}x_2$$ $$h_T = \sum_{k=1}^{T} \bar{A}^{T-k} \bar{B} x_k$$

The output $y_T = C h_T = \sum_k C \bar{A}^{T-k} \bar{B} x_k$.

Define the SSM kernel: $\bar{K} = (C\bar{B},\ C\bar{A}\bar{B},\ C\bar{A}^2\bar{B},\ \ldots,\ C\bar{A}^{T-1}\bar{B})$

Then $y = x * \bar{K}$ — a standard 1D convolution! This can be computed in $O(T \log T)$ with FFT.

Training (parallel):  y = conv(x, K)   → O(T log T)  ← convolve entire sequence at once
Inference (causal):   y_t = C*h_t      → O(1) per step ← just maintain state h

This dual-mode property is the key advantage: fast training via convolution, fast inference via recurrence.

2.5 The S4 Model

S4 (Structured State Space Sequence, Gu et al., 2021) makes SSMs practical:

1. Diagonal-plus-low-rank (DPLR) parameterization of A — can be efficiently computed
2. HIPPO initialization: A initialized to approximate memorizing the input using Legendre polynomial projections
3. Convolutional mode for training on GPUs

The S4 layer processes each feature channel independently with its own SSM:

Input:  x ∈ [B, T, d_model]
                    ↓ (process each feature channel independently)
For each feature f ∈ [0, d_model):
    SSM_f: x[:, :, f] → y[:, :, f]
    (each has its own A_f, B_f, C_f ∈ ℝ^{d_state})
Output: y ∈ [B, T, d_model]

2.6 Mamba: Selective State Spaces

The problem with S4: A, B, C are fixed for all inputs. The model can't selectively remember or forget based on content.

Mamba (Gu & Dao, 2023) makes B, C, and Δ input-dependent — "selective":

$$B_t = \text{Linear}_B(x_t)$$ $$C_t = \text{Linear}C(x_t)$$ $$\Delta_t = \text{softplus}(\text{Linear}\Delta(x_t))$$

Now A is still learned (but fixed at inference), while the way inputs are projected into state and read out changes per-token. This gives the model ability to:

• Focus: high $\Delta_t$ → strong input absorption (large $\bar{B}_t$)
• Forget: small $\Delta_t$ → nearly identity state update ($\bar{A}_t \approx I$)
• Select: B, C can attend to specific content in the input

Comparison to attention:

2.7 The Selective Scan: How Mamba is Actually Computed

At inference time: simple recurrence (O(1) per step)

At training time: naive sequential loop is O(T), but GPUs need parallelism. Mamba uses parallel prefix scan (a.k.a. "scan" or "cumulative operation"):

Parallel prefix scan for associative operation $\oplus$:

• Input: $[a_1, a_2, a_3, a_4]$
• Output: $[a_1, a_1 \oplus a_2, a_1 \oplus a_2 \oplus a_3, a_1 \oplus a_2 \oplus a_3 \oplus a_4]$

For SSM recurrence: $h_t = A_t h_{t-1} + B_t x_t$

Define pair operation: $(A_2, B_2) \oplus (A_1, B_1) = (A_2 A_1,\ A_2 B_1 + B_2)$

This is associative — can be computed in $O(\log T)$ depth with $O(T)$ work in parallel!

Mamba's actual implementation (in Triton/CUDA) uses hardware-aware recomputation during backward pass to avoid storing O(T) states.

2.8 Full Mamba Block Architecture

Input x ∈ [B, T, d_model]
         │
         ├──── in_proj ──→ x_in ∈ [B, T, d_inner]
         │                 │
         │         depthwise conv1d (causal)
         │                 │
         │           silu activation
         │                 │
         │         SelectiveSSM (Δ, B, C input-dependent)
         │                 │
         └──── in_proj ──→ z ∈ [B, T, d_inner]
                           │
                   (SSM_out * silu(z))  ← gating
                           │
                     out_proj ──→ output ∈ [B, T, d_model]

The depthwise conv provides local context (a few tokens of short-range dependency) while the SSM provides global context (entire sequence history in the state).

class MambaBlock(nn.Module):
    def __init__(self, d_model, d_state=16, d_conv=4, expand=2):
        super().__init__()
        d_inner = int(expand * d_model)
        self.d_inner = d_inner
        
        # Expand input to two streams: x_in and z
        self.in_proj = nn.Linear(d_model, 2 * d_inner, bias=False)
        
        # Causal depthwise conv for local context
        # padding = d_conv - 1 ensures output length = input length (causal)
        self.conv1d = nn.Conv1d(d_inner, d_inner, d_conv,
                                padding=d_conv-1, groups=d_inner)
        
        # SSM projections
        self.x_proj = nn.Linear(d_inner, d_state*2 + d_inner, bias=False)  # delta, B, C
        self.dt_proj = nn.Linear(d_inner, d_inner)   # refine delta
        
        # A: diagonal, initialized log-spaced for multi-scale memory
        A_log = torch.log(torch.arange(1, d_state+1).float()).repeat(d_inner, 1)
        self.A_log = nn.Parameter(A_log)
        self.D = nn.Parameter(torch.ones(d_inner))
        
        self.out_proj = nn.Linear(d_inner, d_model, bias=False)
    
    def forward(self, x):
        B, T, _ = x.shape
        
        # Split into two streams
        xz = self.in_proj(x)                # [B, T, 2*d_inner]
        x_in, z = xz.chunk(2, dim=-1)       # each [B, T, d_inner]
        
        # Causal depthwise conv (truncate to T to ensure causality)
        x_conv = self.conv1d(x_in.transpose(1, 2))[:, :, :T]  # [B, d_inner, T]
        x_act  = F.silu(x_conv).transpose(1, 2)               # [B, T, d_inner]
        
        # Selective SSM
        y = self.selective_scan(x_act)       # [B, T, d_inner]
        
        # Gating: multiply by silu(z) branch
        output = self.out_proj(y * F.silu(z))
        return output
    
    def selective_scan(self, x):
        B, T, d = x.shape
        A = -F.softplus(self.A_log)    # [d_inner, d_state], negative
        
        # Project x → delta, B_input, C_input
        delta_BC = self.x_proj(x)      # [B, T, d_state*2 + d_inner]
        delta_raw, B_in, C_in = delta_BC.split([d, self.A_log.shape[1], self.A_log.shape[1]], dim=-1)
        delta = F.softplus(self.dt_proj(delta_raw))  # [B, T, d_inner]
        
        # For each batch item, run sequential scan
        # (actual Mamba uses optimized CUDA parallel scan)
        h = torch.zeros(B, d, self.A_log.shape[1], device=x.device)
        ys = []
        for t in range(T):
            dt = delta[:, t, :, None]          # [B, d, 1]
            A_bar = torch.exp(dt * A[None])    # [B, d, d_state]
            B_bar = (A_bar - 1) / A[None] * B_in[:, t, None, :]  # [B, d, d_state]
            h = A_bar * h + B_bar
            y_t = (C_in[:, t, None, :] * h).sum(-1) + self.D * x[:, t, :]
            ys.append(y_t)
        return torch.stack(ys, dim=1)

2.9 Memory Comparison: O(T²) vs O(T)

Let's make this concrete:

At 32k tokens, attention needs 2 GB for a single head, while the SSM state is 1 MB. This is why Mamba-like architectures are the leading approach for truly long-context models (1M+ tokens).

2.10 Interview Pitfalls

Q: "What's the fundamental trade-off between MoE and SSMs?"

MoE addresses width scaling — more parameters, same compute. SSMs address temporal scaling — longer sequences, same memory. GPT-4 reportedly uses both: MoE blocks for width efficiency, plus the base transformer architecture. Future models will likely combine MoE (for capacity) with SSM or linear attention (for long context).

Q: "Why can't you just use a very large d_state in SSMs to approximate attention?"

Theoretically yes, but: (1) SSMs have a finite-impulse response limitation — with $N$ state dimensions, you can represent at most $N$ distinct "memories". Attention can attend to any of $T$ previous tokens. (2) The selectivity of attention (content-based retrieval) is more powerful than SSM's "smooth decay + input projection" model. In practice, Mamba-2 and similar hybrid models alternate between SSM layers and full attention layers to get the best of both.

Q: "How does Mixtral achieve competitive quality with a dense model 4× larger?"

The key insight is that language modeling doesn't require every parameter to be active on every token. Different "topics" or "styles" of computation naturally cluster to different experts. A math-reasoning token activates different experts than a poetry token. The model effectively partitions its capacity by input domain.

Q: "Explain load balancing loss without looking at the code."

"The naive router collapses to using one expert because that expert gets more gradient signal and becomes better, attracting more traffic — a positive feedback loop. The aux loss breaks this by penalizing the product $f_i \times P_i$: $f_i$ is the fraction of tokens going to expert $i$ (from the hard top-K selection, non-differentiable) and $P_i$ is the mean softmax probability for expert $i$ (differentiable). Their product peaks when load is uneven, so minimizing $\sum f_i P_i$ pushes the router toward uniform load. The neat trick is that even though the routing decision itself has no gradient, $P_i$ does — so we can still backpropagate."

Section 3: Quick Setup Guide

Installing dependencies

pip install torch>=2.1.0
# For actual Mamba (hardware-optimized):
pip install mamba-ssm  # requires CUDA
# Or the pure-PyTorch reference:
pip install mamba-minimal

Running your first MoE forward pass

import torch
from solution import MoEBlock

torch.manual_seed(42)
B, T, d = 2, 16, 128
moe = MoEBlock(d_model=d, d_ff=512, n_experts=4, top_k=2)

x = torch.randn(B, T, d)
out, aux_loss = moe(x)

print(f"Input:    {x.shape}")           # [2, 16, 128]
print(f"Output:   {out.shape}")         # [2, 16, 128]
print(f"Aux loss: {aux_loss.item():.4f}")  # > 1.0 if imbalanced

# Check expert utilization
from solution import compute_expert_utilization
util = compute_expert_utilization(moe.router, x)
print(f"Expert loads: {[f'{l:.2f}' for l in util['expert_load']]}")
print(f"Imbalance:    {util['load_imbalance']:.2f}x")

Running your first SSM forward pass

import torch
from solution import MambaBlock, discretize_zoh, s4_recurrence

# Test S4 recurrence
d = 32
A = -torch.ones(d)              # stable (negative diagonal)
B = torch.randn(d)
C = torch.randn(d)
A_bar, B_bar = discretize_zoh(A, B, delta=0.1)

x_seq = torch.randn(64, d)     # sequence of 64 timesteps
y_seq = s4_recurrence(x_seq, A_bar, B_bar, C)
print(f"S4 output: {y_seq.shape}")  # [64, 32]

# Test MambaBlock
mamba = MambaBlock(d_model=32, d_state=8)
x = torch.randn(1, 64, 32)     # [B, T, d_model]
y = mamba(x)
print(f"Mamba output: {y.shape}")   # [1, 64, 32]

Lab 01 — Transformer Architectural Variants

Phase: 02 — Model Architecture Mastery | Difficulty: ⭐⭐⭐⭐☆ | Time: 3–4 hours

Read ../HITCHHIKERS-GUIDE.md before starting.

What you build

Implement and compare four key attention variants used in production LLMs:

• Multi-Head Attention (MHA) — original Transformer (Vaswani et al.)
• Multi-Query Attention (MQA) — single KV head, multiple Q heads (PaLM, Falcon)
• Grouped-Query Attention (GQA) — G KV groups (Llama 2/3, Mistral)
• RoPE positional embeddings — rotary position encoding used in Llama/GPT-NeoX

Key concepts

Files

Run

pip install -r requirements.txt
python lab.py
pytest test_lab.py -v

Lab 02 — Mixture of Experts (MoE)

Phase: 02 — Model Architecture Mastery | Difficulty: ⭐⭐⭐⭐⭐ | Time: 3–4 hours

This lab implements the MoE layer used in GPT-4, Mixtral, and Switch Transformer.

What you build

• TopKRouter — linear routing with top-K expert selection and Switch Transformer auxiliary load-balancing loss
• ExpertLayer — FFN expert (w1 → GELU → w2)
• MoEBlock — routes each token to K experts, combines weighted outputs
• compute_expert_utilization — detects load imbalance

Key concepts

Files

Run

pip install -r requirements.txt
python lab.py
pytest test_lab.py -v

Lab 03 — State Space Models (S4 / Mamba)

Phase: 02 — Model Architecture Mastery | Difficulty: ⭐⭐⭐⭐⭐ | Time: 4–5 hours

Mamba (Gu & Dao, 2023) achieves linear time/memory in sequence length vs quadratic for attention.

What you build

• discretize_zoh — Zero-Order Hold discretization: Ā = exp(ΔA), B̄ = (Ā-I)/A·B
• s4_recurrence — Sequential SSM scan over a sequence
• SelectiveSSM — Mamba's input-dependent B, C, Δ (selection mechanism)
• MambaBlock — full block: in_proj + depthwise conv + SelectiveSSM + gating
• compare_memory_usage — O(T²) attention vs O(T·d_state) SSM

Key concepts

Files

Run

pip install -r requirements.txt
python lab.py
pytest test_lab.py -v

Lab 04 — ViT + DINO Self-Distillation

Phase: 02 — Model Architecture Mastery | Difficulty: ⭐⭐⭐⭐⭐ | Time: 4–6 hours

DINO (Caron et al., 2021) trains a ViT with no labels and gets attention maps that segment objects. This lab builds both the ViT and the self-distillation machinery from scratch, and makes you prove you understand why it doesn't collapse.

What you build

• PatchEmbed — the conv-as-patch-projection trick (kernel = stride = patch size)
• ViT — CLS token, learned position embeddings, pre-norm encoder blocks, plus get_last_attention for the famous segmentation-without-labels visualization
• DINOHead — MLP → L2-normalized bottleneck → prototype layer (cosine similarity against K learned prototypes)
• DINOLoss — cross-entropy between sharpened teacher and student distributions with EMA centering — the two mechanisms that jointly prevent collapse
• momentum_update — the EMA teacher (a temporal ensemble of past students)
• dino_train_step — multi-crop: teacher sees global views only, student sees all, loss pairs every teacher view with every different student view

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v               # your lab.py
LAB_MODULE=solution pytest test_lab.py -v   # sanity-check the reference

Suggested TODO order

1. PatchEmbed — get the conv-as-projection reshape right first
2. ViT.forward — CLS prepend, pos embed, blocks, norm, return position 0
3. DINOHead — note the L2 normalize before the prototype layer
4. DINOLoss — the detach and the center update are where the learning is
5. momentum_update + the DINO wrapper (teacher frozen!)
6. dino_train_step — pair views correctly, EMA after optimizer step

Success criteria

• All 13 tests pass
• You can explain, without notes, why sharpening alone or centering alone still collapses
• You can state what the teacher would output if you forgot the detach() and trained it by gradient (hint: student and teacher co-collapse to the trivial solution)
• python solution.py style training shows the loss decreasing on repeated identical crops

Full-scale extension (the original spec)

Train ViT-Tiny on CIFAR-10 for 30 epochs (≈1 h single GPU): real augmentation multi-crop, cosine momentum schedule 0.996→1.0, linear-probe accuracy ≥80% without labels, and attention map visualizations showing object segmentation. Add register tokens (DINOv2) and compare.

Interview Q&A

Q: Why does DINO need both centering and sharpening? Centering subtracts a running mean of teacher logits, preventing any prototype from permanently dominating (one-hot collapse) — but pushes toward uniform. Sharpening (low teacher temperature) pushes toward confident outputs — but alone allows one-hot collapse. They balance in opposite directions; remove either and training degenerates.

Q: Why is the teacher an EMA rather than a copy or a separately trained network? EMA gives a temporal ensemble of student checkpoints: smoother and consistently better than the student during training (the paper measures this), providing targets slightly ahead of the student — a bootstrapping effect similar to BYOL's predictor asymmetry.

Q: Why do ViT attention maps segment objects without segmentation labels? The multi-crop objective forces local-crop features to predict global-view prototypes; the CLS token must aggregate from patches that determine object identity, so its attention concentrates on the foreground object. Supervised ViTs show much weaker localization — the objective, not the architecture, produces it.

Q: What changes in DINOv2? Scale (LVD-142M curated data), iBOT-style masked-patch prediction added to the loss, KoLeo regularization, register tokens to fix attention-map artifacts, and efficient implementations — the centering/sharpening core stays.

References

• Emerging Properties in Self-Supervised Vision Transformers (DINO)
• DINOv2: Learning Robust Visual Features without Supervision
• An Image is Worth 16x16 Words (ViT)
• Vision Transformers Need Registers
• Official DINO code

Lab 05 — Multimodal Bridge (CLIP → LLM Projection)

Phase: 02 — Model Architecture Mastery | Difficulty: ⭐⭐⭐⭐⭐ | Time: 4–5 hours

LLaVA's core insight: a 2-layer MLP is all it takes to connect a frozen vision encoder to a frozen LLM. This lab builds that bridge at toy scale, where the architecture, the freezing discipline, and the loss masking are identical to the 7B-parameter version.

What you build

• Projector — the trainable 2-layer MLP mapping vision-feature space → LLM embedding space (the only component that trains; <0.25% of total parameters)
• MultimodalModel — frozen vision tower + frozen causal LM with visual tokens prepended to the text sequence
• loss — next-token cross-entropy on text positions only, with the visual-token causal shift and prompt masking (ignore_index=-100)
• generate — greedy image-conditioned decoding

Frozen TinyVisionEncoder / TinyCausalLM stand-ins are provided with the same interfaces as CLIP ViT-L/14 and TinyLlama, so the tests run in seconds with no downloads; swapping in the real models is the extension.

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v
LAB_MODULE=solution pytest test_lab.py -v

Suggested TODO order

1. Projector — two lines of architecture, but know why MLP > single linear
2. MultimodalModel.__init__ — the freezing; check test_only_projector_is_trainable
3. visual_tokens → forward — concatenation order matters
4. loss — draw the position layout on paper before coding the slice
5. generate

Success criteria

• All 12 tests pass — especially test_loss_alignment_uses_image_conditioned_prediction (the off-by-one catcher) and test_trained_model_uses_the_image (proves information actually flows image → projector → LM)
• You can state the trainable-parameter fraction and why that makes stage-1 training cheap
• You can explain what breaks if you detach() the projector output instead of the vision output (no gradient reaches the projector — silent non-training)

Full-scale extension (the original spec)

Swap stand-ins for openai/clip-vit-large-patch14 (frozen) and TinyLlama-1.1B (frozen); train the projector on the LLaVA-Instruct-150K caption subset; evaluate VQAv2 accuracy (target >50%); compare single-linear vs MLP projector. Then: Q-Former connector (BLIP-2), and measure the NPU-relevant cost — visual tokens inflate the KV cache by N_patches at every layer.

Interview Q&A

Q: Why freeze both towers in stage 1? The projector starts random — its gradient noise would destroy pretrained representations (the LM would "catastrophically adapt" to garbage visual tokens). Train the bridge until visual tokens live in the LM's embedding distribution, then optionally unfreeze for instruction tuning.

Q: Why does LLaVA prepend visual tokens rather than cross-attend (Flamingo-style)? Prepending reuses the LM unchanged (no new cross-attention layers to train), is trivially compatible with KV caching and existing serving stacks, and scales with patch count. Cross-attention isolates modalities better and keeps text-only throughput intact, at the cost of new trained layers — a classic deployment-vs-architecture tradeoff.

Q: What is the inference cost of visual tokens on an edge NPU? Each of the N visual tokens occupies KV cache in every layer (N × layers × 2 × d_kv) and participates in every attention computation — for 576 CLIP tokens this often dominates a short-question VQA prompt. Mitigations: token pruning/merging, pooled visual tokens, Q-Former-style compression to a fixed small budget.

References

• Visual Instruction Tuning (LLaVA)
• Improved Baselines with Visual Instruction Tuning (LLaVA-1.5)
• BLIP-2: Bootstrapping Language-Image Pre-training
• Flamingo: a Visual Language Model for Few-Shot Learning
• CLIP

Phase 03 — Quantization & Model Compression

Difficulty: ⭐⭐⭐⭐⭐
Estimated Time: 4 weeks (90–110 hours)
Roles Supported: Core differentiating skill — every NPU deployment requires quantization; accuracy recovery under aggressive quantization is what Senior Staff engineers own

Why This Phase Exists

This is the skill that distinguishes Qualcomm ML engineers from generic ML engineers. A Qualcomm NPU (HTP) achieves peak performance with INT8/INT4 activations and weights. Moving from FP32 → INT8 gives ~4x memory reduction and ~4–8x compute speedup, but at what accuracy cost? The answer is: it depends entirely on how you quantize.

Naive quantization (round FP32 to INT8 with global min/max clipping) easily degrades MMLU from 72% to 55% for a 7B model. State-of-the-art quantization (GPTQ + activation-aware calibration) can maintain 70.8% — barely above FP16 baseline. The difference between these outcomes is what this phase teaches.

The JD requires "expertise in model training and inference workflows" with "graph optimizations for performance and accuracy." Quantization is the single most important accuracy-vs-performance tradeoff in the NPU deployment pipeline.

Concepts

• Uniform quantization: $Q(x) = \text{round}(x / s) + z$, where scale $s = (x_{max} - x_{min}) / (2^b - 1)$ and zero-point $z$ is the integer offset
• Per-tensor vs per-channel vs per-group quantization: coarser granularity → higher error; finer granularity → more overhead; group quantization (AWQ, GPTQ) balances both
• Symmetric vs asymmetric: symmetric ($z=0$, signed ints) simpler for hardware; asymmetric handles non-zero activation distributions better (important for ReLU outputs)
• Post-Training Quantization (PTQ): quantize after training using a calibration dataset to determine scales; fast but suboptimal
• Calibration strategies: MinMax, Percentile (discard outliers), Entropy/KL-divergence (minimize information loss), MSE minimization
• Quantization-Aware Training (QAT): insert fake-quantization nodes during training; gradients flow through straight-through estimator (STE); best accuracy
• Straight-Through Estimator (STE): $\frac{\partial \text{round}(x)}{\partial x} \approx 1$ in backward pass — enables gradients to flow through the non-differentiable round operation
• GPTQ (Generative Pre-trained Transformer Quantization): layer-wise quantization using second-order (Hessian) information; solves an optimal quantization problem per-layer; 4-bit models at near-FP16 accuracy
• AWQ (Activation-Weighted Quantization): identifies salient weights (high activation magnitude) and applies per-group scaling to protect them; simpler than GPTQ, hardware-friendly
• SmoothQuant: migrates quantization difficulty from activations to weights by multiplying activations by $1/s$ and weights by $s$; activations become easier to quantize
• FP8 quantization: used in H100/H200 and Qualcomm's latest chips; E4M3 or E5M2 format; lower quantization error than INT8 for many distributions
• Knowledge Distillation: train a student to match teacher logits/features; use for accuracy recovery after aggressive quantization
• Magnitude pruning: zero out weights below threshold; reduces model size but needs sparse hardware support for speedup
• Structured pruning: remove entire channels/heads/layers; immediately reduces compute; can be combined with quantization

Labs

Lab 01 — PTQ Pipeline with Accuracy Benchmarking

Lab 02 — QAT from Scratch

Lab 03 — GPTQ: Hessian-Based Optimal Quantization

The shipped lab (lab-03-gptq/) implements GPTQ end-to-end: Hessian computation, column-wise OBQ with Cholesky updates, error propagation, and a full model pipeline with calibration hooks. AWQ is implemented in Phase 10 — Lab 01. The AWQ/SmoothQuant/Pareto sweep described below is the extension target once both labs are done.

Deliverables Checklist

• PTQ pipeline with calibration comparison table on OPT-125M / TinyLlama
• QAT implementation with STE showing accuracy > PTQ on CIFAR-10
• GPTQ + AWQ + SmoothQuant comparison with Pareto curves on LLaMA-3.2-1B
• All test_lab.py suites pass
• Pareto curve plots saved as PNG artifacts

Guides in This Phase

• HITCHHIKERS-GUIDE.md — quantization math, GPTQ, AWQ, SmoothQuant, decision tree, key numbers
• DEEP-DIVE-QAT-GPTQ.md — PTQ, QAT/STE, and GPTQ from first principles (supports Labs 02–03)

Interview Relevance

• "Explain the Straight-Through Estimator and why it works." — after Lab 02, you can draw the backward graph and explain the approximation
• "How does GPTQ minimize quantization error?" — describe the Cholesky-based optimal weight update, the per-column sequential quantization, and the Hessian estimation
• "Given a model that loses 5% accuracy under INT4, what would you try first?" — you have a decision tree: SmoothQuant (free, helps activation outliers), AWQ (salient channels), GPTQ (optimal per-layer), QAT (expensive but best)
• "What is the memory layout consideration for per-group quantization on an NPU?" — group scales must be loaded alongside weights, understand the bandwidth implications

Warmup Guide — Quantization & Model Compression

Zero-to-expert primer for Phase 03. Builds quantization from "what is a number format" through PTQ, QAT, and GPTQ — assuming only basic Python and the Phase 02 architecture vocabulary.

Table of Contents

• Chapter 1: Number Formats — What Precision Actually Is
• Chapter 2: Why Quantize — The Memory Wall
• Chapter 3: Uniform Quantization — The Core Math
• Chapter 4: Granularity — Per-Tensor, Per-Channel, Per-Group
• Chapter 5: Calibration — Choosing the Range
• Chapter 6: PTQ — The Baseline Pipeline
• Chapter 7: QAT and the Straight-Through Estimator
• Chapter 8: GPTQ — Second-Order Quantization
• Chapter 9: The Activation Outlier Problem — AWQ and SmoothQuant in Brief
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: Number Formats — What Precision Actually Is

Zero background: computers store real numbers in fixed-width binary. A floating-point number is sign × mantissa × 2^exponent: the exponent gives range, the mantissa gives precision within that range. Key formats:

The crucial difference: float spacing is relative (numbers near 0 are dense, large numbers sparse); integer spacing is absolute (uniform grid). Quantization is the act of mapping the float world onto an integer grid — and choosing where that grid sits is the entire game.

Production reality: INT8/INT4 are not just smaller — NPU/GPU integer pipelines have 2–8× the throughput of float pipelines, and (more important on edge) memory traffic drops proportionally with bit-width. FP8 (E4M3/E5M2) is the newest middle ground on H100-class and recent NPU hardware.

Chapter 2: Why Quantize — The Memory Wall

A 7B-parameter model in FP16 is 14 GB of weights. To generate one token you must read every weight once — at batch 1, LLM decoding is memory-bandwidth bound, not compute-bound (Phase 07 formalizes this with roofline analysis). So:

$$\text{tokens/sec} \approx \frac{\text{memory bandwidth}}{\text{bytes of weights}}$$

Halve the bytes → roughly double the decode speed, independent of compute. INT4 cuts a 7B model to 3.5 GB — the difference between "doesn't fit on the phone" and "fits with room for the KV cache." That is why quantization is the edge-deployment skill, and why this phase exists.

The price: every weight moves to its nearest grid point. The question of this whole phase is: how much accuracy does that cost, and how do you minimize it? Naive INT8 can take a 7B model's MMLU from 72% to 55%; well-engineered INT4 can hold it within a point. The technique gap, not the bit-width, determines the outcome.

Chapter 3: Uniform Quantization — The Core Math

The affine map. Choose a scale $s$ and zero-point $z$:

$$q = \text{clamp}!\left(\text{round}!\left(\frac{x}{s}\right) + z,; q_{min},; q_{max}\right), \qquad \hat{x} = s,(q - z)$$

• $s$ = step size of the integer grid in real units: $s = (x_{max} - x_{min})/(2^b - 1)$.
• $z$ = which integer represents real 0. Symmetric quantization fixes $z = 0$ (grid centered on zero, signed ints) — simpler hardware, no zero-point multiplies in the matmul inner loop. Asymmetric allows $z \neq 0$ — fits one-sided distributions (post-ReLU activations live in $[0, +)$; symmetric wastes half the grid on negatives).
• Convention: weights symmetric per-channel, activations asymmetric per-tensor is the standard hardware-friendly starting point.

The two error sources:

1. Rounding error: uniform in $[-s/2, s/2]$, variance $s^2/12$. Shrinks as range shrinks.
2. Clipping error: values outside $[x_{min}, x_{max}]$ saturate. Shrinks as range grows.

These trade against each other through $s$ — that tension is calibration (Chapter 5). One outlier weight of 50 in a tensor whose other values live in $[-1, 1]$ forces $s$ 50× larger, multiplying everyone else's rounding error by 50. Outliers are the villain of this entire phase.

Quantized matmul: $Y = XW \approx s_x s_w (Q_x Q_w)$ — the inner product runs in integer arithmetic (accumulating in INT32), with scales applied once at the end. This is what NPU MAC arrays actually execute.

Chapter 4: Granularity — Per-Tensor, Per-Channel, Per-Group

How many values share one $(s, z)$ pair?

• Per-tensor: one scale for the whole weight matrix. Cheapest; one outlier row ruins every row.
• Per-channel: one scale per output row of $W$. Each row's range fits its own distribution — the single biggest free win for weights; supported by essentially all inference hardware.
• Per-group: one scale per contiguous group of (typically 128) weights within a row — the GPTQ/AWQ INT4 standard. Finer adaptation; cost = storing one FP16 scale per 128 INT4 values (≈ +3% memory).

Rule of thumb: INT8 weights → per-channel suffices. INT4 weights → per-group is required. Activations → per-tensor (dynamic per-token helps but costs runtime calibration).

Chapter 5: Calibration — Choosing the Range

Run a few hundred representative samples through the model, observe activation distributions, choose $[x_{min}, x_{max}]$ per quantizer:

• MinMax: exact observed extremes. Zero clipping, maximal rounding error; one outlier ruins it. Right for weights (no sampling noise — you can see all of them).
• Percentile (99.9th): clip the tail outliers deliberately. Usually beats MinMax for activations.
• Entropy / KL-divergence (TensorRT's classic): choose the clip threshold whose quantized distribution diverges least from the original — principled rounding-vs- clipping balance.
• MSE: directly minimize $|x - \hat{x}|^2$ over candidate clip values.

Engineering reality: calibration-set choice matters as much as method — it must cover the deployment input distribution. A model calibrated on English prose mis-ranges on code tokens. Lab 01 makes you implement three methods and measure the differences instead of trusting folklore.

Chapter 6: PTQ — The Baseline Pipeline

Post-Training Quantization = freeze the trained model, calibrate, quantize, ship. No gradients, minutes not days. The pipeline you build in Lab 01:

1. Wrap target layers (QuantizedLinear replacing nn.Linear).
2. Attach observers; run calibration batches; collect ranges.
3. Compute $(s, z)$ per quantizer from the chosen calibration rule.
4. Quantize weights offline; activations quantize at runtime (or simulate with QDQ — quantize-dequantize — nodes for accuracy evaluation in float).
5. Measure: perplexity and task accuracy vs the FP16 baseline, per configuration.

When PTQ suffices: INT8 per-channel on most transformer weights — typically <0.5 PPL degradation. When it fails: INT4 without grouping; activation quantization in the presence of outliers; small models (less redundancy to absorb error); anything where the calibration set misses the deployment distribution.

Chapter 7: QAT and the Straight-Through Estimator

The idea: if quantization will distort the network, let training see the distortion and route around it. Insert fake-quantize nodes — $\hat{x} = s(\text{round}(x/s) - z)$ computed in float — into the forward pass and fine-tune.

The problem: $\text{round}(\cdot)$ has derivative 0 almost everywhere (and undefined at half-integers). Backprop through it kills all gradients.

The Straight-Through Estimator: in the backward pass, pretend round is the identity:

$$\frac{\partial,\text{round}(x)}{\partial x} \overset{\text{STE}}{=} 1 \quad (\text{within the clip range; 0 outside})$$

Mathematically false; empirically it works because round is "identity plus bounded noise" ($|round(x) - x| \le 0.5$), so the identity's gradient points in approximately the right direction, and SGD's stochasticity absorbs the bias. The clip-range masking matters: saturated values get zero gradient (they genuinely can't move the output), which also nudges weights back inside the representable range.

What QAT buys: weights migrate toward grid points and away from clip boundaries during fine-tuning — typically +1–3% accuracy over PTQ at INT8, and it's the main tool that makes INT4/INT2 activations viable. Cost: a training loop, data, and time — which is why PTQ-family methods (GPTQ/AWQ) dominate for LLMs where fine-tuning is expensive.

Lab 02 implements FakeQuant as a torch.autograd.Function with the STE backward and an observer that tracks running ranges — the same architecture as torch.ao.quantization's internals.

Chapter 8: GPTQ — Second-Order Quantization

The question GPTQ answers: when you round one weight, can you adjust the other weights to compensate? Yes — optimally, using curvature information.

Setup: quantize layer-by-layer, minimizing layer output error $|XW - X\hat{W}|^2$ over the calibration activations $X$. The curvature of this objective w.r.t. the weights of one output row is the Hessian $H = 2XX^\top$ (+ small damping $\lambda I$ for stability) — note it depends only on inputs, shared across rows.

The OBQ step, for one weight $w_q$ (from the Optimal Brain Surgeon lineage): quantize it, then update all not-yet-quantized weights in the row by

$$\delta = -\frac{w_q - \text{quant}(w_q)}{[H^{-1}]{qq}} , [H^{-1}]{:,q}$$

— spread the rounding error onto correlated weights, weighted by inverse curvature (move most where the loss is flattest).

GPTQ's algorithmic contributions that make this feasible at LLM scale:

1. Quantize columns in a fixed order for all rows (vectorizes across the matrix).
2. Maintain $H^{-1}$ updates via Cholesky factorization (numerically stable; the naive rank-1 inverse updates accumulate error and produce garbage at 175B scale).
3. Lazy-batch the updates for memory locality.

Result: a 175B model quantized to INT4/3-bit in hours on one GPU, at near-FP16 perplexity. Lab 03 implements the per-layer algorithm: Hessian from calibration hooks, column-wise quantize-and-compensate, and the comparison vs round-to-nearest that shows why the compensation matters (typically 2–5× lower reconstruction error).

Chapter 9: The Activation Outlier Problem — AWQ and SmoothQuant in Brief

Transformer activations past ~6B parameters develop systematic outlier channels — specific hidden dimensions with values 10–100× the median, consistently across tokens (an emergent property tied to attention "sink" behavior). Weights quantize easily; these activations don't.

• SmoothQuant: migrate difficulty from activations to weights with a per-channel rescale: $Y = (X,\text{diag}(s)^{-1})(\text{diag}(s),W)$ — mathematically identical, but $X' = X/s$ has tamed ranges while $W' = sW$ absorbs them (weights have headroom). Choose $s_j = \max|X_j|^\alpha / \max|W_j|^{1-\alpha}$, α ≈ 0.5.
• AWQ: observe that ~1% of weight channels are salient (they multiply the high-magnitude activations); protect them by scaling up before quantization (equivalent trick in the other direction). No Hessian, no backprop — hardware-friendly and fast. You implement AWQ fully in Phase 10 Lab 01.

These two plus GPTQ form the LLM INT4 toolbox; Chapter 8's deep dive file (DEEP-DIVE-QAT-GPTQ.md) covers the math at full depth.

Lab Walkthrough Guidance

Order: Lab 01 (PTQ) → Lab 02 (QAT) → Lab 03 (GPTQ) — each builds on the previous's vocabulary.

• Lab 01: implement quantize/dequantize and verify the round-trip error bound ($\le s/2$ per element) before touching calibration. Then observers, then the three calibration strategies, then the accuracy table. Resist evaluating before all three are in — the comparison is the lab.
• Lab 02: write the STE autograd.Function first and gradcheck your own understanding: the gradient should be 1 inside the clip range, 0 outside. Then the observer, then training. Verify QAT > PTQ on the final table.
• Lab 03: build compute_hessian and check positive-definiteness (damping!). Then single-column quantize-with-compensation on a toy 4×4 case you can verify by hand. Then the full loop, then compare_reconstruction_error — GPTQ must beat round-to-nearest clearly.

Success Criteria

You are ready for Phase 04 when you can, from memory:

1. Write the affine quantization map and derive $s$ from a range and bit-width; state the rounding-vs-clipping tension in one sentence.
2. Explain why weights are symmetric-per-channel and activations asymmetric-per-tensor in standard deployments.
3. State the STE, why it's wrong, and why it works anyway.
4. Write GPTQ's compensation update and explain each factor ($[H^{-1}]_{qq}$ in the denominator = move most where curvature is flattest).
5. Explain the activation-outlier phenomenon and the SmoothQuant rescaling identity.
6. Given "MMLU dropped 5% at INT4", produce the ordered debugging tree: granularity → calibration → SmoothQuant/AWQ → GPTQ → QAT.

Interview Q&A

Q: Why does per-channel weight quantization dominate per-tensor? Output channels of a trained layer have wildly different weight ranges (often 10–50× spread). Per-tensor forces one scale to cover the largest channel, inflating everyone else's step size. Per-channel scales are free at inference (folded into the output rescale) — better accuracy at zero runtime cost.

Q: Decode throughput is memory-bound. Walk me through the INT4 speedup math. Decode reads all weights per token: $t \approx \text{bytes}/\text{BW}$. 7B FP16 = 14 GB; at 100 GB/s that's 140 ms/token ≈ 7 tok/s. INT4 (+group scales ≈ 3.6 GB) → ~28 tok/s. Compute barely changes; we cut bytes 4×. This is also why INT4 weights with FP16 compute (dequantize-on-load) still wins on bandwidth-bound hardware.

Q: When does QAT fail to help? At extreme bit-widths (INT2) where the function space is too coarse for STE's gradient fiction; when the fine-tuning data distribution drifts from deployment; and when the bottleneck is activation outliers (a range problem QAT's weight adaptation can't fix — you need SmoothQuant-style migration first).

Q: Why does GPTQ use the Hessian rather than weight magnitude? Magnitude says how big a weight is, not how much the output changes when you perturb it — that's curvature. $H = 2XX^\top$ measures exactly the output sensitivity under the calibration distribution, so error is spread onto directions the layer is least sensitive to. Magnitude-based rounding is precisely the "naive" baseline GPTQ beats.

Q: NPU memory layout implication of group quantization? Each group's FP16 scale must be co-resident with its INT4 weights during the MAC — layouts interleave scales with weight blocks so a tile fetch brings both; misalignment doubles memory transactions. group_size is chosen to match the NPU's native tile width (often 64 or 128) for exactly this reason.

References

• Jacob et al., Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference (2017) — arXiv:1712.05877
• Nagel et al., A White Paper on Neural Network Quantization (2021) — arXiv:2106.08295 — the single best survey
• Bengio et al., Estimating Gradients Through Stochastic Neurons (STE) (2013) — arXiv:1308.3432
• Frantar et al., GPTQ: Accurate Post-Training Quantization for GPT (2022) — arXiv:2210.17323
• Hassibi & Stork, Optimal Brain Surgeon (1992) — the lineage behind GPTQ
• Dettmers et al., LLM.int8(): 8-bit Matrix Multiplication and Emergent Outliers (2022) — arXiv:2208.07339
• Xiao et al., SmoothQuant (2022) — arXiv:2211.10438
• Lin et al., AWQ: Activation-aware Weight Quantization (2023) — arXiv:2306.00978
• PyTorch quantization docs
• DEEP-DIVE-QAT-GPTQ.md — this phase's companion deep dive

Hitchhiker's Guide — Quantization & Model Compression

"Quantization is not rounding. It is finding the best integer world to map the floating-point world into."

1. Quantization Math — Precise Derivations

1.1 Affine (Asymmetric) Quantization

$$Q(x) = \text{clamp}\left(\text{round}\left(\frac{x}{s}\right) + z, ; q_{\min}, q_{\max}\right)$$

$$x_{\text{dequant}} = s \cdot (Q(x) - z)$$

Scale and zero-point derivation from observed range $[x_{\min}, x_{\max}]$:

$$s = \frac{x_{\max} - x_{\min}}{q_{\max} - q_{\min}}, \qquad z = q_{\min} - \text{round}\left(\frac{x_{\min}}{s}\right)$$

For INT8: $q_{\min} = -128, q_{\max} = 127$ (signed) or $0, 255$ (unsigned).

Symmetric quantization ($z = 0$): $s = \frac{\max(|x|)}{q_{\max}}$. Simpler hardware: multiplication by scale instead of multiply + add.

1.2 Quantization Error

The quantization error (rounding error) for value $x$ is bounded by:

$$|Q(x) - x| \leq \frac{s}{2}$$

Therefore, smaller scale → smaller error. This is why per-channel quantization (separate scale per output channel) outperforms per-tensor (single scale): each channel can have its own optimal scale.

Signal-to-Quantization-Noise Ratio (SQNR):

$$\text{SQNR} \approx 6.02 b + 4.77 \text{ dB}$$ (for Gaussian-distributed values, $b$ = number of bits)

INT8 ≈ 52.9 dB SQNR; INT4 ≈ 28.9 dB SQNR. This is a 24 dB drop — a factor of ~250 in noise power.

1.3 Per-Group Quantization

In group quantization (used by GPTQ, AWQ, llama.cpp), weights are quantized in groups of $G$ elements sharing one (scale, zero-point):

$$Q(W_{i, j:j+G}) = \text{round}(W_{i, j:j+G} / s_{i,j})$$

Trade-off: $G=128$ → 1 extra FP16 scale per 128 INT4 weights = 3.1% overhead, but ~2-3 dB SQNR improvement over per-tensor.

Hardware consideration for Qualcomm HTP: HTP natively supports INT4 with per-group scales at G=32, G=64, G=128. G=32 is most accurate but requires more scale storage; G=128 is most hardware-friendly.

2. GPTQ — Hessian-Based Optimal Quantization

2.1 The Core Idea

GPTQ frames weight quantization as an optimization problem. For a layer with weight $W$ and input activations $X$:

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

subject to $\hat{W}$ being quantized.

The optimal update when quantizing weight $w_q$ to $\hat{w}_q$ is:

$$\delta W = -\frac{w_q - \hat{w}q}{[H^{-1}]{qq}} \cdot H^{-1}_{:,q}$$

where $H = 2X X^T$ is the (approximated) Hessian of the squared error.

Key insight: after quantizing one weight, update all remaining unquantized weights to compensate for the introduced error. This is the same idea as "least squares with constraints" — each quantization decision corrupts the loss function, and GPTQ corrects for it greedily.

2.2 Cholesky Implementation

Direct Hessian inversion is $O(d^3)$ — expensive for large layers. GPTQ uses the Cholesky decomposition trick: process columns in blocked order using triangular solve, achieving $O(d^2)$ amortized.

def gptq_quantize_layer(W, X, bits=4, group_size=128, block_size=128):
    """
    W: (out_features, in_features)
    X: (in_features, n_samples) — calibration activations
    Returns: W_quantized, scales, zeros
    """
    d = W.shape[1]
    H = 2 * (X @ X.T)  # (in_features, in_features)
    
    # Add damping for numerical stability
    H += 0.01 * torch.diag(H).mean() * torch.eye(d, device=H.device)
    
    # Cholesky decomposition for stable inversion
    L = torch.linalg.cholesky(H)
    H_inv = torch.cholesky_inverse(L)
    
    W_q = W.clone()
    scales = []
    zeros = []
    
    for i in range(0, d, block_size):
        block_end = min(i + block_size, d)
        for j in range(i, block_end):
            # Find quantization params for this column group
            g_start = (j // group_size) * group_size
            w_col = W_q[:, j]
            scale, zero = compute_qparams(w_col, bits)
            
            # Quantize this column
            w_q = quantize(w_col, scale, zero, bits)
            
            # Compute error
            err = (w_col - w_q) / H_inv[j, j]
            
            # Update remaining columns to compensate
            W_q[:, j+1:block_end] -= err.unsqueeze(1) * H_inv[j, j+1:block_end].unsqueeze(0)
            W_q[:, j] = w_q
    
    return W_q, scales, zeros

3. AWQ — Activation-Weighted Quantization

3.1 Salient Weight Detection

AWQ's key observation: not all weights are equally important. Weights corresponding to input channels with large activation magnitudes are disproportionately impactful on model accuracy.

$$\text{importance}(j) = \mathbb{E}{x \sim D}[|x_j|] \cdot |W{:,j}|_2$$

For the top-1% most important channels, quantization error multiplies by the (large) activation magnitude. For the bottom 99%, quantization error is scaled down.

3.2 The Scaling Trick

For input channel $j$ with scale $s_j$:

$$y = xW = (x \cdot s) \cdot (W / s) = x' \cdot W'$$

By choosing $s_j > 1$ for salient channels: $W'{:,j} = W{:,j} / s_j$ → smaller magnitude → less quantization error (fits better in the quantization grid). $x'_j = x_j \cdot s_j$ → but this is a multiplication of activations, which can be absorbed into the previous layer's weights (no extra compute).

Critical: the scale $s_j$ is chosen to minimize quantization error of $W'_{:,j}$ on the quantization grid, not based on a fixed formula. AWQ uses grid search: $s_j \in {1, 2, 4, 8, \ldots}$.

4. SmoothQuant — Migrating Difficulty from Activations to Weights

4.1 The Outlier Problem

LLMs have severe activation outliers: certain token positions and embedding dimensions have values 100× larger than typical. These outliers destroy per-tensor activation quantization (the scale becomes too large, compressing most values to a small range of integers).

Example: Activation channel 1053 in OPT-13B might have max value 3847 while 99% of channels are below 4. Per-tensor scale = 3847/127 = 30.3. Channel 0 with value 1.2 → quantizes to round(1.2/30.3) = 0. Massive quantization error.

4.2 The Per-Channel Scale Migration

For matrix multiply $y = xW$, introduce per-input-channel scales $s_j$:

$$y = x \cdot W = \underbrace{(x / s)}{\text{easy to quantize}} \cdot \underbrace{(W \cdot s)}{\text{absorb into weights}}$$

Scale is chosen as:

$$s_j = \frac{\max(|X_{:,j}|)^\alpha}{\max(|W_{j,:}|)^{1-\alpha}}, \quad \alpha = 0.5 \text{ (default)}$$

This migrates the quantization challenge from activations (hard to quantize, dynamic) to weights (easy to quantize, static, absorb offline).

Result: activation quantization improves dramatically; weight quantization barely worsens (weights are already well-distributed).

5. QAT and the Straight-Through Estimator

5.1 Why the Round Function Breaks Gradients

forward:  x → round(x/s) * s → quantized_x
backward: d_loss/d_x = d_loss/d_quantized_x × d_quantized_x/d_x
                                                = 0 (everywhere)

round() has derivative 0 almost everywhere (flat function with unit jumps). Gradient is zero → no learning signal. QAT is stuck.

5.2 STE: The Principled Hack

The Straight-Through Estimator approximates:

$$\frac{\partial \text{round}(x)}{\partial x} \approx \mathbf{1}{q{\min} \leq x \leq q_{\max}}$$

In practice, this means: just pass the gradient through as if quantization didn't happen (but zero it outside the quantization range).

class FakeQuantize(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, scale, zero_point, quant_min, quant_max):
        # Actual quantization (for forward accuracy)
        x_int = torch.round(x / scale) + zero_point
        x_int = torch.clamp(x_int, quant_min, quant_max)
        x_dequant = (x_int - zero_point) * scale
        # Save mask for backward (STE clips gradient outside range)
        ctx.save_for_backward(x, torch.tensor([quant_min * scale]),
                              torch.tensor([quant_max * scale]))
        return x_dequant
    
    @staticmethod
    def backward(ctx, grad_output):
        x, x_min, x_max = ctx.saved_tensors
        # STE: pass gradient through, clip outside quantization range
        mask = (x >= x_min) & (x <= x_max)
        grad_input = grad_output * mask.float()
        return grad_input, None, None, None, None

Why does it work despite being wrong? The approximation error in the gradient is small relative to the overall optimization landscape. The model learns to arrange its weights such that most values are in the quantization-friendly zone — which is exactly the goal.

6. Quantization Decision Tree (Interview Cheatsheet)

Given: model M, target device D, accuracy budget A, latency budget L

1. Can you afford QAT? (requires retraining time + labeled data)
   YES → QAT (best accuracy, especially for INT4)
   NO → go to 2

2. Are activation outliers severe? (check: max(|X|) / median(|X|) > 100)
   YES → Apply SmoothQuant first (free, offline)
   Then → go to 3

3. Target bitwidth?
   INT8 → PTQ with per-channel KL-divergence calibration (usually <0.5 PPL loss)
   INT4 → AWQ if latency matters, GPTQ if accuracy matters, both if you have time

4. Accuracy still insufficient after step 3?
   → Add Knowledge Distillation: quantized model learns from FP16 teacher
   → Try mixed-precision: INT8 for sensitive layers (first/last, embedding), INT4 for others
   → Increase group size (G=32 vs G=128)

5. Deploy to Qualcomm HTP?
   → Use QNN SDK quantization workflow on top of your quantized model
   → Enable HTP-specific optimizations (HTA = Hexagon Tensor Accelerator weight layout)

7. Key Accuracy Numbers to Know

8. Resources

• GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers (Frantar et al., 2022)
• AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration (Lin et al., 2023)
• SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models (Xiao et al., 2022)
• LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale (Dettmers et al., 2022) — the mixed-precision INT8 approach used by bitsandbytes
• Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference (Jacob et al., 2018) — the Google QAT paper
• PyTorch Quantization Documentation
• Qualcomm AI Hub — Quantization Guide

👨🏻 Brother Talk — Phase 03, Off the Record

Quantization — the phase where accuracy goes to die, and where a Senior Staff engineer earns their title by keeping it alive. This is the heart of the Qualcomm job.

Brother, if there's one phase that is the job, it's this one. "Model Accuracy and AI Performance" — quantization is exactly where accuracy and performance collide. The whole reason Qualcomm exists in this space is to run big models on small power budgets, and the only way to do that is to drop precision. Every. Single. Model. Gets quantized before it ships to a phone. So the engineer who can quantize a model to INT8 (or INT4!) without tanking accuracy is the engineer the whole org depends on. This is your bread and butter. Master it.

Here's the first truth that reframes everything: quantization is not rounding. It's finding the best integer world to map the floating-point world into. The scale and zero-point aren't trivia — they ARE the design. Get the range wrong (let one outlier blow up max), and your scale is huge, your steps are coarse, and everything important gets crushed into a few levels. The entire art is: pick the clipping range that minimizes the error on the values that matter. Per-tensor vs per-channel, symmetric vs asymmetric, percentile clipping vs min/max — these are all answers to "what range do I quantize over?" Build the PTQ pipeline and watch what a single outlier does to your accuracy. That's the lesson.

Now the bug I want tattooed on your brain, because I literally just fixed a version of it in this track's QAT lab: the dtype you store the integer in must hold the range you quantize to. 8-bit unsigned quantization lives in [0, 255]. If you cast that to int8 (which is [-128, 127]) before you dequantize, every value above 127 wraps to negative and your weights become garbage — the model's output explodes. It's a one-line mistake and it produces a catastrophic-looking failure that looks like a deep numerical problem but is actually a type-overflow footgun. This is the texture of real quantization work: tiny, mechanical mistakes with huge blast radius. Be paranoid about ranges and dtypes.

The PTQ vs QAT decision is a core Senior Staff judgment call, so internalize it: PTQ (post-training quantization) is fast, needs only a little calibration data, and is your first move — often it's enough, especially at INT8. QAT (quantization-aware training) simulates quantization during training (fake-quant in the forward, straight-through estimator for the gradient) so the model learns to be robust to it — you reach for it when PTQ drops too much accuracy, especially at INT4 or for sensitive layers. The mistake juniors make is jumping to QAT (expensive, needs the training pipeline) when a better PTQ calibration would've done it. The principal move is: try PTQ hard first, reach for QAT when the accuracy budget forces you.

GPTQ is the one that'll make you valuable for LLMs specifically. It's not just "round the weights" — it's a second-order method that quantizes weights column by column while compensating the remaining weights for the error introduced, using the Hessian (approximated from calibration data). That error-compensation is why GPTQ holds accuracy at 4-bit where naive rounding falls apart. Understand the intuition — "quantize, measure the damage, push the correction into the not-yet-quantized weights" — and you understand the whole modern weight-only quantization family (GPTQ, AWQ, which you'll meet again in P10).

The thing nobody tells you about quantization accuracy: it's not uniform across the model. Some layers (often the first, the last, the attention output projections, the layernorms) are sensitive — quantizing them costs you dearly — while others shrug it off. Mixed precision (keep the sensitive layers high, crush the rest) is where the real accuracy/size wins live. The sensitivity analysis — quantize one layer at a time, measure the accuracy hit, rank them — is a technique that feels tedious and is actually the highest-leverage thing you can do. The Pareto frontier of "which layers in what precision" is a Senior Staff deliverable.

And the part specific to this job: hardware quantization is not textbook quantization. The NPU has opinions — it wants per-tensor or per-channel in specific layouts, specific bit-widths, specific symmetric/asymmetric conventions, and it has a fixed-point accumulator that can overflow if you're not careful. The quantization that's mathematically optimal might be the one the hardware can't run, or runs slowly. P06 makes this concrete; here, just plant the seed: always ask "what quantization does the silicon actually want?"

Career truth, brother: quantization expertise is rare and in brutal demand, because it sits at the intersection of ML math, numerical precision, and hardware — most people are strong in at most one. The engineer who can take a model from FP32 to INT4 with a 1% accuracy drop and explain every decision is writing their own ticket. This is the phase to go deep, be paranoid about ranges and dtypes, and build the sensitivity intuition that makes you the person they bring the hard models to.

Build the PTQ pipeline. Break it with an outlier. Fix the QAT overflow yourself if you want to feel it. Then come to P04, where we stop hand-optimizing and let the compiler do it — torch.compile and the FX graph.

— your brother 👨🏻

Deep Dive — PTQ, QAT & GPTQ from First Principles (Phase 03)

"PTQ gets you to INT8 for free. QAT gets you to INT4 with effort. GPTQ gets you to INT4 with math."

Section 1: Why Quantization Exists — The Memory Wall

A 7B parameter model in FP32 = 7 × 10⁹ × 4 bytes = 28 GB VRAM. A single A100 has 80 GB; an iPhone 15 Pro has 8 GB. Without quantization, large models simply don't fit on the hardware where they need to run.

Quantization converts floating-point weights (and optionally activations) to lower-bit integers:

The difference between "0.5%" and "15%" INT4 accuracy loss is entirely in the algorithm. That's why QAT and GPTQ exist.

Section 2: Post-Training Quantization (PTQ) — The Baseline

2.1 Symmetric Quantization Math

Symmetric per-tensor (the simplest form):

Given weights $W \in \mathbb{R}^{m \times n}$ and target bit width $b$:

$$\text{max_int} = 2^{b-1} - 1 \quad (\text{e.g., 127 for INT8, 7 for INT4})$$

$$\text{scale} = \frac{\max(|W|)}{\text{max_int}}$$

$$W_{\text{int}} = \text{round}\left(\frac{W}{\text{scale}}\right) \text{ clamped to } [-\text{max_int}, \text{max_int}]$$

$$W_{\text{deq}} = W_{\text{int}} \times \text{scale}$$

Example with INT4 (max_int = 7):

W = tensor([[0.23, -1.45, 0.07, 2.10],
            [-0.89,  0.51, 1.33, -0.42]])

abs_max = 2.10
scale   = 2.10 / 7 = 0.300

W_int   = round(W / 0.300) = [[1, -5,  0,  7],
                                [-3,  2,  4, -1]]  # clamped to [-7, 7]

W_deq   = W_int * 0.300    = [[0.30, -1.50, 0.00, 2.10],
                                [-0.90,  0.60, 1.20, -0.30]]

Error   = W - W_deq         = [[-0.07, 0.05, 0.07, 0.00],
                                 [0.01, -0.09, 0.13, -0.12]]

The maximum error is bounded by scale/2 = 0.15 (half the quantization step).

2.2 Asymmetric Quantization

Symmetric works when weights are centered around zero. Activations (after ReLU, GELU) are often non-negative — half the INT range would be wasted.

Asymmetric per-tensor:

$$\text{n_levels} = 2^b \quad (\text{e.g., 256 for INT8})$$

$$\text{scale} = \frac{\max(W) - \min(W)}{\text{n_levels} - 1}$$

$$\text{zero_point} = \text{round}\left(\frac{-\min(W)}{\text{scale}}\right) \text{ clamped to } [0, \text{n_levels}-1]$$

$$W_{\text{int}} = \text{round}\left(\frac{W}{\text{scale}} + \text{zero_point}\right) \text{ clamped to } [0, \text{n_levels}-1]$$

$$W_{\text{deq}} = (W_{\text{int}} - \text{zero_point}) \times \text{scale}$$

Example with INT8 and non-negative activations:

Activations after ReLU: min=0.0, max=3.2
scale = (3.2 - 0.0) / 255 = 0.01255
zero_point = round(-0.0 / 0.01255) = 0

For value 1.6:
  W_int = round(1.6 / 0.01255 + 0) = round(127.5) = 128
  W_deq = (128 - 0) * 0.01255 = 1.606
  Error = 0.006  — very small

With symmetric: min=0 would map to -127, wasting half the range.

2.3 Per-Tensor vs Per-Channel

Per-tensor: one scale for entire weight matrix. Simple but inaccurate for weights with very different magnitudes across output channels.

Per-channel: one scale per output channel (row of weight matrix). Much better for weights:

# Weight matrix W: [out_features, in_features]
# Per-channel: compute max abs per row
max_per_channel = W.abs().max(dim=1).values  # [out_features]
scale = max_per_channel / 127.0              # [out_features]

# Quantize row-by-row (broadcast)
W_int = (W / scale.unsqueeze(1)).round().clamp(-127, 127)
W_deq = W_int * scale.unsqueeze(1)

Why this matters: A transformer's linear layers often have output channels spanning 0.01–10.0 in weight magnitude. Per-tensor scale = 10.0/127 = 0.079, so small-magnitude channels are quantized with massive relative error. Per-channel fixes this.

Section 3: Quantization-Aware Training (QAT)

3.1 The Core Problem with PTQ

PTQ quantizes after training is complete. The model was never trained to be robust to quantization noise. This is fine for INT8 but catastrophic for INT4:

INT8 quantization noise ~ 0.5% of weight range → small perturbation
INT4 quantization noise ~ 7%  of weight range → very large perturbation

At INT4, each weight is rounded to one of only 16 values. A model trained in FP32 never experienced this level of discretization and its outputs change dramatically.

Solution: Train the model while simulating quantization. The model learns to be robust.

3.2 The Straight-Through Estimator (STE)

The quantization function $Q(x) = \text{round}(x/s) \times s$ has zero gradient almost everywhere and undefined gradient at step discontinuities. We can't backprop through it.

STE trick (Hinton, Bengio et al.): approximate $\partial Q(x)/\partial x \approx 1$ (pass gradient straight through):

$$\text{FakeQuantize}(x) = x + \underbrace{(Q(x) - x)}{\text{quantization error}} \cdot \underbrace{\mathbf{1}}{\text{stop-gradient}}$$

In PyTorch:

def fake_quantize(x, scale, zero_point, n_levels):
    """Forward: quantize + dequantize. Backward: identity gradient."""
    # Quantize
    x_int = (x / scale + zero_point).round().clamp(0, n_levels - 1)
    # Dequantize
    x_dq  = (x_int - zero_point) * scale
    
    # STE: x_out = x + (x_dq - x).detach()
    # Forward:  x_out = x_dq  (quantized values)
    # Backward: ∂x_out/∂x = ∂x/∂x = 1  (gradient flows through unchanged)
    return x + (x_dq - x).detach()

Why this works: During forward pass, the model sees quantized values — it experiences the quantization error and must produce correct output despite it. During backward pass, the gradient flows as if quantization didn't happen — weights still get useful gradient signal to improve.

Concrete gradient walkthrough:

import torch

x = torch.tensor([1.7], requires_grad=True)
scale = 0.5
zero_point = 0
n_levels = 256

# Forward
x_int = (x / scale + zero_point).round().clamp(0, 255)  # = round(3.4) = 3
x_dq  = (x_int - zero_point) * scale                    # = 3 * 0.5 = 1.5
x_out = x + (x_dq - x).detach()                         # = x + (1.5 - 1.7).detach()
                                                          # = x + (-0.2)  [no grad to -0.2]
# x_out = 1.5 (numerically)

loss = x_out ** 2                                        # = 2.25
loss.backward()

print(x.grad)  # = 2 * x_out * dx_out/dx
               # = 2 * 1.5 * 1   ← STE: dx_out/dx = 1
               # = 3.0
               # NOT: 2 * 1.5 * 0 = 0 (what you'd get from round)

Without STE: x.grad = 0.0 → no learning signal. With STE: x.grad = 3.0 → normal gradient.

3.3 QAT Training Loop

def train_qat(model, train_data, n_epochs=3, lr=1e-3, bits=8):
    """
    Full QAT training loop.
    
    The key: wrap model BEFORE creating optimizer.
    Then train normally — FakeQuantize handles the rest.
    """
    # Step 1: Replace all Linear layers with QATLinear (FakeQuantize on weights)
    qat_model = wrap_model_for_qat(model, bits=bits)
    
    # Step 2: Create optimizer — ALL parameters trainable
    # (the base weight learns to minimize error despite quantization)
    optimizer = torch.optim.AdamW(qat_model.parameters(), lr=lr)
    
    losses = []
    for epoch in range(n_epochs):
        epoch_loss = 0.0
        for x, y_true in train_data:
            optimizer.zero_grad()
            
            # Forward: weights go through FakeQuantize
            #   → model sees INT8-like values
            #   → must produce accurate output
            y_pred = qat_model(x)
            loss   = F.mse_loss(y_pred, y_true)
            
            # Backward: STE allows gradients to flow
            #   → weights update to reduce quantization sensitivity
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        
        losses.append(epoch_loss / len(train_data))
    
    return qat_model, losses

What the model learns during QAT:

1. Weight smoothing: weights tend toward values that quantize well (multiples of the scale)
2. Reduced sensitivity: activations become less sensitive to weight perturbations (the model learns to be more robust)
3. Calibrated ranges: the FakeQuantize scale adapts to actual weight distribution

3.4 Converting QAT Model to INT8

After QAT training, the model's weights are FP32 but trained to be robust to INT8 quantization. To actually save memory, we need to extract the integer weights:

def convert_to_int8(qat_model):
    """
    After QAT training:
    1. For each QATLinear layer:
       a. Read current weight (FP32, but quantization-friendly)
       b. Compute scale/zero_point from actual min/max
       c. Store W_int (INT8 tensor)
    2. Replace QATLinear → nn.Linear with dequantized weights
       (in real deployment: use true INT8 GEMM kernel)
    """
    converted = copy.deepcopy(qat_model)
    quant_info = {}
    
    for name, module in list(converted.named_modules()):
        if isinstance(module, QATLinear):
            W = module.weight.data
            fq = module.fake_q
            
            # Compute quantization parameters
            min_val, max_val = W.min(), W.max()
            scale = (max_val - min_val) / (fq.n_levels - 1)
            zp    = (-min_val / scale).round().clamp(0, fq.n_levels-1).int()
            
            # Quantize to integers
            W_int = (W / scale + zp).round().clamp(0, fq.n_levels-1).to(torch.int8)
            
            # Dequantize (for our CPU reference implementation)
            W_deq = (W_int.float() - zp) * scale
            
            # Save quantization metadata
            quant_info[name] = QuantizedWeight(W_int, scale.item(), zp.item(), fq.bits)
            
            # Replace module with standard Linear using dequantized weights
            # (in production: use torch.ao.nn.quantized.Linear with INT8 GEMM)
            new_linear = nn.Linear(W.shape[1], W.shape[0], bias=module.bias is not None)
            new_linear.weight = nn.Parameter(W_deq)
            # ... set on parent module
    
    return converted, quant_info

Memory layout in true INT8 deployment:

FP32 Linear weight: [4096, 4096] × 4 bytes = 64 MB
INT8 Linear weight: [4096, 4096] × 1 byte  = 16 MB  ← 4× smaller
Scale: [4096] × 4 bytes = 16 KB (per-channel scale, negligible)

During inference:
  x: FP16  → INT8 quantize
  W: INT8  (stored)
  GEMM:    INT8 × INT8 → INT32 accumulator  (fast on Qualcomm HTP/CUDA INT8 cores)
  Output:  INT32 → dequantize → FP16

3.5 QAT vs PTQ: When to Use Each

Section 4: GPTQ — Second-Order Quantization

4.1 The Core Insight

PTQ quantizes each weight independently: $W_q[i,j] = \text{round}(W[i,j] / s) \times s$.

This ignores correlations. But the reconstruction objective is:

$$\min_{W_q} |W \cdot X - W_q \cdot X|_F^2$$

We want the output of the layer to be close, not the weights themselves. Some weight perturbations cancel out, others amplify. GPTQ uses the Hessian of this objective to make smarter quantization decisions.

4.2 The Hessian of Layer Reconstruction

For the layer output $WX$ (W is the weight matrix, X is the input data), the reconstruction loss is:

$$L(W_q) = |WX - W_q X|_F^2$$

The gradient: $\nabla_{W_q} L = -2(W - W_q)X^T$
The Hessian: $H = 2XX^T$

This Hessian captures how much the output changes when we perturb each weight. Large eigenvalues = that direction of weight space strongly affects output. Small eigenvalues = we can change those weights without hurting accuracy.

def compute_hessian(X):
    """
    X: [N, d_in]  — N examples of layer inputs
    Returns H: [d_in, d_in]
    """
    # H = 2 * X^T @ X  (in the form X is [d_in, N] convention)
    Xt = X.T.float()       # [d_in, N]
    H  = 2.0 * Xt @ Xt.T  # [d_in, d_in]
    
    # Numerical stabilization: add damping
    # dampening = λ_mean * 0.01 prevents near-singular H
    damp = 0.01 * H.diag().mean()
    H += damp * torch.eye(H.shape[0])
    
    return H

# Example:
# Network input x with batch=100 samples, d_in=32
X = torch.randn(100, 32)
H = compute_hessian(X)
# H is 32×32 positive definite symmetric matrix
# H[i,i] ≈ 2 * sum_k x_k[i]^2 = variance of input dimension i

Intuition: $H_{ii}$ tells us how sensitive the reconstruction loss is to weight column $i$. If the $i$-th input feature is always near-zero (small $H_{ii}$), quantizing that column's weights doesn't matter much. If $H_{ii}$ is large (that feature is active and important), quantization error there hurts a lot.

4.3 Optimal Brain Quantizer (OBQ) — The Algorithm

Problem: Given weight matrix $W \in \mathbb{R}^{d_{out} \times d_{in}}$ and Hessian $H$, find quantized $W_q$ minimizing $|WX - W_q X|_F^2$.

OBQ idea (one column at a time):

For column $j$ (processing in order $j = 0, 1, \ldots, d_{in}-1$):

1. Quantize column $j$: $w_q^{(j)} = \text{round}(W[:, j] / s) \times s$
2. Compute quantization error: $\delta^{(j)} = W[:, j] - w_q^{(j)}$
3. Update remaining unquantized columns to compensate: $$W[:, j'] \leftarrow W[:, j'] - \delta^{(j)} \cdot \frac{H_{j, j'}^{-1}}{H_{jj}^{-1}}$$

The update step propagates the error to remaining columns in a way that minimizes the total reconstruction loss. This is the "brain surgery" analogy from OBS (Optimal Brain Surgeon): removing (quantizing) one weight, then adjusting others to compensate.

4.4 GPTQ: OBQ Made Practical

OBQ processes columns in arbitrary order (choosing the column with minimum quantization error each time) — $O(d_{in}^2)$ work per column, $O(d_{in}^3)$ total. Too slow for LLMs.

GPTQ simplification:

1. Process columns left-to-right (no selection overhead)
2. Use Cholesky decomposition of $H^{-1}$ for numerically stable updates

def gptq_quantize_layer(W, H, bits=4):
    """
    W: [d_out, d_in]  weight matrix
    H: [d_in, d_in]   Hessian from calibration data
    """
    W = W.clone().float()
    d_out, d_in = W.shape
    
    # Step 1: Compute H^{-1} via Cholesky decomposition
    # torch.linalg.cholesky gives L such that H = L @ L^T (lower triangular)
    try:
        L = torch.linalg.cholesky(H)
        I = torch.eye(d_in, device=H.device)
        H_inv = torch.cholesky_solve(I, L)  # H^{-1} = L^{-T} L^{-1}
    except:
        H_inv = torch.linalg.inv(H)  # fallback
    
    W_q = W.clone()
    
    # Step 2: Process columns left-to-right
    for j in range(d_in):
        # Quantize column j
        w_col   = W[:, j]                    # [d_out]
        w_q_col = quantize_symmetric(w_col, bits)
        W_q[:, j] = w_q_col
        
        error   = w_col - w_q_col            # [d_out]  quantization error
        h_inv_j = H_inv[j, j+1:]             # [d_in - j - 1]  off-diagonal
        h_inv_jj = H_inv[j, j].clamp(min=1e-8)  # diagonal element
        
        # Propagate error to remaining columns:
        # W[:, j+1:] -= error.outer(h_inv_j) / h_inv_jj
        if j + 1 < d_in:
            W[:, j+1:] -= (error.unsqueeze(1) * h_inv_j.unsqueeze(0)) / h_inv_jj
    
    return W_q

Step-by-step walkthrough with 1D output (d_out=1) and 4 input features:

W     = [0.5, 1.2, -0.3, 0.8]        # original weights
H_inv = [[0.4, -0.1,  0.0,  0.2],    # inverse Hessian
          [-0.1, 0.3,  0.1, -0.1],
          [0.0,  0.1,  0.5,  0.1],
          [0.2, -0.1,  0.1,  0.6]]

bits=4, scale = max(|W|)/7 = 1.2/7 ≈ 0.171

Step j=0:
  w[0] = 0.5
  w_q[0] = round(0.5/0.171)*0.171 = round(2.92)*0.171 = 3*0.171 = 0.513
  error = 0.5 - 0.513 = -0.013
  H_inv[0, 1:] = [-0.1, 0.0, 0.2],  H_inv[0,0] = 0.4
  Update: W[1:] -= (-0.013) * [-0.1, 0.0, 0.2] / 0.4
        W[1]  += (-0.013) * (-0.1) / 0.4 = +0.00325  → 1.203
        W[2]  += (-0.013) * (0.0)  / 0.4 = 0         → -0.300
        W[3]  += (-0.013) * (0.2)  / 0.4 = -0.0065   → 0.7935

Step j=1:
  w[1] = 1.203 (already adjusted!)
  w_q[1] = round(1.203/0.171)*0.171 = 7*0.171 = 1.197  (now 1.197 ≈ 1.203)
  error = 1.203 - 1.197 = 0.006  ← much smaller error after adjustment!
  ... (continue)

The key: after compensating, the next column's quantization error is smaller — errors are propagated forward and partially cancelled.

4.5 Why Cholesky?

Direct matrix inversion is numerically unstable and $O(n^3)$ every time. Cholesky gives:

1. Numerical stability: Cholesky only works for positive-definite matrices (which $H$ is, after damping), and the triangular solve is well-conditioned
2. Efficiency: compute $H = LL^T$ once ($O(n^3)$), then triangular solves are $O(n^2)$
3. Block processing: GPTQ processes weights in blocks of 128 columns, recomputing a local Cholesky for each block — this reduces numerical drift

# Why this works:
# H @ H_inv = I
# cholesky(H) = L  (lower triangular, L @ L.T = H)
# cholesky_solve(I, L) solves:  H @ X = I  →  X = H^{-1}
# 
# This is more stable than: H_inv = torch.linalg.inv(H)
# because inv() suffers from floating-point catastrophic cancellation
# for near-singular H

L = torch.linalg.cholesky(H)
H_inv = torch.cholesky_solve(torch.eye(n), L)

4.6 GPTQ vs AWQ vs Naive Quantization: Benchmark

For Llama-2-7B at INT4, on WikiText-2 perplexity:

GPTQ's advantage: purely post-training, needs only ~128 calibration samples, takes ~4 hours for 7B model on a single A100. No gradient computation needed.

4.7 Full Pipeline: Calibration + GPTQ + Deployment

import torch, copy
from transformers import AutoModelForCausalLM, AutoTokenizer

# 1. Load model
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# 2. Collect calibration data
def get_calibration_data(tokenizer, n_samples=128, seq_len=512):
    """Use a diverse sample of real text for calibration."""
    # In practice: C4, WikiText, or domain-specific data
    texts = ["The quick brown fox..."] * n_samples
    inputs = tokenizer(texts, return_tensors="pt", max_length=seq_len,
                      truncation=True, padding=True)
    return inputs["input_ids"]

cal_data = get_calibration_data(tokenizer)

# 3. Run GPTQ layer-by-layer
def quantize_model_gptq(model, cal_data, bits=4):
    """Collect inputs via hooks, then quantize each layer."""
    quantized_model = copy.deepcopy(model)
    
    for name, module in quantized_model.named_modules():
        if isinstance(module, nn.Linear):
            # Register hook to capture inputs during calibration pass
            inputs_collected = []
            hook = module.register_forward_hook(
                lambda m, inp, out: inputs_collected.append(inp[0].detach()))
            
            # Run calibration
            with torch.no_grad():
                model(cal_data)
            hook.remove()
            
            # Stack calibration inputs and compute Hessian
            X_cal = torch.cat([x.reshape(-1, x.shape[-1]) for x in inputs_collected])
            H = compute_hessian(X_cal)
            
            # Quantize this layer's weights
            W_q = gptq_quantize_layer(module.weight.data, H, bits=bits)
            module.weight.data = W_q
    
    return quantized_model

quantized_model = quantize_model_gptq(model, cal_data, bits=4)

# 4. Evaluate
def evaluate_perplexity(model, tokenizer, text):
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        loss = model(**inputs, labels=inputs["input_ids"]).loss
    return torch.exp(loss).item()

4.8 Interview Pitfalls

Q: "Explain STE without jargon to a PM."

"We train the model as if the weights were quantized to 8-bit integers — every forward pass uses the rounded values so the model learns to deal with it. But for the backward pass (updating weights), we pretend the quantization step wasn't there and use the normal gradient as if the weights were still floating-point. This trick — called Straight-Through Estimator — lets us train the model to be robust to quantization noise even though we can't actually differentiate through the rounding operation."

Q: "Why does GPTQ need calibration data? Can't you just quantize without it?"

"GPTQ uses the Hessian $H = 2XX^T$ to determine which weights are sensitive to quantization (large $H_{ii}$ = that input dimension is highly active, so quantizing that column's weights hurts more). Without calibration data, we don't know which directions in weight space are important. With 128 random samples, we get a good approximation of the typical input distribution. Without it, we'd have to treat all weights equally — that's naive round-to-nearest, which loses 5–15% accuracy at INT4."

Q: "QAT vs GPTQ — when would you choose each at Qualcomm?"

"For a new on-device model being trained from scratch or fine-tuned: QAT, because you have training data and compute budget, and it gives the best accuracy. For quantizing a large pretrained model (Llama-3, Gemma-2) for NPU deployment without access to the training data or compute: GPTQ (or AWQ), because you only need 128 calibration samples and a few GPU-hours. GPTQ is the practical choice for the productionization of public models."

Q: "What is the relationship between PTQ, QAT, and GPTQ?"

"PTQ: quantize post-hoc with no training. Cheap, works for INT8, fails at INT4.
QAT: simulate quantization during training with STE. Best accuracy, expensive.
GPTQ: second-order PTQ — use Hessian to propagate quantization errors across columns, compensating for each quantized column before moving to the next. Near-QAT accuracy, no training needed."

Section 5: Quick Setup

pip install torch>=2.1.0
# For running GPTQ on real LLMs:
pip install auto-gptq transformers

QAT quick test:

import torch, torch.nn as nn
from phase_03_fundamentals.lab_02_qat.solution import FakeQuantize, wrap_model_for_qat

model = nn.Sequential(nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 16))

# Check STE works
fq = FakeQuantize(bits=8)
x = torch.randn(4, 64, requires_grad=True)
y = fq(x)
y.sum().backward()
print(f"STE gradient norm: {x.grad.norm():.4f}")  # Should be non-zero!

# QAT training
qat_model = wrap_model_for_qat(model, bits=8)
print("QAT layers:", [name for name, m in qat_model.named_modules() if 'QAT' in type(m).__name__])

GPTQ quick test:

from phase_03_fundamentals.lab_03_gptq.solution import compute_hessian, gptq_quantize_layer, compare_reconstruction_error

d_out, d_in, N = 32, 64, 128
X = torch.randn(N, d_in)        # calibration inputs
W = torch.randn(d_out, d_in)    # pretrained weight

H = compute_hessian(X)
result = compare_reconstruction_error(W, X, bits=4)
print(f"GPTQ error:  {result['gptq_error']:.4f}")
print(f"Naive error: {result['naive_error']:.4f}")
print(f"Improvement: {result['improvement_ratio']:.1f}x")  # GPTQ should be better

Lab 01 — PTQ Pipeline

Phase: 03 — Quantization & Model Compression | Difficulty: ⭐⭐⭐⭐☆ | Time: 2–3 hours

Post-Training Quantization (PTQ) converts a trained FP32 model to INT8/INT4 without retraining.

What you build

• Symmetric and asymmetric per-tensor/per-channel quantization
• Calibration-based scale/zero-point computation
• Full PTQ pipeline for a linear model
• Accuracy evaluation: cosine similarity and relative error

Key concepts

Files

Run

pip install -r requirements.txt
python lab.py
pytest test_lab.py -v

Lab 02 — Quantization-Aware Training (QAT)

Phase: 03 — Quantization & Model Compression | Difficulty: ⭐⭐⭐⭐☆ | Time: 3–4 hours

QAT trains with simulated quantization noise — the model learns to be robust to INT8/INT4 precision.

What you build

• FakeQuantize — asymmetric per-tensor fake-quant with Straight-Through Estimator (STE)
• QATLinear — Linear layer with fake-quantized weights
• wrap_model_for_qat — replace all Linear layers with QATLinear (deepcopy safe)
• train_qat — AdamW training loop on QAT model
• convert_to_int8 — fold scale/zero_point into INT8 weights post-training

Key concepts

Files

Run

pip install -r requirements.txt
python lab.py
pytest test_lab.py -v

Lab 03 — GPTQ: Optimal Brain Quantization

Phase: 03 — Quantization & Model Compression | Difficulty: ⭐⭐⭐⭐⭐ | Time: 4–5 hours

GPTQ (Frantar et al., 2022) quantizes LLMs to INT4 with minimal accuracy loss using second-order weight optimization.

What you build

• compute_hessian — H = 2·X·X^T + damping·I from calibration activations
• symmetric_quantize — INT4/INT8 symmetric per-tensor quantization
• gptq_quantize_layer — column-wise OBQ with Cholesky-based Hessian inverse update
• compare_reconstruction_error — GPTQ vs naive round-to-nearest
• GPTQLayer — inference layer with W_q as buffer (no grad)
• gptq_quantize_model — full pipeline with forward hooks for calibration

Key concepts

Files

Run

pip install -r requirements.txt
python lab.py
pytest test_lab.py -v

Phase 04 — torch.compile & FX Graph Optimization

Difficulty: ⭐⭐⭐⭐⭐
Estimated Time: 3 weeks (70–90 hours)
Roles Supported: Framework engineers, NPU backend engineers — understanding torch.compile is mandatory for writing ops that work at production scale

Why This Phase Exists

torch.compile is the future of PyTorch performance. Since PyTorch 2.0, every major model deployment at Google, Meta, and NVIDIA goes through some form of graph compilation. At Qualcomm, the ML framework team's job is to make their NPU backend a first-class torch.compile target.

This requires deep understanding of three systems: TorchDynamo (how Python bytecode becomes an FX graph), AOTAutograd (how the joint forward+backward graph is constructed), and TorchInductor (how FX graphs are lowered to CUDA/C++/Triton kernels). If you can write a custom torch.compile backend that correctly handles 90% of LLaMA's operations, you are the kind of engineer Qualcomm is hiring.

Concepts

• TorchDynamo: Python bytecode analyzer; installs a frame evaluation callback; traces tensor-producing operations into an FX graph; inserts guards for assumptions made during tracing
• Graph breaks: when TorchDynamo encounters unsupported Python constructs (dynamic control flow on tensor values, non-PyTorch library calls), it splits the computation into multiple FX graphs separated by Python execution
• FX Graph IR: a DAG of Node objects with ops: placeholder, call_function, call_module, call_method, output; serializable, manipulable, inspectable
• AOTAutograd: ahead-of-time autograd — computes the joint forward+backward graph before execution; enables compiling the backward pass alongside the forward
• TorchInductor: default backend; generates Triton kernels (GPU) or C++/OpenMP (CPU); performs operator fusion, memory layout optimization, loop tiling
• FX graph transformations: Interpreter (rerun with custom op dispatch), Transformer (node replacement), pattern matching via torch.fx.subgraph_rewriter
• Operator fusion: combining multiple ops into a single kernel (e.g., conv+BN, LayerNorm+dropout, matmul+bias+ReLU) reduces memory bandwidth pressure
• Dead code elimination (DCE): remove graph nodes whose output is never used
• Constant folding: evaluate subgraphs with only constant inputs at compile time
• Shape specialization vs dynamic shapes: torch.compile specializes on input shapes by default (fast) or can use dynamic=True for shape-agnostic compilation (slower, more general)
• Custom backends: any callable (fx.GraphModule, example_inputs) → callable can be registered as a torch.compile backend
• Triton: Python-embedded GPU kernel language; used by TorchInductor; enables writing performance-critical kernels without CUDA C++

Labs

Lab 01 — FX Graph Transformation Passes

Lab 02 — Custom torch.compile Backend

Lab 03 — Triton Fused Kernel + Inductor Lowering

Deliverables Checklist

• 4 FX graph transformation passes with benchmark results
• Custom torch.compile backend with op coverage report
• Triton fused kernel registered as Inductor lowering
• All test_lab.py suites pass

Interview Relevance

• "Walk me through what happens when I call torch.compile(model)(x)." — you can describe the full stack from Python bytecode → FX → AOTAutograd → Inductor → CUDA/Triton
• "How would you write a torch.compile backend for the Qualcomm HTP?" — Lab 02 gives you the complete template
• "What is a graph break and why does it hurt performance?" — you've encountered them in Lab 02 during op coverage analysis
• "Write a Triton kernel for softmax." — Lab 03 gives you the mental model for tiling and warp-level reduction

Warmup Guide — torch.compile, FX Graph & Custom Backends

Zero-to-expert primer for Phase 04. Builds the compilation stack from "why graphs at all" through TorchDynamo's bytecode capture, FX IR, guards, backends, Inductor, and Triton — assuming Phase 01's dispatcher knowledge.

Table of Contents

• Chapter 1: Eager vs Graph — Why Capture Exists
• Chapter 2: torch.fx — The IR You Will Actually Touch
• Chapter 3: Graph Passes — Reading, Matching, Rewriting
• Chapter 4: TorchDynamo — Capture Without Tracing's Lies
• Chapter 5: Guards and Graph Breaks
• Chapter 6: AOTAutograd and Decompositions
• Chapter 7: Backends — Where Your Code Plugs In
• Chapter 8: TorchInductor and Triton
• Chapter 9: Kernel Fusion — The Point of All This
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: Eager vs Graph — Why Capture Exists

Eager mode (default PyTorch): each Python line dispatches one kernel immediately. Maximum flexibility — arbitrary control flow, prints, debugger — but the GPU sees one op at a time. Costs: per-op Python/dispatcher overhead (microseconds each, fatal for small ops), and no cross-op optimization — every op writes its result to DRAM and the next op reads it back.

Graph mode: first capture the computation as a data structure, then optimize it as a whole (fuse ops, plan memory, specialize on shapes), then execute the optimized artifact. The entire phase is about the capture-optimize-execute pipeline and where you, the NPU/compiler engineer, insert your own stages.

The capture problem: Python is too dynamic to compile directly. Three historical strategies, all of which you should be able to compare:

1. Trace (torch.jit.trace, fx.symbolic_trace): run the function with example inputs, record ops. Fast, but control flow is baked in — the trace lies if a branch depends on data.
2. Script (TorchScript): parse Python source into a static language subset. Honest about control flow, but the subset is painful and the project is effectively frozen.
3. Bytecode interception (TorchDynamo, Chapter 4): hook CPython frame evaluation, symbolically execute bytecode, and fall back to Python wherever capture fails. This is the approach that finally worked.

Chapter 2: torch.fx — The IR You Will Actually Touch

FX is PyTorch's Python-level intermediate representation: a GraphModule wrapping a Graph of Nodes. Each node has:

• op: one of placeholder (input), get_attr (parameter fetch), call_function (torch.add, operator.mul...), call_method (x.view), call_module (self.linear), output.
• target: the thing called (function object, attr name, module path).
• args / kwargs: inputs, which may be other Nodes — this linking is the dataflow graph.
• meta: a dict where shape/dtype propagation and your own passes stash information (node.meta["tensor_meta"] after ShapeProp).

How symbolic tracing works: call the module with Proxy objects instead of tensors; every operation on a Proxy records a node and returns a new Proxy. This is also exactly why it fails on data-dependent control flow — if proxy > 0: needs a concrete bool that doesn't exist. (Dynamo solves this differently; FX remains the IR.)

Why FX matters to you: it is the interchange format. Dynamo produces FX graphs; quantization workflows annotate FX graphs; your Lab 02 NPU backend consumes an FX graph. It's Python objects all the way down — printable (graph.print_tabular()), mutable, and testable.

Chapter 3: Graph Passes — Reading, Matching, Rewriting

A pass is a function GraphModule → GraphModule. The three core skills (all in Lab 01):

1. Analysis: walk graph.nodes (topologically ordered), classify, count, propagate shapes (torch.fx.passes.shape_prop.ShapeProp executes the graph with FakeTensors to fill meta).

2. Pattern matching: find subgraphs like conv → batch_norm. The robust way is structural: for each batch_norm node, check node.args[0] is a conv with exactly one user (you!). Single-user checks matter — if the conv output feeds elsewhere, fusing it away is wrong. This single-user discipline is where most homegrown passes have bugs.

3. Rewriting: create the replacement (e.g., fold BN's $\gamma/\sqrt{\sigma^2+\epsilon}$ scale into conv weight, $\beta - \mu\gamma/\sqrt{\sigma^2+\epsilon}$ into bias), insert with graph.inserting_after(node), reroute consumers via node.replace_all_uses_with(new), graph.erase_node the dead ones, then graph.lint() and gm.recompile().

Conv-BN folding math (Lab 01's centerpiece): BN at inference is an affine transform per channel; an affine transform of a conv output is absorbable into the conv's weights and bias:

$$W' = \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}} \cdot W, \qquad b' = \frac{\gamma,(b - \mu)}{\sqrt{\sigma^2+\epsilon}} + \beta$$

One kernel instead of two, and — critical for Phase 03 — quantization observers see the folded distribution, which is the one that ships.

Chapter 4: TorchDynamo — Capture Without Tracing's Lies

Mechanism: PEP 523 lets a C extension replace CPython's frame evaluation. When a torch.compile'd function runs, Dynamo intercepts the frame before execution and symbolically executes its bytecode: tensor operations append FX nodes; Python-level operations (list appends, attribute reads) are tracked symbolically; anything un-traceable triggers a graph break (Chapter 5).

The output per frame: (1) an FX graph of the tensor compute, (2) a guard set — runtime predicates under which this graph is valid, (3) rewritten bytecode that calls the compiled graph and resumes Python where capture stopped.

Why bytecode, not source or tracing: bytecode is what actually executes (no parser divergence); symbolic execution sees real Python semantics (closures, globals, builtins); and fallback is natural — any unsupported construct just ends the graph and lets CPython continue. Tracing can't fall back; scripting can't handle full Python. Dynamo's bet: capture most of the program honestly, rather than all of it wrongly.

Chapter 5: Guards and Graph Breaks

Guards make specialization safe. The captured graph assumed things: x.shape == (8, 512), x.dtype == float16, self.training == False, that global flag's value. Each assumption becomes a guard — a cheap predicate checked on every subsequent call. All pass → run the compiled artifact. Any fail → recompile for the new situation (each function holds a cache of (guards → compiled graph) entries; exceeding torch._dynamo.config.cache_size_limit falls back to eager, your classic silent perf bug).

Dynamic shapes: by default Dynamo specializes on exact shapes; on the second shape it sees, it generalizes that dimension to symbolic (s0) with SymInt arithmetic — guards become constraints like s0 > 1 rather than s0 == 8. dynamic=True requests this up front.

Graph breaks happen at: data-dependent control flow on tensor values (if x.sum() > 0), .item()/.tolist() (forces materialization), unsupported builtins/C extensions, prints. The cost is not correctness — it's fragmentation: each fragment compiles separately, killing cross-fragment fusion, with eager Python between. Diagnose with torch._dynamo.explain(fn)(args) or TORCH_LOGS=graph_breaks; the fix is usually rewriting the offending line in tensor terms (torch.where instead of if).

Chapter 6: AOTAutograd and Decompositions

The FX graph Dynamo captures is forward only and in terms of high-level ops. AOTAutograd traces through the autograd machinery ahead of time, producing a joint forward+backward graph, then partitions it into separate forward/backward graphs — choosing what to save vs recompute between them (this is where activation recomputation decisions live). Your backend may therefore receive a backward graph too.

Decompositions: PyTorch has ~2000 operators; no backend wants to implement them all. The _decomp registry rewrites composite ops into a small primitive set (prims / core ATen IR — a few hundred ops). Your Lab 02 backend declares which ops it handles natively and lets decompositions lower the rest. Functionalization also runs here: in-place ops (add_) become pure ops plus explicit copies, giving backends a side-effect-free graph.

Chapter 7: Backends — Where Your Code Plugs In

A backend is just a callable:

def my_backend(gm: torch.fx.GraphModule, example_inputs):
    # inspect/compile gm.graph however you like...
    return compiled_callable        # same signature as gm.forward

torch.compile(model, backend=my_backend). Whatever you return gets called with real tensors; returning gm.forward unchanged is a valid (null) backend — start there.

The real-world pattern your Lab 02 implements is partitioning: walk the graph, mark nodes your NPU supports, group maximal supported subgraphs, compile those to your ISA, and leave the rest to eager. This is exactly how QNN/TensorRT/CoreML integrations work, and the partition boundaries (with their device-transfer costs) are where deployment performance is won or lost — few large islands beat many small ones, every time.

Chapter 8: TorchInductor and Triton

Inductor is the default backend: FX → its own loop-level IR (define-by-run, scheduling fusion decisions) → Triton source for GPU (or C++/OpenMP for CPU) → compiled and cached. You can read everything it generates: TORCH_LOGS=output_code.

Triton is a Python DSL for writing GPU kernels at the block level: you write the program one block of threads executes (tl.load a tile, compute, tl.store), and the compiler handles the within-block parallelism, memory coalescing, and software pipelining that CUDA makes you do by hand. The mental model shift from CUDA: think "tile program with masks", not "one thread's scalar program". Lab 03 has you write the fused kernel by hand that Inductor would generate, then benchmark both.

Chapter 9: Kernel Fusion — The Point of All This

Why does any of this matter? Memory traffic. Unfused y = relu(x + b):

• kernel 1 reads $x, b$ (2n), writes $t$ (n)
• kernel 2 reads $t$ (n), writes $y$ (n) → 5n DRAM traffic + 2 launches

Fused: read $x, b$, write $y$ → 3n and 1 launch. For elementwise chains (bias, norm, activation, dropout, residual — i.e., half a transformer layer), fusion is a 1.5–3× end-to-end win because these ops are bandwidth-bound (Phase 07's roofline makes this quantitative: fusion raises arithmetic intensity). Matmuls themselves are usually left to cuBLAS/vendor kernels; the win is fusing the elementwise epilogue into them (Inductor does both).

This is also the NPU story (Phase 06): NPU compilers live and die by how long an op chain they can keep inside on-chip SRAM without a DRAM round-trip.

Lab Walkthrough Guidance

Order: Lab 01 (FX passes) → Lab 02 (backend) → Lab 03 (Triton).

• Lab 01: print the graph first (print_tabular()); write the analysis pass before the rewrites; for conv-BN folding, verify numerics against the unfused module on random inputs to 1e-6 — then check the single-user edge case (a conv feeding two consumers must not fold).
• Lab 02: start with the null backend (return gm.forward) to see what Dynamo hands you and how often (add a counter — watch recompiles happen as you vary shapes). Then op classification, then partitioning, then the mock-ISA codegen. Keep a supported_ops set as data, not code — that's how real integrations stay maintainable.
• Lab 03: get the unfused PyTorch reference and its timing first. Write the Triton kernel with masks for the ragged tail; verify exactness, then benchmark across sizes — small sizes show launch overhead, large show bandwidth; explain both regimes with Chapter 9's math.

Success Criteria

You are ready for Phase 05 when you can, from memory:

1. Name the six FX node ops and walk a printed graph aloud.
2. Write conv-BN folding's weight/bias formulas and state the single-user precondition.
3. Explain how Dynamo captures (bytecode symbolic execution), what a guard is, and the recompile-vs-graph-break distinction.
4. List four graph-break triggers and the tensor-level rewrite for each.
5. Describe what AOTAutograd and decompositions transform, and why backends want core ATen IR rather than all ~2000 ops.
6. Compute the 5n→3n fusion traffic example and generalize it to a longer chain.

Interview Q&A

Q: torch.compile made my model slower. Diagnose. In order: (1) compile-time amortization — measure steady state, not first calls; (2) graph breaks fragmenting the model — torch._dynamo.explain, count fragments; (3) recompile storms from varying shapes/Python values — TORCH_LOGS=recompiles, fix with dynamic=True or by passing values as tensors; (4) cache-limit fallback to eager; (5) tiny workload where launch overhead dominated and CUDA graphs (mode= "reduce-overhead") is the actual fix.

Q: Why did Dynamo succeed where TorchScript failed? TorchScript demanded the whole program fit a static subset — one unsupported construct and you rewrite your codebase. Dynamo captures incrementally with graceful fallback: unsupported code just breaks the graph and runs in Python. Adoption cost went from "rewrite" to "decorate", with a perf gradient (fix breaks as you find them) instead of a correctness cliff.

Q: Your NPU supports conv/matmul/elementwise but not softmax. What does your backend do? Partition: maximal supported subgraphs run on NPU; softmax stays on host. Each boundary costs a device round-trip, so fuse aggressively within islands and consider whether a decomposition of softmax (exp/sum/div — all supported elementwise ops!) lets you keep it on-device — that judgment call is exactly Lab 02's design space.

Q: What does a Meta/Fake kernel do and why does compile need it? It computes output shapes and dtypes without touching data. Dynamo/AOTAutograd trace with FakeTensors — every op must propagate metadata or capture fails; it's also how shape specialization and memory planning happen before any real kernel runs. (You wrote one for warp_softmax in Phase 01 Lab 02.)

References

• torch.fx documentation
• Reed et al., torch.fx: Practical Program Capture and Transformation for DL in Python (2021) — arXiv:2112.08429
• TorchDynamo deep-dive (PyTorch dev forum)
• Ansel et al., PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation (ASPLOS 2024)
• torch.compile troubleshooting guide
• Tillet et al., Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations (2019)
• Triton tutorials — work through vector-add and fused-softmax
• PEP 523 — Adding a frame evaluation API to CPython

The Hitchhiker's Guide — Phase 04: torch.compile, FX Graph, Custom Backends

"The first time you look at what torch.compile actually does to your model, you will have one of two reactions: either 'this is incredibly clever' or 'this is incredibly frightening.' Both reactions are correct." — anonymous PyTorch contributor

Section 1: torch.compile Architecture — Three Layers

torch.compile is not a single tool. It is a three-layer compiler pipeline:

Python code
     ↓
[Layer 1] TorchDynamo — bytecode rewriting, Python frame capture
     ↓  
FX Graph (Aten/Prims ops only, no Python control flow)
     ↓
[Layer 2] AOTAutograd — ahead-of-time differentiation, lifts grad into the graph
     ↓
Joint forward+backward FX graph
     ↓
[Layer 3] Backend — inductor (default), nvfuser, custom backend
     ↓
Optimized code / Triton kernels

Layer 1: TorchDynamo — Bytecode Rewriting

TorchDynamo intercepts Python execution at the CPython bytecode level. It does not use tracing in the classical sense — it evaluates Python code and builds an FX graph from the operations it observes.

The key mechanism is torch._dynamo.eval_frame.set_eval_frame(), which installs a callback that intercepts CPython's frame.f_back evaluation. When a guarded region encounters Python control flow (an if statement, a for loop with dynamic bounds), Dynamo breaks the graph at that point, compiling everything before the break into one subgraph and resuming Python execution.

Graph breaks are the primary performance enemy of torch.compile. Every graph break means:

1. A Triton/C++ dispatch overhead
2. Loss of fusion opportunities across the break
3. Potential synchronization

Common graph break causes:

# 1. data-dependent control flow
if x.sum() > 0:  # .sum() returns a Python scalar → break
    ...

# 2. non-PyTorch library calls
import numpy as np
y = np.array(x.numpy())  # leaves PyTorch graph → break

# 3. custom __torch_function__ with Python fallback
class MyTensor(torch.Tensor):
    def __add__(self, other):
        return super().__add__(other)  # may break depending on implementation

# 4. print statements
print(x.shape)  # accesses tensor metadata but causes Python escape → break

# 5. in-place ops on tensors that require grad
x += 1  # may cause aliasing issues → break

How to debug graph breaks:

import torch._dynamo
torch._dynamo.config.verbose = True

# Or use explain():
explanation = torch._dynamo.explain(model)(x)
print(explanation.graphs)       # compiled subgraphs
print(explanation.break_reasons) # why each break happened

Guards: for each compiled subgraph, Dynamo generates guard conditions — Python predicates that must be true for the compiled version to be valid. If guards fail (e.g., input shape changed), Dynamo recompiles (within the cache_size_limit).

Example guard for a linear layer:

# Dynamo generates something like:
def guard_fn(input):
    return (
        input.dtype == torch.float32 and
        input.device.type == 'cuda' and
        input.shape[1] == 512  # static shape guard
    )

Dynamic shapes (torch.compile(dynamic=True)) replace static shape guards with symbolic shape guards using torch.fx.experimental.symbolic_shapes.

Section 2: FX — Functional Transformation of PyTorch Graphs

torch.fx is the intermediate representation used throughout torch.compile. It represents a PyTorch program as a directed acyclic graph of nodes.

FX Node Types

Capturing an FX Graph

import torch
import torch.fx as fx

model = torch.nn.Sequential(
    torch.nn.Linear(512, 256),
    torch.nn.ReLU(),
    torch.nn.Linear(256, 10)
)

# Symbolic tracing (static, follows Python code)
tracer = fx.symbolic_trace(model)
print(tracer.graph)

# torch.export (newer, handles more cases, produces aten-level graph)
from torch.export import export
exported = export(model, args=(torch.randn(1, 512),))
print(exported.graph)

Writing FX Passes

An FX pass is a function that takes a GraphModule, transforms its graph, and returns the modified GraphModule.

Pattern: Replace all ReLU with LeakyReLU

import torch.fx as fx

def replace_relu_with_leaky(gm: fx.GraphModule) -> fx.GraphModule:
    for node in gm.graph.nodes:
        if node.op == 'call_function' and node.target == torch.nn.functional.relu:
            with gm.graph.inserting_after(node):
                new_node = gm.graph.call_function(
                    torch.nn.functional.leaky_relu,
                    args=node.args,
                    kwargs={'negative_slope': 0.01}
                )
                node.replace_all_uses_with(new_node)
            gm.graph.erase_node(node)
    gm.recompile()
    return gm

Pattern: Fuse Linear + ReLU → LinearReLU

def fuse_linear_relu(gm: fx.GraphModule) -> fx.GraphModule:
    graph = gm.graph
    for node in list(graph.nodes):
        # Match: relu(linear(x))
        if (node.op == 'call_module' and 
            isinstance(gm.get_submodule(node.target), torch.nn.ReLU)):
            linear_node = node.args[0]
            if (linear_node.op == 'call_module' and 
                isinstance(gm.get_submodule(linear_node.target), torch.nn.Linear)):
                # Create fused module
                linear = gm.get_submodule(linear_node.target)
                fused = torch.nn.utils.fusion.fuse_linear_bn_eval(linear, None)
                # ... (register fused module, replace both nodes with one)
    gm.recompile()
    return gm

Pattern: Constant folding

def constant_fold(gm: fx.GraphModule) -> fx.GraphModule:
    """Evaluate ops with constant inputs at compile time."""
    graph = gm.graph
    for node in list(graph.nodes):
        if node.op == 'call_function':
            # Check if all inputs are constants (graph inputs are not)
            if all(isinstance(arg, (int, float, bool)) or 
                   (isinstance(arg, fx.Node) and arg.op == 'get_attr')
                   for arg in node.args):
                try:
                    # Execute the op with concrete values
                    const_value = node.target(*[
                        gm.get_parameter(a.target) if isinstance(a, fx.Node) else a
                        for a in node.args
                    ])
                    # Replace node with constant
                    with graph.inserting_before(node):
                        const_node = graph.get_attr(f'_const_{node.name}')
                    gm.register_buffer(f'_const_{node.name}', const_value)
                    node.replace_all_uses_with(const_node)
                    graph.erase_node(node)
                except Exception:
                    pass  # can't fold — skip
    gm.recompile()
    return gm

Section 3: Registering a Custom torch.compile Backend

A custom backend receives an FX GraphModule and returns a callable (typically a Python function wrapping compiled code).

from torch._dynamo.backends.common import aot_autograd
import torch._dynamo as dynamo

def my_custom_backend(gm: torch.fx.GraphModule, example_inputs):
    """
    gm: FX GraphModule (after AOTAutograd, contains aten-level ops)
    example_inputs: list of example input tensors (for shape/dtype info)
    
    Must return: callable with same signature as gm.forward
    """
    print(f"Custom backend received graph with {len(list(gm.graph.nodes))} nodes")
    
    # Option 1: Apply transformations and return gm itself
    gm = my_optimization_pass(gm)
    return gm
    
    # Option 2: JIT compile the graph (e.g., via TorchScript)
    scripted = torch.jit.script(gm)
    return scripted
    
    # Option 3: Lower to custom ISA and return simulation
    npu_program = lower_to_npu(gm, example_inputs)
    return npu_program.simulate  # callable that simulates execution

# Register the backend
@dynamo.register_backend
def my_backend(gm, example_inputs):
    return my_custom_backend(gm, example_inputs)

# Use it:
model = MyModel()
compiled = torch.compile(model, backend="my_backend")
output = compiled(x)

Using AOT Autograd with a custom backend:

from torch._dynamo.backends.common import aot_autograd

def my_inference_backend(gm, example_inputs):
    # Receives joint graph (fwd+bwd combined)
    # Lower to hardware
    return lower_to_hardware(gm)

my_backend_with_aot = aot_autograd(fw_compiler=my_inference_backend)
compiled = torch.compile(model, backend=my_backend_with_aot)

Section 4: Triton Kernels — When and How

Triton is a Python-based GPU kernel language that compiles to PTX. It is the backend for TorchInductor's generated kernels.

Key concepts:

• Programs: Triton launches many programs (like CUDA blocks); each program handles a tile of data
• tl.program_id(axis): which tile this program handles
• tl.load/tl.store: vectorized memory ops with optional masking
• tl.dot: matrix multiply of tiles (uses tensor cores)
• tl.constexpr: compile-time constant (like CUDA constexpr template param)

Softmax kernel pattern:

import triton
import triton.language as tl

@triton.jit
def softmax_kernel(
    input_ptr, output_ptr,
    n_rows, n_cols,
    BLOCK_SIZE: tl.constexpr  # must be power of 2, ≥ n_cols
):
    row_idx = tl.program_id(0)
    row_start = input_ptr + row_idx * n_cols
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < n_cols
    
    # Load row (with masking for out-of-bounds)
    row = tl.load(row_start + cols, mask=mask, other=float('-inf'))
    
    # Numerically stable softmax
    row_max = tl.max(row, axis=0)
    row -= row_max  # broadcast
    row_exp = tl.exp(row)
    row_sum = tl.sum(row_exp, axis=0)
    softmax_output = row_exp / row_sum
    
    # Store result
    output_start = output_ptr + row_idx * n_cols
    tl.store(output_start + cols, softmax_output, mask=mask)

Autotuning:

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE': 128}),
        triton.Config({'BLOCK_SIZE': 256}),
        triton.Config({'BLOCK_SIZE': 512}),
        triton.Config({'BLOCK_SIZE': 1024}),
    ],
    key=['n_cols']  # re-tune when n_cols changes
)
@triton.jit
def softmax_kernel(..., BLOCK_SIZE: tl.constexpr):
    ...

Integrating Triton kernel as a custom torch.op:

class TritonSoftmax(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        n_rows, n_cols = x.shape
        output = torch.empty_like(x)
        BLOCK_SIZE = triton.next_power_of_2(n_cols)
        softmax_kernel[(n_rows,)](
            x, output, n_rows, n_cols, BLOCK_SIZE=BLOCK_SIZE
        )
        ctx.save_for_backward(output)
        return output
    
    @staticmethod
    def backward(ctx, grad_output):
        output, = ctx.saved_tensors
        # Softmax backward: grad_in = output * (grad_out - (grad_out * output).sum(dim=-1, keepdim=True))
        ...

triton_softmax = TritonSoftmax.apply

Section 5: FX Graph Optimization Passes — Reference List

These are the passes TorchInductor applies (and you can apply in a custom backend):

Section 6: Interview Pitfalls

Section 7: Resources

1. PyTorch torch.compile tutorial — https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html
2. TorchDynamo overview — Jason Ansel et al., "TorchDynamo: Towards Principled Python JIT Compilation" (Meta AI, 2022)
3. TorchInductor — blog post on the default backend architecture: https://dev-discuss.pytorch.org/t/torchinductor-a-pytorch-native-compiler-with-define-by-run-ir-and-symbolic-shapes
4. Triton paper — Philippe Tillet et al., "Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations" (MLSys 2019)
5. FX documentation — https://pytorch.org/docs/stable/fx.html
6. torch.export documentation — https://pytorch.org/docs/stable/export.html
7. AOTAutograd internals — https://github.com/pytorch/pytorch/blob/main/torch/_functorch/aot_autograd.py (read the docstring)

👨🏻 Brother Talk — Phase 04, Off the Record

torch.compile & FX — the JD names this explicitly ("model level optimization using techniques like torch compile"). This is the bridge from eager PyTorch to the compiler world.

Brother, the JD calls out torch.compile by name — that's rare and it tells you this is a must-have, not a nice-to-have. And here's why it matters for Qualcomm specifically: a custom hardware backend (like Qualcomm's QNN) hooks into PyTorch through exactly this machinery. When you write a torch.compile backend, you're doing the same thing the QNN integration does — capture the graph, transform it for your hardware, hand back something runnable. So this phase isn't "learn a speedup flag." It's "learn how to put your hardware into PyTorch." That's the job.

Let me demystify the stack, because the names intimidate people. torch.compile is three pieces: Dynamo (captures your Python bytecode into an FX graph — the hard part, because Python is dynamic), the FX graph (a clean, inspectable IR — nodes and edges you can rewrite), and Inductor (the default backend that lowers FX to Triton/C++ kernels). The magic and the pain both live in Dynamo: when it can capture your whole function, you get one fast graph; when it hits something it can't trace (data-dependent control flow, an unsupported op, a .item() call), it graph-breaks — splits into multiple graphs with eager glue between them. Graph breaks are the #1 reason torch.compile underdelivers, and finding/fixing them is a real skill.

The FX graph is where you'll spend your time, and here's the thing that will trip you, because it just tripped this track's lab and I had to fix it: the same model produces different FX graphs depending on how you capture it. symbolic_trace on nn.Linear/nn.GELU modules gives you call_module nodes (target = the submodule name). Dynamo gives you a flattened graph of function calls. AOT/make_fx gives you aten-level ops (aten.linear.default). A graph pass written to match call_function aten ops will silently do nothing on a symbolic_trace graph of modules — no error, just no fusion, no offload. The bug looks like "my optimization isn't working" and the cause is "I'm pattern-matching the wrong graph dialect." A real backend must handle the dialect it actually receives. Build the fusion pass, then test it on a module-based graph and an aten graph, and feel the difference.

That fusion lab (fuse Linear+GELU into one node) is the core skill in miniature: pattern-match a subgraph, verify the match is safe to fuse, rewrite the graph, prove it's numerically identical. The "safe to fuse" part is where seniors separate from juniors — you can only fuse Linear→GELU if the Linear's output feeds only the GELU (single user). Fuse it when something else also reads that output, and you've changed the computation. The single-user check is the whole correctness argument. This pattern — match, check safety, rewrite, verify — is every graph optimization you'll ever write, including the ones that target the NPU.

The custom-backend lab is the payoff: register a backend, classify which nodes your "NPU" can run vs which fall back to CPU, transform the eligible ones, return a callable. This is exactly the shape of a real hardware backend. And the classification bug I fixed (matching aten strings when the graph had function-identity targets) is the real-world trap: your op-support table has to match the graph dialect, or you classify zero nodes as NPU-eligible and silently run everything on CPU — a "why is my accelerator doing nothing?" incident. CPU fallback is supposed to be a safety net, not the default because your matcher is broken.

The thing nobody tells you about torch.compile: correctness first, speed second, and verify both. A backend that's fast but wrong is worse than useless — it ships silent accuracy bugs. So every backend lab here ends with a numerical-equivalence check against eager. Internalize that discipline: every graph transformation must be proven to preserve the output. At Qualcomm, "the compiled model is 3× faster" is meaningless if you can't follow it with "and bit-exact to eager within tolerance." The verification is not optional ceremony; it's the deliverable.

Career truth, brother: torch.compile / FX fluency is a differentiator because most ML engineers treat the compiler as a black box they flip on. The person who can read an FX graph, find the graph breaks, write a custom pass, and build a backend is operating at the layer Qualcomm hires for — the layer where models meet silicon. And it's deeply transferable: the same graph-rewriting mindset applies to ONNX (P05), to TVM, to MLIR, to any IR. Learn it once here, recognize it everywhere.

Build the fusion pass. Build the backend. Make it classify NPU nodes correctly across graph dialects, and prove every transform is numerically equivalent. Then come to P05, where we go broader — the whole ML-compiler / IR landscape (ONNX, TVM, MLIR) that this is one instance of.

— your brother 👨🏻

Lab 01 — FX Graph Passes

Phase: 04 — torch.compile, FX Graph & Custom Backends | Difficulty: ⭐⭐⭐⭐☆ | Time: 3–4 hours

FX graph passes are how torch.compile transforms your model before lowering to a backend.

What you build

• symbolic_trace a model to capture its FX graph
• Write a graph pass that fuses linear → relu into a single node
• Write a graph pass that eliminates redundant operations (dead code)
• Verify output numerical equivalence before and after transformation

Key concepts

Files

Run

pip install -r requirements.txt
python lab.py
pytest test_lab.py -v

Lab 02 — Custom torch.compile Backend

Phase: 04 — torch.compile, FX Graph & Custom Backends | Difficulty: ⭐⭐⭐⭐⭐ | Time: 4–5 hours

Implement a custom backend that classifies ops as NPU-eligible or CPU-fallback and substitutes them in the FX graph.

What you build

• NPU op registry with @register_npu_op decorator
• classify_nodes — analyze FX graph for NPU vs CPU ops
• transform_graph — replace ATen ops with registered NPU ops via FX node substitution
• npu_backend — torch.compile(model, backend=npu_backend) callable
• verify_backend_correctness — compare eager vs compiled numerics

Key concepts

Files

Run

pip install -r requirements.txt
python lab.py
pytest test_lab.py -v

Lab 03 — Triton Fused Kernel

Phase: 04 — torch.compile, FX Graph & Custom Backends | Difficulty: ⭐⭐⭐⭐⭐ | Time: 4–5 hours

Write a GPU kernel in Triton that fuses multiple operations into a single pass over memory.

What you build

• Triton kernel for fused LayerNorm + Linear in a single kernel
• Tile-based computation for cache efficiency
• Memory bandwidth analysis: fused vs unfused
• Benchmarking with triton.testing.Benchmark

Key concepts

Files

Run

pip install -r requirements.txt  # requires CUDA GPU
python lab.py
pytest test_lab.py -v

Phase 05 — ML Compilers & IRs

Difficulty: ⭐⭐⭐⭐⭐
Estimated Time: 3 weeks (70–90 hours)
Roles Supported: Framework & compiler engineers — required for writing NPU backend compilers, ONNX export pipelines, and cross-framework optimization toolchains

Why This Phase Exists

The ML compiler stack is the path from a PyTorch model to a running binary on a Qualcomm NPU. The chain is:

PyTorch Model
  → torch.export / torch.onnx.export
  → ONNX / MLIR / StableHLO  (portable IRs)
  → SNPE/QNN compiler  (Qualcomm-specific)
  → .dlc or .bin  (NPU binary)
  → Hexagon runtime

Understanding each IR is not optional if you want to debug why a model produces wrong results after compilation, or why 3 ops are missing from the NPU coverage report, or why the compiled model is slower than expected despite the op coverage being 95%.

This phase builds working knowledge of ONNX, TVM, and MLIR — the three dominant ML compiler IRs in production systems. ONNX is what most pipelines export to. TVM is a complete compiler that demonstrates autotuning and schedule optimization. MLIR is the foundation of many new compilers including IREE (used for Qualcomm AI Engine).

Concepts

• ONNX (Open Neural Network Exchange): protobuf-based IR; nodes are NodeProto with op_type, inputs, outputs, attributes; opset versioning (opset 17 = standard); computation graph is a topological list of nodes
• ONNX opset: each op version specifies exact semantics; mismatches between export opset and runtime opset cause failures
• onnxoptimizer: graph-level rewrites for ONNX graphs (constant folding, reshape fusion, conv+BN fusion); similar to FX passes but for ONNX
• onnxruntime: Microsoft's ONNX execution engine; multiple execution providers (CPU, CUDA, CoreML, QNN)
• TVM: complete ML compiler stack; Relay (high-level IR) → TIR (tensor IR) → codegen → machine code; includes AutoTVM (search-based schedule tuning) and Ansor (auto-scheduler)
• Relay IR: functional, typed, higher-order; supports data types, shapes, control flow; analogous to FX but more formal
• TIR (Tensor IR): imperative, loop-based; describes memory access patterns, loop tiling, vectorization; where most performance optimization happens
• Schedule in TVM: explicit loop transformation specification (tile, vectorize, parallelize, reorder, unroll); AutoTVM/Ansor searches the schedule space automatically
• MLIR (Multi-Level IR): modular IR framework; dialects define ops and types; transformation passes lower between dialects; core of LLVM 17+ and many ML compilers
• MLIR dialects: linalg (named mathematical operations), affine (affine maps on memrefs), vector (SIMD), memref (memory regions), arith (scalar arithmetic), func (functions), tosa (Tensor Operator Set Architecture — used by TFLite and Qualcomm)
• TOSA (Tensor Operator Set Architecture): small, well-defined op set targeting hardware accelerators; used by TFLite converter, IREE, and Qualcomm's pipeline
• IREE (Intermediate Representation Execution Environment): end-to-end ML compiler; ingests MLIR/StableHLO, outputs code for CPU/GPU/NPU; used for Qualcomm AI Engine deployments
• StableHLO: XLA's HLO dialect stabilized as a cross-framework standard; JAX/TF/PyTorch can all emit StableHLO

Labs

Lab 01 — ONNX Export Pipeline + Optimization Passes