LLM-QP: Query Planning for Large Language Model Inference

Plan Equivalence

Two inference plans $P_1, P_2$ are equivalent if they produce identical token sequences under identical model and constraint semantics.

Definition ($\epsilon$-Equivalence)

Two plans are $\epsilon$-equivalent if logits differ by at most $\epsilon$ over the valid candidate set and produce identical decoded outputs.

Lemma: Dense and Sparse Equivalence

Let $\mathcal{A}(s)$ denote the valid token set. Dense masked scoring:

Sparse scoring:

Then:

Roofline Regime Analysis

Dense head compute:

Sparse head compute:

Sparse dominates when:

Amortized Query Hypothesis

We model decoder state as:

Full transformer recomputation may provide limited incremental information when margins are large. Define margin:

Routing decision — refine if:

Otherwise use amortized scoring.

Planner Routing as a Contextual Bandit

At decoding step $t$, the planner observes a context vector $x_t \in \mathcal{X}$ and selects an execution plan $a_t \in \mathcal{A}$. We define a scalar loss:

and seek to minimize $\sum_{t=1}^T \ell_t(a_t)$.

Oracle and Regret

Define the per-step oracle action:

and dynamic regret:

For stationary regimes, define the best fixed plan:

with static regret:

Linear Contextual Bandit

Assume realizability: $\mathbb{E}[\ell_t(a) \mid x_t] = x_t^\top \theta_a$, with $\|x_t\| \le 1$. LinUCB/linear Thompson sampling achieves sublinear regret of the form:

where $d = \dim(x_t)$ and $\tilde{O}$ hides logarithmic factors.

Nonstationarity and Variation Budget

To capture policy churn, define a variation budget:

Sliding-window or restart-based bandits can achieve dynamic regret bounds of the form:

under standard boundedness assumptions.

Connection to Margin-Based Routing

The threshold router $\pi_\tau$ from the amortization section is a finite policy class when $\tau$ is discretized. This enables expert-style learning over $\{\pi_{\tau_i}\}$ with regret $O(\\sqrt{T \\log |\\Pi|})$ for $|\Pi|$ candidate thresholds.

Bandit loop for LLM-QP. The planner maps runtime context to a plan, observes latency and quality proxies (optionally via shadow evaluation), and updates the routing policy online.

Compiler Integration: Plan Selection in MLIR / StableHLO

Modern inference stacks already compile model graphs through MLIR and StableHLO. LLM-QP can therefore be implemented as a compiler pass that introduces multiple equivalent execution plans and selects among them using a cost model.

Logical vs Physical Plans

We define a logical operator: DecodeStep(query, constraint_state) which expresses the semantics of a single constrained decoding step. Physical implementations include:

• Dense projection head
• Sparse adjacency scoring
• Amortized query update
• Amortized update with rerank
• Full recomputation (refinement)

Plan Expansion Pass

During compilation the planner expands the logical node into candidate physical plans:

DecodeStep
  -> DenseHead
  -> SparseHead
  -> AmortizedHead
  -> AmortizedHead + Rerank

Cost Model

Each plan is annotated with an estimated cost covering kernel latency estimates, memory bandwidth usage, refinement probability, and device utilization. The cost function follows:

Rewrite Rule

if cost(SparseHead) < cost(DenseHead):
    use SparseHead
else:
    use DenseHead

In practice this is implemented using MLIR pattern rewrites or XLA custom calls.

Runtime Adaptation

Runtime telemetry (margins, branching factor, cache state) feeds back into the cost model. The compiler therefore emits a small routing kernel that chooses the implementation at runtime. This hybrid compile-time / runtime planning architecture allows LLM-QP to reuse existing compiler infrastructure while enabling adaptive execution.

Visualizations

Amortized vs full recompute routing tradeoff.

Routing cost vs margin threshold.

Synthetic margin vs branching factor.

Planner Pseudocode

LinUCB for Plan Selection

Let $A = |\mathcal{A}|$ and feature dimension $d$. Maintain per-action matrices $V_a \in \mathbb{R}^{d \times d}$ and vectors $b_a \in \mathbb{R}^d$ initialized as $V_a = \alpha I$, $b_a = 0$. At each step $t$:

1. Observe context $x_t$.
2. For each action $a$, estimate $\hat{\theta}_a = V_a^{-1} b_a$ and compute:
   $$\mathrm{ucb}_a = x_t^\top \hat{\theta}_a + \beta \sqrt{x_t^\top V_a^{-1} x_t}$$
3. Choose $a_t = \arg\min_a \mathrm{ucb}_a$ (loss-minimization).
4. Execute plan $a_t$, observe loss $\ell_t(a_t)$.
5. Update $V_{a_t} \leftarrow V_{a_t} + x_t x_t^\top$ and $b_{a_t} \leftarrow b_{a_t} + x_t \ell_t(a_t)$.

Expert-Over-Thresholds Routing

Discretize thresholds $\{\tau_1, \dots, \tau_N\}$ and define policies $\pi_i = \pi_{\tau_i}$. Maintain weights $w_i$ initialized uniformly. At each step $t$:

1. Observe context $x_t$ and compute each expert’s suggested action $a_t^{(i)} = \pi_i(x_t)$.
2. Sample an expert index $i$ proportional to weights (or pick $\arg\max w_i$ for greedy).
3. Execute $a_t = a_t^{(i)}$, observe loss $\ell_t(a_t)$.
4. Update weights using exponential weighting:
   $$w_i \leftarrow w_i \exp(-\eta \hat{\ell}_{t,i})$$
   where $\hat{\ell}_{t,i}$ is an unbiased loss estimate.

This realizes regret $O(\sqrt{T \log N})$ against the best threshold policy in hindsight under standard assumptions.

End-to-End System Architecture

End-to-end architecture of LLM-QP. A constrained decoding query is transformed into a constraint graph, analyzed by the planner, routed through the execution operator space, and refined via telemetry feedback.

The pipeline operates as follows:

1. Query embedding produces a representation of the decoding state.
2. The constraint graph defines valid token transitions.
3. The planner selects an execution plan using contextual signals.
4. The chosen operator executes the step.
5. Telemetry feeds back to update routing decisions.

Minimal Reproducible Benchmarks

We provide a portable benchmark harness (synthetic, deterministic) that instantiates the cost models and routing policies used throughout the paper. The suite outputs: (i) dense vs sparse crossover vs $K$, (ii) amortized vs full recomputation crossover vs decode length, and (iii) router regret relative to an oracle policy (threshold router).

Dense vs sparse crossover (synthetic roofline + refinement).

Cumulative amortization crossover over sequence length.

Cumulative router regret vs oracle (threshold policy).