Joint compilation speeds up paired ai model decoding on gpus
Conjugate Compilation: Self-Improving Agentic Co-Compilation for Speculative Decoding
Machine LearningDistributed, Parallel, and Cluster ComputingProgramming Languages
Summary
Running two AI models together can be slow, especially when one model guesses answers and the other double-checks them. The authors propose combining these two models into a single optimized process, rather than handling them separately. This method rearranges how the computer handles the models on the graphics card to make them faster and use resources better. They use a smart trial-and-error system guided by an AI to find the best way to do this. Their approach speeds up the combined process by up to about 60% and improves how well the graphics card is used.
What this means in practice
- •For machine learning engineers: Optimize deployment of paired AI models to improve inference speed and GPU efficiency in real-time applications.
- •For graphics hardware developers: Design GPU architectures and compilers that better support joint execution of multiple AI workloads for enhanced performance.
Authors
Annabelle Sujun Tang, Rohan Mahapatra, Christopher Priebe, Hanyang Xu, Manya Ghobadi, Christopher Rossbach, Hadi Esmaeilzadeh
Abstract
As LLM adoption surges, scaling model serving requires unlocking stranded GPU capacity alongside deploying more GPUs. Speculative decoding brings two models into the serving process: a small draft model proposes tokens that a larger target model verifies. Yet compiling these models individually leaves their shared execution outside the compiler's optimization boundary. We introduce CONJUGATE COMPILATION, which reframes speculative decoding as a cross-model compilation problem and jointly optimizes the draft and target models within their shared execution rather than treating each in isolation. Our key insight is that cross-model compilation should be informed by the measured performance of the co-located shared execution it produces: an optimization that improves a kernel in isolation need not improve the models when they execute together. To expand cross-model opportunities, we devise operand layouts that interleave blocks from both models within common tiles, co-locating GEMM and non-GEMM work while preserving their mathematics. We realize CONJUGATE COMPILATION through a self-improving agentic search system in which an LLM edits cross-model optimization plans using the compiler and GPU execution as its environment. Each proposed plan is compiled and evaluated through its co-located shared execution, whose measured score, together with kernel-level verdicts and compiler feedback, informs subsequent proposals, while Monte Carlo tree search accumulates rewards to guide exploration. Across six single-GPU speculative-decoding configurations on NVIDIA B300 and RTX 5090 GPUs, CONJUGATE COMPILATION achieves a geometric-mean per-round speedup of 1.37x over conventional execution. Interleaved kernels raise this speedup to 1.45x, reaching 1.59x in the best case. CONJUGATE COMPILATION also improves GPU utilization on average by 1.46x on RTX 5090 and 1.50x on B300.