Adaptive Self-improvement LLM Agentic System for ML Library Development
Authors: Genghan Zhang
Weixin Liang
Olivia Hsu
Kunle Olukotun
Venue: ICML 2025 poster (2025)
Keywords: LLM agents
Self-improvement learning
Machine learning library
Primary area: applications
TL;DR: How can we use LLMs to improve the efficiency of themselves? We introduce an LLM agentic system with self-improvement for ML library development using hardware architecture specific language, automatically implementing 25 of 26 key LLM operators.

Abstract: ML libraries, often written in architecture-specific programming languages (ASPLs) that target domain-specific architectures, are key to efficient ML systems. However, writing these high-performance ML libraries is challenging because it requires expert knowledge of both ML algorithms and the ASPL. Large language models (LLMs), on the other hand, have shown general coding capabilities. However, challenges remain when using LLMs for generating ML libraries using ASPLs because 1) this task is complicated even for human experts and 2) there are limited code examples due to the esoteric and evolving nature of ASPLs. We present an adaptive self-improvement agentic system that enables LLMs to perform such complex reasoning under limited data by iteratively improving their capability through self-generated experience. In order to evaluate the effectiveness of our system, we construct a benchmark of a typical ML library and generate ASPL code with both open and closed-source LLMs on this benchmark. Our results show improvements of up to $3.9\times$ over a baseline single LLM.

Adaptive Self-improvement LLM Agentic System for ML Library Development
Genghan Zhang'
Weixin Liang"
Olivia Hsu'
Kunle Olukotun
Abstract
b. Adaptive self-improvement learning
ML libraries, often written in architecture-specffic programming languages (ASPLs) that target domain-specific architectures, are key to efîi cient ML systems. However, writing these high- performance MI. libraries is challenging because it requires expert knowledge of both ML algo- rithms and the ASPL. Large language models (LLMs), the other hand, have shown general coding capabilities. However, challenges remain when using LLMs for generating ML libraries using ASPLs because 1) this task is complicated even for human experts and 2) there are limited code examples due to the esoteric and evolving nature of ASPLs. We present an adaptive self- improvement agentic system that enables LLMs to perform such complex reasoning under lim- ited data by iteratively improving their capability through self-generated experience. In order to evaluate the effectiveness of our system, we con- struct a benchmark of typical ML library and generate ASPI. code with both open and closed- source LLMs on this benchmark. Our results show improvements of up 3.9x over baseline single LLM
Select
Multi-level experiences
Filter
Demonstrations
Stratify
Task
LLM
Verifier
Answer
Agents
Agentic system organization
Parallel sampling
Figure I. We propose an adaptive self self-improveeett LLM agentic system. LLM agents start from their base knowledge and ac- cumulate experiences through parallel sampling. Our adaptive self-improvement learning algorithm filters high-quality answers, stratifies the carned experiences by difficulty, and adaptively se- lccts demonstrations enhance LLM agents.
generations of DSAs evolve because new DSAs introduce specialized functions to these existing ASPL. (Choquette, 2023). Efficiently utilizing these new functions requires fundamentally different programming styles and thus new ASPLs (Thakkar et al., 2023; Hagedorn et al., 2023).
Each DSA needs corresponding ML library, collection of ML operators written in its ASPL, before programmers can effeetively use the DSA accelerate ML applications. ML library development is challenging it requires expertise in both ML algorithms and the target ASPL. Essentially, library development is a generation process that composes low- level ASPL primitives into high-level ML operators (Dong et al., 2024; Ye al., 2025).
1. Introduction
With the ending of Dennard Scaling and Moore's Law, computer architectures are specializing in domain appli- cations to achieve greater performance and efficiency and will continue do so (Hennessy & Patterson, 2019). New domain-specific architectures (DSA) typically come with new architecture-specffc programming languages (ASPL), such as CUDA for NVIDIA GPUs (NVIDIA, 2025), HIP for AMD GPUs (AMD, 2025), and Pallas for Google TPUs (Google, 2025). Even existing ASPLs change as
Furthermore, MI. library development using ASPLs requires complex reasoning while minimizing data requirements. ML libraries development often occurs in parallel with hard- ware manufacturing to meet production deadlines (Villa et al., 2021). This time-constrained library-chip co-design process leaves limited code examples. Moreover, this task complicated even for expert human programmers. For exam- ple, the publication of FlashAttentinn-3 (Shah al., 2024) lagged behind the release of the H100 by two years. The challenge is further intensified by the need for MI libraries to co-evolve with new hardware to sustain performance. For instance, directly adapting FlashAttention-2 (Dao, 2023) from the A100 to H100 GPU witnessed 47% performance drop (Spector al., 2024).
'Department of Computer Science, Stanford University, USA. Correspondence to: Genghan Zhang <zgh23 @tta@fsrnndfddu
Proceedings of the 4nnd International Conference on Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
The example code is public https:://// /ttp//gitttttt com/ zhang677 w7t.677771111111 lite

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
The challenges in ML library development call for more automatic solutions. Furthermore, these automatic solutions need to self-improve to perform complex reasoning start- ing from simple and limited data. Large language models (LLMs) have demonstrated emerging capabilities in code generation (Kaplan et al., 2020; Wei et al., 2022a). More- over, empirical evidence implies that LLMs already have the base knowledge of ML algorithms (Ouyang et al., 2024). Therefore, explore the use of LLM agents develop ML libraries with emerging ASPLs.
next-generation reconfigurable dataflow architectures (Prab- hakar et al., 2017), a family of DSAs for Al (Prabhakar & Jairath, 2021; Chen et al., 2023; Prabhakar et al., 2024). The only public document of STeP is non non-archival three-page page workshop publication (Sohn et al., 2024a), which defines its semantics without any code examples or execution en- vironments. Therefore, STeP programs do not exist the training corpus of any LLM.
Current self-improvement methods for LLM agents fall short because of limited exploration or low data efficiency. LLM agents can enhance their performance by synthesiz- ing semantically similar data (Yu et al., 2023; Shinn al., 2024; Zhao et al., 2024). Although these methods are ef- fective for local exploration (Chen et al., 2024), they are insufficient for tasks that require substantial cognitive ef- fort (Huang et al., 2023). Self-improvement learning can significantly improve reasoning ability through reinforce- ment learning (Cobbe et al., 2021; Bai et al., 2022; Singh al., 2023). This approach, however, currently requires hundreds of effective trajectories sampled from LLM agents for each problem (Wang et al., 2024), making it unsuitable for complex scenarios with limited data availability.
Performance Comparison Across Models and Methods
Single
1.3x
Agent Self-improved
0.8
Passinn
0.6
3.5x
0.4
3.6x
3.9x
0 Claude 3.5 Sonnet
GPT-4o
Llama 3.1-4058 DeepSeek-V3 Models
Portion of completed tasks (Pass@n) across models us- ing single LLM, agentic system, and adaptive self- self-improvement agentic system, highlighting performance improvement.
Figure 2.
To address these limitations, we design an adaptive self- improvement learning algorithm integrated with an agentic system organization. This approach not only produces a self-iimpoving agentic system to assist humans but also generates qualityaity ML operators that can be leveraged by other systems. We show our system in Figure 1. Simi- lar to human experiential learning (Kolb, 2014), our tech- niques create self- soff-improvement cycle: LLM agents evolve through earned experiences and these evolved agents can earn more experiences. This self-improvement cycle is fully automated, involving no human experts beyond the ASPL designers themselves, who initially tell the models how to the ASPL primitives.
Putting all these together, our system solves up to 96% of the tasks in our benchmark and achieves up to a 3.9x improvement over a baseline single LLM, as shown in Fig- ure 2. The contributions of this paper are: (1) an adap- tive self-improvement learning algorithm that enables LLM agents to continuously construct MI. libraries through adap- tive experience-dree -diive evolution; (2) an end-to-end agentic system that uses adaptive self- improvemenn to develop an ML library for STeP, an ASPL for next-generation AI accelerator; (3) complete evaluation of the adaptive self- improvement learning algorithm and the integrated agentic system on realistic benchmark constructed from common MI operators.
Our algorithm prioritizes hard-earned experiences gained from completing challenging tasks, inspired by curriculum learning (Bengio et al., 2009). When these hard-earned ex- periences are exhausted, the algorithm adaptively increases the number of demonstrations by incoporating experiences from less challenging tasks. Section 6. shows that hard- earned experiences improve LLM agents more efficiently than mixed ones. While mixed experiences may dilute demonstrations, they can help agents overcome learning obstacles and complete more tasks. As a byproduct, the algo- rithm adaptively increases test-time compute on challenging tasks until they are complete or the data exhausted.
2. Background
In this section, we provide background on how DSAs are programmed through their ASPLs, describe how these AS- PLs are used to create end-user MI libraries, and identify key challenges of this library generation process. We also establish ST&P as the target ASPL to explore LLM tech- niques for ML library development. Key concepts related to the background are listed in Table for reference.
To emulate the library-chip co-design process, choose Streaming Tensor Programs (STeP) target ASPL for library generation. STeP is an emerging ASPL designed for
2

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
alternative current automation techniques that focus simplifying the learning curve for ASPLs. Main- stream automation techniques compromise by optimizing ML operators whose performance is most significantly af- fected by the ASPI. update (Tillet et al., 2019; Thakkar et al., 2023), Only focusing on certain ML operators does not fully incorporate some ASPL updates, such as memory optimiza- tions, which could potentially accelerate any ML operator. Meanwhile, new ML operators being proposed (Gu & Dao, 2023; Sun et al., 2024). Given these factors, we need better automation to improve the productivity of MI. library development using APSLs.
End-to-end of Mixture-of-Experr (4)
MoE
Merger
Expert execution (3)
MLP;
MLP.
Index generation (4)
Router
Norm
FlashAttention with variants (3)
Attention"
BMM
Softmax
Batch matrix-matrix product (3)
BMM
Rotary position embedding (3)
RoPE
RoPE
1.,,,.00,00,11000
Copy
Matrix Multiplication (3)
@Wa
@Wk
ew,
Map frivi MLPW111.
Zip
(60.r))
Partition N
Merge x
Outsuts
Normalization (3)
Norm
Map MLPM11 myy)
(a)
Partition
Expert 0
MU
cu
MU
cu
MU
Figure benchmarked eight groups ML operators within a common LLM layer (Jiang et al., 2024). These operators are categorized as dynamic, static matrix, and static vector, with paren- theses indicating the number of tasks associated with each group.
Si Ripyrfor
cu
CU
MJ
cu
MU
CU
Expert
Merge
2.1. Architecture-speeifie programming languages
MU
GU
MU
GU
MU
ASPLs describe the low-level programming interface of DSA using primitives and specialized functions. Primitives model the basic execution pattern similar to general-purpose programming language constructs, and specialized functions control specialized accelerator units that are optimized for domain applications on the DSA. Unlike domain- speccffcc languages (DSLs), which are top-down distillation of the domain algorithms, ASPLs refer to bottom-up abstraction of the underlying chip architecture.
(b)
(c)
Figure 4. STeP is ASPL for cext-eennration RDAs that models streaming dataflow execution. (a) STeP example program for a sim- plified MoE module. More details can be found in Appendix A.1. (b) An illustration of the streaming dataflow exccution of (a) on an RDA of memory units (MU), and compute units (CU). (c) The SN40L, deployed RDA chip.
2.2. ML library development using ASPLs
2.3. STeP for next-generation RDA
ML libraries developed in ASPLs face portability challenges because ASPLs rapidly evolve align with DSA updates in new generations to meet the demands of growing ML workloads. For example, the matrix multiplication units NVIDIA GPUs and their corresponding MMA instructions have been updated every generation of GPU since their in- troduction (NVIDIA, 2017). Consequently, every library function that uses MMA instructions must be rewritten new ASPL (e.g. CUDA with new instructions) per genera- tion. Moreover, ML libraries need to be shipped at the same time as the chip. In this case, library development costs are no longer negligible.
We chose STeP target ASPL due its potential for better efficiency and status research prototype ASPL. STeP's efficiency potential stems from its role ASPL for wext-generation RDAs, which are promising alterna- tive to GPUs. The SN40L., a deployed RDA implementation shown in Figure 4(c), demonstrates record- booudd-beeakigg infer- ence speeds for Llama 3.1 405B model (SambaNova, 2024). Although STeP does not yet have a path a fabri- cated chip, developing ML libraries in STeP still presents similar challenges as other ASPLs. Writing STeP programs requires complex reasoning about streaming dataflow exe- cution, and our work began without any existing executable STeP programs to reference.
To solve these challenges, we propose to enhance users' learning capabilities for a given ASPL. This approach offers
Similar to other ASPL., primitives and specialized functions
3

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
compose form operational semantics in STeP. Specifically, STeP primitives describe different stream token manipu- lation strategies, and STeP specialized functions express different configurations of RDA units. Inspired by parallel patterns and array programming (Hsu et al., 2023; Rucker et al., 2024), the streaming dataflow execution model in STeP treats data as streams of tokens flowing between com- putational units in space. STeP extends the conventional dataflow execution model of RDAs by introducing stream- ing semantics, as shown in Figure 4(b). This approach uni- fies data values and control signals into stream tokens, em- bedding control flow directly into the data to enable greater dynamism. A stream can be consumed by at most one primi- tive because of the queueing nature of dataflow (Zhang et al., 2021), which is called an affine type constraint in program- ming language theory (Wikipedia, 2024). The affine type constraint is global property of the program since it counts the usage of variable in the whole program.
3. Adaptive self-improvement learning
Our adaptive self-improvement learning evolves LLM agen- tic systems with data generated by the systems themselves. This algorithm samples the agentic system in parallel for cor- rect answers with success rates, filters high-quality correct answers, stratifies the earned experiences, and adaptively updates demonstrations until all the tasks are solved or the demonstrations are exhausted. The complete algorithm for adaptive self- self-improvement learning is shown in Algorithm I. As a byproduct, the algorithm adaptively assigns more test- time compute to harder tasks by excluding tasks from the task set after completion. Figure 5 illustrates running example in which only the unfinished tasks are fed to the agentic system. Additionally, this algorithm is independent of the specific organization of the agentic system, as shown in Figure 1.
3.1. Filtering high-quality answers
STeP primitives are categorized as either arithmetic or shape manipulation. Arithmetic primitives apply computations and control flow to stream tokens. Shape manipulation primitives reshape the data within the stream by changing the control tokens. Figure 4(a) and Figure 17 are example STeP programs that contain 5 arithmetic primitives and only shape manipulation primitives, respectively.
The filter function, a, collects earned experience 'D by fil- tering one answer for each newly solved task in S, and records the success rate of this answer in C.. a first groups the correct answers B, by the isomorphic abstract syntax tree (Knuth, 1968). Then, a groups equívalent answers using AST isomorphism and randomly selects one repre- sentative from each group. Among these representatives, a chooses the one with the minimal length of pure code (excluding comments and empty lines) as the final answer for the task in S,, o also stores the success rate from C, for B. The minimal length selection follows the Minimum Description Length principle for higher information den- sity (Rissanen, 1978). On the other hand, shorter text might lose chain-of-thought comments (Wei et al., 2022b). We purposefully introduce randomness to the selection of rep- resentative answers for each isomorphic group to balance these two contradicting intuitions.
Algorithm/ 1 Adaptive self- self-improvement learning
input X:: task set, m: adaptive granularity
Require: 0: LLM agentic system, r: reward from verifier, a: filter function, 3: selection function
1: D 0
2: t ++0
> iteration
3: repeat
4:
E E+ +-B(D.m) m)
stratification
5:
for for E & do
6:
t + --|,,, eI..e]]
>> selection
3.2. Stratification and selection
7:
Parallle sampling
8:
C. +
(f) v)] 24 E X}
The selection function 3 stratifies the earned experiences D by binning them into m levels of difficulty and demon- strations d, are selected incrementally from the stratified experiences E. We define the difficulty as the opposite of the success rate following (Lightman al., 2023). / sorts D in ascending order of success rate and then bins the tasks as evenly as possible to get the boundaries of each bin. Then tasks are rebinned using these boundary values. This selec- tion strategy can cause repetitive steps as exemplified by the dashed-line circles (Iter 4&5 in Figure 7(a)) when the newly finished tasks are easier than the demonstrations. Our algorithm keeps these repetitive steps instead of avoiding them because the tasks with low success rates have a better chance of getting one correct answer with the number of samples doubled. For example, Iter 4 performs better than
9:
B. + (I., | g) = 1, E x
10:
S. + | > 0,c, E C.)
11:
if B, # then
12:
D DUD(B, C1, St)
filtering
13:
X x+x1 S,
14:
1+1+1
15:
break
16:
end if
17:
t t +-t+1
18:
end for
until X = V (d,- = D B. =
output Solutions: D

---

Adaptive Seff-improvement LLM Agentic System for ML Library Development
Iter
Iter
Iter
Iter
12
Agentic system
Agentic system
Agentic systemm
Agentic system
Agentic system
Agentic system
Correct samples
9
1
3
2 3
2
Hard
Hard+Madium
Hard
Hard+Medium
Hard-MedilmmEaay
Passøn:
Pass@n:
Pass@n:
Figure
Running example Algorithm with m = | = 5. Orange circles are demonstrations d, for current iteration.
Green circles finished tasks and white ones not. means 3 stratifies demonstrations into 3 levels: hard, medium, and easy. Iter I and 3 consider hard-only examples, and the other iterations consider mixed examples.
Domain expert
Structural IR
Host language (Python)
Base Prompt
Single
Agent 1
Agent 2
Pytest
LLM
Task
Proposer
Gate
Corrector
Code generator
STeP Program
Input
Prompt Composer
API Query
Output
Type Checker
Configuration
User
Agentic system
|... User
--.
(a) A single LLM
Agentic system organization
Figure Details our agentic system organization. (a) shows single LLM. (b) shows system components and their representations. The user is either human or seff-improvement learning process. The filled colors align with the text colors in Section 4.
Iter 3 in Figure 7(b) with the same demonstrations. This method can also cause later iterations to have fewer tokens but with higher quality than previous iterations (Iter 2&3 in Figure 7(a)) when the boundary value crosses two bins and the newly finished tasks are easier.
4.1. LLM agents
As shown in Figure 6(a), each single LLM is designed for specific purpose assigned by the domain expert through the base prompt. The base prompt contains the task description, base knowledge, and demonstrations. The prompt composer chains the base prompt and task susk-specific input, which is then fed into a configured LLM serving API.
3.3. Discussion
Algorithm can also extend to new tasks that typically in- volve new ML operators and new hardware specialized func- tions. These can be incorporated into the initial task pool X' and handled using the same adaptive self-improvement learning process by selectively sampling only the new tasks. Each example takes about 500 tokens, so cutting-ed, -dgee LLMs can handle hundreds more tasks. If the number exceeds the context length, better stratification and selection functions are needed to preserve experience quality within the context window limit. If the task X' finite, as our case, the algorithm will terminate.
Specifically, we design two agents, proposer and guardian. As explained in Section 2.3, the affine type con- straint in STeP is global property that requires thinking back and forth beyond step-by-step reasoning, which is chal- lenging for the causal generation of LLMs. Therefore, we design a guardian agent to check and correct the affine type error globally. We exclude the demonstration tasks in the base prompt of the agents from the benchmark to avoid the model directly copying the answer.
The proposer agent generates candidate STeP implemen- tation whose base prompt comes from ASPL designers. The base prompt is composed of STeP references and usage patterns. Figure 14 shows the reference for the Accum prim- itive, and Figure 17 exemplifies one of the usage patterns for shape manipulation.
4. Agentic system organization
In this section, we introduce a specific agentic system or- ganization tailored for ML library development using STeP as shown Figure 6. The agentic system comprises LLM agents, a code generator, and verifiers, with struetural in- termediate representation the interface between users and system components.
The guardian agent decides whether the output of the pro- poser violates the affine type constraint and corrects the implementation when necessary. The guardian agent con- sists a fused gate and corrector. Figure 20 shows the base
5

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
prompt for the guardian agent, which provides input and out- put examples of variable reuse where the variable is reused zero, one, or two times.
Vasilache al., 2022) and use structural IR unify the interfaces of agentic system. Specifically, this structural IR is the interface between users and the agentic system and between LLM agents and the code generator within the system.
4.2. Code generator
After the LLM agents, a code generator takes in the gener- ated implementation and outputs self-contained pytestable Python script. With this code generator, the LLM agents only need to output implementations without other helper code. We embed the STeP specifications written in natural language to Python (see Figure 21 and Figure 22). Beneath this Python frontend is a functional simulator that calcu- lates the result of STeP programs. Since the essence of each ASPL is the programming abstraction (semantics) instead of its syntax (Liskov, 1), we choose to prompt LLMs with their familiar Python syntax. In this way. LLMs can focus more on reasoning about the STeP programming abstraction without being distracted by less familiar syntax.
Our structural IR encodes necessary information using a data serialization language. It externalizes and condenses programs instead of simply formatting the prompt without changing the content. Comparing the structural IR in Fig- ure 18 with the equivalent bare Python in Figure 19 for the same task, users only need to state two things: the ML operator to implement and the specialized functions as in Figure 15 without any redundant glue string. Moreover, structural IR saves tokens by reducing redundant prompts, allowing for more demonstrations.
5. Benchmark
We construct a set of tasks to measure the adaptive self- improvement agentic system proposed in Section and Sec- tion 4. This benchmark should cover a diverse set of popular ML operators and specialized functions. In total, we collect 26 tasks covering 8 groups of ML operators in common LLM model architectures, as shown in Figure 3.
4.3. Verifier
Fast verification is vital because it bottlenecks adaptive self- improvement learning. ASPLs further increase this complex- co-design
ity by requiring simulator for the library- process. Simulating STeP general dataflow system in Python would be slow because of its high dynamism (Zhang et al., 2024b). Since we only focus on functional correctness in this work, we meticulously limit the level of dynamism to the degree that ML. operators require. Consequently, our simulator emulates stream execution using tensor computa- tion with the necessary control flow.
5.1. Metric
We choose pass@k (Chen al., 2021) as the metric for task completion. Pass@k is calculated as Equation (1) given T tasks, n samples of the agentic system, and e; correct responses for each task i. It is useful for us to analyze the metric at two extremes: pass@ @ and pass@n. Pass@1 is the expectation of the success rate across tasks. @ass@n is the expectation of the portion of tasks that can be solved given all samples.
Our system organization contains two verifiers, and the re- ward function for our system in Algorithm 1 is r(z, y) for task if answer passes these two verifiers and r(z,)) v) = 0 otherwise. One verifier checks for functional correctness. Users program PyTorch to express their ML operators, which elicits L.LMs' base knowledge of MI. al- gorithms. The verifier compares the execution results of our simulator with the result tensors of the corresponding PyTorch program a single set of shapes with random input values. The fidelity of our unit test method builds on practice (Jia et al., 2019) and theory (Gulwani & Necula, 2003). The other verifier checks the affine type constraint by performing static analysis on the abstraction syntax tree of the STeP program with the Python ast module (Ronacher, 2008).
(n-c))
pass @k
(I)
5.2. Benchmark Construction
We construct the benchmark from first principles and do not favor any kind of task. Firstly, the number of tasks in each group is nearly the same as shown in Figure 3. Secondly, the benchmark has an even distribution of difficulties. Table 2 shows that both tasks requiring shape and arithmetic primi- tives and tasks with and without reused variables distribute fairly evenly.
4.4. Structual intermediate representation
We provide reference implementation for each task. These oracle implementations ensure that each task has at least one correct answer. Each task of one type has either differ- ent specialized functions for the same operator or different operators with different specialized functions. More details on the benchmark are in Appendix A.2.
Good interfaces can improve the performance of agentic coding systems (Yang al., 2024; Wei et al., 2024). There- fore, we are inspired by the intermediate representation (IR) technique from compiler literature (I (Lattner et al., 2021;
6

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
6. Experiments
Success diversity
Failure diversity
Overall diversity
Method
@
Detailed experimental settings are in Appendix A.3, We also benchmark the tasks with OpenAI-ol in Table but do not include it in the following experiments control for test- time compute. As described in Section 4.4, all prompts are formatted in YAML because structural prompts generally benefit (He et al., 2024).
Single- w/o-IR
0.32
0.47
0.41
0.46
Single Agent 0.34
0.27
0.64
0.50
0.62
0.68
0.52
0.85
Table I. Analysis of the correlation between semantic diversity answers and the performance. "-w/o-IR" means "without structural IR". Higher values indicate higher semantic diversity.
6.1. Analysis of adaptive seff-improvement learning
This section uses the agentic system organization described in Section 4.1 for the best possible base learning capability. Our evaluation provides the following insights:
Mode
Metrie
Method
Once
Reuse
Avg
Single
0.685
0.232
0.444
The hard-only examples improve performance more ef- fectively than examples mixed with easier ones. As shown in Figure the Pareto optimal is composed of hard exam- ples (denoted by "H"" for all three models. Moreover, hard examples bring the most significant improvement along the learning curve. In some cases, fewer hard examples may perform better than more examples mixed with easier exam- ples. For example, Iter 3 has better performance than Iter 2 for gpt-4o.
Pass @
Agent
0.663
0.455
0.552
Arith
Single
7/7
5/8
12/15
Pass n
Agent
7/7
8/8
15/15
Single
0.145
0.004
0.094
Pass @
Agent
0.167
0.051
0.125
Shape
Single
3/7
1/4
4/11
Pass @ n
Agent
4/7
3/4
7/11
Single
0.415
0.156
0.296
Pass @
Agent
0.415
0.320
0.371
Mixed examples are required to generate better hard- only examples. In DeepSeck-V3, although HM, performs the same as H₃ and HME, performs worse than H2 while taking more tokens, H2 would not be discovered without HM, and HME, Although the new hard-only examples do not necessarily improve the performance, they can save input tokens, as exemplified by Hs of gpt-4o and HM&&H, of llama.
Avg
Single
10/14
6/12
16/26
Pass @n
Agent
11/14
11/12
22/26
Table 2. Analysis of improvement brought by the agentic method. "Arith" and "Shape" mean the oracle implementation only involves arithmetic primitives and involves shape manipulation primitives introduced in Section 2.3, respectively. "Once" and "Reuse" mean all the streams are used once and more than once, respectively.
Adaptive granularity affects token efficiency and peak performance. If the adaptive granularity m is too small, then easy examples might dilute the difficulty of training data. If m is too large, then exploration steps might be too conservative and thus waste input tokens. Therefore, there a sweet spot that balances the training data difficulty and cost of input tokens. As shown in Figure 8, 3 is 3 saves 1.07x tokens over maintaining performance. Additionally, = 3 improves the performance by 1.5x at a similar input token cost when compared to m 1.
provided in Figure 16, which is considered challenging for ordinary programmers.
Our structual IR improves performance by increasing the sample diversity. Figure 9 shows the efficacy of our structural IR and agentic method. We calculate the semantic diversity of sampled answers. group of answers is con- sidered have the same semantics if their abstract syntax trees are isomorphic. We calculate success diversity the number of semantically different correct answers divided by the total number of correct answers, averaged across all tasks. Failure diversity calculated in a similar way but for wrong answers. Overall, diversity combines both metrics. As shown in Table 1, Pass @ has positive correlation with failure, overall diversity, and the complexity of the methods. However, structural IR can hurt success diversity.
that spot.
4 while
6.2. Ablation study of agentic system organization
The experiments below study the base learning capability of agentic systems without the self-improvement process. The agentic system can discover non-trivial STeP programs. Surprisingly, the agentic system composes attention oper- ators with over 50% Pass@1
shown in Figure 13. That means the agentic system can discover online softmax (Mi- lakov & Gimelshein, 2018) 8) and memory-free streaming attention (Sohn et al., 2024b) with specialized functions
The guardian agent can correct affine type errors but might corrupt the correct answers output by the pro- poser agent. As shown in Table 2, the guardian agent effectively corrects the proposer agent, solving 5 extra reuse tasks (Pass@n @n of Avg-Reuse from 6/12 to 11/12). Notably,

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
GPT-4o
Liama 3.1-405B
DeepSeek-V3
1.0
1.0
1.0
HMMe
0.8
0.8
0: HMM
HMESs
HMs
H:M
0.6
2:
Pessigi
PessiBin
Hb
Ho
2: Hi
0.2
intermediate Pareto optimal
D.None
intermediate Peretoo optimel
None
Intermed
Parelo optimel
8000
6000
8000
10000
12000
2000
6000
8000 Tokens
10000
12000
4c0o
6000
8000 (6)
10000
12000
Irput Toklers
(A)
Figure Adaptive self- eeff-improeement learning improves the agentic system with data generated by itself. The input tokens are averaged across tasks. A task takes increasing input tokens until success or data is used up. "9: HME,, means 9-th iteration, 5-th cycle of adaptive sampling, demonstrations contain hard (H), medium (M), and easy eEa-nd experiences. "None" means no examples iteration. For Iter if Pass@n;-1 > Pass mass@n, then new cycle of adaptive sampling starts from "H"" Otherwise, the current cycle continues the order HM HMM., The dash lines are connected the itcration order. Claude 3.5 Sonnet result Figure 11.
first
P Input Tokens
Pass@k Method
1.0
1.0
1.07x
Agent
0.8
0.8
Single
Single-w/o-IR
Passigin
0.6
Pass@ik
0.4
0.4
m=1
0.2
m-2
0.2
0.0
4000
6000
8000
10000
12000
10 Number of samples (k)
Input Tokens
Figure Hyperparameter tuning of adaptive granularity m on GPT-4o. This supports using m=3 for Figure? 7.
Figure 9. Pass@k against the number of samples for single LLM without structural IR, baseline single LLM, and agentic system.
7.1. Self-improvement learning for LLMs
all the Arith-Reuse tasks can be solved by the agentic system (Pass@n of Arith-Reuse from 5/8 to 8/8). Surprisingly, the agentic system also helps with non-reuse tasks (Pass@nn of Shape-One from 3/7 4/7). However, the agentic system reduces the Pass @ 1 of Arith-Once from 0.685 0.663, im- plying that the guardian agent might corrupt the proposer's output. The agentic system can compensate for such cor- ruption by finishing more tasks, resulting in an unchanged @ass@1 I of Avg-Once.
Self-improvement learning for LLMs typically involves two stages: scoring generated samples (trajectories) and incor- porating those samples to enhance the model. Scoring can be achieved through human labeling (Cobbe et al., 2021; Lightman al., 2023) through automated methods such as verifiers and heuristics (Wang et al., 2024; Singh et al., 2023). Our method stands out in this context by utilizing AST analysis, offering more interpretable approach to scor- ing. When it comes to incorporating samples, models may rely on retrieval (Zhao et al., 2024; Park al., 2023), reflec- tion (Shinn et al., 2024; Liu al., 2023: Yao al., 2023), or reward feedback (Opsahl-Ong et al., 2024; Fernando et al., 2023). Our method introduces new mechanism in this stage by adaptively extending and prioritizing high-scoring samples.
7. Related Work
This work introduces self-in inffiipproeeent agentic system for ASPL code generation, enhancing L.M effectiveness through an adaptive learning algorithm and deliberate agen- tic system design. Accordingly, we focus the related work on approaches that leverage LLMs for self celf-improvement and agentic systems on specialized tasks.
Our method shares similar reinforcement principles with self-improvement learning at the post-training stage (Bai
8

---

Adaptive Seff-improvement LLM Agentic System for ML Library Development
al., 2022; Gulcehre et al., 2023; Tian al., 2024) but is rewarded task-level granularity instead of token-level. Specifically, each action is program implementation in- stead of token prediction and the state is defined by earned experiences rather than generated sequences.
tial implications warrant ongoing discussion and careful
consideration the technology develops.
Acknowledgement
We thank Vishnu Sarukkai, Haochen Shi, Xinhao Li, Gina Sohn, Nathan Zhang, Tian Zhao, Yu Sun, Anjiang Wei, Chun Deng, Qizheng Zhang, Rubens Lacouture, Fredrik Kjolstad and the anonymous reviewers for their feedback this paper. This research supported in part by the Pervasive Parallelism Lab Affiliates Program Fund. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the aforementioned funding agencies.
7.2. Agentic system organization for specialized tasks
Task-specific organization has proven effective in enhancing the performance of agentic systems across diverse coding tasks (Zhang et al., 2024c; Fang et al., 2024; Guan et al., 2024). We adopt an agentic system organization specifically for ML library development using an ASPL. Such domain knowledge can be further augmented with automatic agentic system design tools (Khattab et al., 2023; Hu et al., 2024; Zhang et al., 2024a). Furthermore, well well-designed interfaces between agents, tools, and other agents have been shown to improve performance (Schick et al., 2023; Yang et al., 2024; Wu et al., 2023), and a structural IR enables these interfaces to be highly task-aligned in our system.
References
AMD. Amd hip documentation, 2025.
URL
https://roam, docs, amd. com/projects/ HIP hie//nnlateest/ Accessed: 2025-01-12.
8. Conclusion
Bai, Y., Kadavath, S., Kundu, S., Askell, A., Kernion, J.,
ML library development using ASPLs is a critical compo- nent of the ML ecosystem, but it remains poorly automated. To address this limitation, we co-design the learning process and agentic system around central objective: enabling complex reasoning with limited data. Our methods simulta- neously implement non- non-trivial ML operators and produce self-improving agent. Consequently, our system not only supports experts in developing ML libraries but also of- fers valuable resources for other systems. We recognize that the challenges of complex reasoning under limited data extend beyond this domain and envision applying similar self-imppovement techniques with LLM agents to a broad class of tasks.
Jones, A., Chen, A., Goldie, A., Mirhoseini, A., McKin- non, C., et al. Constitutional ai: Harmlessness from ai feedback. arXiv preprint arXiv:2212.08073 2022.
Bengio, Y., Louradour, J., Collobert, R., and Weston, J. Curriculum learning. In Proceedings of the 26th annual international conference on machine learning, pp. 41-48, 2009.
Chen, M., Tworek, J., Jun, H., Yuan, Q.. Pinto, H. P.D.D. O., Kaplan, J., Edwards, H., Burda, Joseph, N., Brockman. G., et al. Evaluating large language models trained on code. arXiv preprint arXiv: 2021.
Chen, 7., Huang, D., Wang, M., Yang. B., Shin, J. L.., Hu, C.,
Li, B., Prabhakar, R., Deng, G., Sheng, Y., et al. Ai soc design challenges in the foundation model era. In 2023 IEEE Custom Integrated Circuits Conference (CICC), pp. 1-8. IEEE, 2023.
Impact Statement
Our work on adaptive self-improvement learning LLM agen- tic systems for automating ML library development using ASPLs has several potential societal implications. The pri- mary impact is the potential significantly enhance the productivity of ML library developers, which could accel- erate the development of more efficient ML systems. This advancement could democratize access to ML development tools and reduce the technical barriers to entry in the field. Additionally, our technique offers an approach for deploy- ing I.IM agents in scenarios requiring complex reasoning with limited data availability. While these developments primarily aim to advance the field of Machine Learning, acknowledge that increased automation in software de- velopment could impact the nature of programming work and skills required in the field. We believe these poten-
Chen, Z., Deng, Y., Yuan, H., Ji, K., and Gu, Q. Self-play fine-tuning converts weak language models to strong lan- guage models. arXiv preprint arXiv:2401.01335, 2024.
Choquette, J. Nvidia hopper h100 gpu: Scaling performance. IEEE Micro, 43(3):9-17, 2023.
Cobbe, K.,, Kosaraju. V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10. 14168, 2021.
Dao, T. Flashattention-22 Faster attention with bet- ter parallelism and work partitioning. arXiv preprint urXiv:2307.0869/ 2023.

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. Flashat- tention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Process- ing Systems, 35:163444100 16359, 2022.
Hennessy, J. L. and Patterson, D. A. A new golden age for computer architecture. Communications of the ACM, 62 (2):48-60, 2019.
Hsu, O., Strange, M., Sharma, R., Won, J., Olukotun, K., Emer, S., Horowitz, M. A., and Kjølstad, The sparse abstract machine. In Proceedings of the 28th ACM In- ternational Conference on Architectural Support for Pro- gramming Languages and Operating Systems, Volume 3, pp. 710-726, 2023.
Dong, J., Feng, B., Guessous, D.. Liang, Y., and He, H. Flex attention: programming model for gen- erating optimized attention kernels. arXiv preprint arXiv:2412.05596, 2024.
Fang. W., Li, M.. Li, M., Yan, Z., Liu, S., Zhang, H., and Xie, Z. Assertllm: Generating hardware verification as- sertions from design specifications via multi-llms. In 2024 IEEE LLM Aided Design Workshop (LAD), pp. 1-1. IEEE, 2024.
Hu, S., Lu, C., and Clune, Automated design of agentic systems. arXiv preprint 2rXiv:2408.08435, 2024.
Huang, J., Chen, X., Mishra, S., Zheng, S., Yu, W., Song, X., and Zhou, Large language mod- els cannot self-correct reasoning yet. arXiv preprint arXiv:2310.0199,, 2023.
Fernando, C., Banarse, D., Michalewski, H., Osindero, S., and Rocktäschel, T. Promptbreeder: Self-referential self-improvement via prompt evolution. arXiv preprint arXiv:2309.16797, 2023.
Jia, Z., Padon, O., Thomas, J., Warszawski, T., Zaharia, M., and Aiken, A. Taso: optimizing deep learning computa- tion with automatic generation of graph substitutions. In Proceedings of the 27th ACM Symposium Operating Systems Principles, pp. 47-62, 2019.
Google. Pallas: jax kernel language, 2025, URL latest/ Accessed: 2025-01-12.
https://jax. readthedocs hes 1o/ en pallas/tpu/
Jiang, Q. Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, S., Casas. D. d. 1., Hanna, E. B., Bressand, F., al. Mixtral of experts. arXiv preprint arXiv:2401.04088 2024.
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:23/2.00755: 2023.
Guan, Y., Wang, D., Chu, Z., Wang, S., Ni, F., Song, R., and Zhuang, C. Intelligent agents with Ilm-based process automation. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 5018-5027, 2024.
Kaplan, J., McCandlish, S., Henighan, T., Brown, B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836// 2020.
Khattab, O., Singhvi, A., Maheshwari, P., Zhang, Z., San- thanam, K., Vardhamanan, S., Haq. S., Sharma, A.. Joshi, T. T., Moazam, H., al. Dspy: Compiling declarative language model calls into self-improving pipelines. arXiv preprint arXiv:2310.03711, 2023.
Gulcehre, C., Paine, T. L., Srinivasan, S., Konyushkova, K., Weerts, L., Sharma, A., Siddhant, A., Ahern, A., Wang, M., Gu, C., et al. Reinforced self-training (rest) for language modeling. arXiv preprint arXiv:2308.08990 2023.
Knuth, D. E. Semantics of context-free languages. Mathe-
Gulwani, S. and Necula, G. C. Discovering affine equali- ties using random interpretation. In Proceedings of the 30th ACM SIGPLAN-SIGACT symposium on Principles of programming languages, pp. 74-84, 2003.
matical systems theory. 2(2):127-145, 1968.
Kolb, D. A. Experiential learning: Experience as the source of learning and development. FT press, 2014.
Hagedorn, B., Fan, B., Chen, H., Cecka, C., Garland, M., and Grover, V. Graphene: An ir for optimized tensor computations on gpus. Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3, pp. 302-313, 2023.
Lattner, C., Amini, M., Bondhugula, U., Cohen, A., Davis, A., Pienaar, J., Riddle, R., Shpeisman, T., Vasilache, N., and Zinenko, O. Mlir: Scaling compiler infrastructure for domain specific computation. International Symposium on Code Generation and Opti- mization (CGO), pp. 2-14. IEEE, 2021.
2021 IEEE/ACM
He, J., Rungta, M., Koleczek, D., Sekhon, A., Wang, F. X., and Hasan, S. Does prompt formatting have any impact on I1m performance? arXiv preprint arX1v22411.10541, 2024.
Lightman, H., Kosaraju, V., Burda, Y., Edwards, H., Baker, B., Lee, T., Leike, J., Schulman, J., Sutskever, I., and Cobbe, K. Let's verify step by step. arXiv preprint arXiv:2305. 20050, 2023.
10

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
Liskov, B. The Power of Abstraction, pp. 2008. Association for Computing Machinery, New York, NY, USA, 2011. ISBN 9781450310499, 1145/1283920 0115511283920.192222.. 1962422.
Puschel, M., Moura, J. M., Johnson, R., Padua, D., Veloso,
M. M., Singer, B. W., Xiong, J., Franchetti, F,, Gacic, A., Voronenko, Y., et al. Spiral: Code generation for dsp transforms. Proceedings of the IEEE, 93(2):232-275, 2005.
Liu, Z., Bahety, A., and Song, S. Reflect: Summarizing robot experiences for failure explanation and correction. arXiv preprint arXiv:2306. 1...6.553 15724, 2023.
Raposo, D., Ritter, S., Richards, B., Lillicrap. T. Humphreys, P. C., and Santoro, A. Mixture-of- Dynamically allocating compute in transformer-based lan- guage models. arXiv preprinl urXiv:2404.02258, 2024.
Milakov, M. and Gimelshein, N. Online normalizer cal- culation for softmax. arXiv preprint arXiv:1805.02800 2018.
Rissanen, Modeling by shortest data description. Auto- matica, 14(5):465-471, 1978.
NVIDIA. Nvidia tesla v100 gpu architecture, the world's most advanced data center gpu, 2017. URL https: htt./mmaaes.vviii volta-archhe sewpittecture-whitepee pdf. Whitepaper.
Ronacher, A. Python ast module, 2008 URL https: t/docs..pthhaeee www.hh...ggg.............. waarrg///ibaaaar /ast. aat.attttt html.
Rucker, A. C., Sundram, S., Smith, C., Vilim, M. Prabhakar, R., Kjølstad, F., and Olukotun, K. Revet: language and compiler for dataflow threads. 2024 IEEE In- ternational Symposium on High-Performant Computer Architecture (HPCA), pp. 1-14. IEEE, 2024.
volta
NVIDIA. Cuda c++ programming guide, 2025. URL /ttps://docs..o nvidia com/ cuda/ cedda-c-rrogaamiin
index. html.
Accessed: 2025-01-12.
SambaNova. Rdu: The gpu alternative, 2024. URL ai-chip..
Opsahl-Ong, K., Ryan, M. J., Purtell, J., Broman, D., Potts, C., Zaharia, M., and Khattab, O. Optimizing instruc- tions and demonstrations for multi-stage language model programs. arXiv preprint arXiv:2406. 1695, 2024.
Schick, T., Dwivedii Yu, J., Dessi, R., Raileanu, R., Lomeli, M., Hambro, E., Zettlemoyer, L., Cancedda, N., and Scialom, T. Toolformer: Language models can teach themselves to use tools. Advances Neural Information Processing Systems, 36:68539 000668559668551,10 23551,2023
Ouyang, A., Guo, S., and Mirhoseini, A. Kernl- bench: Can Iims write gpu kernels?, 2024. URL https:/ wew.caaiinnneelleennn.oe.e stanford. edu/blogs//ernellenn
Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., and Dao, T. Flashattention-3: Fast and accurate atten- tion with asynchrony and low-precision. arXiv preprint arXh::2407.08660 2024.
Park, J. S., O' Brien, J., Cai, C. Morris, M. R., Liang, P., and Bernstein, M. S. Generative agents: Interactive simulacra of human behavior. In Proceedings of the 36th annual acm symposium on user interfuce software and technology, pp. 1-22, 2023.
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q. Hinton, G., and Dean, J. Outrageously large neural networks: The sparsely-gated mixture-of-cxperts layer. arXiv preprint arXiv: 1701. 00100000.658,. 2017.
Prabhakar, R. and Jairath, S. Sambanova 8010 rdu: Accel- erating software 2.0 with dataflow. In 2021 IEEE Hot Chips 33 Symposium (HCS), pp. 1-37. IEEE, 2021.
Shinn, N., Cassano, F., Gopinath, A., Narasimhan, K., and Yao, S. Reflexion: Language agents with verbal rein- forcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Prabhakar, R., Zhang, Y., Koeplinger, D.,, Feldman, M., Zhao, T., Hadjis. S., Pedram, A., Kozyrakis, C., and Olukotun, K. Plasticine: A reconfigurable architecture for parallel paterns. In Proceedings of the 44th Annual International Symposium on Computer Architecture, pp. 389-402, 2017.
Singh, A., Co-Reyes, J. D., Agarwal, R., Anand, A., Patil, P., Garcia, X., Liu, P. J., Harrison, J., Lee, J., Xu, K., et al. Beyond human data: Scaling self-training for problem-solving with language models. arXiv preprint arXiv:2312.06555, 2023.
Prabhakar, R., Sivaramakrishnan, R., Gandhi, D.,. Du, Y., Wang, M., Song, X., Zhang, K., Gao, T., Wang, A., Li, X., et al. Sambanova sn401: Scaling the ai memory wall with dataflow and composition of experts. In 2024 57th JEEE/ACM International Symposium on Microarchitee-- ture (MICRO), pp. 1353-1366. IEEE, 2024.
Snell, C., Lee, J., Xu, K., and Kumar, A. Scaling Ilm test- time compute optimally can be more effective than seal- ing model parameters. arXiv preprint ABt::2008.03314 2024.

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
Sohn, G., Gyurgyik, C., Zhang, G., Velury, S., Mure, P., Zhang, N., and Olukotun, K. Streaming tensor programs: A programming abstraction for streaming dataflow accel- erators. In ASPLOS Young Architect Workshop (YArch), 2024a.
Wang, P., Li, L.,, Shao, Z., Xu, R., Dai, D., Li, Y., Chen, D., Wu, Y., and Sui, Z. Math-shepherd: Verify and reinforce Ilms step-by-step without human annotations. In Proceed- ings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 9426-9433 2024.
Sohn. G., Zhang, N., and Olukotun, K. Implementing and optimizing the scaled dot-product attention on streaming dataflow. arXiv preprint arXiv:2404. 16629, 2024b.
Wei, A., Nie, A., Teixeira, T. S., Yadav, R., Lee, W., Wang, K., and Aiken, A. Improving parallel program perfor- mance through dsl-driven code generation with Ilm opti- mizers. arXiv preprini arXiv:2410. 15625, 2024.
Spector, B. F., Arora, S., Singhal, A., Fu, D. Y., and Ré, C. Thunderkittens: Simple, fast, and adorable kernels. arXiv preprint arXiv:2410.2000 200:2999..0099 2024.
Wei, J., Tay, Y., Bommasani, R., Raffel, C., Zoph, B., Borgeaud, S., Yogatama, D., Bosma, M., Zhou, D., Met- zler, D., et al. Emergent abilities of large language models. arXiv preprint arXiv:2206. 07682, 2022a.
Sun, Y., Li, X., Dalal, K., Xu, J., Vikram, A., Zhang, G., Dubois, Y., Chen, X., Wang, X., Koyejo, S., et al. Lear- ing to (learn at test time): Rons with expressive hidden states. arXiv preprint arXiv:2407.04620. 2024.
Wei, J., Wang, X., Schuurmans, D., Bosma, M., Xia, F., Chi, E., Le, Q. V., Zhou, D., et al. Chain-of of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24837. 2022b.
Thakkar, V., Ramani, P., Cecka, C., Shivam, A., Lu, H., Yan, E., Kosaian, J., Hoemmen, M., Wu, H., Kerr, A., Nicely, M., Merril,, D., Blasig, D., Qiao, F., Majcher, P., Springer, P. Hohnerbach, M., Wang, J., and Gupta, M. Cut- lass. attps://githuie wwt.///i///////./... cutlas,, 2023. URL https: http::://githbb.. wow.coo////o/// NVIDIA/ cutlass.
Wikipedia. Substructural type system Wikipedia, the free encyclopedia, 2024. URL https: // wikipedia. org
t_ttp_______________ _erpe____semtt Affine MEfine__tye__syteeee Accessed: 2025-01-12.
Tian, Y., Peng, B., Song, L., Jin, L., Yu, D., Mi, H., and Yu, D. Toward self-improvement of Iims via imagination, searching, and criticizing. arXiv preprint arXiv:2404. 12253, 2024.
Wu, Q., Bansal, G., Zhang, J., Wu, Y., Zhang, S., Zhu, E., Li, B., Jing,, L., Zhang, X., and Wang, C. Autogen: Enabling next-gen Ilm applications via framework. arXiv preprint arXiv:2308.08133 2023.
conversation
Tillet, P., Kung, H.-T., and Cox, D. Triton: intermediate language and compiler for tiled neural network computa- tions. In Proceedings of the 3rd ACM SIGPLAN Interna- tional Workshop on Machine Learning and Programming Languages, pp. 10-19, 2019.
Yang, J., Jimenez, C. E., Wettig, A., Lieret, K., Yao, S., Narasimhan, K., and Press, O. Swe-agent: Agent- computer interfaces enable automated software engineer- ing. arXiv preprint arXiv: /5793, 2024.
Yao, S., Yu, D., Zhao, J., Shafran, I., Griffiths, T., Cao, Y., and Narasimhan, K. Tree of thoughts: Deliberate problem solving with large language models. Advances in neural information processing systems, 36: 1 1808-8198822 2023.
Vasilache, N., Zinenko, O., Bik, A. J., Ravishankar, M., Raoux, T., Belyaev, A., Springer, M., Gysi, T., Caballero, D., Herhut, S., et al. Composable and modular code generation in mlir: A structured and retargetable ap- proach to tensor compiler construction. arXiv preprint arXiv:2202.0322 03290, 2022.
Ye, Z., Chen, L., Lai, R., Lin, W., Zhang, Y., Wang, S., Chen, T., Kasikei, B., Grover, V., Krishnamurthy, A., et al. Flash- infer: Efficient and customizable attention engine for I1m inference serving. arXiv preprint 2XXi:2500100000 2025.
Villa, O., Lustig, D., Yan, Z., Bolotin, E., Fu, Y., Chat- terjee, N., Jiang, N., and Nellans, D. Need for speed: Experiences building a trustworthy system level gpu sim- ulator. In 2021 IEEE International Symposium on High- Performance Computer Architecture (HPCA), pp. 868- 880. IEEE, 2021.
Yu, L., Jiang W. Shi, H., Yu, Liu, Z., Zhang, Y., Kwok, J. T., Li, Z., Weller, A., and Liu, W. Metamath: Boot- strap your own mathematical questions for large language models. arXiv preprint arXiv:2309. 12284, 2023.
VLLMM
Rotary embedding implementation. wwtp///wwwnooo/e/////eeooooooe/e
Zhang, J., Xiang, J., Yu, Z., Teng, F., Chen, X., Chen, J., Zhuge, M., Cheng, X., Hong, S., Wang, J., et al. Aflow: Automating agentic workflow generation. arXiv preprint 1xXiv2440.100700 2024a.
blob/main/vl1mmnoodll httpa/////i/nneoooeeeeeeeeeeottr
eotarr_embedding. py, 2023.
Accessed: 2025-01-
12.
12

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
Zhang, N., Feldman, M., and Olukotun, K. High perfor- mance lattice regression on fpgas via high level hard- ware description language. In 2021 International Con- ference on Field-Proo Solldd-Porraammbll Technology (ICFPT), Pp. 1-10. IEFF, 2021.
Zhang. N., Lacouture, R., Sohn, G., Mure, P., Zhang, Q.. Kjolstad, F., and Olukotun, K. The dataflow abstract machine simulator framework. In 2024 ACM/IEEE 51st Annual International Symposium on Computer Architec- ture (ISCA), pp. 532-547. IEEE, 2024b.
Zhang, Q. Imran, A., Bardhi, E.. Swamy, T., Zhang, N., Shahbaz, M., and Olukotun, K. Caravan: Practical online learning of In-Network ML mod- els with labeling agents. In 18th USENIX Sympo- sium on Operating Systems Design and Implementa- tion (OSDI 24), pp. 325 325-345, 345, Santa Clara, CA, July 2024c. USENIX Association. ISBN 978-1-99393133 URL httpppps/wwww...........o//oooea osdi24
Zhao, A., Huang. D., Xu, Q., Lin, M., Liu, Y.J., and Huang, G. Expel: Llm agents experiential learners. In Pro- ceedings of the AAAI Conference Artificial Intelli- gence, volume 38, pp. 19632-19642, 2024.
Zhao, J., Li, B., Nie, W., Geng, Z., Zhang, R., Gao, X., Cheng, B., Wu, C., Cheng, Y., Li, Z., al. Akg: auto- matic kernel generation for neural processing units using polyhedral transformations. In Proceedings of the 42nd ACM SIGPLAN International Conference on Program- ming Language Design and Implementation, pp. 1233- 1248, 2021.
Zheng, L.,, Jia, C., Sun, M., Wu, Z., Yu, C. H., Haj-Ali, A., Wang, Y., Yang. J., Zhuo, D., Sen, K., et al. Ansor: Gen- erating [High-Performance) tensor programs for deep learning. In 14th USENIX symposium on operating sys- tems design and implementation (OSDI 20), pp. 863-879, 2020.
Zhou, Y., Lei, T., Liu, H., Du, N., Huang, Y., Zhao, V., Dai, A. M., Le, Q. V., Laudon, J., et al. Mixture-of. experts with expert choice routing. Advances in Neural Information Processing Systems, 35:7103-7114 2022.
13

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
A. Appendix
A.I. STeP introduction
G gelu(W;X) with N; I[, 0]
(2)
Algorithm STeP for a simplified MoE module
input X: [m,n] of Buffer(k), N: [m,n] of Multihot(e), G:
[m,n] of Buffer(e)
output [m.n] of Buffer(k)
param Wo: [k,d], W1: [k,d]
func weightedsum
external function
type: Buffer(k) (Buffer(k), Scalar) Buffer(k)
fn (s,v) = Up v1
init: =
fune expertu
type: (Buffer(k), Buffer(e)) > (Buffer(k), Scalar)
(v) = gel= (0])
func expert
type: (Buffer(k). Buffer(e)) (Buffer(k), Scalar)
fn (v) ge))( (1(1)
S. Zip(X, G)
2: S/. S (N)
3: S2 Partition(2, So. SD)
4: 53 [Map(experto. S2(0)), Map(expert, S2(1)]
S4 Merge(weightedssum S3, S/)
Figure Algorithm is STeP program example Equation (2). type signature follows the Haskell style where + connects sequence of argument types with one return type. Three funes are external functions provided by the hardware.
As shown in Figure 10, Copy duplicates stream for the affine type constraint. Zip combines two streams of values into stream of tuples. Map applies the function (fn) each input value. Partition routes tokens of the data stream to experts assigned by the index stream (SP). Merge accumulates tokens from experts. The accumulation of Merge is parameterized by fn and init where init initializes the state and fn updates the state with the input value. Partition and Merge are used in pairs, sharing the same index stream. () represents the Tuple type of value tokens. Buffer, Multihot, and Scalar are also types of value tokens, parametrized by the generic data type like float and half, which omitted for simplicity in the algorithm. Buffer and Multihot types are further parametrized by the shape in the paratheses. Shape manipulation primitives include Promote, Repeat and Repeat Ref. Figure 4(a) also shows two streams. S1 is a control token signaling the end of rank-1. 01 is value token of multihot vector type served in index streams. (v0 v1) ) is value token of type Tuple(Scalar, Reference) because the weight and input streams are composed of scalar and reference values, respectively.
14

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
A.2. Benchmark details
As shown in Figure 3, Group 4 tasks have the same operator: so softmar(S) V where equals QK They differ in external functions shown in Figure 16. Group can RDA's on-chip fusion to compose scale-dot- dot--rodupppttttco attention with Group 5. Group also differs in external functions. Group 5 contains three dataflow orders: inner-product(" poadrrrrouutmmmnkk mnnk mdk - mnd"), row-wise("mnn mmnk, mkd - mnd"), and pnter-product("mbaa mk,n, mkd - mnd"), Group 6 contains GptJ and NeoX styles which differ in pairing even and odd or the first and second half positions (vLLM, 2023). Group 8 contains LayerNorm and RMSNorm.
The last three Table index, expert, and all come from MoE. MoE contains token (Shazeer et al., 2017) Equation (3) and expert choice routing (Zhou et al., 2022) Equation (4).
S soflmax(X X W Ra
(3)
G, I = (S)
Along expert dimension
= softmar (X W2), ga
(4)
G, = Topk
Along token dimension
The expert choice routing can have another MLP auxiliary predictor for causal inference when is binary choice (Raposo al., 2024) Equation (5).
S = [[gell((X Wo)) W ,,,, gaxt
expert=2
(5)
G 0.5
Group
Description
Count
Softmax(S)@V V part
gemm
Matrix multiplication with expanded reduction dimension
3
bmm
Batch matrix-matrix product
3
norm
RMSNorm and LayerNorm (without bias and gain)
3
rope
GptJ and NeoX style of RoPE
index
Index generation MoE router
expert
Expert execution MoE
End-to-end of MoE module
4
Table 3. Description of each type of tasks. Matmul short for matrix multiplication.
15

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
A.3. Experiment settings
In Section 6. I sample 64 times for each temperature of 0.4, 0.7, and 1.0, recording the best result. In Section 6.2, sample 64 times at temperature 0.7 Claude 3.5 Sonnet to control variables. Four models are: --side------omme--222224 1-20 of OpenAI API (GPT-4o), deepseek-chat of DeepSeck API of TogetherAl API (Llama 3.1 1-405B). Maximum output tokens
of Anthropic API (Claude 3.5 Sonnet), gpt-44-2024-1 2024-11-20
(DeepSeek- V3), and 1att---laa---1144000 1-405BB-Instruct are set as 1024 and the seed for GPT-4o 42.
A.4. Additional experiments
Claude 3.5 Sonnet
1.0
Pass@n
0.6
0.4
0.2
Intermediate steps Pareto optimal
0,0
4000
6000
8000 Input Tokens
10000
12000
Figure Self-improvement learning curve with m=3. Claude 3.5 Sonnet consumes much fewer tokens than other models because the input tokens counted by the least necessary number of tokens averaged across tasks. A task can take increasing input tokens but still fail. Sonnet only has ved tasks left. Therefore, although has most example tokens, the average number of input tokens across tasks is still less than others.
(a) Pass@1
attn
gemm
bmm
norm
rope
index
expert
etoe
Avg
3wen2.5-Coder-328 0.01
0.01
0.002
Llama3-405B
0.06
0.007
DeepSeek-V3
10.06
0.056
GPT-4o
0.14
0..6
0.099
Claude-3-5-sonnet
0.13
0,380
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
(b) Pass@n
attn
gemm
bmm
norm
rope
index
expert
Avg
0,083
Ulama3-45BB
0.125
DeepSeek-v3
0.146
GPT-4o
0.146
Claude-3-S-sonnet
0.844
1.0
1.0
0
1.0
0
1.0
1.0
1.0
0
1.0
0
1.0
Figure 12, Result models temperature 0.7.
As shown in Table 4, OpenAl-ol achieves similar performance at the cost of more tokens than single This observation aligns with previous findings that scaling pretraining is preferable over inference for challenging tasks (Snell et al., 2024). Meanwhile, a single Claude-3-5-Sonne proposer can finish more tasks than the OpenAl-ol using fewer tokens. Table 3 contains abbreviations for all tasks.
Our system completed each task under 10 minutes on average. In contrast, during our pilot study, domain expert was
16

---

Adaptive Seff-improvement LLM Agentic System for ML Library Development
(a) Pass@1
attn
gemm
bmm
norm
rope
index
expert
etoe
avg
Single
0.01
0.19
Single-Structure
0.30
Agent-Structure
0.37
1.0
1.0
1.0
0
1.0
1.0
0
1.0
1.0
1.0
0
1.0
(b) Pass@n
attn
gemm
bmm
norm
rope
index
expert
etoe
avg
Single
0.46
Snrute-Stuctture
0.62
Agent-Structure
0.85
1.0
1.0
1.0
0
1.0
1.0
0
1.0
1.0
1.0
0
1.0
Figure 13. Result of three methods on Claude 3.5 Sonnet temperature 0.7.
Model
Metric (1)
attn
gemm
bmm
norm
rope
index
expert
ctoc
avg
Pass@1
0.010
0.052
0.004
0
0.008
Qwen2.5-Coder-32B
Pass@n
0.33
0
0.33
0.25
0
0.115
Pass@1 0.010
0
0
0.057
0.089
0.016
0
0.020
Llama3-405B
P
0.33
1.00
0.67
0.25
0
0.269
Pass@1 0.438
0
0
0.113
0
0.068
DeepSeek-Vi V3
P 1.00
0
0
0.75
0
0.231
Pass@] 0.021
0.016
0.214
0.094
0.258
0.005
0.080
GPT-4o
@ass@n 0.33
0.67
0.33
0.67
0.5
0.33
0.346
Pass
0.620
0.229
0.146
0.208
0.526
0.676
0.688
0.324
0.433
Claude-3-5-some
Pass n
1.00
1.00
0.33
1.00
1.00
0.75
1.00
1.00
0.885
Pass@1 0.208
0,042
0.083
0.343
0.583
0.159
OpenAI-ol (n=8)
Pass@n
0.67
0.33
0.67
0.5
1.00
0.385
Table 4. The performance of self-improvement agentic system across models.
unable to write single program within 48 hours, as they had to do riill-nn-erne and accumulate experience sequentially. Our system, by comparison, can perform these explorations in parallel. In the future, collaboration with HCI researchers will enable more extensive experiments comparing time and effort required by human programmers versus the system, providing quantitative data on usability and cognitive load.
Differences in tool access and computational load might have influenced outcomes. Therefore, we conducted an experiment that aligned both aspects.
For computation fairness, matched the token count of the single model with the agent and self-improved models by resampling. All model variants (single, agent, self-improved) have access to the same verifier so it is fair; the difference lies in how the verifier is leveraged. Self-improved models incorporate it throughout the process, while others use it only at the end as a final judge. We chose Souue----5-onnee and GPT-4o as base models. As shown in Table 5, our agentic systems still perform better under this fair setting.
Claude Sonnet Single
Claude Sonnet
GPT-4o Single
GPT-4o Agent & Self Self-improved
Pass @n
Agent Self-im self-improved
From
0.73
0.73
0.23
0.23
To
0.77
0.96
0.38
0.81
Table 5. Performance comparison of single models vs agent & self-improved models with token count matching.
17

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
We conducted supervised finetuning (SFT) using GPT-4o. Table shows that SFT can improve performance, but less than self-improvement approach.
Since we do not know the exact SFT algorithm of OpenAl service for FLOPs matching, we tried our best to favor the SFT method. We began with the same 133 correct samples from all completed tasks used in the first iteration of self-improvement Different from aeff-improvement which only picks correct program per completed task, we picked all 133 programs to form training dataset of SFT. created three SFT datasets with varying prompt compositions:
133 (base promppt-qquesiion-aasseee
133 (question+answer)
17 (question+answers deduplicated via AST)
Each dataset was used to train separate SFT model. After that, we sampled each model on all the uncompleted tasks.
Pass @n
Finetuned
Self-improved
From
0.35
0.35
To 0.62
0.81
Table 6. Performance comparison between finetuned and self- improved Models
assessed code maintainability using two metrics: maintenance index without comments (MIwoc) and with comments (MI) 2. The comment weight (MIew) defined as MI MIwoc and falls in [0, 49); MI 85 indicates good maintainability. Using all correct programs from our best model (self-improved agentic Claude Sonnet), we recorded the top MIwoc and MI per task. The mean MIwoc is 102, MI is 149, and Mlew is 47 -nniniicai well-commented, maintainable code.
As shown in Table 7, we used the same set of programs as the maintainability statistics to measure complexity. We measured these three metrics:
Lines of code: Counted via primitive calls (excluding comments/blank lines)
Shape transformations: Counted by use of Promote, Repeat, RepeatRef, and Flatten primitives
Specialized instructions: Counted as the number specialized functions task descriptions
Metric
Min
Max
Mean
Lines of code
4
17
8.67
Shape transformations
0
6
1.13
Specialized instructions
2
7
3.68
Table 7. The assessment of complexity across all completed tasks.
A.5. Generalization
The proposed system can be extended to other scenarios that require complex reasoning with limited example data and well-defined evaluation metrics. We outline the general recipe below.
As shown in Figure 1, the agentic system organization is constructed in three main steps. First, system designers define the format of both the task and its expected output. Once the format specified, the next step is to build verifier for the task. With the format and verifier in place, designers can either use single LLM design LLM agents tailored the
2htpss//www.erfffaft..omml maintainability.hml
18

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
domain-similar how handle the type constraints of the STeP language. After completing these three steps, the task can be handed over system, which will automatically carry adaptive self-improvement learning.
The adaptive self- infrommprveeent learning system also exposes several tunable hyperparameters which will be helpful when the results are not satisfactory. The most direct control is the number of parallel sampling. Users can also adjust the adaptive granularity parameter m for experience stratification. Additionally, domain-specific filtering heuristics-suce as the minimal code length heuristic used by us-can be incorporated to further guide the learning process.
Our approach has two parts: adaptive iaff-improvement learning and agentic system organization. The learning process broadly applicable other programming languages; the challenges lie in tailoring agentic systems other languages.
Mainstream languages like CUDA, HIP, and CPU vector intrinsics exhibit global properties such as arbitrary memory access, data layout sensitivity, and side effects. Similarly, STeP enforces a global affine type constraint. Our framework addressed this using a guardian agent that detects and corrects affine type violations. This concept generalizes: domain-specing guardian agents can monitor and enforce global properties of various languages, adapting the STeP solution more broadly.
second challenge is that LLMs may lean toward surface-level patterns in mainstream languages due their existence in training data, potentially missing more optimal or novel transformations. As shown in Section 6.2, our structural IR can increase sample diversity and thus boost the LLM agentic system performance. Extending this, structural IRs and tailored code generators can guide LLMs toward more creative solutions beyond conventional patterns.
We also conducted an experiment on the AIME-2024 dataset which contains 30 challenging problems from the American Invitational Mathematics Examination (AIME) 2024. We applied our adaptive self- celf-imppoveeent learning to the Claude-3-5. Sonnet base model and increased Pass@n from 0.50 to 0.67. This demonstrates the potential capabilities of our system on other tasks.
A.6. Additional explanation
Abbv.
Description
DSA
Domain Specific Architectures
DSL
Domain Specific Language
ASPL
Architecture Specific Programming Language
RDA
Reconfigurable Dataflow Architecture, DSA for AI
STeP
Streaming Tensor Program, an ASPI, for next- next-generation RDA
Table 8, Explanation of the abbreviations.
Existing auto-tuning code generation systems like Spiral (Puschel et al., 2005), Ansor (Zheng al., 2020), and AKG (Zhao al., 2021) rely on predefined optimization strategies and cost models. On the contrary, we employ LLMs to generate and refine code for machine learning libraries, particularly targeting ASPLs. We explore the shift towards utilizing the reasoning capabilities of LL for code generation and optimization, enabling the system to adapt and improve over time without explicit human intervention.
hhnn://hnggnnaaccooddssss/eesses 2024
19

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
A.7. Code and prompt examples
name: Accum
desc:
primitive operation applies function
The function is applied to the first element of the stream and the initial state to produce the first output
The function is then applied to the second element of the stream and the output of the previous application to produce second output slement, and on. The state is initialized at rank "b' of the input stream. The output stream's shape is the ingut stream' shape excluding the first 'b' dimensions.
examples:
imputs:
name: E0
dtype: fp32
dims: [M, N]
den:_ tordh.rand
name: Sum
apply:
return (state [6] ingut [0]]
init: [0]
input_dtype: fp32
atype: fp32
outputs:
name:
dtype: fp32
[N]
teta_transform:
torch. sum (input_ E'E0'], 1, keepdin-False)
impl:
step. Accum sum,
apply (ED)
return E1
Figure The reference Acoum that contains the definition and an example. Each example in the examples field is composed of task description and implementation.
20

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
inputs:
name: E0
dtype: Ep32
dims: [M,
data_ gon: torch.rand
dtype: Buffer [D)
[M,
data_gen: tardh.rand
fns:
name: MaxSum
apply:
t,t, 1,t, a.t state t scalar, scalar, []]
ttt, v.t input scalar, [D]
m_next torch, torch.max(mLE (m_t, 1.primtt torch. expch im_t a_next) 1_t
scald:
torch. expis_t
next = p.t. 1.prim_t o_next prim_t
p.t v.t
return im [m_next, nexex o_next]
init: [- inf,
0]
input_dtypes [fp32, "Buffer [DI""]
output output_dtype dtype: (fp32, fp32, "Buffer (fp32, [D]]"] func name: fn
name:
apply:
return [input [2]]
input_dtype: [fp32, fp32, "Buffer :fp32, [D])*]
output_ dtype: Buffer (fp32, [D))
func name: fn getthird
outputs:
name: 50
dtype: Ep32
dims: [D, tata.transform
corch.bmm/nn (torch. softmax (aapuut_ddtt('EO0 1). unsqueeze (1), daputt_data('E 1) ..quueze (1)
impl:
Figure 15. Task description of atin task our structural IR, where the LLM needs to complete impl.
21

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
Cask
McxSum
ExpMExx1.E
ExpMaxDiEE
apply (state, input):
apply (state, input): state de. input
apply istate, input:
4,.00 state
fc. state
6c, Ut input
8x. input
m++1 (m,,,
m4+1 maz(m,,
m4+1 maz
Am = m. mil1
Am = mi 18111
= =
fr+1 elor-mees) am
6--1 = =
4+1 = P +1
4:11 c
0c+1 = (12 0. p(
return (Me41:
return (mp-1) Fe41.
return (14411,642,1111 0111)
def init return (-00,0,0)
def init
return (-20,0,0)
def init !): return (-00,0,0)
F DivSum
betSenondThrid
apply (state, input):
apply linput);
GetThrid
Be. ct. de input
return input input [2]
def apply (input):
1, 0z state
return input [2]
KeighheddBnminggle
6+1
def apply istate, inputi: input ni state (ry de + ee)
Pr+1 1,s1 (1/
return (6411
def init(!:
init (): return 0
raturn (0.0)
Get.Secondthrid
weighhedSumbouble def apply istate, input:
apply (imput):
return
[2]
ve, input return (state d, + + Uz)
Get.Second
def apply (input)
def init (): raturn
input [1]
Div
apply limput): rt. = input
Figure 16. Inner functions for 3 tasks of attn. Task 0 encapsulates the whole innermost loop body of FlashAttention (Dao et al., 2022) in MaxSun function. Task splits the MaxSum into ExpMaxDit and DivSum. Task postpones the division of summation as FlashAttention2 (Dao, 2023). The bold symbols are streams with type ID Buffer, and the symbols are streams with type Scalar.
Stashing dimension
desc:
inserted pair can wrap the primitives to adjust the dimension.
the pritmives require
dimension, Bufferizze:Stremmity
This pattern is useful for Repeat and RepeatRef primitives. examples:
inputs: E0
dtype: fp32
dims:
data_gen: forch.rand
outputs: 80
dtype: fp32
dims: N. x
traassormmo
imput data E00 Junquuezze (1) .ropeat (L, D. D_value, 1, 11
impl:
El Bufferrie (a=2) apply (EO) gls (deype: Buffer (fp32, [t,, N)), dims: [8]
B2 step.Repaas Repoat (n=D) apply IE1 E2: [dsype: Buffer (fp32, [M, M)), dima: [D, K]
ctop.strramity.. apply
[dsypes 1p32, dimss (N,
return B3
Figure 17, An example of usage pattern that contains 3 shape manipulation primitives: Bufferize, Repeat, and Streamify.
22

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
inputs:
name: E0
dtype: fp32
dims: IM, N
data_gen: torch.rand
name: E1
dtype: (fp32, [D))
dims: [M,
gen: forch.rand
fns:
apply:
t, 1,t, o.t
scalar, scolar, [D]
e_t, v_t input scalar, [D]
m_next torch.max max (mtt, s_t) scalar
torch. m_next)
next)
p_t torch, expis_t
p.t 1_prim_t
o_next pprim_ o.t / next + p.t v_t / 1 next [m, next, next, o next
init: [-inf,
input_dtypa:
dtype: [fp32, "Buffer (fp32, [D]"]
output dtype: [*p32, fp32, "Buffer (dp32, [D]]"] Junc_names fn_maxsum
name: GetThird
apply: return [input
input_dtypa: output
input dtype: [fp32, fp32, "Buffer (fp32, [D))*]
dtype: output_dtype:
[D])
Junc_name fn_getthird
outputs:
name:
dtype: Ep32
dims: [D, N]
data_traneform:
torch.
(input Aianpuu_ddttaaE data('E0'), unsuueeze (1), input_data data('E1') sq.qqueeze (1)
impl:
E3 step. Zepp. .apply (EO, E1:)
24 Atep.Accum (fa=ff__maxuumm maxsum, b-1) applyy (R3)
Accum
step.Map (fn=EEn gee .apply
Z2
Sep..Stremiffy () apply (ES)
Figure 18. The complete implementation of atin task 0 written in structural IR.
23

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
import step
sympy import Symbol
import torch
Symbol (*M)
N Symbol (
8ymbol (*x*)
D Symbol ("D")
value
N_value
X_value
value 12
ctx .M: M. M_vale 8: N N_valu K: K. value, D: D_value:
dtppe_dtype~ ('80':
Scalar 'Ei': step.Buffee. istep.scalar ('float"), [D))
input__de data = E'E0': ranch.rand rand (M value, v_vaee,, 'El'; torch. rand (N value, M value, D_ D_value)
class MaxSum (ansuumettp..F...
def init iself, input, output ("Max5um", input, output!
super
def getInit (self):
return (torch.tensor (float
corch.tensor (0), torch.zeros value)
def apply iself, state, input):
m.t, Lt,, o.t state scalas, scaldi, [D] v_t input scalar, [V]
nanext torch max im_t, s_t) sealor
p exp
next:
p_t exxchhexpla..
m_next;
next =p. p_prim
o_next 1_prim_t o_t 1_next
1_next
return (m_next, nexxt, o_next]
fn
MaxSumister sta.ummitepp.Suppeeee.. (step. ("float"), Step.Buffer. (step. (step.Scalar Scalar ("float"), (D)))),
stop.STuple (step. ("floatt), Scalar ("float"), step.Buffer Sstap..calar ""cloatt"flat"",, (D))))
class GetThird :n):
def init (self, input, output!:
super init ("GetThird", input, output)
def (self, imput):
return [input[2]]
fn gatthird catThird (step. sTuple (stap. aatee..calle.r..eeeeeeaaeaaaeeee Scaler ") [D))
latep..calar Scalar ("float"), Scalar Btfp. Scalar Scalar("float" ("float"),
[D])),, Buffer (step.
prepare ():
E0 Sttep.Stremm"' "EO", step.Scalar ("float"), 1, [M, M]!
g0.ctx = ctx
20.data [input data (EE)
E1 Step..tream "E1", step. Buffer (step. Scalar ("float"), [D)), 1, MM, N))
gl.data linput data (EE11]
return BO, E1
def check shape (80)
shape
assert S0, dtype == step. Scalar("flost
data def check_data($0) (S0):
ttorch.sottma..... softmax (input. data ('E0' 1) unsqueaze(1)
(1), (1), input_data ('E1 aquuezz
torch.testing assert__sossee close data [O], 80. _0_
test
E0, E1 prepare body (ED, check shape (80)
check_ check_data data
daf body (80, Bl)
step. Zip() apply
E4 Step.Accum(far (wt-:fffnnnaaeeeee b=1) apply (E3)
step.Map (fn=fn getthird) apply (E4)
E2 atp.Strmraii apply (E5)
return
Figure 19. The unit test of the implementation of attn task 0 in Python produced by the code generator from structural IR shown in Figure 18.
24

---

input implementation. need copy any Adaptive Self-improvement LLM Agentic System for ML Library Development
desc:
Streaming Tansor Programs (STeP) provides a higher-level. abstraction for dataflow systems,
only of the input stream when necessary.
Copy primitives
that copy
examples:
imput_impl:
B2 step. Partition IN=E value) applyy (E0, 21))
E3 sstep..M. Map fn=fni asply for In, a in zip (matmul fns, E2:
step. Merge
apply (E3,
return E4
output_immp
E1_0, 0, 21_ step. aapply (E1)
B2 step. Partition (N=E value) apply: (E0,
R3 stte. Map fn=fni asply for In, a in xip (matmul fns, E2:
step. Merge
apply (E3, E1_1))
return E4
explanation:
Stream El is consumed twice in the input implementation. To ensure that the stream is consumed only once, we copy of the El and the copy in the second step.
input_ impl:
El step. Kap (fn=fn, prreiicredict)...... apply (20)
step. Map (fn=fn router) apply (B1)
E3 step. Map (fn=fn_ htt.aaffiiitt...l. apply (E0)
B4 step. () .pplyy (20, B3)
output impl:
во 0, 20 1 = step. Copy 1) apply (E0)
2, R0 step. Copy 1) apply
El step. (fn=fn_predict) .ppply (E0_1)
B2 step. Kap (fa=fn router) apply (B1)
Map (fn=fn_ affinity) apply (EO_2)
B4 step. () applyy [E0_3, E3)
explanation:
Stream EO consumed 3 times the input implementation. that all streams consumed only we create copy of the stream Eo and use the copy in the subsequent steps.
input impl:
stop.Buffeeee Bufferize (a=1) apply (EO)
E2 step. Map (fn=fn, sonuff_gatte_._ gate) .apply (B1)
E3 step, Map (fn=fn_top2) .pply (82)
return E3, E2
output_imm impl:
E1 step. Bufferize (a=1) .apply (E0)
step. Map (fn=fn gate) apply (B1)
E3 step. Map (fn=fn_toog (op2)) .apply (B2)
return E3, E2
explanation:
only
Figure 20. Base prompt for the guardian agent.
25

---

Adaptive Self-improvement LLM Agentic System for ML Library Development
Accumulate the lower "b" dimensions in "Stream-A, az into single value of type 'B" Accumm will continue to dequeue
calling the accumulation function (' >B") it sees '.Sb' in the input stream. Then, it will emit the accumllated value of type 'B" into the output accumulator with given
type
Accum-A, B, ; Ph B, Fa!) B, Stream-A, Stream-B, b>
(accumulate) (initialize)
Precondition: 0 c=
We can think of 'b' as the minimum stop token level --Accum*. has to see before emitting the accumulated values. according type
be found examples.
<details>
<summary>
Examples
</ s/summary>
.-Examplel; R
<br/>>
Goal: B. 8 -> [B (Reduce over the inner-most dim) Accum-A=f32,
a=3, Fn (f32, f32) ->332, Fn()->f32, Stream::32, F32,
Precondition: 0 (=1) (m3)
call given (max) on every dequoue
accumulated
is in this example because have to see the whole vector to obtain the reduced value.
<be/>
Figure 21. The specification of Accum primitive in the STeP document.
class Accum (OpBase):
def init iself, v.kwargs:: ('Accum", kwwaggs)
super
def apply iself, input: Mase.Stream name=*");
b welf.config. "b"]
fn: base. self. self.config "fn"
isinstance base.Fn), f"Accum should
provided fns input, but (typa(fa))
assert fn. input == input. dtype, I"Accum should take (fn. input/ as input, but get (Anput.dtype) input. rank,
take a positive integer than equal of the input, but get b: [b] and input rank: rankk
result base. !tream
self. gatNama (name), fn. output, Iaput.rank b, input, shape <hapa/b:) [b:]
if input.dad data is not None:
result. input.cet
TODO: construct general application Tuset here
output_indiee indices get full indices (base. subsoutrrhhe (result. shape, result. output iuput_indices . dtype, base,Element) or isinstance (result. trasult.dtype, base. basa.Buffer]. if isinstance /result.
fofffullindioo (base. subsouterShape (input. shapel:bl input.ctt .txt))ct)
output_shappe shapes [base. subsFullShage (result. dtype, elif isinstance result, dtype, base STuple!
shape,
[+:-]]]
output output_ehapes [base. subsFullShape result. shape,
[tt-1] dtspel.dtype]
elae:
raise ValueError ('Invaid dtype")
result. data zeros (shape) for shape output shapes for idx in output indicas:
fn.getInit()
for i in imput_indices.
Periialddats [d[full_idx (...,)] for in intut.data] data] state fn. apply (state, partial.date result,
for desult.data in] [idx
(state):
result
Figure 22. The definition of Accum primitive in the Python frontend.
26


Reviews:
[{"review_id": "M3xUKEEiFP", "reviewer": "Reviewer_qrAB", "summary": "This paper introduces a novel task: utilizing LLM Agents that adaptively evolve to develop architecture-specific programming languages, addressing the challenges faced by human engineers in developing corresponding languages for rapidly evolving hardware. The experimental results appear very promising, and the proposed adaptive method is highly innovative, making this a noteworthy paper.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer qrAB for positive comments and helpful feedback on our work. We are encouraged to hear the reviewer found the task and method to be innovative and the experimental results to be promising. \n\n## Larger scale evaluation potential \n> *”The benchmark dataset, consisting of only 26 tasks across 8 groups, is relatively small and could benefit from a larger scale evaluation.”*\n\nWe thank the reviewer for highlighting the potential benefits of large-scale evaluation. We briefly addressed this point in Section 3.3 of the paper, and we will expand the discussion in the revised version. New tasks typically involve new ML operators and new hardware specialized functions. These can be incorporated into the existing task pool and handled using the same adaptive self-improvement learning process by selectively sampling only the new tasks. In our current experiments, the longest prompt is approximately 14k tokens, with each example averaging around 0.5k tokens. Given Claude Sonnet’s 200k-token context window, there is capacity to include hundreds of additional tasks.\n\n## Related work\n> *”The paper could strengthen its literature review by incorporating more recent work on agent self-improvement, such as ADAS[1], AFLOW[2] to better position its contributions.”*\n\nWe thank the reviewer for pointing out two relevant works—ADAS [1] and AFLOW [2]—that can enhance our literature review. The revised version will cite both papers.\n\n## Human programmer comparisons\n> *”The inclusion of human programmer comparisons would provide valuable context and better demonstrate the practical significance of the system's achievements.”*\n\nWe appreciate the reviewer’s suggestion. Our system completed each task in under **10 minutes** on average. In contrast, during our pilot study, a domain expert was unable to write a single program within **48 hours**, as they had to do trial-and-error and accumulate experience sequentially. Our system, by comparison, can perform these explorations in parallel. We agree with the reviewer that the comparison to human programmers offers valuable insight, and we will include these results in the revised version. \n\nIn the future, we can also collaborate with HCI researchers to conduct more extensive experiments on the time and effort required by human programmers versus our system, aiming to better understand usability and cognitive load.\n\n## References\n\n[1] Hu S, Lu C, Clune J. Automated design of agentic systems[J]. arXiv preprint arXiv:2408.08435, 2024.\n\n[2] Zhang J, Xiang J, Yu Z, et al. Aflow: Automating agentic workflow generation[J]. arXiv preprint arXiv:2410.10762, 2024.", "claims_and_evidence": "Yes.\n\nAll major claims of the paper, including adaptive self-improvement learning, curriculum-based example stratification, structured intermediate representation, and complex program discovery, are well supported by comprehensive experimental results with up to 3.9× improvement over baselines and 96% task completion rate. The authors provide detailed ablation studies and cross-model validations that demonstrate the effectiveness of their approach across different model architectures, with clear empirical evidence showing the superiority of hard-example training and the benefits of their structured intermediate representation design.", "methods_and_evaluation_criteria": "The Task is novel, so the paper establish a comprehensive benchmark consisting of 8 groups with 26 ML operator tasks for evaluation.\nAlthought it's a newly constructed dataset, the paper employs solid evaluation metrics and semantic diversity analysis, effectively demonstrating the system's capabilities and make the experimental validation convincing and meaningful.", "theoretical_claims": "The paper does not include extensive theoretical analysis.", "experimental_designs_or_analyses": "Although the evaluation metrics are reasonable, the small size of the dataset needs to be noted, and experimenting with larger-scale datasets would likely further demonstrate the value of this paper's contributions.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "There aren't.", "other_strengths_and_weaknesses": "Strengths:\n\n1. The paper introduces a novel and significant task that addresses a critical need in developing ASPL for emerging hardware.\n2. The proposed self-improvement methodology through adaptive curriculum learning and experience stratification is innovative and well-designed.\n3. The experimental results demonstrate impressive performance improvements, achieving up to 3.9× enhancement over baselines and completing 96% of benchmark tasks.\n4. The paper is well-written and clearly structured, effectively presenting complex concepts (such as those mentioned in background) and experimental validations.\n\nLimitations:\n\n1. The benchmark dataset, consisting of only 26 tasks across 8 groups, is relatively small and could benefit from a larger scale evaluation.\n2. The paper could strengthen its literature review by incorporating more recent work on agent self-improvement, such as ADAS[1], AFLOW[2] to better position its contributions.\n3. The inclusion of human programmer comparisons would provide valuable context and better demonstrate the practical significance of the system's achievements.\n\n[1] Hu S, Lu C, Clune J. Automated design of agentic systems[J]. arXiv preprint arXiv:2408.08435, 2024.\n\n[2] Zhang J, Xiang J, Yu Z, et al. Aflow: Automating agentic workflow generation[J]. arXiv preprint arXiv:2410.10762, 2024.", "other_comments_or_suggestions": "Please see the weaknesses above", "questions_for_authors": "Please see the weaknesses above", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3arGqip4ww", "reviewer": "Reviewer_hNb4", "summary": "This paper proposes an adaptive self-improving agent system to unleash the ability of LLM to perform complex reasoning using limited data. It aims to automate the ML library development process using ASPL. To evaluate this, this paper builds a benchmark to conduct experiments to demonstrate the effectiveness of the proposed approach.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer hNb4 for the positive comments and helpful feedback. We are encouraged to hear the reviewer found our experiments and demonstrations clear and convincing. We also appreciate the reviewer’s careful review of the entire appendix content for prompt details.\n\n## Other agentic methods for ML library development using ASPLs\n> *”…, are there other existing Agentic systems/workflows that can be applied to the current task?”*\n\nWe thank the reviewer for raising this important point. As discussed in the paper, the tight timeline of the library-chip co-design process highlights the need for better automation. The community has begun to explore agentic solutions to this challenge in parallel [1, 2, 3]. Our adaptive self-improvement learning can enhance these efforts by making better use of correct samples. Additionally, existing systems often struggle with evolving language features, whereas our method, designed for new languages, adapts naturally to such changes.\n \n## Other fields and scenarios potential\n> *”Can some experiments be designed in the future to prove the reliability of the system in other fields and scenarios?”*\n\n> *”I hope the authors can discuss the potential of the proposed system in other fields and scenarios.”*\n\nThe proposed system can be extended to other scenarios that require complex reasoning with limited example data and well-defined evaluation metrics. We outline the general recipe below.\n\nAs shown in Figure 1 of our paper, the agentic system organization is constructed in three main steps. First, system designers define the format of both the task and its expected output. Once the format is specified, the next step is to build a verifier for the task. With the format and verifier in place, designers can either use a single LLM or design LLM agents tailored to the domain—similar to how we handle the type constraints of the STeP language. After completing these three steps, the task can be handed over to our system, which will automatically carry out adaptive self-improvement learning.\n\nThe adaptive self-improvement learning system also exposes several tunable hyperparameters which will be helpful when the results are not satisfactory. The most direct control is the number of parallel sampling. Users can also adjust the adaptive granularity parameter `m` for experience stratification. Additionally, domain-specific filtering heuristics—such as the minimal code length heuristic used by us—can be incorporated to further guide the learning process.\n\nWe also conducted an experiment on the [AIME-2024 dataset](https://huggingface.co/datasets/Maxwell-Jia/AIME_2024) which contains 30 challenging problems from the American Invitational Mathematics Examination (AIME) 2024. We applied our adaptive self-improvement learning to the Claude Sonnet base model and increased Pass@n from **0.50** to **0.67**. This demonstrates the potential capabilities of our system on other tasks. \n\nSince we agree with the reviewer that demonstrating potentials in other domains can benefit the community, we will add the recipe and results in the revised version.\n\n## References\n\n[1] https://developer.nvidia.com/blog/automating-gpu-kernel-generation-with-deepseek-r1-and-inference-time-scaling/\n\n[2] Lange, R.T., Prasad, A., Sun, Q., Faldor, M., Tang, Y. and Ha, D., 2025. The AI CUDA Engineer: Agentic CUDA Kernel Discovery, Optimization and Composition.\n\n[3] Ouyang, A., Guo, S., Arora, S., Zhang, A.L., Hu, W., Ré, C. and Mirhoseini, A., 2025. KernelBench: Can LLMs Write Efficient GPU Kernels?. arXiv preprint arXiv:2502.10517.", "claims_and_evidence": "Yes, the claims made in the paper are supported by clear and convincing evidence. Especially, multiple clear flowcharts and algorithms demonstrate the operating principles and processes of the system.", "methods_and_evaluation_criteria": "The proposed method is reasonable for the current problem. The benchmark simulates the library-chip co-design process, which is close to the real scenario and can verify the potential of the system if it is applied.", "theoretical_claims": "This paper does not involve theoretical proof.", "experimental_designs_or_analyses": "I think the experiments in this paper have fully demonstrated the effectiveness of the various parts of the proposed system. Although this paper is about ML library development using an ASPL, are there other existing Agentic systems/workflows that can be applied to the current task?", "relation_to_broader_scientific_literature": "The key contributions of the paper are related to self-improvement learning for LLMs and designing ML library using ASPLs.", "essential_references_not_discussed": "No, the paper is well-cited and covers the essential references.", "other_strengths_and_weaknesses": "Since its main goal is to achieve complex reasoning with limited data, I think the proposed system should not be limited to the field of machine learning library design. Can some experiments be designed in the future to prove the reliability of the system in other fields and scenarios?", "other_comments_or_suggestions": "NA", "questions_for_authors": "I hope the authors can discuss the potential of the proposed system in other fields and scenarios.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6KnFy3OtTk", "reviewer": "Reviewer_6GnU", "summary": "The paper suggests an (agentic) system based on LLMs that self-improves using sampling to learn programming for (architecture) specific languages. It claims that this is a complex task for which little data is available therefore necessitating the need for a reasoning system.\n\n## update after rebuttal\n\nI acknowledge the effort the authors put into the response. However, I don't intend to update my score.\n\n> While Claude Sonnet achieves 70% on our benchmark, \n\nIf the baseline is good on your benchmark, it is more or less easy. Analogies are not good arguments..\n\n> ”Also the generality of the framework is unclear - is it only for that particular language?” \n\nYou need to evaluate on more tasks... Otherwise don't call it framework.\n\nGenerally, the idea of a rebuttal is not to fix the paper within a few days, e.g., adding essential experiments - there is intentionally no possibility to upload a revised paper version, which would be needed to properly assess major changes. It is more for clarifications or pointing out misunderstandings. Thus, do not expect that doing so will be seen as a fix to major issues in the paper that leads to a better score, though no doubt sooner or later you should do so.\n\nReviews and rebuttal read. Thank you. The paper has merit and if not ICML, it will still make its way. No update to score was done.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer 6GnU for the constructive comments and helpful feedback. We are encouraged that the reviewer found our improvement non-trivial. Below, we respond to the raised concerns.\n\n## Challenge of this programming task\n> *”It is not clear why this programming task should be so challenging.”*\n\nWe appreciate the question. While Claude Sonnet achieves 70% on our benchmark, this does not imply that ML library development using ASPLs is easy. As a real-world example, as we pointed out in the paper, it took the community two years post-H100 release to optimize attention—a key LLM operator—to ~70% peak performance. As an analogy, although Claude Sonnet scores 81.7% Pass@1 on HumanEval-Mul (Table 6 in [1]), generating code from language instructions remains a hard problem. Our benchmark represents only a subset of the broader challenge: implementing transformer operators using a Python-embedded, side-effect-free ASPL. Many other ML operators and ASPLs involve more complex semantics. We plan to expand the benchmark with more difficult tasks to push system capabilities further.\n\n## Generality of the framework\n> *”Also the generality of the framework is unclear - is it only for that particular language?”*\n\nWe thank the reviewer for raising this important point. While our evaluation focuses on one language, STeP, we argue that the framework is general. As discussed in Section 2.1, we identify two essential features of ASPLs—primitives and specialized functions—and show in Section 2.3 how STeP embodies them. Due to space constraints, we refer the reviewer to the \"Generality of the framework\" section of our response to Reviewer 4wDa for the challenges and solutions of applying our framework to other languages.\n\n## Comparison fairness\n> *”… This seems unfair as their agentic systems perform a lot of extra computation and have access to tools (like the verifier).”*\n\nAs the reviewer pointed out, differences in tool access and computational load might have influenced the outcomes. Therefore, we conducted an experiment that aligned both aspects. \n\nFor computation fairness, we matched the token count of the single model with the agent and self-improved models by resampling. All model variants (single, agent, self-improved) have access to the same verifier so it is fair; the difference lies in how the verifier is leveraged. Self-improved models incorporate it throughout the process, while others use it only at the end as a final judge. We chose Claude Sonnet and GPT-4o as base models. Below is the result:\n\n| Pass@n | Claude Sonnet Single | Claude Sonnet Agent & Self-improved | GPT-4o Single | GPT-4o Agent & Self-improved |\n|-----|----|----|----|----|\n| From | 0.73 | 0.73 | 0.23 | 0.23 |\n| To | 0.77 | **0.96** | 0.38 | **0.81** |\n\nTherefore, our agentic systems still perform better under this fair setting. Since we also agree aligned comparisons can provide a more comprehensive view, we will add these results in the revised version.\n\n\n## Finetuning\n> *”…, it is unclear, if a system fine-tuned say on samples that got filtered by the verifier or in some other way, would not outperform the proposed system”*\n\nWe conducted supervised finetuning (SFT) using GPT-4o and found it improved performance, but less than our self-improvement approach.\n\nSince we do not know the exact SFT algorithm of OpenAI service for FLOPs matching, we tried our best to favor the SFT method. We began with the same 133 correct samples from all completed tasks used in the first iteration of self-improvement. Different from self-improvement which only picks 1 correct program per completed task, we picked all 133 programs to form the training dataset of SFT. We created three SFT datasets with varying prompt compositions:\n- 133 (base prompt+question+answer)\n- 133 (question+answer)\n- 17 (question+answers deduplicated via AST)\n\nEach dataset was used to train a separate SFT model. After that, we sampled each model on all the uncompleted tasks. Below is the result:\n\n| Pass@n | Finetuned | Self-improved |\n|-----|----|----|\n| From | 0.35 | 0.35 | \n| To | 0.62 | **0.81** |\n\nWe appreciate the reviewer’s suggestion and will include these results in the revised version.\n\n\n## Paper organization\n> *”The paper should more clearly carve out early on what the contribution to the ML field are”*\n\n> *”The claims like \"we do human style learning with some ref\" are too brief and vague… “*\n\nWe thank the reviewer for these helpful suggestions. In the revised version, we will emphasize our contributions to ML more clearly in the introduction and better define human-style learning in the discussion.\n\n## Reference\n\n[1] Liu, A., Feng, B., Xue, B., Wang, B., Wu, B., Lu, C., Zhao, C., Deng, C., Zhang, C., Ruan, C. and Dai, D., 2024. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437.", "claims_and_evidence": "It is not clear, why this programming task should be so challenging (even for experienced programmers) as claimed in the intro - also in the light that Claude Sonnet achieves already 70%.\nAlso the generality of the framework are unclear - is it only for that particular language (judging from the evaluation it is, as there is just one dataset constructed toward that language).", "methods_and_evaluation_criteria": "The benchmark is self-constructed and consists of just few tasks. This limits severely the generalizability.", "theoretical_claims": "no theory", "experimental_designs_or_analyses": "The comparison against other models is not fully clear. It appears that they are comparing against raw base models, e.g. GPT4o .This seems unfair as their agentic systems performs a lot of extra computation and has access to tools (like the verifier). Thus, while the improvement is still non-trivial, it is unclear, if a system fine-tuned say on samples that got filtered by the verifier or in some other way, would not outperform the proposed system.", "relation_to_broader_scientific_literature": "Agentic AI is a hot topic.", "essential_references_not_discussed": "---", "other_strengths_and_weaknesses": "The paper should more clearly carve out early on what the contribution to the ML field are. It focuses too much on the domain-specific problem.\n\nMinor comment: The claims like \"we do human style learning with some ref\" are too brief and vague but still appearing multiple times. If important discuss properly, otherwise maybe just mention it in the discussion or", "other_comments_or_suggestions": "None", "questions_for_authors": "None - but see uncertainties above", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Rz905Or0YQ", "reviewer": "Reviewer_4wDa", "summary": "The authors propose an agentic system with adaptive self-improvement capabilities, specifically designed for synthesizing high-performance ML libraries. The proposed synthesis algorithm targets architecture-specific programming languages (ASPLs), with experiments conducted on Streaming Tensor Programs. The primary motivation is that domain-specific accelerators often change drastically with each new hardware generation, creating a pressing need for the rapid development of ML libraries in low-level specialized languages, often without access to large corpora of examples.\n\nThe paper presents an iterative approach that employs LLMs to generate new code solutions, filters out high-quality solutions, and then leverages these examples as demonstrations for increasingly complex tasks.\n\nThe proposed approach is evaluated on a suite of 26 tasks (curated by the authors from first principles), covering eight types of common operators (e.g., matrix multiplication, attention blocks, mixture-of-experts). The findings indicate that this method achieves higher pass rates (up to 96%) and a 3.9× improvement over a single-LLM baseline.\n\n\n## Update After Rebuttal:\nI find this paper very interesting and recommend its acceptance. Please make sure to incorporate the changes discussed during rebuttal in the final version of the paper.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer 4wDa for the positive comments and helpful feedback. We were encouraged that the reviewer enjoyed reading the paper, found our ideas novel, and took the time to review the appendix. We will include all the discussions and results below in the revised version.\n \n## Run-time performance measurement\n> *”Q1: Have you measured the actual run-time performance of any of these automatically generated kernels on real or simulated hardware?”*\n\nWe manually translated the generated implementation of Task 2 in Figure 16 of our paper to a simulator built on top of [DAM-RS](https://github.com/stanford-ppl/DAM-RS), which models the streaming behavior of each STeP primitive and assumes every operation and specialized function takes one cycle. This task implements *softmax(S)@V* where S and V are of shape [n,m] and [m], respectively; since S and V are streamed sequentially, the cycle count should scale with n*m. Simulation results match this expectation. Detailed result: https://anonymous.4open.science/r/ICML2025-rebuttal-4D6B/fig.png \n\n## Generality of the framework\n> *”Q2: Do you anticipate any unique challenges in applying this approach to more mainstream, well-established languages like CUDA/HIP or to CPU vector intrinsics?”*\n \nOur approach has two parts: adaptive self-improvement learning and agentic system organization. The learning process is broadly applicable; the challenges lie in tailoring agentic systems to other languages.\n\nMainstream languages like CUDA, HIP, and CPU vector intrinsics exhibit global properties such as arbitrary memory access, data layout sensitivity, and side effects. Similarly, STeP enforces a global affine type constraint. Our framework addressed this using a *guardian* agent that detects and corrects affine type violations. This concept generalizes: domain-specific guardian agents can monitor and enforce global properties of various languages, adapting the STeP solution more broadly.\n\nA second challenge is that LLMs may lean toward surface-level patterns in mainstream languages due to their existence in training data, potentially missing more optimal or novel transformations. As shown in Section 6.2, our structural IR can increase sample diversity and thus boost the LLM agentic system performance. Extending this, structural IRs and tailored code generators can guide LLMs toward more creative solutions beyond conventional patterns.\n\n## Code maintainability and complexity\n> *”Q3: Have you evaluated the readability or maintainability of the final code solutions?”*\n\n> *”Metrics such as lines of code, the number of shape transformations, or specialized instructions would help quantify the difficulty.”*\n\n**Maintainability statistics**. We assessed code maintainability using two metrics: maintenance index without comments (MIwoc) and with comments (MI) [1]. The comment weight (MIcw) is defined as MI - MIwoc and falls in [0, 49); MI > 85 indicates good maintainability. Using all correct programs from our best model (self-improved agentic Claude Sonnet), we recorded the top MIwoc and MI per task. The mean MIwoc is 102, MI is 149, and MIcw is 47—indicating well-commented, maintainable code.\n\n**Complexity statistics**. We used the same set of programs as the maintainability statistics to measure complexity. We measured all three metrics the reviewer suggested:\n- Lines of code: Counted via primitive calls (excluding comments/blank lines)\n- Shape transformations: Counted by use of Promote, Repeat, RepeatRef, and Flatten primitives\n- Specialized instructions: Counted as the number of specialized functions in task descriptions\n\nAcross all completed tasks:\n| Metric | Min | Max | Mean |\n|-----|----|----|----|\n| Lines of code | 4 | 17 | 8.67 |\n| Shape transformations | 0 | 6 | 1.13 |\n| Specialized instructions | 2 | 7 | 3.68 |\n\nDetailed result: https://anonymous.4open.science/r/ICML2025-rebuttal-4D6B/tab.md  \n\n## Scalability\n> *”…it is unclear how the system would scale on libraries that are much bigger or that require 50–100 times more operators.”*\n\nAs discussed in the “Larger scale evaluation potential” section of our response to Reviewer qrAB, current LLMs still have the capacity to self-improve over hundreds more tasks. If the number exceeds the context length, better stratification and selection functions are needed to preserve experience quality within the context window limit.\n\n## Related work\n> *”... TVM and Spiral, and more recent papers on end-to-end auto-tuning code generators, may be relevant. Additionally, a broader set of reflection-based LLM coding pipelines, such as Reflexion, and Tree-of-Thought could be cited or compared.”*\n\nWe thank the reviewer for providing more relevant work. We will incorporate these papers into related work and give a more thorough discussion of how our work improves the results of them.\n\n## Reference\n[1] https://www.verifysoft.com/en_maintainability.html", "claims_and_evidence": "Overall, the paper's main claim, i.e. that adaptive self-improvement leads to higher pass rates and higher code-correctness coverage, is well supported by pass@k results on a specialized but well-motivated benchmark. None the less, it is unclear how the system would scale on libraries that are much bigger or that require 50–100 times more operators.", "methods_and_evaluation_criteria": "The primary metric for evaluating the proposed agentic system is functional correctness, measured by pass@k across 26 tasks. Pass@n is also reported, indicating how many tasks are eventually solved by at least one attempt. The authors further analyze the number of input tokens consumed per attempt and whether more complex examples yield better outcomes. \n\nThese criteria are well-suited for code-generation tasks, and the emphasis on pass@k aligns with standard practices in LLM-based coding research. However, the authors do not provide timing data or real hardware evaluations of the generated STeP code, leaving the claim of \"high-performance\" implementations unvalidated. The experimental design primarily assesses correctness rather than raw performance.", "theoretical_claims": "The authors do not provide formal mathematical proofs or new theorems.", "experimental_designs_or_analyses": "The experiments are conducted using 8 systematically created categories of ML operators, such as those involving shape manipulation and advanced arithmetic. Each category consists of multiple tasks that differ in certain details, such as whether partial streams are reused. The evaluation is quite thorough for a custom suite of 26 tasks.", "relation_to_broader_scientific_literature": "The paper's approach relates to self-refining, multi-agent code generation systems, specifically in the context of LLM agentic methods for self-play and self-improvement. It provides sufficient citations and discussions of the broader scientific literature in this domain. The idea of using IR to represent partially structured code with a dedicated compiler reminds me of work using MLIR to improve the efficiency of Tensor Compiler Construction [1].\n\n[1] VASILACHE, NICOLAS, et al. \"Composable and Modular Code Generation in MLIR.\" arXiv preprint arXiv:2202.03293 (2022).", "essential_references_not_discussed": "Overall, the related works section is well-written and provides sufficient background on prior research related to self-improving LLM agents and their design for specialized code generation tasks. Some prior efforts, such as TVM and Spiral, and more recent papers on end-to-end auto-tuning code generators, may be relevant. Additionally, a broader set of reflection-based LLM coding pipelines, such as Reflexion, and Tree-of-Thought could be cited or compared.", "other_strengths_and_weaknesses": "Strengths:\n- The paper is well-written, provides sufficient background, and introduces the method in detail with appropriate examples. I enjoyed reading this paper.\n- It presents novel ideas as well as interesting instantiations of well-studied techniques in a specialized domain. I particularly found the idea of a “guardian” agent for checking a global type-theoretic property to be a clever application of multi-agent prompting.\n\nWeaknesses:\n- The paper’s evaluation focuses almost entirely on the correctness of relatively small tasks. There is no direct measurement of the speed, efficiency, or memory overhead of the generated STeP libraries.\n- The authors do not provide strong evidence of how the approach scales beyond these 26 tasks or how maintainable and comprehensible the generated solutions will be in practice.", "other_comments_or_suggestions": "The paper could be strengthened by providing concrete results on the size or complexity of the final STeP programs, beyond just pass@k metrics. Metrics such as lines of code, the number of shape transformations, or specialized instructions would help quantify the difficulty.", "questions_for_authors": "- Q1: Have you measured the actual run-time performance of any of these automatically generated kernels on real or simulated hardware?\n- Q2: Do you anticipate any unique challenges in applying this approach to more mainstream, well-established languages like CUDA/HIP or to CPU vector intrinsics?\n- Q3: Have you evaluated the readability or maintainability of the final code solutions?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes an adaptive self-improvement agentic framework that leverages LLMs to automate the generation of ML libraries using architecture-specific programming languages (ASPLs), demonstrating strong results in synthesizing 25 out of 26 key LLM operators.

The strengths of the paper include:

- Novelty and relevance of the task: The paper introduces a timely and challenging problem of synthesizing ML libraries in ASPLs using LLMs, addressing an underexplored but impactful area (Reviewers qrAB, 4wDa).

- Innovative self-improvement framework: The adaptive curriculum-based learning loop, stratified experience sampling, and structured IR contribute meaningfully to agentic system design (Reviewers qrAB, hNb4).

- Strong empirical performance: The system achieves up to 3.9× improvements over baselines and 96% task completion, with detailed ablation studies and fair comparisons across model variants (Reviewers qrAB, 4wDa).

- Well-structured presentation and thoughtful design: The paper clearly articulates its design choices and motivations, making a compelling case for its adaptive methodology (Reviewers hNb4, 4wDa).

Before the rebuttal, common concerns raised by reviewers included the limited dataset size and benchmark scope (Reviewers qrAB, 6GnU, 4wDa), questions around generalizability beyond a single ASPL (STeP) (Reviewers 6GnU, hNb4), and lack of real hardware evaluation or human expert comparison (Reviewers qrAB, 4wDa). The authors effectively addressed these concerns in the rebuttal, offering clarifications, pilot simulation results, maintainability metrics, and concrete plans for broader evaluations.

Despite a mixed score from Reviewer 6GnU due to skepticism about benchmark difficulty and generality, the overall reviewer consensus supports the novelty, rigor, and potential impact of the work.

Therefore, AC recommend acceptance of this paper.
Decision: Accept (poster)