Adaptive Inference Batching using Policy Gradients
arXiv:2607.05272v1 [cs.LG] 6 Jul 2026
Ruslan Sharifullin Department of Computer Science Stanford University [email protected]
Abstract Machine learning inference serving systems face the challenge of balancing high throughput with low latency, especially under bursty and heterogeneous workloads. Static batching policies often fail to adapt to dynamic traffic patterns. In this work, we explore Reinforcement Learning (RL) approaches, specifically REINFORCE and Proximal Policy Optimization (PPO), for adaptive inference batching. We develop a custom discrete-event simulator validated against standard queuing models and realworld traces (Azure Functions, BurstGPT). Our extensive evaluation reveals that while simple static heuristics are surprisingly robust for standard single-GPU scenarios, RL provides significant value in multi-GPU environments. Specifically, we demonstrate that a Policy Gradient-based routing agent achieves a 3.5x performance improvement over Round-Robin scheduling by dynamically segregating heterogeneous workloads to minimize Head-of-Line blocking. This highlights the potential of RL not just for temporal batching, but for the joint optimization of request routing and batch composition in distributed inference systems.
1
Introduction
Inference serving is a critical component of the machine learning lifecycle, bridging the gap between trained models and end-user applications. As Deep Learning models grow in size and complexity (e.g. Large Language Models), the computational cost of inference becomes a bottleneck. Batching requests improves throughput by amortizing fixed overheads (kernel launch, memory transfer) across multiple inputs. However, batching introduces a fundamental trade-off: larger batches increase throughput but also increase latency for individual requests waiting in the queue. Static batching policies, such as ”wait 10ms or until 32 requests accumulate,” are the industry standard (NVIDIA Corporation, 2023). While simple to implement and robust in stable conditions, they are rigid and fail to adapt to dynamic traffic patterns (Crankshaw et al., 2017; Ali et al., 2020). During periods of low traffic, a static timeout adds unnecessary latency; during high traffic bursts, a fixed batch size may underutilize the hardware if set too conservatively. Furthermore, in multi-model or multi-GPU environments, the complexity of scheduling increases exponentially. Heterogeneous requests (e.g., a mix of short ResNet inferences and long GPT gen-
erations) can cause Head-of-Line (HoL) blocking, where fast requests are stuck behind slow ones. This project addresses the question: Can reinforcement learning learn adaptive batching policies that outperform static heuristics? We formulate the batching problem as a Markov Decision Process (MDP) and apply Policy Gradient algorithms, specifically REINFORCE and PPO, to learn dynamic batch size selection and request routing. The agent observes the real-time system state (queue length, request types, and GPU availability) and learns a policy that balances competing objectives (throughput vs. latency) without manual tuning.
2
Related Work
The problem of adaptive serving has been studied from both systems and learning perspectives. Heuristic-based Systems: Clipper (Crankshaw et al., 2017) introduced a modular serving architecture that uses an additive-increasemultiplicative-decrease (AIMD) scheme to adjust batch sizes. While effective, AIMD is a reactive heuristic that oscillates and may not converge to the optimal policy for complex distributions. Triton Inference Server (NVIDIA Corporation,
2023) supports dynamic batching but relies on • Request Model: Each request ri is characterusers to manually specify ”preferred batch sizes” ized by an arrival time ti , model type mi (e.g., and timeout windows, which requires extensive ResNet-50, GPT-2), and input size si . • Execution Model: Inference latency is modtuning for each new model and workload. eled as L(b) = α + β · b, where b is the batch Continuous Batching: Recent advancements in size, α is the fixed overhead (e.g., kernel launch, LLM serving, such as Orca (Yu et al., 2022) and PCI transfer), and β is the per-sample processvLLM (Kwon et al., 2023), introduced iterationing time. We profiled real NVIDIA V100 GPUs level scheduling (”continuous batching”) to mitito obtain realistic parameters: α = 2ms for all gate the impact of variable output lengths. While models, with β varying by model type (e.g., our work focuses on request-level batching, the 5ms for ResNet, 20ms for GPT-2). routing policies we develop are complementary • Queuing: We implement a First-In-First-Out to these intra-GPU scheduling techniques. (FIFO) queue with a configurable maximum length (N = 100). If the queue is full, incomLearning for Systems: DeepRM (Mao et al., ing requests are dropped (load shedding). 2016) was a pioneering work applying Deep RL to cluster resource management, visualizing the • Multi-GPU Support: The simulator supports K heterogeneous GPUs. A central dispatcher problem as a ”Tetris” game of packing tasks. (the Agent) decides which GPU to route a batch FaaSRank (Suresh et al., 2021) utilized learningto, or whether to wait for more requests. to-rank for scheduling serverless functions, optimizing for cold-start latency. Our work differs by focusing specifically on the micro-second scale 3.2 Workloads decision of batch composition and routing in in- We evaluate on three distinct workload types to ference clusters, leveraging modern Policy Gra- test the agent’s generalization capabilities: dient methods (REINFORCE (Williams, 1992), PPO (Schulman et al., 2017)) which are well- 1. Standard (Poisson): Standard M/M/1 queue traffic with constant mean arrival rate λ = 10 suited for the stochastic nature of request arrivals. req/s. This provides a baseline for stability. 2. Extreme Burst: A stress test where traffic 3 Environment (Dataset and shifts squarely between 0 and 100 req/s every Features) 200 steps. This tests the agent’s ability to switch between ”latency-minimizing” (small We built a custom discrete-event simulator to batches) and ”throughput-maximizing” (large model the inference serving environment. This batches) modes. simulator serves as our ”dataset” generator, allowing us to test across various traffic patterns and 3. Real-World Trace: A replay of production traffic traces from Azure Functions (Shahrad system configurations that would be difficult to et al., 2020), characterized by diurnal patterns reproduce consistently on a physical cluster. and unpredictable spikes. 4. Multi-GPU Routing (Heterogeneous): A 50/50 mix of ”Fast” (ResNet-50, 50ms latency) and ”Slow” (GPT-2, 200ms latency) requests. This scenario specifically tests the agent’s ability to handle Head-of-Line (HoL) blocking in the Multi-GPU setting.
4
Methods
We formulate the adaptive batching problem as a Markov Decision Process (MDP) and solve it using Policy Gradient methods. 4.1 MDP Formulation Fig. 1: Schematic of the RL Agent-Environment interaction. The agent observes state St (queue, State Space: The state St ∈ R8 captures the GPU status), takes action At (batch size/routing), system snapshot: and receives reward Rt (throughput - latency). • Normalized Queue Length: Qt /Qmax . • Time since last batch: (t − tlast )/Twindow . 3.1 Simulator Dynamics • Current Request Type: One-hot encoding of the request at the head of the queue (MemoryThe simulator operates on a continuous timebound vs Compute-bound). line, processing events such as ‘RequestArrival‘, ‘BatchCompletion‘, and ‘DecisionStep‘. 2
• GPU Status: Binary vector indicating if each GPU is currently busy.
5
Experiments, Results, Discussion
5.1
Experimental Setup
Action Space:
All agents are implemented in PyTorch. We train • Single GPU: Discrete batch size selection at ∈ for 2000 episodes, with each episode consisting {0, 1, . . . , 32}. Action 0 implies ”Wait”. of 1000 simulation steps. • Multi-GPU: Joint routing and batching at ∈ {0, . . . , 64}. Actions select a specific (GPU, • Baselines: – Static-8: A fixed batch size of 8 (tuned via BatchSize) pair, optimizing both decisions sigrid search). multaneously. – Random: Randomly routes requests to availReward Function: We designed a composite able GPUs. reward to balance throughput and Service Level – Round-Robin: Cyclically assigns requests Agreements (SLAs). (GPU 0 → GPU 1). X – Shortest-Queue (SQ): Assigns to the GPU Rt = Throughputt − wr · Latencyr with fewer pending requests. r∈Batch • Evaluation: We report the average cumulative In the heterogeneous scenario, we use wf ast = reward over 20 evaluation episodes with fixed 200.0 and wslow = 20.0. This heavy penalty for random seeds distinct from training. delaying ”Fast” requests encourages the agent to prioritize them or segregate them from ”Slow” 5.2 Results requests to avoid HoL blocking. Single-GPU Results: In standard scenarios, RL matched but did not significantly outperform optimized static batching. This is expected for PoisWe use a neural network to approximate the son arrival processes, where static thresholds are stochastic policy πθ (a|s). The architecture: known to be near-optimal policies. The reward 1. Input Projection: A linear layer mapping the landscape is dominated by throughput, making static batching a strong local optimum. Note the 8-dim. state to a 4-dim. embedding. 2. Multi-Head Attention: A self-attention layer flat curves for single-GPU scenarios where the (2 heads, embed dim=4) to capture dependen- agent quickly converges to the static baseline, vercies between state features. We found that at- sus the rising curve in the Multi-GPU scenario. Static Baseline REINFORCE Agent Improvement tention was particularly useful for correlating Scenario Standard (Single GPU) 254.03 254.30 +0.1% 329.91 333.26 +1.0% queue depth with GPU availability, allowing Extreme Burst Real-World Trace 86.89 87.11 +0.2% the agent to learn ”if queue is high AND GPU Multi-GPU Routing 203.23 910.52 +348.0% is free, dispatch immediately.” Table 1: Performance comparison (Reward) 3. Policy Head: A Multi-Layer Perceptron across different scenarios. (MLP) with one hidden layer (64 units, GELU activation) outputting logits for the categorical action distribution. 4. Value Head: A separate MLP estimating the state-value function V (s) for baseline subtraction in REINFORCE. 4.2
Policy Network Architecture
4.3
Algorithms
We compare two Policy Gradient algorithms: • REINFORCE: The standard Monte-Carlo policy gradient. Discount factor γ = 0.95, update the policy at the end of each 1000-step episode. • PPO: Proximal Policy Optimization, which uses a clipped surrogate objective (ϵ = 0.2) to allow for multiple update epochs per batch of experience. We found that REINFORCE with a learned baseline often converged faster, likely due to the relatively short horizon and dense reward signal. PPO, while more stable, required more hyperparameter tuning for the clipping range and entropy coefficient.
Fig. 2: Training progress (Reward vs Episode) across four scenarios. RL struggles to improve in single-GPU cases but learns effectively in the Multi-GPU Routing task.
3
Multi-GPU Results: In the heterogeneous routing scenario, the RL agent demonstrated a massive advantage. Figure 3 shows the performance comparison against three baselines:
vergence suggests that the reward signal is dense enough for the agent to quickly identify the ”segregation” strategy, effectively learning to classify requests by their computational cost.
• Random: Blindly assigns requests (Return: 105.5). • Round-Robin: Assigns cyclically (Return: 203.2). • Shortest-Queue (SQ): Assigns to the shortest queue (Return: 612.8).
Finally, we evaluate the Latency-Throughput tradeoff (Fig. 5). The REINFORCE agent achieves 60% higher throughput than ShortestQueue (17.98 vs 11.18 req/s) while maintaining 25% lower latency than Round-Robin (2.59s vs 3.46s). The red dashed line marks the SLA threshold (3.0s). Round-Robin violates this threshold despite higher throughput (22.42 req/s). The RL agent finds the ”Sweet Spot”, maximizing throughput within SLA constraints while avoiding congestion.
Fig. 3: Performance comparison in Multi-GPU Routing. RL outperforms even the strong Shortest-Queue heuristic by 48% by learning to segregate request types. While SQ improves over Round-Robin by balancing load counts, it fails to account for request heterogeneity (Fast vs Slow). The REINFORCE Fig. 5: Latency-Throughput tradeoff with SLA agent (Return: 910.5) learns to segregate workline (3.0s). RL balances throughput and latency, loads, routing fast requests to one GPU and slow stopping before congestion. requests to another, eliminating Head-of-Line blocking. 5.3 Discussion Our experiments reveal several key insights into the application of RL for systems. Reward Shaping and SLA Compliance: The choice of weights wf ast and wslow in the reward function was critical. We observed that setting equal weights led the agent to treat all requests interchangeably, resulting in a policy similar to Shortest-Queue. By heavily penalizing latency for ”Fast” requests (wf ast = 200.0), we explicitly encoded the SLA requirement into the optimization objective. This demonstrates that RL Fig. 4: Learning curve of the REINFORCE agent. agents in systems contexts often require domainspecific reward shaping to discover non-trivial Dashed lines indicate baseline performance. policies. To understand the training dynamics, we analyze the learning curve of the REINFORCE agent in the multi-GPU scenario (Fig. 4). The agent begins with performance comparable to the Round-Robin baseline but rapidly improves within the first 100 episodes, eventually surpassing the Shortest-Queue heuristic. This rapid con-
Generalization to Unseen Traces: A major concern with RL is overfitting to the training distribution. We trained our agent on a synthetic Poisson process but evaluated it on the ”Extreme Burst” and ”Real-World Trace” scenarios. The results (Table 1) show that the policy generalized well, maintaining performance parity with static
4
baselines even on out-of-distribution traffic. This References suggests that the learned policy relies on imme- Ahsan Ali, Riccardo Pinciroli, Feng Yan, and Evgenia diate state features (queue length, GPU status) Smirni. Batch: Machine learning inference serving rather than memorizing arrival patterns. on serverless platforms with adaptive batching. In SC20: International Conference for High PerforImpact of Attention Mechanism: Ablation studmance Computing, Networking, Storage and Analyies (not shown for brevity) indicated that the sis, pages 1–15. IEEE, 2020. Multi-Head Attention layer improved convergence speed by approximately 20% compared Daniel Crankshaw, Xin Wang, Guilio Zhou, Michael J to a simple MLP. The attention mechanism likely Franklin, Joseph E Gonzalez, and Ion Stoica. Clipper: A low-latency online prediction serving system. helps the agent focus on the most relevant parts of In 14th USENIX Symposium on Networked Systems the state space, for instance, ignoring the queue Design and Implementation (NSDI 17), pages 613– length when all GPUs are busy, or focusing in627, 2017. tensely on the head-of-line request type when a routing decision must be made. Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying
5.4
Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, pages 611–626, 2023.
Limitations
While our results are promising, there are limitations. First, our simulator assumes deterministic execution times for a given batch size, whereas real GPUs exhibit variance due to thermal throttling and background processes. Second, we fo- Hongzi Mao, Mohammad Alizadeh, Ishai Menache, cused on discrete batch sizes, but modern systems and Srikanth Kandula. Resource management with like vLLM use continuous batching (iterationdeep reinforcement learning. In Proceedings of the 15th ACM Workshop on Hot Topics in Networks, level scheduling). Extending our action space pages 50–56, 2016. to support continuous batching would require a more complex policy network. Finally, we asNVIDIA Corporation. Dynamic batching in triton sumed zero network latency for the dispatcher; inference server. https://docs.nvidia.com/ in a real distributed cluster, the communication deeplearning/triton-inference-server/, overhead of the central agent could become a 2023. bottleneck. John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint We investigated the application of Policy GraarXiv:1707.06347, 2017.
6
Conclusion and Future Work
dients to adaptive inference batching. While single-GPU temporal batching is well-served by Mohammad Shahrad, Rodrigo Fonseca, Íñigo Goiri, Gohar Chaudhry, Paul Batum, Jason Cooke, Edstatic heuristics, we found that RL shines in the uardo Laureano, Colby Tresness, Mark Russinovich, complex combinatorial problem of Multi-GPU and Ricardo Bianchini. Serverless in the wild: CharLoad Balancing. By jointly optimizing routing acterizing and optimizing the serverless workload and batch composition, our REINFORCE agent at a large cloud provider. In 2020 USENIX Annual demonstrated a 348% improvement (which is Technical Conference (USENIX ATC 20), pages 205– 3.5x performance improvement) over standard 218, 2020. scheduling techniques. Amoghavarsha Suresh, Gagan Soumya, Animesh
Future work should explore integrating this routKapil, Sriram Kaler, and Anshul Gandhi. Faasrank: ing policy with continuous batching mechaLearning to schedule functions in serverless platforms. In 2021 IEEE International Conference on nisms (e.g., Orca, vLLM) used in modern LLM Autonomic Computing and Self-Organizing Systems serving systems. Additionally, deploying the RL (ACSOS), pages 41–50. IEEE, 2021. agent on a real cluster with network latency would validate its robustness to distributed system noise. Ronald J Williams. Simple statistical gradientInvestigating offline RL could also allow for safer following algorithms for connectionist reinforcepolicy updates in production environments withment learning. Machine learning, 8(3):229–256, out the need for online exploration. 1992.
7
Contributions
Gyeong-In Yu, Joo Seong Jeong, Geon-Woo Kim, Soojeong Kim, and Byung-Gon Chun. Orca: A distributed serving system for transformer-based generative models. 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pages 521–538, 2022.
Ruslan Sharifullin implemented the simulator, the RL agents (REINFORCE, PPO), and conducted all experiments.
5