用大模型指导生成高效CUDA代码,加速Transformer推理。
AgentCompile: An LLM-Guided Compiler for Direct CUDA Inference

- 大模型提供编译器优化建议,约束生成范围提升可靠性。
- 实测推理速度比PyTorch快2.23到6.98倍,比vLLM快1.04到1.16倍。
- 适合追求极致推理性能的模型部署开发者使用。
Transformer推理越来越依赖专用编译器和运行时支持,而当前大模型已能生成复杂CUDA内核。然而,无约束生成难以保证正确性与性能。我们提出\textsc{AgentCompile},一种由大模型引导的CUDA推理编译器,融合两种互补的大模型使用方式:首先,大模型为编译器生成的区域摘要和候选空间提供咨询元数据;编译器随后基于模板生成候选内核,验证正确性,通过实测延迟选择最优实现,并在特殊化不支持或不划算时回退。其次,在编译器定义的契约下,大模型直接生成五类解码关键内核,提示基于提炼的优化原则。\textsc{AgentCompile}将这些内核集成到支持分页键值缓存、连续批处理、抢占、分块预填充和桶式全步CUDA图重放的推理运行时中。在六个模型家族上的评估显示,\textsc{AgentCompile}在单请求生成上相较PyTorch eager获得2.23–6.98×加速,在单请求和多请求服务上相较vLLM获得1.04–1.16×加速。代码已公开于https://github.com/veneno1213822/AgentCompile。
原文摘要 · Abstract (English)
Transformer inference increasingly relies on specialized compiler and runtime support, while recent LLMs can generate nontrivial CUDA kernels. However, unconstrained generation guarantees neither correctness nor performance. We present \textsc{AgentCompile}, an LLM-guided CUDA inference compiler that combines two complementary uses of LLMs. First, the LLM provides advisory metadata for compiler-derived region summaries and bounded candidate spaces. The compiler then instantiates template-based CUDA candidates, validates correctness, selects implementations by measured latency, and falls back when specialization is unsupported or unprofitable. Second, under compiler-defined contracts, the LLM directly generates five classes of decode-critical kernels to accelerate inference, prompted by distilled optimization principles. \textsc{AgentCompile} integrates these kernels into a serving runtime with paged KV cache, continuous batching, preemption, chunked prefill, and bucketed full-step CUDA Graph replay. Across six evaluated model families, \textsc{AgentCompile} achieves speedups of \textbf{2.23--6.98$\times$} over PyTorch eager for single-request generation, and \textbf{1.04--1.16$\times$} over vLLM for both single-request generation and multi-request serving. Our code is publicly available at https://github.com/veneno1213822/AgentCompile.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。