Resource-Aware Joint Kernel Fusion for Ahead-of-Time Inference

Real-time inference of learned models in environments with low latency and memory constraints is becoming challenging with their growing complexity. Hardware architectures such as GPUs are actively utilized for developing efficient serving systems, leveraging their massive parallelism benefits towards computing heavy operations for throughput improvements. Kernel fusion, being a method to harness parallelism in GPUs, is typically applied along a single axis, either horizontal (co-launching independent operators to maximize occupancy) or vertical (chaining dependent operators to avoid DRAM round-trips), with each axis unaware of the other's decisions, and sized against the hardware's physical ceiling rather than what the application is willing to spend. We present Esraj (Efficient Scheduler of Resource-Aware kernels Jointly) a joint fusion scheduler that builds vertical fusion candidates and pairs them with horizontal grouping, considering mutual benefits under a user-declared resource budget. We evaluate across three parallel-branch transformer building blocks, chosen such that every block contains both a horizontal and a vertical fusion opportunity. We compare against internal single-axis ablations and against policies approximating related work, reporting inference latency across batch sizes and resource budgets.