Language model inference on edge devices is fundamentally limited by thermals before compute. A Jetson system operates within a 2- to 50-watt budget under a 100°C ceiling, frequently without active cooling. Standard frequency governors handle these constraints poorly: they drive the accelerator near its maximum until throttling intervenes, and they remain blind to the distinct compute characteristics of prefill and decode phases. EnergyLLM replaces these governors with a dynamic runtime scheduler. Using a reinforcement learning agent trained via Group Relative Policy Optimization (GRPO), it independently sets prefill and decode GPU frequencies, request batch sizes, and the number of transformer layers to execute, conditioning every decision on real-time thermal and power headroom.
Scheduling is formulated as a Markov decision process. The state space captures normalized prefill and decode frequency bins, batch size, and remaining thermal headroom, mapping to an action space of several thousand discrete configurations. The reward function balances normalized throughput, energy efficiency, and the energy-delay product. By applying exponential penalties on temperature, latency, power, and exit entropy, the system effectively converts a constrained optimization problem into an unconstrained formulation. The agent’s behavior divides into two regimes: it maximizes throughput when headroom is ample, but shifts to preserving energy once effective headroom falls below 33% (roughly 80°C junction temperature). Training proceeds in two phases: supervised pre-training against reward-weighted grid search data, followed by online fine-tuning on hardware using group-relative advantages.
The scheduler controls four primary dimensions. First, independent frequency scaling exploits the asymmetry between the compute-bound prefill phase and the memory-bound decode phase. Second, batch sizing trades per-request energy costs against token throughput. Third, dynamic early exit recovers the throughput lost to larger batches by terminating the forward pass before the final layer. This utilizes a Llama 3.2 1B LayerSkip model, allowing the agent to dynamically select between layers 9 and 16. Unlike quantization or pruning, this requires no specialized hardware or weight reformatting, and unused layers are never transferred to the GPU. To regulate this, an entropy penalty prevents low-confidence early exits when thermal headroom permits deeper execution. Finally, a factorized action head reduces the output layer from 123,904 to 3,872 parameters. This cuts the memory footprint by 96.8% and policy update costs by over 48%, all while maintaining less than a 3% accuracy loss.
Evaluation was conducted on a Jetson AGX Orin using 1,000 prompts from the Natural Questions dataset. The built-in fan was disabled to model passive cooling, and power was sampled via tegrastats at one-millisecond intervals. Compared to a DVFS baseline pinned to maximum frequency and batch size, which quickly hits the thermal limit and crashes, EnergyLLM stabilizes at 92°C against the baseline’s 99°C peak. This reduces end-to-end latency by 87.5% and power draw by 47%, holding latency near 4 seconds against a 15-second Service Level Objective (SLO) that the baseline eventually violates. Furthermore, dynamic early exit drives a 29.3% reduction in total request time relative to the best static configuration, and a 33.1% reduction relative to DVFS at 99°C. Agent overhead remains marginal: action selection consumes under 0.1% of total latency and energy, while online policy updates require under 3%. Compared to an exhaustive grid search, the policy converges on the same optimal configuration while reducing search energy, power draw, and latency by roughly 80%.