Capstone 03 — Custom torch.compile NPU Backend

Phase: 11 — Capstone | Difficulty: ⭐⭐⭐⭐⭐ | Time: 5–6 hours

Build a torch.compile backend that routes eligible ops to a simulated NPU and falls back to CPU for unsupported ops.

What you build

  • _NPU_OPS registry + @register_npu_op decorator
  • npu_linear, npu_relu, npu_gelu — simulated NPU implementations
  • classify_nodes — partition FX graph into NPU-eligible vs CPU-fallback
  • transform_graph — replace ATen ops with NPU ops via FX node substitution
  • npu_backend — torch.compile(model, backend=npu_backend) callable
  • verify_backend_correctness — max absolute diff < 1e-4

Integrates

Phase 04 FX graph passes + Phase 04 custom backend + Phase 06 NPU op support concepts

Run

pip install -r requirements.txt
python lab.py
pytest test_lab.py -v