Learning to Formulate: Agentic Hardware-Software Co-Design of AI Infrastructure through Progressive Optimization and Verification
Organizations: Shenzhen Loop Area Institute (SLAI) · National University of Singapore
Abstract
AI infrastructure, the chips and software that train and serve AI models, is designed in layers. Layers make the system manageable, but they also fix which choices an optimizer can consider. Gains that need the chip and its software to change together are excluded before optimization begins: the formulation leaves them out. Language-model agents make it possible to revisit these boundaries, but an agent's proposal cannot show by itself that a design is correct or optimal. We present Formulaic, a flow that divides the work. The agent writes the formulation, a mathematical solver optimizes within it, and verification and measurement check its assumptions and results. The agent refines the problem step by step, from the computation graph toward the chip. Each step is a template: a family of programs with the conditions under which they are correct. Wall relaxation searches the chip by pricing the hardware limits that bind the software. This paper takes a first step, through separate studies rather than one end-to-end run. On NVIDIA GB10, templates the agent wrote run complete Qwen3-8B layers faster than SGLang. In a simulated Ventus GPU study, joint design cuts the worst-scenario cycles of Qwen2.5-0.5B to 0.428 of the reference, against 0.934 for software alone. On a five-step GEMM refinement chain, reusing proofs makes checking faster.
Figures & tables
| Structure | Example on a GPU | What the search can drop |
|---|---|---|
| What to optimize | ||
| Bottleneck dominance | In decode at batch one, a projection reads each weight once and does about one operation per byte read, so HBM bandwidth alone sets its time. | a full latency model: rank designs by HBM bytes |
| Concavity | A larger tile reuses more data until the kernel becomes compute-bound. Past that point it gains nothing and only uses more shared memory. | every tile size except the few near the knee |
| Monotonicity | Giving each block more shared memory never forces a slower tiling, because every tile that fit before still fits. | regions whose valid lower bounds exceed the incumbent |
| What is independent | ||
| Near-decomposability | The tile staged from HBM into shared memory and the blocking inside the registers interact only through the size of the shared-memory tile. | one three-way search: two smaller ones remain |
| Level | Data transactions at this level | Software-stack structure introduced | Chip structure introduced | What must be proved |
|---|---|---|---|---|
| Workload specification | Required results: what is computed, not how | Required computation, allowed inputs, operator semantics | — | Independent reference contract |
| 0. Computation graph | Tensor values flowing along graph edges | Operators, edges, parameter and state versions | — | Matches the workload specification |
| 1. Operator transformation | Reductions and partial results inside each operator | Fusion, decomposition, reduction partitioning, algorithmic families (e.g. materialized vs. streaming attention) | — | Operator semantics and dependencies preserved |
| 2. Tensor execution mapping | Tile copies between memory levels; values held in storage slots | Tiles, layouts, memory placement, processor mapping | Register and on-chip memory capacity, number of processing units | Coverage, exact-once reductions, address maps, value lifetime |
| 3. Instruction scheduling | Issued and completed transfers; synchronized reads | Instruction selection, transfers, synchronization | Instruction interface, transfer mechanisms, software–hardware boundary | Operand readiness, visibility, progress |
| 4. System architecture | Transfers across units, memory channels and interconnect | Placement of work and data across units | Execution resources, memory hierarchy, interconnect | Resource compatibility, communication contracts |
| Step | Coarse transaction finer transactions | The optimizer gains | The verifier must prove |
|---|---|---|---|
| 1. Reduction partition | one sum of products partial sums over groups, then a combine | the assignment of each reduction index to a group | exact-once: every product contributes exactly once |
| 2. Physical storage | a logical partial result or operand an occupancy of a reusable buffer | buffer assignments and the order of values sharing a buffer | lifetime: a value survives until its last read |
| 3. Asynchronous transfer | an atomic write an issue event and a later completion | issue and completion times; the transfer mechanism (software, instruction or controller) | readiness: an operand is read only after its transfer completes |
| 4. Double buffering | one buffer per value generations rotating through slots, loaded tiles ahead | slots , prefetch distance , trailing barrier , and the shared-memory capacity that bounds them | generation correctness: a consumer reads the right generation; lifetime becomes (Listing 2.2 ) |
| 5. Shared resources | an untimed transfer or access requests to banks, ports and bounded queues | bindings of requests to banks and ports; queue depths and port counts on the chip | conflict freedom, no deadlock, and progress under the declared arbitration contract |
| Source | What goes wrong | Repair |
|---|---|---|
| Representation | The template cannot express the best design | A new or revised template |
| Correctness constraints | The preconditions exclude safe designs | A less restrictive condition that is still sound |
| Optimization | The solver stops short of the model optimum | A better encoding or search |
| Prediction | The cost model ranks designs wrongly | A better model, fitted to measurements |
| Study | Evidence | Who supplied what | What it does not show |
|---|---|---|---|
| GB10 layers ( 5.2 ) | Device latency of complete layers against SGLang; numerical checks on tested inputs | agent wrote program families and cost expressions; production components inherited; solver selects the instance | No formal proof of the GPU program families |
| White vs. black box ( 5.3 ) | One construction run per method from the same eager code under equal caps | agent wrote both programs; host-side help in diagnosis | No demonstrated selection gain over copying; construction-run variability |
| Template reuse / construction ( 5.4 ) | Fixed-family transfer; three construction runs per knowledge condition | Fixed families transferred; agent constructs new families; solver selects instances | Total construction cost including inherited knowledge |
| Ventus co-design ( 5.5 ) | Cycle simulation, calibrated aggregation and extrapolation | Study menus and references set by the authors; solver selects | Cross-study comparisons under a common reference |
| Wall relaxation ( 5.5 ) | Exact inner solves on a 12-chip capacity domain; purchases confirmed by simulation | Author-defined menus; solver-based outer search | A general search-efficiency advantage; random search finds the same optimum |
| Gemmini ( 5.5 ) | Analytical co-design on DoSA’s BERT GEMM list with Timeloop/Accelergy | agent derived costs from the Gemmini source and wrote the enumeration | Automated template construction/reuse loop; complete Transformer inference |
| Fixed-hardware software | Hardware–software co-design | |
|---|---|---|
| Platform | NVIDIA GB10; TileLang 0.1.15 [ 22 ] | Ventus, parameterized RISC-V GPU |
| Workload | Qwen3-8B and Qwen3.5-35B-A3B [ 23 , 24 ] ; BF16, batch one | Qwen2.5-0.5B |
| Input scenarios | Prefill 128/512/2048/4096; decode history 512/2048/4096/8192 | Prefill 128/512; 16 decode steps from history 128/512 |
| Design freedom | Graph structure, kernels, layouts, schedules | Same software freedoms within the study menus, plus chip resources |
| Primary metric | Complete-scope device latency; service TTFT/TPOT | Worst-scenario normalized cycles at a fixed area budget |
| Efficiency | Recorded search/build cost and reuse controls | Formulation/search cost and incremental verification cost |
| Speedup over | |||||
|---|---|---|---|---|---|
| Scenario | Eager | SGLang | Formulaic | Eager | SGLang |
| Prefill 128 | 2211.4 | 1920.7 | 1796.3 | ||
| Prefill 512 | 4797.8 | 2851.8 | 2659.7 | ||
| Prefill 2048 | 36476.3 | 11121.9 | 10087.9 | ||
| Prefill 4096 | 153917.9 | 22799.9 | 20969.3 | ||
| Decode 512 | 2255.1 | 1811.8 | 1680.3 | ||
| Scenario | SGLang | Formulaic | Speedup |
|---|---|---|---|
| Prefill 128 | 22411.5 | 24938.8 | |
| Prefill 512 | 31366.2 | 29034.8 | |
| Prefill 2048 | 53139.5 | 50752.8 | |
| Prefill 4096 | 83451.5 | 79748.2 | |
| Decode 512 | 2843.7 | 2784.0 | |
| Decode 2048 | 2873.9 | 2808.5 |
| History | Eager | Direct | Formulaic | Copied | FI TC | FI SIMT |
|---|---|---|---|---|---|---|
| 512 | 46.14 | 8.44 | 7.78 | 5.67 | 6.78 | 12.02 |
| 2048 | 374.22 | 17.61 | 11.00 | 10.16 | 11.32 | 15.56 |
| 4096 | 783.15 | 30.01 | 15.74 | 15.74 | 15.90 | 20.34 |
| 8192 | 1597.65 | 141.98 | 138.16 | 138.16 | 137.05 | 138.41 |
| KV length | Task only | + Executable template |
|---|---|---|
| 1025 | 0.835 | 1.488 |
| 6145 | 0.942 | 1.158 |
| 24577 | 1.038 | 1.031 |
| Objective | SW only | HW only | Joint | Speedup over SW |
|---|---|---|---|---|
| Prefill | 0.758 | 0.440 | 0.377 | 2.01 |
| Decode | 0.934 | 0.454 | 0.428 | 2.18 |
| All scenarios | 0.934 | 0.454 | 0.428 | 2.18 |
| Prefill | Decode | Mixed | ||||
|---|---|---|---|---|---|---|
| Method | 1024/128 | 2048/256 | 512/256 | 2048/256 | 512/256 | 2048/256 |
| Fixed hardware | .642 (1) | .231 (1) | .431 (1) | .269 (1) | .642 (1) | .269 (1) |
| Sequential | .642 (1) | .231 (1) | .431 (1) | .269 (1) | .642 (1) | .269 (1) |
| Random + exact inner | .231 (12) | .231 (12) | .269 (12) | .269 (12) | .269 (12) | .269 (12) |
| Wall relaxation | .231 (12) | .231 (3) | .269 (8) | .269 (3) | .269 (11) | .269 (3) |
| Coupled wall | .231 (16) | .231 (4) | .269 (11) | .269 (4) | .269 (16) | .269 (4) |
| Method | Energy ( J) | Latency (cycles) | EDP ( J cycles) | Area (mm 2 ) |
|---|---|---|---|---|
| Default chip + CoSA | 2631.85 | 2,963,456 | 7.7994 | 0.708 |
| DoSA | 1983.04 | 917,508 | 1.8195 | 14.088 |
| Formulaic | 1655.24 | 917,508 | 1.5187 | 10.043 |
| Added refinement | Full (ms) | Progressive (ms) | Queries, full/ours | Speedup |
|---|---|---|---|---|
| Reduction partition | 25.5 | 26.3 | 8/8 | 0.97 |
| Storage | 73.8 | 48.7 | 23/15 | 1.51 |
| Asynchronous transfer | 289.0 | 200.6 | 43/20 | 1.50 |
| Double buffering | 528.1 | 218.9 | 65/22 | 2.41 |
| Shared resources | 958.4 | 380.6 | 92/27 | 2.51 |
| Sequence | Full | Node reuse | Query reuse | Full/Query |
|---|---|---|---|---|
| 7 agent proposals | 10.74 | 10.56 | 8.56 | 1.25 |
| 59-version history | 83.04 | 81.66 | 72.17 | 1.15 |
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
| Slots | Trailing barrier | Safe prefetch distances |
|---|---|---|
| 1 | yes | 0 |
| 1 | no | none |
| 2 | yes | 0, 1 |
| 2 | no | 0 |
| 3 | yes | 0, 1, 2 |
| 3 | no | 0, 1 |
| GEMM | Shared memory | Modeled time | |
|---|---|---|---|
| 72 KiB | 1826.9 s, memory-bound | ||
| 24 KiB | 75.1 s, memory-bound | ||
| , ragged tails | 96 KiB, at capacity | 53.1 s, compute-bound |
| Precondition given to the solver | |||
|---|---|---|---|
| Exact, | optimal, passes | optimal, passes | optimal, passes |
| Too strong: always keep the trailing barrier | +0.0%, passes | +0.0%, passes | +1.34% , passes |
| Too weak: , barrier ignored | rejected by verifier | rejected by verifier | rejected by verifier |
| Variant | Failing condition | Replayed defect |
|---|---|---|
| Floor-divided grid (tail uncovered) | coverage | missing reduction terms |
| Unguarded tail | no phantom | write outside ; padding term added |
| Overlapping tiles (stride ) | exact-once | receives twice |
| Single buffer with prefetch one | lifetime | missing term |
| Double buffer, , no trailing barrier | lifetime | the race of Figure 7 |
| Missing leading barrier | readiness, lifetime | reads of uninitialized slot elements |
| M | Solver total | Build/pack | SGLang | Reused selection | Speedup |
|---|---|---|---|---|---|
| 128 | 19.060 | 2.878 | 1349.65 | 1264.97 | 1.067 |
| 512 | 19.162 | 2.851 | 1984.91 | 1505.20 | 1.319 |
| 2048 | 16.782 | 2.575 | 6980.16 | 5973.93 | 1.168 |
| 4096 | 17.081 | 2.499 | 13736.33 | 12096.62 | 1.136 |
| History KV | Split | Online ( ) | Formulaic ( ) | Paired speedup |
|---|---|---|---|---|
| 3072 | 4 | 66.464 | 20.043 | 3.296–3.383 |
| 6144 | 4 | 131.915 | 67.573 | 1.942–2.200 |
| 1536 | 4 | 34.882 | 11.442 | 3.047–3.052 |
| 5120 | 16 | 108.400 | 29.826 | 3.637–3.644 |
| Task | Condition | Latency ( ) | vs. FlashInfer | Active cost (s) |
|---|---|---|---|---|
| Task only | 5.598 [5.326, 6.400] | 0.835 | 2225 | |
| Executable template | 3.142 [3.137, 3.443] | 1.488 | 2069 | |
| Task only | 71.725 [58.424, 72.786] | 0.942 | 2089 | |
| Executable template | 58.362 [56.626, 61.149] | 1.158 | 2083 | |
| Task only | 422.976 [419.722, 423.445] | 1.038 | 1826 | |
| Executable template | 425.649 [424.966, 468.482] | 1.031 | 2075 |
| Sequence | Versions | Full (s) | Node (s) | Query (s) | Full/Node | Full/Query |
|---|---|---|---|---|---|---|
| Controlled, 1 shape | 16 | 4.06 | 2.34 | – | 1.72 | – |
| Controlled, 2 shapes | 16 | 6.48 | 3.94 | – | 1.66 | – |
| Controlled, 4 shapes | 16 | 11.74 | 7.24 | – | 1.62 | – |
| Controlled, 8 shapes | 16 | 25.19 | 15.87 | – | 1.59 | – |
| agent proposals | 7 | 10.74 | 10.56 | 8.56 | 1.02 | 1.25 |
| Complete execution history | 59 | 83.04 | 81.66 | 72.17 | 1.02 | 1.15 |
Hardware-Software Co-Design
Hardware-software co-design jointly shapes algorithms, software, and computing hardware to meet performance, energy, cost, and deployment constraints.
See all 70 papers