Energy-Aware Routing to Large Reasoning Models

Technical Summary: Energy-Aware Routing to Large Reasoning Models

1. Problem Statement

The paper addresses the challenge of managing energy consumption in "AI factories"—large data centers hosting Large Reasoning Models (LRMs). These systems face a multi-dimensional optimization problem:

• Heterogeneous Energy Costs: Different LRMs have varying inference energy requirements based on their size and the amount of "thinking time" (compute) allocated to a task.
• Stochastic Energy Supply: Data centers increasingly rely on renewable energy (solar, wind), which is highly variable and often misaligned with workload demands.
• Task Constraints: Each task arrives with specific requirements, including a latency deadline ($\lambda$) and an error tolerance ($\epsilon$).
• The Objective: To design a routing policy that dispatches tasks to specific models and allocates sufficient compute to satisfy accuracy and deadline constraints while minimizing the use of auxiliary (non-renewable) energy to cover deficits.

2. Methodology

The authors develop a mathematically principled framework to characterize the limits of energy-aware routing.

• System Modeling: The system is modeled as a stochastic process where tasks arrive in a queue and energy availability is governed by a baseline (renewable) source and an auxiliary source. The authors use Brownian motion and Donsker’s invariance principle to approximate the battery's energy trajectory.
• First- and Second-Order Characterization:
  • First-order (Drift): Analyzes the mean energy deficit caused by the mismatch between average energy supply and average consumption.
  • Second-order (Fluctuation/Volatility): Analyzes how stochastic variability in both energy harvesting and task arrivals necessitates energy reserves, even when the mean supply meets the mean demand.
• Scaling Laws Integration: To make the model practical, the authors incorporate empirical training-compute scaling laws (to model model capability) and inference-compute scaling laws (to model how energy and latency scale with token generation and context length).
• Myopic Policy Analysis: They evaluate a "myopic" baseline policy—one that dispatches tasks to the most energy-efficient model available at the moment of arrival—to establish a benchmark for performance.

3. Key Contributions

• Mathematical Formulation: A formal definition of the energy-aware routing problem that treats routing decisions as a means to manage energy volatility.
• Regime Identification: The discovery of three distinct operating regimes based on the relationship between energy supply ($R$) and consumption ($C$):
  1. Persistent Deficit ($\mu < 0$): Auxiliary energy grows linearly with time.
  2. Persistent Surplus ($\mu > 0$): The system is stable, but energy is systematically wasted.
  3. Critical Regime ($\mu = 0$): The "sweet spot" where neither energy is wasted, but the system is volatility-limited, meaning the energy deficit scales with $\sqrt{T}$.
• Routing Error Analysis: A characterization of how "routing error" (suboptimal model selection) shifts the system from a fluctuation-dominated regime ($\sqrt{T}$ scaling) to a drift-dominated regime (linear scaling).
• Scaling Law Mapping: A bridge between high-level routing decisions and low-level transformer physics (parameter memory access, FLOPs, and attention mechanisms).

4. Results

• Deficit Scaling: The authors prove that in the critical regime (where mean supply equals mean demand), the expected auxiliary energy required scales as $O(\sqrt{T})$.
• The Impact of Accuracy: Numerical simulations demonstrate that as prediction error ($E$) increases, the system quickly transitions from a manageable square-root scaling of energy reserves to a much more expensive linear scaling.
• Model Trade-offs: Simulations show that for easy tasks, small models are superior. However, as task difficulty increases, small models require disproportionately more tokens (and thus more energy/time) to meet error tolerances, eventually making large models more efficient.

5. Significance

This work shifts the focus of AI efficiency from merely "making models smaller" to "intelligent dispatching." It provides a theoretical foundation for building "variance-aware" routers. Instead of just optimizing for the average case, future AI factory controllers can use these insights to balance model selection and compute allocation against the real-time volatility of renewable energy grids. This is crucial for the sustainable scaling of large-scale AI infrastructure.