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

FieldValue
GoalWrite 4 FX graph transformation passes: (1) Conv2d + BatchNorm fusion, (2) constant folding, (3) dead node elimination, (4) Op replacement (swap relu for hardswish for NPU compatibility)
Conceptstorch.fx.symbolic_trace, torch.fx.Interpreter, torch.fx.Transformer, node pattern matching, torch.nn.utils.fusion.fuse_conv_bn_eval, GraphModule.recompile()
Steps1. Implement ConvBNFuser(fx.Transformer) that fuses sequential Conv2d+BN; 2. Implement ConstantFolder(fx.Interpreter) that evaluates constant subgraphs; 3. Implement DCEPass that removes nodes with zero users; 4. Implement OpReplacer for configurable op substitution; 5. Compose passes into a PassManager with before/after accuracy validation; 6. Benchmark ResNet-50: passes applied vs baseline (latency + accuracy)
StackPyTorch 2.3+, torchvision
DatasetsImageNet val subset (1000 images)
OutputPassManager with all 4 passes; ResNet-50 benchmark: latency reduction with <0.1% accuracy drop; pass unit tests
How to Testpytest test_lab.py — checks: (1) Conv+BN fusion reduces node count, (2) fused model accuracy matches original within 1e-5, (3) constant subgraphs are correctly evaluated, (4) dead nodes are eliminated
Talking PointsWhat is the mathematical equivalence behind conv+BN fusion? Why can constant folding improve latency? What makes DCE safe (when might DCE be incorrect)?
Resume BulletWrote 4 FX graph optimization passes (Conv-BN fusion, constant folding, DCE, op replacement); demonstrated 12% latency improvement on ResNet-50 with <0.1% accuracy delta
ExtensionsAdd quantization-aware fusion (conv+bn+relu → quantized fused op); add symbolic shape analysis; implement algebraic simplification (x + 0 = x)

Lab 02 — Custom torch.compile Backend

FieldValue
GoalWrite a custom torch.compile backend (npu_simulator) that lowers an FX graph to a mock NPU instruction set (a simple JSON-based IR), with fallback to PyTorch CPU for unsupported ops
Concepts@register_backend, fx.GraphModule, example_inputs, lowering, fallback partitioning, torch.fx.passes.operator_support, BackendPatternBase
Steps1. Define 15-op NPU ISA in JSON (matmul, add, relu, softmax, layernorm, etc.); 2. Implement FX→NPU IR lowering (supported ops → NPU IR nodes, unsupported → CPU fallback); 3. Implement NpuSimulator that executes the JSON IR; 4. Register backend via @register_backend; 5. Test with torch.compile(model, backend="npu_simulator"); 6. Measure partition quality: % ops handled by NPU vs CPU fallback on LLaMA-tiny
StackPyTorch 2.3+, torch._dynamo, torch.fx
DatasetsSynthetic (random inputs)
OutputRegistered npu_simulator backend; LLaMA-tiny running under torch.compile(backend="npu_simulator") with >85% ops handled by NPU; op coverage report
How to Testpytest test_lab.py — checks: backend registered correctly, output matches torch.compile(backend="eager") within 1e-4, partition report generated, fallback rate < 15% for standard models
Talking PointsHow does torch.compile determine which ops go to which backend? What is operator_support.OperatorSupport? When should you prefer a graph break over a fallback?
Resume BulletWrote custom torch.compile backend targeting mock NPU ISA; achieved 87% op coverage for LLaMA-style models with automatic CPU fallback for unsupported ops
ExtensionsAdd profiling mode (timing each NPU op); implement operator decomposition (break unsupported ops into supported primitives); add memory layout optimization (NHWC vs NCHW)

Lab 03 — Triton Fused Kernel + Inductor Lowering

FieldValue
GoalWrite a fused LayerNorm+GeLU Triton kernel and register it as a TorchInductor lowering so that torch.compile automatically uses your kernel whenever it encounters this pattern
ConceptsTriton programming model, blocked tile computation, SRAM vs DRAM, Inductor lowering via @torch._inductor.lowering.register_lowering, FX pattern matching, performance benchmarking
Steps1. Implement layernorm_gelu_kernel in Triton: fused computation that reads x once, computes LayerNorm, applies GeLU, writes output (avoids two DRAM passes); 2. Verify correctness vs F.layer_norm + F.gelu within 1e-4; 3. Benchmark vs sequential and vs torch.compile default; 4. Register as Inductor lowering; 5. Verify torch.compile(model) now uses your kernel automatically via Triton's @triton.jit
StackTriton 2.3+, PyTorch 2.3+, CUDA GPU required
DatasetsSynthetic: torch.randn(batch, seq, d_model)
Outputlayernorm_gelu_kernel achieving ~1.4× speedup over sequential on batch=32,seq=2048,d=4096; registered as Inductor lowering with correctness verification
How to Testpytest test_lab.py — checks: Triton output matches PyTorch within 1e-4, kernel handles edge cases (d_model not divisible by BLOCK_SIZE), perf > 1.2x on target shape, Inductor lowering fires correctly
Talking PointsHow does Triton's blocked programming model map to GPU warp/block hierarchy? What determines BLOCK_SIZE choice? What is memory coalescing and why does your kernel achieve it?
Resume BulletWrote fused LayerNorm+GeLU Triton kernel registered as TorchInductor lowering; achieved 1.43× speedup over sequential ops at transformer-scale dimensions (seq=2048, d=4096)
ExtensionsAdd FP8 input support; implement fused MHA (Q,K,V projection + attention in one kernel); contribute to PyTorch's Triton op library

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