Today's AI workloads depend on optimized GPU kernels for efficient inference of a model's forward pass. The task of partitioning work into a minimal set of fused kernels is typically handled by machine learning compilers. However, existing compilers are too restrictive; they rely on static fusion rules, often fall back to opaque third-party kernels that cannot be fused, enforce strict and suboptimal constraints on kernel launch parameters, and are not compatible with state-of-the-art inference engines.

We show that achieving the fastest forward pass for a model necessitates using AI agents to discover the optimal block of fused kernels. We present Together Compiler (TKCC), a compiler for AI inference optimization that leverages agents for joint optimization of a model's forward pass. We redesign compiler infrastructure around maximizing iteration velocity for agents around three key ideas, all focused on optimizing a block of kernels from an inference engine.

First, we introduce the Block Capture, which produces a detailed, kernel-level representation of a model's forward pass, exactly as it was executed in the engine.

Second, we explore the search space via the Block Optimizer, which leverages independent parallel strategies and agents that balance exploration and exploitation to expand the frontier of optimized kernels.

Third, we build the Block Persistent Engine, an atomized instance of the inference engine that allows rapid benchmarking for measurement and correctness validation and catches numerical issues that single-kernel harnesses miss.