We present TritorX, an agentic AI system designed to generate functionally correct Triton PyTorch ATen kernels at scale for emerging accelerator platforms. TritorX integrates open-source large language models with a custom linter, JIT compilation, and a PyTorch OpInfo-based test harness. This pipeline is compatible with both real Meta Training and Inference Accelerator (MTIA) silicon and in hardware simulation environments for next-generation devices. In contrast to previous kernel-generation approaches that prioritize performance for a limited set of high-usage kernels, TritorX prioritizes coverage. Our system emphasizes correctness and generality across the entire operator set, including diverse data types, shapes, and argument patterns. In our experiments, TritorX successfully generated kernels and wrappers for 481 unique ATen operators that pass all corresponding PyTorch OpInfo tests (over 20,000 in total). TritorX paves the way for overnight generation of complete PyTorch ATen backends for new accelerator platforms.


翻译:本文提出TritorX——一种智能AI系统,旨在为新兴加速器平台大规模生成功能正确的Triton PyTorch ATen内核。TritorX将开源大语言模型与定制化代码检查器、即时编译技术以及基于PyTorch OpInfo的测试框架相集成。该流水线兼容真实的Meta训练与推理加速器(MTIA)硬件芯片,同时支持下一代设备的硬件仿真环境。相较于以往优先为有限高频使用内核优化性能的核函数生成方案,TritorX更注重覆盖广度。我们的系统强调在整个算子集合中实现正确性与通用性,涵盖多样化数据类型、张量形状及参数模式。实验表明,TritorX成功为481个独特ATen算子生成可通过全部对应PyTorch OpInfo测试(总计超20,000项)的内核与封装层。TritorX为新一代加速器平台实现‘隔夜生成’完整PyTorch ATen后端的能力开辟了道路。

0
下载
关闭预览

相关内容

预知未来——Gluon 时间序列工具包(GluonTS)
ApacheMXNet
24+阅读 · 2019年6月25日
深度学习人脸识别系统DFace
深度学习
17+阅读 · 2018年2月14日
大数据分析研究组开源Easy Machine Learning系统
中国科学院网络数据重点实验室
17+阅读 · 2017年6月13日
国家自然科学基金
4+阅读 · 2015年12月31日
国家自然科学基金
46+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
7+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2014年12月31日
Arxiv
0+阅读 · 11月18日
VIP会员
相关资讯
预知未来——Gluon 时间序列工具包(GluonTS)
ApacheMXNet
24+阅读 · 2019年6月25日
深度学习人脸识别系统DFace
深度学习
17+阅读 · 2018年2月14日
大数据分析研究组开源Easy Machine Learning系统
中国科学院网络数据重点实验室
17+阅读 · 2017年6月13日
相关基金
国家自然科学基金
4+阅读 · 2015年12月31日
国家自然科学基金
46+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
7+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2014年12月31日
Top
微信扫码咨询专知VIP会员