自动构建无需手写代码的高效推理内核，支持多GPU型号自适应。
AutoMegaKernel: A Statically-Checked Agent Harness for Self-Retargeting Megakernel Synthesis
- 用静态验证机制自动合成跨GPU型号的协同计算内核。
- 在10个模型上实现零错误生成，实测性能超传统cuBLAS方案。
- 适合追求极致推理效率与自动化部署的研究者和工程师。
AutoMegaKernel（AMK）将HuggingFace Llama系列模型编译为单一持久协作的CUDA内核，一次性完成前向传播，无需针对每模型手动编写CUDA代码。其核心贡献在于系统设计：通过冻结的调度IR验证器进行静态图检查，确保无死锁与竞态，7,160个对抗性调度中零误接受，360个真实降级全部通过。同一源码可适配sm_80/sm_90/sm_120，自动生成10个受支持模型的正确巨内核；在SmolLM2-135M真实检查点上，逐令牌重现HuggingFace贪婪解码（困惑度匹配2.5e-7）。无人值守的代理驱动自研循环可自我优化内核，相较基线提升1.25-1.72倍。搜索得到的int8（W8A16）内核在批量1解码时超越CUDA图化cuBLAS：L4最高1.33倍，L40S达1.25-1.27倍，A10G在规模下最高1.08倍，消费级RTX 5090达1.19-1.23倍。性能排序不依赖带宽（如864 GB/s的L40S优于600 GB/s的A10G），差异源于推理类与训练类芯片定位。在高带宽训练类卡（A100/H100）上，该框架因局域跨SM同步瓶颈而落后于cuBLAS，差距已明确披露。本比较为解码位置0的精度非对称（W8A16 vs bf16），最大真实检查点为TinyLlama-1.1B。代码与框架见：https://github.com/RightNow-AI/AutoMegaKernel
原文摘要 · Abstract (English)
AutoMegaKernel (AMK) compiles a HuggingFace Llama-family model into a single persistent cooperative CUDA kernel that runs the whole forward pass in one launch, with no per-model hand-written CUDA. The contribution is the system, not raw speed. A frozen schedule-IR validator statically certifies deadlock-freedom and race-freedom via static graph checks (not a mechanized proof), so an unsafe agent-proposed schedule is rejected before launch: across 7,160 adversarial schedules (6,091 unsafe) it had zero false-accepts and accepted all 360 real lowerings. The same source retargets sm_80/sm_90/sm_120 from one codebase, auto-generates correct megakernels for 10 of 10 supported models, and on a real SmolLM2-135M checkpoint reproduces HuggingFace greedy decode token-for-token (perplexity match 2.5e-7). An unattended, agent-drivable autoresearch loop self-improves the megakernel over its own baseline (1.25-1.72x). A search-found int8 (W8A16) megakernel beats CUDA-graphed cuBLAS bf16 at batch-1 decode across NVIDIA's datacenter inference fleet: L4 up to 1.33x, the current-gen L40S 1.25-1.27x, A10G up to 1.08x at scale, and the consumer RTX 5090 1.19-1.23x. The ordering is not a clean function of bandwidth (the 864 GB/s L40S beats the 600 GB/s A10G); the divide is inference-class vs training-class. AMK trails cuBLAS on the high-bandwidth training-class A100/H100, where the harness localizes the cross-SM-sync bottleneck; we report the gap plainly. This is a precision-asymmetric (W8A16 vs bf16) comparison at decode position 0; the largest real checkpoint is TinyLlama-1.1B. Code and the harness: https://github.com/RightNow-AI/AutoMegaKernel
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。