Queueing-Aware Optimization of Reasoning Tokens for Accuracy-Latency Trade-offs in LLM Servers

TL;DR

This paper studies token-budget optimization for a single FIFO LLM server handling heterogeneous query types arriving as a Poisson process. It models per-type reasoning budgets $\ell_k$ to shape both service time and accuracy, and defines a system-wide objective $J(\boldsymbol{\ell})$ that rewards weighted accuracy while penalizing mean queueing time under the stability constraint $\lambda\mathbb{E}[S]<1$. The authors prove $J(\boldsymbol{\ell})$ is strictly concave on the stability region, ensuring a unique optimum, and develop a contracted projected fixed-point method via a Lagrangian/KKT formulation along with a robust projected gradient ascent approach for broader convergence. They also provide a practical integer-budget projection by rounding continuous solutions and quantify the associated performance loss. Numerical experiments with six heterogeneous task types demonstrate substantial gains from task-aware token budgeting over uniform allocations and reveal the sensitivity of the optimal budget to workloads, illustrating the method’s applicability to real LLM serving systems.

Abstract

We consider a single large language model (LLM) server that serves a heterogeneous stream of queries belonging to $N$ distinct task types. Queries arrive according to a Poisson process, and each type occurs with a known prior probability. For each task type, the server allocates a fixed number of internal thinking tokens, which determines the computational effort devoted to that query. The token allocation induces an accuracy-latency trade-off: the service time follows an approximately affine function of the allocated tokens, while the probability of a correct response exhibits diminishing returns. Under a first-in, first-out (FIFO) service discipline, the system operates as an $M/G/1$ queue, and the mean system time depends on the first and second moments of the resulting service-time distribution. We formulate a constrained optimization problem that maximizes a weighted average accuracy objective penalized by the mean system time, subject to architectural token-budget constraints and queue-stability conditions. The objective function is shown to be strictly concave over the stability region, which ensures existence and uniqueness of the optimal token allocation. The first-order optimality conditions yield a coupled projected fixed-point characterization of the optimum, together with an iterative solution and an explicit sufficient condition for contraction. Moreover, a projected gradient method with a computable global step-size bound is developed to guarantee convergence beyond the contractive regime. Finally, integer-valued token allocations are attained via rounding of the continuous solution, and the resulting performance loss is evaluated in simulation results.

Figures (4)

• Figure 1: System model of a single LLM server processing $N$ different heterogeneous query types.
• Figure 2: Empirical accuracy as a function of the enforced reasoning-token budget $\ell$ for each task type. Markers denote measured accuracies; solid curves show the fitted model $p_k(\ell)=A_k(1-e^{-b_k\ell})+D_k$.
• Figure 3: Objective value $J(\boldsymbol{\ell})$ under uniform token allocations ($\ell_k \in \{0,100,500\}$) and the proposed optimal heterogeneous allocation $\boldsymbol{\ell}^\star$.
• Figure 4: Objective value $J(\boldsymbol{\ell})$ as a function of the GSM8K reasoning-token budget $\ell_{\mathrm{GSM8K}}$, with all other budgets fixed to their optimal values.

Theorems & Definitions (3)

• Lemma 1
• Lemma 2
• Lemma 3