发表机构
Alibaba Group（阿里巴巴集团）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
研究大语言模型生成GPU内核用于生产的可行性，提出Atrex-Bench基准测试，发现现有模型表现不佳。为此发布Atrex-Kernel-Agent优化代理，结合多种技术，经案例研究能将回退转换为匹配或超越生产基线的内核。
AI 中文摘要
现有的GPU内核生成基准测试的问题来源于合成或精心策划的来源，与实际部署的工作负载不同。我们提出了Atrex-Bench基准测试，其30个运算符和440种形状直接从计算受限、内存丰富的GPU的全集群生产推理追踪中采样。每个问题都有一个重要性权重，通过应用卡小时加权，并针对其运行的服务阶段单独计算，还有每个问题的屋顶线上限。评估六个前沿编码代理表明，即使是最好的原始模型在生产运算符上也只能达到硬件屋顶线的约10%。为了缩小差距，我们共同发布了Atrex-Kernel-Agent（AKA），它结合了迭代测量-修正搜索、用于避免搜索上下文停滞的优化随机失活，以及分层的GPU优化知识库。在一个受控案例研究中，该代理将零FlyDSL回退转换为与手工调整的生产基线匹配或超过的实际内核。
英文摘要
Existing GPU kernel generation benchmarks draw problems from synthetic or curated sources that diverge from deployed workloads. We present Atrex-Bench, a benchmark whose 30 operators and 440 shapes are sampled directly from full-cluster production inference traces of compute-limited, memory-rich GPUs. Each problem carries an importance weight derived from its share of observed GPU time, weighted by application card-hours and computed separately for the serving phases in which it runs, together with a per-problem roofline ceiling, so the aggregate score emphasizes the kernels that consume the most serving time. Evaluating six frontier coding agents on Atrex-Bench shows that even the best vanilla model reaches only ${\sim}10\%$ of the hardware roofline on production operators; and correctness alone overstates capability, since much of the apparent pass rate comes from PyTorch fallbacks rather than kernels the model wrote. To close this gap, we co-release Atrex-Kernel-Agent (AKA), a profile-driven kernel-optimization agent that combines iterative measure-revise search, optimization dropout for escaping stalled search contexts, and a layered GPU-optimization knowledge base (298 reference-kernel files and 244 optimization-knowledge documents, plus external upstream reference projects for API/ISA lookup). In a controlled case study, the agent converts zero-FlyDSL fallbacks into real kernels that match or exceed hand-tuned production baselines.
CommentsBoth artifacts are released as open source: Atrex-Bench (https://github.com/alibaba/atrex-bench) and Atrex-Kernel-Agent (https://github.com/alibaba/atrex-kernel-agent)
