Source-linked AI summary
Design Docs Are All You Need: An AI-native Machine-Learning Performance Tool
Samuel Kushnir, Kimia Noorbakhsh, Kavya Sreedhar, Liqun Cheng, Ming Liu, Parthasarathy Ranganathan, Mohammad Alizadeh, Fred Kjolstad, Suvinay Subramanian
TL;DR
ML performance-modeling frameworks face continual churn from evolving architectures and hardware, while incremental updates accumulate technical debt. SMART makes natural-language design documents the durable source of truth, regenerates implementations through an agent-orchestrated DAG, and uses worked examples with a minimal symbolic IR. The resulting workflow makes full regeneration practical and supports design docs as the durable artifact for ML-systems co-design.
Problem
Evolving ML architectures and hardware continually invalidate performance-modeling abstractions, while incremental code updates accumulate technical debt and context-blind revisions degrade structural coherence.
Method
SMART stores self-contained natural-language design documents in a dependency DAG, regenerates code with sub-agents, and combines worked examples with a minimal symbolic IR for symbolic cost modeling.
Results
Full clean-slate regeneration takes 1.5–3 hours and costs around 100 USD, making recurring library rebuilds economically viable.
Takeaways & Limitations
Design docs can serve as the durable artifact for ML-systems co-design when worked examples, a machine-discovered dependency DAG, and a minimal symbolic IR support regeneration.
Abstract
from arXiv · showhide
Machine-learning performance modeling is a uniquely hostile terrain for long-lived software: the assumptions baked into today's abstractions are invalidated by tomorrow's models and systems, forcing perpetual refactoring of performance-modeling frameworks. Meanwhile, AI coding agents have become fast and capable enough that regenerating an entire library is cheaper than paying down the tech debt of incrementally patching it. We describe SMART, a rigorous symbolic performance-modeling library for ML systems whose main branch contains almost no code: the repository is a DAG of self-contained natural-language design docs, coding sub-agents regenerate the implementation from only the docs on new version updates, and every human change is a natural-language edit to a doc--self-documenting by construction. Two ingredients make regeneration reliable: (i) a design-doc style built around step-by-step worked examples that act as in-context demonstrations for the generating agents, and (ii) a minimal, recursively defined operator IR with symbolic (SymPy) cost expressions, a fast analytical roll-up mode for large sweeps, and a slow modulo-scheduling mode for fine-grained schedule studies. Regenerated implementations reproduce hand-audited reference models--including DeepSeek-V3 serving on a TPU pod slice--to round-off precision, suggesting that design docs--not code--can be the durable artifact for ML-systems co-design tools.
1 Introduction
SMART addresses the instability of ML performance-modeling abstractions and the technical debt of incremental code updates by making natural-language design documents the durable artifact and regenerating code from them.
- Rapidly changing model architectures and hardware systems continually invalidate performance-modeling abstractions, forcing repeated framework refactoring.
- AI coding agents add new pressure because finite context windows encourage fragmented revisions that miss global invariants and architectural intent.
- Incremental generation debt measures the distance between a repeatedly patched implementation and one regenerated from the current specification alone.
- The regeneration workflow orders document-specific agents through a dependency DAG and reconciles the generated implementation against hand-built references before replacement.
- SMART makes design docs the durable artifact and treats code as a regenerable build product, driving incremental-patching debt to zero through clean-slate regeneration.
2 Design docs as the source of truth
SMART organizes its source of truth as a dependency-ordered DAG of self-contained design documents, using specialized agents and worked examples to make full regeneration practical and reliable.
- SMART’s main branch contains almost no code; self-contained Markdown design documents form a DAG whose dependencies determine generation order.
- Read-only agents infer document dependencies, then an orchestrator assigns coding sub-agents to implement each document in topological order.
- Per-document generation bounds context windows, exposes difficult documents for targeted prose refinement, and supports routing complex documents to stronger models.
- A complete clean-slate regeneration takes 1.5–3 hours and costs around 100 USD, making recurring full-library rebuilds economically viable.
- Step-by-step worked examples specify intermediate shapes, values, and exact closed-form costs, providing concrete traces that stabilize independent agent generations.
3 A minimal symbolic IR for performance co-design
SMART uses a minimal recursive operator IR that separates algorithm composition from system pricing, propagates symbolic costs, and supports both fast analytical sweeps and detailed scheduling studies.
- The IR is deliberately minimal, orthogonal, and stable under architecture churn, with fast and slow modes for different performance-modeling tasks.
- An Op is recursively an interior loop or subgraph node or a leaf carrying system-specific resource reservations and cost expressions.
- Leaves connect algorithm structure to TPU resources through MXU matmuls, VMEM loads, and ICI collectives, while system and algorithm changes remain separately documented.
- The builder DSL traces decorated blocks and loops into symbolic subgraphs, allowing one nest to serve both prefill and flash-decoding through asymmetric query and key-value lengths.
- Fast roll-up scales leaf costs by enclosing trip counts and models overlap analytically for sweeps over thousands of design points.
- Symbolic propagation produces closed-form SymPy expressions in design variables, with numeric substitution deferred until each design point is evaluated.
4 Conclusion
SMART concludes that design docs can replace code as the durable artifact for ML-systems co-design when worked examples, a dependency DAG, and a minimal symbolic IR support regeneration.
- SMART comprises 50 design docs and approximately 9,000 lines of specification prose spanning TPU systems, cost models, numerics, schedulers, and frontier model families.
- The library is regenerated with each version change, reducing the main branch to documents plus a few utilities and driving incremental-patching debt to zero.
- Worked-example documents, a machine-discovered dependency DAG, and a minimal symbolic IR are presented as the enabling ingredients for this workflow.