Command Palette
Search for a command to run...
자율 진화 탐색을 위한 행위자 기반 변이 연산자
자율 진화 탐색을 위한 행위자 기반 변이 연산자
초록
행위자 기반 변이 연산자(AVO)는 기존 진화 탐색의 고정된 변이, 교차, 수작업 휴리스틱을 자율 코딩 행위자로 대체하는 새로운 진화 변이 연산자 계열이다. 언어 모델을 사전 정의된 파이프라인 내 후보 생성에 국한하지 않고, AVO는 변이를 자기 주도적 행위자 루프로 구현하여 현재 계통, 도메인 특화 지식 베이스, 실행 피드백을 참조하며 구현 편집을 제안, 수정, 비판, 검증할 수 있도록 한다. 우리는 AI 분야에서 가장 공격적으로 최적화된 커널 대상 중 하나인 어텐션을 NVIDIA Blackwell(B200) GPU 상에서 평가한다. 멀티헤드 어텐션에 대한 7일간의 지속적인 자율 진화를 통해, AVO는 평가된 구성 전반에서 cuDNN 대비 최대 3.5%, FlashAttention-4 대비 최대 10.5% 성능이 우수한 커널을 발견한다. 발견된 최적화는 그룹화 쿼리 어텐션으로 쉽게 전이되어 단 30분의 추가 자율 적응만으로 cuDNN 대비 최대 7.0%, FlashAttention-4 대비 최대 9.3%의 성능 향상을 달성한다. 이러한 결과는 행위자 기반 변이 연산자가 행위자를 후보 생성기에서 변이 연산자로 격상시킴으로써 기존의 LLM-인-더-루프 진화 파이프라인을 넘어서며, 오늘날 가장 진보된 GPU 하드웨어에서 최첨단 전문가 설계 어텐션 구현을 능가하는 커널을 만들어내는 성능 결정적 마이크로아키텍처 최적화를 발견할 수 있음을 보여준다.
One-sentence Summary
NVIDIA researchers propose Agentic Variation Operators (AVO), autonomous coding agents that replace fixed evolutionary operators by self-directing variation through a loop consulting lineage, a knowledge base, and execution feedback to propose, repair, critique, and verify implementation edits, and over 7 days of continuous evolution on multi-head attention discover kernels on Blackwell B200 GPUs outperforming cuDNN by up to 3.5% and FlashAttention-4 by up to 10.5%, with rapid transfer to grouped-query attention requiring only 30 minutes of adaptation to yield up to 7.0% and 9.3% gains respectively.
Key Contributions
- Agentic Variation Operators (AVO) replace fixed evolutionary variation heuristics with autonomous coding agents that subsume sampling, generation, and evaluation into a single self-directed loop.
- Applied to multi-head attention on NVIDIA Blackwell GPUs, AVO discovers kernels that outperform cuDNN by up to 3.5% and FlashAttention-4 by up to 10.5% after seven days of continuous evolution.
- The discovered optimizations transfer to grouped-query attention with 30 minutes of additional adaptation, yielding improvements of up to 7.0% over cuDNN and 9.3% over FlashAttention-4.
Introduction
The authors address the challenge of using large language models in evolutionary code optimization, where prior systems confine the LLM to a single-turn candidate generator within a fixed pipeline. This limits the ability to perform the iterative, multi-step engineering required to surpass highly tuned implementations like attention kernels on GPUs, which demand studying hardware documentation, profiling, testing, and revising. To overcome this, they introduce Agentic Variation Operators (AVO), which replace the entire variation step with a self-directed coding agent that autonomously consults knowledge, evaluates candidates, and refines its approach over extended time horizons. Applied to multi-head attention on NVIDIA Blackwell B200 GPUs, AVO discovers micro-architectural optimizations that outperform cuDNN by up to 3.5% and FlashAttention-4 by up to 10.5%, with the learned techniques transferring to grouped-query attention in just 30 minutes of additional autonomous adaptation.
Method
The authors begin by detailing the complexity of attention computation on modern GPUs. Given query, key, and value matrices Q,K,V, attention computes O=softmax(QK⊤/d)V. While FlashAttention avoids materializing the full score matrix by processing key blocks sequentially, state-of-the-art kernels like FA4 on NVIDIA's Blackwell architecture employ warp specialization. Different warp groups handle matrix multiplications, online softmax, output rescaling, and data movement concurrently. Optimizing such highly tuned kernels requires deep hardware expertise and extensive profiling.
To automate this optimization, the authors propose Agentic Variation Operators (AVO). Traditional evolutionary search decomposes the variation operator into separate sampling and generation steps, confining the LLM to a fixed pipeline. AVO replaces this decomposition with a single autonomous agent run:
Vary(Pt)=Agent(Pt,K,f)where Pt is the full lineage of solutions and their scores, K is a domain-specific knowledge base, and f is the scoring function.
The input to the AVO system consists of three main components. The population Pt contains the lineage of prior CUDA kernel implementations and their scores. The knowledge base K includes CUDA and PTX documentation, Blackwell architecture specifications, and reference GPU kernel codebases. The scoring function f evaluates candidates on numerical correctness and throughput in TFLOPS; any candidate failing correctness checks receives a zero score regardless of throughput.
The core of the method is the AVO Main Agent Loop, which produces a new solution xt+1 from the current lineage Pt. This autonomous loop involves several interconnected phases. During Planning, the agent proposes edits using the knowledge base and prior implementations. In Implementation, it applies code modifications. The Evaluation phase uses the scoring function f to measure performance. If issues arise, the Bug-Fixing phase allows the agent to diagnose, repair, and adapt its plan. The agent frequently examines multiple prior implementations to identify bottlenecks and consults documentation to understand hardware constraints before implementing optimizations. This edit-evaluate-diagnose cycle repeats until a satisfactory xt+1 is committed. A new version is persisted only when it passes correctness checks and matches or improves the benchmark score relative to the best committed version.
To ensure forward progress during long-running autonomous optimization, the authors incorporate an AVO Supervisor Agent. This supervisor monitors the main agent for stagnation, such as exhausting a line of exploration or entering unproductive edit cycles. When triggered, the supervisor reviews the overall evolutionary trajectory and provides conditional intervention to steer the search toward new candidate optimization directions.
The system operates as a continuous evolution loop. Each committed version is persisted as a git commit with its score, maintaining full state continuity across the entire process. In their evaluation, a 7-day run produced a final multi-head attention kernel spanning 40 successive versions, with the main agent autonomously deciding when to attempt new optimizations or shift strategies, while the supervisor intervened during periods of stagnation.
Experiment
The evaluation uses an agentic variation operator (AVO) to evolve CUDA attention kernels on NVIDIA B200 GPUs, benchmarking against cuDNN and FlashAttention-4 across multi-head and grouped-query attention configurations. AVO autonomously discovers optimizations that yield consistent throughput gains over both baselines, with the largest improvements coming from architectural changes like branchless rescaling, pipeline overlap, and register rebalancing, while later versions provide fine-grained tuning. The agent successfully adapts its evolved MHA kernel to GQA in under an hour, demonstrating that the discovered optimizations generalize beyond the original search space. Overall, the experiments show that agent-driven evolution can produce expert-level kernel optimizations by jointly reasoning about multiple hardware subsystems through iterative profiling and code modification.
Three agent-discovered kernel optimizations each improved throughput, with gains varying by attention type. Branchless accumulator rescaling delivered a large 8.1% improvement on non-causal attention but only 1.6% on causal, because the branchless path applies only to unmasked key blocks. Pipeline overlap and register rebalancing added further gains, primarily for non-causal scenarios where the correction warp is on the critical path. Branchless accumulator rescaling provides a substantial 8.1% geomean throughput gain on non-causal attention, but only 1.6% on causal attention, as the optimization is limited to fully unmasked key-block iterations. Register rebalancing across warp groups yields a 2.1% improvement on non-causal attention and negligible impact on causal, because the correction warp runs concurrently with the second GEMM only in the non-causal pipeline.
Agent-discovered kernel optimizations improved throughput, with gains concentrated in non-causal attention. The branchless accumulator rescaling was highly effective for non-causal scenarios but offered limited benefit on causal attention because it only applies to unmasked key blocks. Additional pipeline overlap and register rebalancing further boosted non-causal performance, where the correction warp lies on the critical path, while causal attention saw minimal improvement.