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
cs.PL, cs.AI
2026-09-05
SMART rebuilds the ML performance library from ~50 design docs in 1.5–3 hours and ~$100, matching DeepSeek-V3 TPU refs to round-off.
ML performance modeling sits between two stacks that refuse to sit still. Model architectures pick up MoE routing, latent attention, and split prefill/decode phases on a months-long cadence. Accelerators and interconnects turn over on their own clock. A framework that assumed every transformer layer looks the same is already lying. The usual response is a stream of refactors and local patches.
Patches compound. Humans almost never delete the old implementation and start from the current spec. Each revision is G(new spec, old code), and the distance to a clean G(new spec) is what the paper calls incremental generation debt. Coding agents add a second failure mode: the context window cannot hold a mature library, so the agent is fed local snippets and produces edits that are locally plausible and globally worse. The paper calls this context-window myopia.
SMART's bet is that regenerating the library is now cheaper than servicing that debt. Humans edit natural-language design docs. Code is a build product.
The main branch holds almost no code. What is checked in is a DAG of markdown design docs, a directed acyclic graph whose edges say which module must be generated first. Those edges are inferred by read-only agents, not maintained by hand. An orchestrator walks the graph in topological order and assigns one coding sub-agent per doc. The result has to pass reconciliation against hand-built references, parameter guards, and unit tests; failures loop back into regeneration, and only a passing build replaces the previous one. Humans edit prose only.
One doc, one agent buys three things. Each generation task fits in a bounded context, which raises the chance it is correct. The orchestrator log records which docs the agents struggle with and which topological waves inject bugs, so the next human pass knows which paragraphs to clarify. Foundational DSL docs can be routed to a stronger model; downstream catalog docs can use a cheaper one.
The docs themselves are written in a specific style. Tests and high-level constitutions are the usual top-down control surface. SMART adds bottom-up worked examples: walk a piece of pseudocode on a concrete input, write down intermediate shapes, intermediate values, and the closed-form cost that should come out. Every number-bearing doc ends in a reconciliation anchor, a small preset whose expected outputs are stated exactly and enforced by generated tests. A typical vignette: on a 2×2×2 torus with wraparound the per-node link count is 3, not 6, so an all-gather of V bytes costs ...
The artifact being specified has to be small and stable too. The IR is one recursive Op. Tensors carry symbolic shapes. Cost is a bundle of SymPy expressions for compute, memory, and communication, kept closed-form until numeric substitution. An RRT (resource reservation table) has hardware resources on rows, cycles on columns, and occupancy in the cells. Params are a loop nest, a subgraph, or a leaf. Leaves are the software/system boundary: the system side prices a leaf with an RRT and an OpCost. Current leaves are TPU-shaped, MXU matmul tiles (the TPU matrix unit), VMEM loads (on-chip vector memory), and ICI collectives (chip-to-chip interconnect). The algorithm side composes leaves into nests. A new attention variant or a new interconnect generation touches only its own docs.
Models are authored in a thin Python tracing DSL. @smartloop is a reduction with a carried accumulator; @smartmaploop is a parallel map. Every dimension is a SymPy symbol, so one trace covers the design space. The flash-attention core uses the same nest for prefill (query length Tq and KV length Tkv both equal sequence length T) and flash-decoding (Tq = 1, Tkv equals context length); the (B, H, Tq, Tkv) score matrix never leaves VMEM. Distribution is sharding annotations. Each tensor dimension names the mesh axis it is split on, and a sharded-einsum wrapper infers AllGather and ReduceScatter from operand and output shardings. Only layout-moving collectives are written by hand; the dispatch all-to-all in the DeepSeekMoE block, which moves the expert-parallel axis from the token-group dimension onto the expert dimension, is the example.
Two roll-up modes turn the Op tree into wall-clock time. Fast mode scales each leaf cost by enclosing trip counts and applies roofline-style analytical overlap (weight pre-collection, all-to-all hiding), using compute and bandwidth ceilings as a lower bound on time. Evaluation is closed-form and intended for sweeps of thousands of points. Slow mode modulo-schedules each loop into its RRT, software-pipelining the body against per-resource capacity, and rolls the initiation interval (how many cycles before the next iteration can issue) up the tree. Numeric substitution happens only at the edge. A design doc can even state the exact expected expression for a collective's cost, and generated tests assert it.
The library is 50 design docs, about 9,000 lines of specification prose, covering TPU topology, collective costs, numerics, schedulers, and a catalog of dense, MoE, latent-attention, and robotics/VLA families. Master is the docs plus a handful of leaf utilities.
A clean-slate regeneration with Claude Code takes 1.5 to 3 hours and about 100 USD in API spend, roughly 20% of a weekly Claude Max budget. Regenerated implementations match hand-audited reference models to round-off precision. The named example is DeepSeek-V3 serving on a TPU pod slice.
| Item | Number | Against |
| Design docs | 50, 9,000 lines | Almost no code on main |
| Full regeneration | 1.5–3 hours | Claude Code |
| API cost | 100 USD per rebuild | 20% of weekly Claude Max |
| DeepSeek-V3 TPU serving | Round-off match | Hand-audited reference models |
There is no table of predicted latency versus measured TPU wall-clock, and no comparison against existing estimators. Round-off agreement is regenerated code versus handwritten reference models.
For systems groups whose specs churn faster than the code can absorb them, treating design docs as the durable artifact and the implementation as a build product is now a runnable workflow. The debt in equation (1) is the distance between an incrementally patched system and a from-scratch generation of the current spec. Full regeneration drives that distance to zero by construction. Vague intent also surfaces earlier, because a guess has to be written into a doc to survive the next rebuild.
The portable piece is the documentation discipline: one doc per agent, worked examples, reconciliation tests on every number-bearing doc, orchestrator logs as a quality gauge. A 100-dollar, two-hour rebuild puts this in the range of routine CI for a library of this size.
The IR is familiar roofline plus modulo scheduling. The contribution is the software lifecycle, not a new performance theory.
The paper is five pages and has no experiments section. The DeepSeek-V3 round-off claim appears in the abstract and the conclusion. There is no error table, no count of failed regenerations, and no report of how many doc edits it took to pass reconciliation. The references are hand-audited models; the reader cannot tell how far those models sit from real TPU measurements. Driving generation debt to zero by construction assumes the generator actually implements the docs. If an agent misreads a doc and the tests are generated from the same docs, reconciliation can fail together. The real external anchor is those handwritten references, and their coverage is not described.
Leaves are TPU-shaped. Porting to GPUs means rewriting the system-side docs, and that work is not shown. Regeneration is reported with Claude Code only. Humans did not leave the loop; they now maintain prose that has to include worked examples. Vague docs show up as orchestrator-log failures.