Measured Joules, Learned Routes: Learning to Route for Energy-Efficient LLM Serving

For modern AI model serving, it is hard to apply one LLM model fit all inputs. This is due to the fact that it becomes challenging to strike a balance between being accurate on hard items and staying efficient on easy queries. For instance, a small, open-weight model generates few tokens and spends little energy, but it can fail queries which need more model and reasoning capacity. Whereas a large, reasoning-based model does the reverse, because they may overthink, and can waste significant amount of joules on items a smaller model would have easily answered (Zhao et al., 2026).

To this end, we study a serving system that selects one answer LLM for each incoming query. Let $\mathcal{M}=\{m_{1},\ldots,m_{K}\}$ denote a fixed pool of $K$ candidate LLMs, and let $\mathcal{A}=\{1,\ldots,K\}$ index the available routes. A query context $x$ contains the question and any metadata provided to the router. The routing policy $\pi_{\theta}(a\mid x)$ selects a candidate index $a$, after which model $m_{a}$ generates an answer $y$. The router makes the model-selection decision; the selected candidate performs the task. Each evaluated query involves one routing decision and one answer-model outcome. Multi-turn recovery and tool use could also be applied to our framework with proper extensions.

The pool of candidate LLMs is heterogeneous by design, and the model card is illustrated in Table 1. Candidates differ in model family, parameter count, and reasoning style, from sub-billion Qwen2.5 drafts through mid-size models such as Gemma2-2B, Qwen2.5-7B, and Llama 3.1 8B, up to DeepSeek-R1 8B and 32B. That spread is what makes routing meaningful. These differences create opportunities for query-dependent selection, since a model that is effective for one task need not offer the best accuracy–energy trade-off for another. For example, Figure 2 shows a GSM8K query answered correctly by Qwen2.5-0.5B using $82\,\mathrm{J}$ and by DeepSeek-R1-8B using $2864\,\mathrm{J}$, while Llama 3.1-8B answers incorrectly using $638\,\mathrm{J}$. For this question, though Llama model has larger size, its overthinking and extra steps around equation modeling causes both inefficient and inaccurate inference. Such an example motivates measuring outcomes for each query-energy-model triplet rather than inferring suitability from model size alone. We assume all answer models are open-weight checkpoints served locally through Ollama on that device, so each candidate can be invoked under the same power logger rather than billed as a separate API price.

For query $i$ and candidate $a$, let $y_{ia}$ be the recorded answer and $y_{i}^{\star}$ its reference. We denote correctness by

where $g_{b(i)}$ is the evaluator for benchmark $b(i)$. This notation covers task-specific answer matching and code-test evaluation without requiring exact and literal string equality. Let $E_{ia}$, $L_{ia}$, and $\bar{P}_{ia}$ denote the each answer call $i$’s incremental GPU energy in joules, wall-clock duration in seconds, and mean incremental GPU power in watts, respectively.

The candidate models are served locally through Ollama on NVIDIA A100-SXM4-80GB GPUs with tensor parallelism. During profiling, the logger polls board power through NVML using pynvml at a $200\,\mathrm{ms}$ interval and subtracts an idle-power baseline. We adopt nvidia-ml-py bindings, which read the same device counters as nvidia-smi. Let $P_{ia}^{\mathrm{board}}(t)$ denote the board power trace and $P_{ia}^{\mathrm{idle}}$ the baseline associated with the query call $i$, the measurement quantities of power, energy, and duration are modeled as

The energy integral is estimated from the sampled trace. Note these quantities describe idle-subtracted GPU consumption within the measured query call interval, which is the core energy associated with AI inference. They do not represent total server or data-center-level energy consumption.

We first profile AI model’s energy and performance through an offline tournament of mild size, which records the outcomes of the candidate models on the collected questions. With given query input $x_{i}$, the per-query record is profiled as

These are realized outcomes of the underlying LLM inference generations. Replaying a record reproduces that recorded outcome, rather than estimating the answer model’s stochastic behavior or measurement variation. The complete outcome vector provides feedback for all candidates on a training query, although the learned policy selects only one candidate at inference time.

In this work, we explore the possibilities of meeting LLM service objectives at low energy cost. On a fixed logged dataset $\mathcal{D}$ of $N$ queries, the expected accuracy and answer energy of a feasible routing policy are

For deterministic routing, these expressions reduce to averages of the selected candidates’ outcomes. For sampled decisions, the corresponding sample averages estimate these expectations. Handling of invalid controller outputs is discussed in Section 3.1.

One design to handle the desired accuracy-energy trade-off is

where $B$ is an average answer-energy budget in joules per query. This constraint expresses a service goal, not a guarantee made by the training algorithm. The reported method instead optimizes the gated rank-based surrogate defined in Section 3.2 and evaluates its resulting accuracy and measured answer energy. Problem (7) can be treated as a single-step contextual decision problem. A contextual-bandit formulation accommodates either a classical or a text-conditioned neural policy. In addition, our logged training data provide detailed outcomes for every candidate. We therefore distinguish the routing problem, the controller’s representation, and the feedback used to train it.

We note with the addition of the router policy $\pi_{\theta}$, the complete serving energy also depends on the controller and deployment costs, including model loading and residency under the chosen serving configuration. Those costs is measured use the consistent metric to establish end-to-end savings, and they are not supplied by the lookup of $E_{ia}$ alone. More importantly, as LLM’s consumed energy are strongly dependent on input queries and output contexts, it requires dedicated design and investigation of the router.

To fully release the routing controller’s effectiveness, we propose to utilize a lightweight, finetuned LLM in a text-conditioned approach. Its role is to interpret a query and select a candidate from the fixed pool, using a smaller LLM than most of the candidate answer models. The controller prompt contains the question like in Figure 9, a fixed legend mapping labels to candidate indices, and an optional difficulty field when available. Recorded correctness and energy are regarded as training feedback rather than inputs supplied to the controller when selecting a model.

Let $p_{\theta}(o\mid x)$ denote the probability of a controller completion $o$ given query context $x$, and let $f(o)$ parse its route label into an index in $\mathcal{A}$. The induced route probabilities are

This separates probabilities over generated texts from probabilities over candidate models. The reported variants use either a short rationale with a Choice: <letter> field or a choice-only completion. Both formats produce the same model-selection action. An un-parseable output is an invalid action and receives a negative training reward. If multi-turn routing is enabled, the router also sees any prior user-assistant turns on the same item, the current turn number, the model chosen last time, and whether that attempt was correct. The router can output Stop when a satisfactory answer already exists and does not require to route to candidate LLMs. In all reported experiments we use a single routing turn per question, so each item produces one letter choice. Probability on invalid outputs is separate from the valid-route probabilities in Eq. (8).

The tournament contains approximately $3,000$ collected questions from MMLU, GSM8K, BBH, ARC Challenge, HellaSwag, MATH-500, and HumanEval. The reported training and held-out subsets contain $2050$ and $913$ questions, respectively; benchmark counts are given in Table 5 in Appendix. The training subset supplies oracle labels and policy rewards. We use Qwen2.5-1.5B-Instruct as a text-conditioned routing controller, and find it achieve desired post-training behaviors with respect to routing decisions. During policy optimization, the controller generates new completions, but answer-model outcomes are retrieved from the fixed training records. Thus, router sampling is on-policy with respect to the rollout controller, while its environmental feedback is replayed.

To ensure in the RL post-training, the router can effectively evaluate the input queries’ contexts and make decisions, we come up with a ranking-based reward for energy, which is in the same scale as of performance reward. Such a reward combines task correctness with preferences over inference costs (e.g., energy, task completion time and etc). For query $i$ and candidate $a$, $C_{ia}$ denotes correctness, $L_{ia}$ denotes answer latency in seconds, and $\bar{P}_{ia}$ denotes mean incremental GPU power in watts. The measured answer energy $E_{ia}$ is expressed in joules.

We derive a bounded energy penalty from the ascending rank $k_{ia}$ of $E_{ia}$ among the $K$ candidates for the same query:

With distinct energies, zero corresponds to the cheapest candidate and one to the most expensive. For example, energies of $85$, $95$, $280$, and $3800\,\mathrm{J}$ produce penalties of $0$, $1/3$, $2/3$, and $1$. $z_{ia}$ denotes rank throughout. It is distinct from energy $E_{ia}$. Ranking bounds the training penalty but discards the magnitude of differences in joules.

Let $\boldsymbol{\lambda}=(\lambda_{t},\lambda_{e},\lambda_{p})$ denote nonnegative cost coefficients associated with session duration, energy rank, and average GPU power. The general gated reward is

Latency, energy, and power describe related but different aspects of a session call. The same energy can be consumed over different durations and at different mean power levels. Thus, separate latency and power coefficients can express preferences over those quantities, but an average-power penalty alone does not enforce an instantaneous device or rack power limit. We distinguish this general reward form from the configuration used in the reported experiments as follows.

For each training query, the oracle target is selected from

where $\lambda_{e}^{\mathrm{SFT}}$ is the coefficient used during target construction. Under the reported positive coefficients $0<\lambda_{e}^{\mathrm{SFT}}\leq 1$, every correct route receives a larger reward than every incorrect route. The oracle therefore selects a lowest-energy correct candidate whenever one exists. Changing the positive coefficient within this range does not change that hard target if the logs and tie-breaking rule are fixed. For $\lambda_{e}^{\mathrm{SFT}}=0$, correct candidates tie, and the draft’s accuracy-only configuration selects the largest correct model. When all candidates are wrong, every valid route ties under the reward.

Let $o_{i}^{\star}$ be the target completion encoding the selected route in the chosen output format. Supervised fine-tuning initializes the controller by minimizing the completion negative log likelihood

where $\mathcal{I}_{\mathrm{tr}}$ indexes the $N_{\mathrm{tr}}$ training queries. We can include two variations of the target during SFT. In the choice-only variant, the target contains the route field; in the rationale variant, it also contains an explanation. The resulting checkpoint is evaluated as an SFT baseline and initializes policy optimization. SFT learns to imitate target completions, whereas the next stage uses the rewards of the controller’s sampled decisions.

We fine-tune the controller using GRPO (Shao et al., 2024), which is based on groupwise comparisons instead of using a separate critic model. For each training query $i$, a rollout policy $p_{\theta_{\mathrm{old}}}$ samples $G$ completions $\{o_{ig}\}_{g=1}^{G}$. Each completion is parsed and scored using its selected candidate’s logged outcome and coefficient $\lambda_{e}^{\mathrm{RL}}$, or the invalid-action penalty. Denote these completion rewards by $r_{ig}$. Their group statistics and advantages are

where $\delta>0$ stabilizes the denominator. A completion receives positive advantage when its reward exceeds the group mean. If all group rewards are equal, all advantages are zero and that group supplies no reward-driven update.

At the completion level, the clipped policy surrogate is expressed using importance sampling ratio

with the resulting loss function

where $\varepsilon$ controls ratio clipping and is distinct from the advantage stabilizer $\delta$. This surrogate favors higher-reward completions while limiting the incentive for large policy-ratio changes.

The KL-regularized variant additionally anchors the controller to the SFT reference distribution $p_{\mathrm{ref}}$. At the distribution level its regularized objective is

with regularization coefficient $\beta\geq 0$. This expression concerns controller completions rather than the marginal distribution over model indices in Eq. (8). Note KL regularization discourages departure from the reference policy. It does not by itself guarantee diverse routing or prevent concentration on a single candidate. Some configurations center rewards by subtracting a query-specific mean candidate reward. If the same scalar is subtracted from every reward before the group normalization in Eq. (16), it cancels from the advantages. Such centering is therefore distinct from any additional reward scaling or clipping used by the trainer.

In this paper, we unify and use TRL to conduct main SFT, GRPO, and KL-GRPO post-training (von Werra et al., 2020). Policy optimization starts from the corresponding SFT checkpoint, and samples multiple controller completions per training prompt. The main reported configuration uses $\lambda_{e}^{\mathrm{SFT}}=0.5$ and $\lambda_{e}^{\mathrm{RL}}=0.4$ as we find these hyperparameters give the best overall SFT and RL performances. The experiments also examine matched coefficients and alternative checkpoints. These coefficients identify the oracle-construction and policy-reward settings, respectively. Their effects on hard targets, sampled rewards, and learned policies should be distinguished. The output format is also part of each configuration, where choice-only and rationale-generating controllers have different generation costs, and should be compared with their matching SFT initializations.

On held-out queries, the router can be treated as the lightweight add-on to inference engines, where it selects a LLM model and the evaluator retrieves that candidate’s stored outcome. We report accuracy, mean answer-model energy, latency, power, and model-selection frequencies. Reward comparisons require a common evaluation coefficient: reward values computed with different $\lambda_{e}$ are different metrics even when the policies share an SFT initialization. Accuracy and measured joules provide the primary comparison across configurations and baselines.

The replay protocol isolates model-selection behavior on a fixed outcome table. It does not measure live changes in answer generation, model-loading cost, queueing, or concurrent serving. Controller measurements and answer-model measurements must therefore be distinguished when discussing complete serving costs. The evaluation addresses single-turn routing on the profiled pool. Adaptation to new deployment conditions and multi-turn agent workflows would require additional experiments.

We evaluate routing policies on a mixed held-out set drawn from seven standard LLM benchmarks that span knowledge, reasoning, math, and code: MMLU (Hendrycks et al., 2020), GSM8K (Cobbe et al., 2021), BBH (Suzgun et al., 2023), ARC Challenge (Clark et al., 2018), HellaSwag (Zellers et al., 2019), MATH-500 (Lightman et al., 2024), and HumanEval (Chen et al., 2021). These tasks stress different answer-model costs, from short multiple-choice items to long generative math and code solutions, which is what makes energy-aware routing nontrivial on the shared nine-model pool.

Power measurements are based on NVML¹¹ 1 https://docs.nvidia.com/deploy/nvml-api/. To keep energy and performance comparable across runs, we place all router training and inference on instances of 1x NVIDIA A100-SXM4-$80$,GB GPU. Answer models are served locally with Ollama on that device. The routing controller is trained and run with Hugging Face Transformers and TRL on the same GPU. Offline tournament collection used a separate machine with the same A100-SXM4-$80$,GB specification. During each answer-model generation we poll board power at a $200,\mathrm{ms}$ interval via pynvml, subtract a short idle baseline, and record instantaneous effective power $P(t)$ in watts. Effective energy for the call is $E_{\mathrm{call}}$ as defined in Eq. (3), and latency is the wall-clock duration of the same generation. Table 2 reports mean accuracy together with measured duration, energy, and power summaries.

All methods in Table 2 are scored on the same 913 held-out test questions covering a spectrum of difficulty and categories. Scoring is an offline lookup from pre-logged tournament outcomes, so no live answer models are called at evaluation time.

K-NN estimates per-model error from embedding neighbours on the train split. UniRoute matches the prompt to a train K-means cluster and uses per-cluster error plus a cost term. Smoothie is label-free and scores candidates from embeddings of $[\mathrm{prompt},\mathrm{output}]$. Unlike the others, Smoothie uses all logged generations on the test item to decide, so we treat it as an output-informed diagnostic rather than a deployable pre-generation router. We still charge only the chosen model’s energy. GreenServ is a local LinUCB adaptation using task, semantic-cluster, and complexity features. Its reward is $(1-\lambda)\mathrm{Acc}-\lambda\cdot\mathrm{energy}$ and is warmed on train and frozen before test. It is not the paper’s online serving stack. RouteLLM is a binary router over two fixed answer models, not a classifier over the full pool. We set the strong arm to DeepSeek-R1-32B and the weak arm to the smallest pool model, train a prompt classifier for $P(\text{strong wins}\mid q)$ on train outcomes, and route to the strong arm when $P\geq\alpha$. Fixed-model baselines always route every question to one predetermined answer model (always-smallest, always-medium, or always-large).

Table 2’s methods and results can be split into three categories. High-accuracy part is K-NN and always-medium, near $79\%$ at about $3.7$–$4.1\,\mathrm{kJ}$. Always-large and RouteLLM spend even more and do worse. The mid-accuracy band is UniRoute, Smoothie, and our $\lambda_{e}{=}0.5{\to}0.4$ GRPO runs, about $58$ to $62\%$ at $0.8$ to $2.2\,\mathrm{kJ}$. The cheap accuracy band is GreenServ $\lambda{=}0.5$ and matched GRPO $\lambda_{e}{=}0.3$, about $57\%$ at about $280$–$300\,\mathrm{J}$. SFT sits below that mid accuracy band. Always-smallest is cheap in energy but achieves poor accuracy.