Paper deep dive
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
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 90%
Last extracted: 9/8/2026, 2:18:33 AM
Summary
The paper introduces SMART, an AI-native machine-learning performance modeling library where natural-language design docs serve as the durable source of truth instead of code. The system uses a directed acyclic graph (DAG) of self-contained design documents, which are processed by coding sub-agents to regenerate the implementation library. This approach eliminates incremental technical debt by rebuilding the code from scratch on every update, leveraging worked examples in the docs for reliable regeneration and a minimal symbolic intermediate representation (IR) for precise performance cost modeling.
Entities (9)
Relation Signals (7)
Symbolic IR → uses → sympy
confidence 95% · symbolic (SymPy) cost expressions
SMART → uses → Design Docs
confidence 95% · the repository is a DAG of self-contained natural-language design docs
Coding Sub-agents → generates → SMART
confidence 92% · coding sub-agents regenerate the implementation from only the docs
Design Docs → contains → Worked Examples
confidence 90% · a design-doc style built around step-by-step worked examples that act as in-context demonstrations
SMART → implements → Symbolic IR
confidence 90% · a minimal, recursively defined operator IR with symbolic (SymPy) cost expressions
SMART → targets → TPU
confidence 90% · We target TPUs
SMART → validatesagainst → DeepSeek-V3
confidence 88% · Regenerated implementations reproduce hand-audited reference models--including DeepSeek-V3
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract: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.
Tags
Links
- Source: https://arxiv.org/abs/2609.05364v1
- Canonical: https://arxiv.org/abs/2609.05364v1
PDF not stored locally. Use the link above to view on the source site.
Full Text
16,523 characters extracted from source content.
Expand or collapse full text
Design Docs Are All You Need: An AI-native Machine-Learning Performance Tool Samuel Kushnir Kimia Noorbakhsh Affiliation: MIT Kavya Sreedhar Affiliation: Google Liqun Cheng Affiliation: Google Ming Liu Affiliation: Google Parthasarathy Ranganathan Affiliation: Google Mohammad Alizadeh Affiliation: MIT Fred Kjolstad Affiliation: Stanford Suvinay Subramanian Google DeepMind Abstract 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 (i) 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 Machine-learning performance modeling is a punishing environment for software abstractions, because the assumptions of current architectures and systems are constantly evolving. An abstraction that seemed reasonable a year ago—for instance, that every transformer layer looks the same—is suddenly false in the face of mixture-of-experts routing, latent attention, and heterogeneous inference phases. Performance-modeling frameworks sit at the intersection of the two fastest-moving parts of the stack—rapidly evolving model architectures on top, and underlying hardware accelerators and interconnects on the bottom—so they absorb this churn from both directions; the result is a never-ending stream of refactorings and hacks to absorb each new architecture or system change into an aging framework. At the same time, AI coding tools have become remarkably capable and fast. Together, these two forces produce failure modes that compound technical debt: 1. Incremental generation debt (predates AI). Let StS_t denote the specification of a library at time t and let G denote a code generator—a human engineer or an agent. A greenfield build computes Ct=G(St)C_t=G(S_t). In practice, however, the generator at time t+1t+1 is handed the previous implementation as an extra argument and computes Ct+1=G(St+1,Ct)C_t+1=G\! (S_t+1,\,C_t ). We can define the technical debt of this process as Dt+1=∥G(St+1,Ct)−G(St+1)∥,D_t+1\;=\; \,G(S_t+1,\,C_t)\;-\;G(S_t+1)\, , (1) the distance between the incrementally patched system and the one that would have been built from the current spec alone. Because starting over—deleting all of one’s code—is mentally challenging and time-consuming for humans, the technical debt is in practice much larger than zero: incremental patches inevitably introduce compromises that compound with every spec revision. 2. Context-window myopia AI coding agents are constrained by finite context windows, making it impossible to pass an entire mature codebase into a single prompt. Consequently, developers must feed the agent fragmented, localized code snippets when requesting revisions. Because the agent cannot see the whole picture—missing global invariants, cross-module dependencies, and the broader architectural intent—it frequently generates code that is locally plausible but globally sub-optimal. Over time, these piecemeal, context-blind updates degrade the structural coherence of the framework. To solve these problems while exploiting the speed of AI coding agents, we make natural-language design docs the durable artifact and treat code as a regenerable build product (Section 2). Reliable regeneration imposes real demands on how docs are written (Section 2) and how the modeled system is factored into abstractions (Section 3). Because Ct=G(St)C_t=G(S_t) is recomputed from scratch on a regular cadence, Equation 1 is driven to zero by construction. This complete regeneration workflow is illustrated in Figure 1. docdocdocdocdocdesign-doc DAG(checked into master)agentic regenerationone sub-agent per doc,topological ordergenerated librarysymbolic cost model(build product)validationreference-model reconciliation,parameter guards, unit testsrepairhuman edits: natural language only Figure 1: The regeneration workflow: sub-agents regenerate the implementation from the design-doc DAG in dependency order, and the result must pass reconciliation against hand-built references before it replaces the previous build. Humans only ever edit the docs. 2 Design docs as the source of truth Our main branch contains almost no code. Instead, it consists of a folder structure of design document markdown files that compose into a directed acyclic graph (DAG). Some modules must be generated before others—for example, hardware topology and numerics precede the collective-cost models, which in turn precede the model catalog. The edges of this DAG are machine-discovered rather than hand-maintained. We utilize a distributed process where read-only agents analyze the design documents and infer the dependency edges between them. Once the graph is resolved, an orchestrator agent walks the DAG in topological order, assigning a dedicated coding sub-agent to implement each self-contained doc. Throughout this process, the orchestrator maintains a central log file recording areas where sub-agents struggled to interpret the prose, as well as bugs uncovered in the output of agents from earlier topological waves. Factoring the codebase into self-contained documents and deploying a distinct sub-agent for each yields three distinct benefits: 1. Bounded context windows. By restricting the scope of each generation step to a single document, we limit the context window and task length for each individual LLM. This avoids overwhelming the model, significantly increasing the probability of correct code generation for each isolated task. 2. Targeted human iteration. The orchestrator’s log provides fine-grained visibility into the “hardness” of each design doc. By observing where agents struggle the most and introduce the most bugs, human engineers know exactly which documents require prose refinement and clarification for the next version. 3. Dynamic model routing. The orchestrator can dynamically allocate varying levels of intelligence based on document complexity. While foundational documents (such as the core DSL design) might be routed to a larger, more capable LLM, many downstream design docs can be successfully translated into working code using smaller, more cost-effective models. Cost and iteration speed. In practice, a full clean-slate regeneration of the library takes between 1.5 and 3 hours. The API cost for a complete rebuild using Claude Code is around 100 USD and under a standard high-tier usage plan (e.g., Claude Max), this represents roughly 20% of a weekly usage budget, making continuous, full-library regeneration both practical and economically viable. An opinionated way of writing design docs Conventionally, in the age of AI, practitioners have focused on writing the right tests and high-level rules to ensure code correctness and adherence to project conventions. This is, in essence, a top-down approach—a constitution the generator must obey. We believe a complementary, bottom-up ingredient is crucial: worked examples. Writing out how a piece of pseudo-code executes on a given input, step by step—intermediate shapes, intermediate values, the exact closed-form cost expression that should result—is what maintains consistency across independent agentic code-generations. Like a human learner, in-context learning benefits from a tangible walk-through (Brown et al., 2020; Dong et al., 2022): a concrete trace pins down semantics that prose alone leaves ambiguous, directly attacking failure mode 2 of Section 1. Our docs therefore favor executable-in-your-head vignettes (“on a 2×2×22×2×2 torus with wraparound, the per-node link count is 3, not 6; the all-gather of V bytes therefore costs …”) and every number-bearing doc ends with a reconciliation anchor: a small preset whose expected outputs are stated exactly and enforced by generated tests. 3 A minimal symbolic IR for performance co-design Reliable regeneration also constrains the artifact being specified: the abstractions must be few, orthogonal, and stable under architecture churn. We use a flexible and minimal IR with both a fast mode for large sweeps and a slow mode for careful scheduling studies. The Op abstraction. A model is defined by a single recursively defined operation: ⬇ Op: inputs: List[Tensor] # symbolic shapes outputs: List[Tensor] cost: OpCost # SymPy exprs per key cost (compute, memory, comm) rrt: RRT # resource reservation table: rows = resources, # cols = cycles, cell = units used, # e.g. ("MXU", cycle 3) -> 1 params: Union[InnerLoop(n_iter, body: Graph), # loop nest GraphParams(graph), # subgraph LeafParams(...)] # leaf op An Op is therefore either an interior node—a loop with a trip count and a body graph, or a plain subgraph—or a leaf. The leaf nodes are where the software and system sides meet: a leaf is an op for which the system specifies an RRT and an OpCost. We target TPUs (Jouppi et al., 2023), so the leaves are TPU-shaped: an MXU matmul tile (mxu_op), a VMEM tile load (load_tile_to_vmem), or an ICI collective (allgather). The algorithm side composes leaves into loop nests; the system side prices them. Swapping either side—a new attention variant, a new interconnect generation—touches only its own docs. The builder DSL. Models are authored in a thin Python-embedded tracing DSL and never construct Op nodes by hand: a decorated block traces into a named subgraph, decorated loops become InnerLoop nodes (@smart_loop is a true reduction with a carried accumulator, @smart_map_loop a parallel map), and builder calls emit system-priced leaves. Every dimension is a SymPy symbol, so trip counts like Tq/qblkT_q/q_blk stay symbolic and one trace serves the whole design space. Listing 1 shows the (lightly condensed) flash-attention core: the (B,H,Tq,Tkv)(B,H,T_q,T_kv) score matrix never leaves VMEM, and the asymmetric TqT_q/TkvT_kv make the same nest serve prefill (Tq=Tkv=T_q=T_kv=T) and flash-decoding (Tq=1T_q=1, Tkv=TctxT_kv=T_ctx). ⬇ def flash_attention_core(b, q, k, v, *, B, H, T_q, T_kv, hd_qk, hd_v, q_blk, kv_blk): @smart_map_loop(b, n_iter=T_q/q_blk, name="q_blocks", gather_axis=2) def q_block(): # parallel MAP over query tiles q_i = b.load_tile(q, (B, H, q_blk, hd_qk)) acc0 = b.load_tile(v, (B, H, q_blk, hd_v)) @smart_loop(b, n_iter=T_kv/kv_blk, name="kv_blocks") def kv_block(o_acc): # REDUCTION over kv tiles k_t = b.load_tile(k, (B, H, kv_blk, hd_qk)) v_t = b.load_tile(v, (B, H, kv_blk, hd_v)) s = b.mxu_op("bhte,bhse->bhts", q_i, k_t, name="scores") pv = b.mxu_op("bhts,bhse->bhte", s, v_t, name="attn_v") return b.tensor_add(o_acc, pv, name="accumulate") return b.store_tile(kv_block(acc0)) # this tile’s (B,H,q_blk,hd_v) return q_block() # gathered (B, H, T_q, hd_v) Listing 1: the flash-attention core in the builder DSL. Distribution is expressed as sharding annotations, not hand-placed collectives: tensors name the mesh axes each dimension is sharded on, and a sharded-einsum wrapper infers the collectives from the operand/output shardings—a just-in-time AllGather of a sharded contracting weight, a ReduceScatter when an output is reduced over a sharded dimension. Only layout-moving collectives are explicit: in the DeepSeekMoE block, the dispatch all_to_all moves the expert-parallel axis from the token-group dimension onto the expert dimension (the combine moves it back); every other collective in the block is inferred. Rolling up the Op hierarchy. Two modes turn a tree of per-op costs into wall-clock time. In fast mode, loops are rolled up coarsely: each leaf’s cost is scaled by the product of enclosing trip counts, and simple analytical schedulers model communication/computation overlap (weight pre-collection, all-to-all hiding) as composable transforms in the style of a roofline bound (Williams et al., 2009); evaluation is closed-form, fast enough for sweeps over thousands of design points. In slow mode, each loop is modulo-scheduled (Rau, 1994) into its resource reservation table—software-pipelining the body against per-resource capacity—and the achieved initiation interval rolls up recursively up the tree, yielding dependency- and resource-aware schedules for the design points that sweeps flag as interesting. Symbolic propagation. All cost formulas are propagated upward symbolically and every roll-up produces a closed-form SymPy (Meurer et al., 2017) expression in the free variables of the design space (batch, sequence length, bandwidths, mesh axes, datatype widths). Numeric binding happens only at the edge—one substitution per design point—so a single symbolic build serves an entire sweep. A design doc can even state the exact expected expression for a collective’s cost, and the generated tests assert it. 4 Conclusion Developers routinely use agents to modify code. We believe the smart alternative approach of generating whole systems from design documentations has a lot of benefits and should increasingly be considered. smart today comprises 50 design docs (∼ 9,000 lines of specification prose) spanning TPU topology, collective cost models, numerics, schedulers, and a catalog of frontier model families (dense, MoE, latent-attention, and robotics/VLA variants). We regenerate with every new version change: master is reduced to the docs plus a handful of leaf utilities, and the library is rebuilt by sub-agent orchestration. Treating design docs—rather than code—as the durable artifact pays the incremental-patching debt of Equation 1 down to zero at every regeneration and surfaces vague intent early, because a guess must be written into a doc to survive. The enablers—worked-example docs, a machine-discovered dependency DAG, and a minimal symbolic IR—should generalize wherever specs churn faster than software absorbs them. References Brown et al. [2020] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, 33:1877–1901, 2020. Dong et al. [2022] Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Jingyuan Ma, Rui Li, Heming Xia, Jingjing Xu, Zhiyong Wu, Tianyu Liu, et al. A survey on in-context learning. arXiv preprint arXiv:2301.00234, 2022. Jouppi et al. [2023] Norman P Jouppi, George Kurian, Sheng Li, Peter Ma, Rahul Nagarajan, Lifeng Nai, Nishant Patil, Suvinay Subramanian, Andy Swing, Brian Towles, et al. TPU v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings. In Proceedings of the 50th Annual International Symposium on Computer Architecture, pages 1–14, 2023. Meurer et al. [2017] Aaron Meurer, Christopher P Smith, Mateusz Paprocki, Ondřej Čertík, Sergey B Kirpichev, Matthew Rocklin, AMiT Kumar, Sergiu Ivanov, Jason K Moore, Sartaj Singh, et al. SymPy: symbolic computing in Python. PeerJ Computer Science, 3:e103, 2017. Rau [1994] B Ramakrishna Rau. Iterative modulo scheduling: An algorithm for software pipelining loops. In Proceedings of the 27th Annual International Symposium on Microarchitecture, pages 63–74, 1994. Williams et al. [2009] Samuel Williams, Andrew Waterman, and David Patterson. Roofline: an insightful visual performance model for multicore architectures. Communications of the ACM, 52(4):65–76, 2009.