Adaptive Inference Batching using Policy Gradients
2026-07-06 • Machine Learning
Machine LearningArtificial IntelligenceDistributed, Parallel, and Cluster ComputingPerformance
AI summaryⓘ
The authors studied how to improve systems that handle AI requests by using reinforcement learning (RL) to decide how to batch and route tasks. They found that for simple setups with one GPU, classic batching methods work almost as well as RL. However, in more complex cases with multiple GPUs and mixed task types, RL learned smart ways to avoid delays and boosted performance a lot. Their RL approach also worked well on new, unpredictable traffic even though it was trained on simple simulated data. This shows RL is most helpful when juggling many resources at once rather than just scheduling tasks over time.
reinforcement learningbatching policiesinference serving systemsmulti-GPU routingHead-of-Line blockingMDP (Markov Decision Process)throughputlatencyPoisson trafficSLA (Service Level Agreement)
Authors
Ruslan Sharifullin
Abstract
Inference serving systems must balance throughput and latency under bursty, heterogeneous workloads, yet the industry standard remains static batching policies that require manual tuning and cannot adapt to shifting traffic. We investigate whether reinforcement learning (RL) can learn adaptive batching and routing policies that outperform these heuristics, training REINFORCE and PPO agents on a discrete-event simulator validated against queuing theory and production traces (Azure Functions, BurstGPT). We formulate the problem as an MDP over queue state, request type and GPU availability, evaluating across standard Poisson traffic, extreme bursts, real-world traces and heterogeneous multi-GPU routing. Our central finding is a clear boundary condition for RL's value in systems problems. In single-GPU settings, a well-tuned static batching policy is already near-optimal under Poisson-like arrivals and RL offers only marginal gains (+0.1% to +1.0%). In multi-GPU heterogeneous routing, however, where fast and slow requests compete for shared resources, the agent discovers a workload-segregation policy that eliminates Head-of-Line blocking, yielding a 3.5x (348%) improvement over Round-Robin and a 48% improvement over the strongest heuristic baseline (Shortest-Queue), with 60% higher throughput and 25% lower latency while respecting SLA constraints. The policy generalizes to unseen bursty and real-world traffic despite training only on synthetic Poisson arrivals and an attention-augmented policy network converges roughly 20% faster than an MLP baseline. These results suggest RL's advantage over engineered heuristics concentrates in combinatorial, multi-resource decisions rather than single-resource temporal scheduling, a practical distinction for deciding where learned policies justify their engineering cost in production inference infrastructure.