PreSTO: Predictive Subtree Prefetching for Fast LLM Power Sampling

Nan Jiang; Panagiotis Theodoropoulos; Wei Deng

cs.CLOct 9, 2026

PreSTO: Predictive Subtree Prefetching for Fast LLM Power Sampling

Authors: Nan Jiang, Panagiotis Theodoropoulos, Wei Deng

Organizations: University of Texas at El Paso · Georgia Institute of Technology · Morgan Stanley

Abstract

Power sampling improves large language model (LLM) reasoning at test time by sampling from a sharpened version of the base model's distribution. Its sampler, PowerMH, is slow because its Metropolis--Hastings (MH) updates are sequential: each accept/reject decision sets the state from which the next proposal is generated. Prior work reduces this cost by changing the sampler or the model (e.g., sequential Monte Carlo, approximate lookahead, distillation). We propose Predictive SubTree Prefetching (PreSTO), which accelerates PowerMH by prefetching future proposals and provably leaves its MH transition kernel unchanged. PreSTO views future MH outcomes as a tree of accept/reject decisions: in each model call, it generates the proposals of a prefetched subtree of this tree in one batch, then follows the realized path with the original MH acceptance rule. We also observe that MH reaches only a sparse, nearly unary part of each prefetched subtree, we further design a predictive traversal rule that chooses the subtree to maximize the predicted number of completed transitions. Across seven datasets and five base LLMs, PreSTO achieves median speedups of $1.71\times$ , $1.91\times$ , and $2.33\times$ on PowerMH and its follow-ups EntropyCut and MultiTryMH, respectively; on the first two, it makes a median of $3.9\times$ fewer model calls, each only $2.1\times$ longer. On LCB V6 with Qwen3.5-9B, the predictive traversal rule is a further $1.25\times$ -- $1.37\times$ faster than the default breadth-first traversal.

Figures & tables

Figure 1: Across 105 sampler–LLM–dataset settings, PreSTO (ours) speeds up PowerMH, EntropyCut, and MultiTryMH by a median of

1.71\times

1.91\times

, and

2.33\times

, respectively. Each point is one setting, and faster settings lie toward the upper right.

Figure 2: Our subtree prefetching (PreSTO) accelerates PowerMH for transitions

\mathbf{x}\to\mathbf{y}\to\mathbf{y}

. (a) PowerMH performs two sequential rounds of proposal generation with

q_{\tau}

followed by an MH accept/reject decision. (b) Our method collects feasible requests by traversing the prefetching tree, generates proposals

(\mathbf{y},\mathbf{z})

in one batched call, and then samples a path over the subtree.

Figure 3: The

\epsilon

-rounded subtree (orange) is sparse and mostly has one child per node within the full prefetched subtree. Blue edges trace a sampled MH path.

Figure 4: Most edges in prefetched subtrees are

\epsilon

-certain. For PreSTO-PowerMH with Qwen3.5-9B and the default BFS traversal,

67\%

–

95\%

of edges are

\epsilon

-certain across all seven datasets, with

\epsilon=0.01

and

N_{\mathrm{pf}}=20

(a) PowerMH vs. PreSTO-PowerMH
Base LLM	MATH500	AIME 24&25	MBPP	GPQA	HumanEval	MMLU	LCB V6
Qwen3-4B	$2.11\times$	$2.12\times$	$1.10\times$	$1.63\times$	$1.77\times$	$1.59\times$	$2.04\times$
Qwen3-8B	$2.23\times$	$2.29\times$	$1.00\times$	$2.12\times$	$1.81\times$	$1.50\times$	$2.05\times$
Qwen3.5-4B	$1.49\times$	$1.66\times$	$1.67\times$	$1.69\times$	$1.59\times$	$1.76\times$	$2.02\times$
Qwen3.5-9B	$2.06\times$	$2.08\times$	$1.61\times$	$2.05\times$	$1.21\times$	$1.18\times$	$1.44\times$
Gemma4-12B-it	$1.74\times$	$1.76\times$	$1.44\times$	$1.67\times$	$1.71\times$	$1.52\times$	$1.95\times$

Table 1: Speedup of PreSTO over sequential samplers: the mean cumulative model-call time of the sampler divided by that of its PreSTO variant ( 20 ), on the first 20 prompts with

K=100

and

N_{\mathrm{pf}}=20

Figure 5: Under LCB V6 and Qwen3.5-9B, PreSTO (

N_{\mathrm{pf}}=20

, default BFS traversal) speeds up PowerMH by

1.44\times

in mean cumulative model-call time: it makes

4.92\times

fewer model calls (a ratio of means), each taking

3.42\times

as long, while the per-prompt transitions per call average

4.96

(Appendix D.2 ).

	PowerMH	PreSTO-PowerMH, prefetch budget $N_{\mathrm{pf}}$
	PowerMH	8	10	12	16	20
MH transitions per model call	1.00	3.59	3.94	4.27	4.63	4.91
Time per model call (s)	5.39	16.15	16.56	17.58	18.15	18.43
Speedup over PowerMH	1.00 $\times$	1.20 $\times$	1.28 $\times$	1.31 $\times$	1.38 $\times$	1.44 $\times$

Table 2: Under dataset LCB V6 and base LLM Qwen3.5-9B, raising the prefetch budget

N_{\mathrm{pf}}

from

8

20

increases the speedup of PreSTO-PowerMH over PowerMH from

1.20\times

1.44\times

. Transitions per call is

K

divided by the mean number of calls, and time per call is the total call time divided by the total number of calls.

Figure 6: Under base LLM Qwen3.5-4B across five datasets, PreSTO-PowerMH and PowerMH produce closely overlapping distributions of base-model log-likelihood and confidence, and the paired TOST finds them equivalent in 7 of the 10 comparisons (shaded).

Figure 7: Predictive traversal rule against BFS for PreSTO-PowerMH on LCB V6 with Qwen3.5-9B. (a–c) The predictive rule is

1.25\times

–

1.37\times

faster at every prefetch budget because it realizes

19\%

–

33\%

more MH transitions per call; points in (a, b) are ratios of totals over the 20 prompts, and error bars in (c) show

\pm 1

SD over prompts. (d) First prefetched subtree of prompt 0, a favorable case: 3 transitions against 2 at every budget (prompts 1–3 in Appendix E.5 ). Label MH

0\rightarrow k

: the sampled path advances the chain from step 0 to

k

. Solid left edges: acceptance; dashed right edges: rejection; orange:

\epsilon

-rounded subtree (

\epsilon=0.01

); blue: sampled path (endpoint double-bordered); gray: other prefetched branches.

Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 8 — Figure 8: Two batched suffix-resampling requests, their prefetching tree, and their shared KV cache. (a) Both requests reuse the cached prefix $\mathbf{x}_{0:c_{2}}$ , and the request for $\mathbf{y}$ also reuses $\mathbf{x}_{c_{2}+1:c_{1}}$ . (b) Each acceptance (left) edge carries the MH acceptance probability computed from the pair of sequences it connects in (a). (c) Radix-tree representation of the shared KV-cache segments.

Panel of Figure 9 — Figure 9: Visualization of ( 13 ): (a) certain rejection, (b) certain acceptance, and (c) an uncertain rounded decision. A ✗ marks a pruned edge. The proposal at $v^{+}$ is required in every case.

Panel of Figure 10 — Figure 10: Budget decomposition visualization for Proposition 2 .

Category	Dataset	Source (Hugging Face)	Instances
Mathematical reasoning	MATH500	HuggingFaceH4/MATH-500	500
Mathematical reasoning	AIME 24&25	Maxwell-Jia/AIME_2024 MathArena/aime_2025	45
Code generation	HumanEval	openai/openai_humaneval	164
	MBPP	google-research-datasets/mbpp	257
	LCB V6	livecodebench/code_generation_lite	175
Scientific and general knowledge	GPQA	Idavidrein/gpqa	198

Appendix

Table 3: Datasets and their full sizes; every run uses only the first 20 prompts of each dataset.

Family	Base LLM	Source (Hugging Face)	Released
Qwen3	Qwen3-4B	Qwen/Qwen3-4B	Apr 2025
Qwen3	Qwen3-8B	Qwen/Qwen3-8B	Apr 2025
Qwen3.5	Qwen3.5-4B	Qwen/Qwen3.5-4B	Feb 2026
Qwen3.5	Qwen3.5-9B	Qwen/Qwen3.5-9B	Feb 2026
Gemma4	Gemma4-12B-it	google/gemma-4-12B-it	May 2026

Appendix

Table 4: Base LLMs and their sources.

Software	Version / configuration
Python	3.13.0
PyTorch	2.13.0
vLLM	0.27.1 (V1 engine) with custom engine patches
Transformers	5.14.1
(a) Software configuration
Hardware	Specification

Appendix

Table 5: Software and hardware configuration.

Parameter	Setting
Response length	$n=1024$ tokens (fixed; no early stop at EOS)
Block size	a single block, $B=n=1024$
Power exponent	$\alpha=4.0$
Proposal LLM	$q_{\tau}=$ base model sampled at temperature $\tau$
Proposal temperature	$\tau=1/\alpha=0.25$
Realized MH transitions	$K=100$

Appendix

Table 6: Hyperparameter configurations.

Figure 11: PowerMH and PreSTO-PowerMH on MATH500.

Dataset	Base LLM	Total Time (s) $\downarrow$	Transitions per call $\uparrow$	Time/call (s) $\downarrow$	Speedup
MATH500 (Fig. 12 )	Qwen3-4B	$277.41_{\pm 45.48}/157.85_{\pm 20.51}$	$1/4.26_{\pm 0.54}$	$2.77_{\pm 0.45}/6.62_{\pm 0.23}$	$1.76\times$
	Qwen3-8B	$336.64_{\pm 58.87}/192.00_{\pm 24.30}$	$1/4.34_{\pm 0.47}$	$3.37_{\pm 0.59}/8.23_{\pm 0.27}$	$1.75\times$
	Qwen3.5-4B	$365.78_{\pm 80.61}/233.47_{\pm 73.77}$	$1/3.81_{\pm 0.99}$	$3.66_{\pm 0.81}/8.22_{\pm 0.47}$	$1.57\times$
	Qwen3.5-9B	$478.02_{\pm 74.36}/272.29_{\pm 59.03}$	$1/3.84_{\pm 0.65}$	$4.78_{\pm 0.74}/10.11_{\pm 0.52}$	$1.76\times$
	Gemma4-12B-it	$980.09_{\pm 240.16}/431.85_{\pm 175.94}$	$1/4.07_{\pm 2.00}$	$9.80_{\pm 2.40}/15.06_{\pm 1.37}$	$2.27\times$
AIME 24&25	Qwen3-4B	$255.04_{\pm 37.63}/150.54_{\pm 16.85}$	$1/4.57_{\pm 0.49}$	$2.55_{\pm 0.38}/6.81_{\pm 0.18}$	$1.69\times$

Appendix

Table 8: EntropyCut execution metrics across base LLMs.

Figure 12: EntropyCut and PreSTO-EntropyCut on MATH500.

Dataset	Base LLM	Total Time (s) $\downarrow$	Transitions per call $\uparrow$	Time/call (s) $\downarrow$	Speedup
MATH500	Qwen3-4B	$741.55_{\pm 35.32}/322.90_{\pm 17.49}$	$1/2.02_{\pm 0.11}$	$7.42_{\pm 0.35}/6.49_{\pm 0.17}$	$2.30\times$
	Qwen3-8B	$865.84_{\pm 42.90}/417.08_{\pm 21.21}$	$1/2.04_{\pm 0.10}$	$8.66_{\pm 0.43}/8.50_{\pm 0.18}$	$2.08\times$
	Qwen3.5-4B	$2174.72_{\pm 142.35}/375.23_{\pm 14.43}$	$1/2.05_{\pm 0.08}$	$21.75_{\pm 1.42}/7.68_{\pm 0.16}$	$5.80\times$
	Qwen3.5-9B	$3246.66_{\pm 262.85}/463.26_{\pm 25.28}$	$1/2.05_{\pm 0.11}$	$32.47_{\pm 2.63}/9.48_{\pm 0.15}$	$7.01\times$
	Gemma4-12B-it	$1558.70_{\pm 79.33}/686.49_{\pm 45.49}$	$1/1.86_{\pm 0.13}$	$15.59_{\pm 0.79}/12.75_{\pm 0.20}$	$2.27\times$
AIME 24&25	Qwen3-4B	$633.31_{\pm 31.36}/313.16_{\pm 17.48}$	$1/2.03_{\pm 0.11}$	$6.33_{\pm 0.31}/6.33_{\pm 0.18}$	$2.02\times$

Appendix

Table 9: MultiTryMH execution metrics across base LLMs.

Figure 13: Per-transition traces across prefetch budgets on LCB V6 with Qwen3.5-9B. Solid curves show the PreSTO mean for each budget, from light (small

N_{\mathrm{pf}}

) to dark (large

N_{\mathrm{pf}}

). Dashed curves show PowerMH, and shaded bands show

\pm 1

SD over traces.

Figure 14: Likelihood and confidence distributions with and without PreSTO. Boxes show the per-prompt token-normalized log-likelihood (a, c) and confidence (b, d) of terminal traces under

p_{0}

for the eight PowerMH pairs (top) and the eleven EntropyCut pairs (bottom), each against its PreSTO variant. Shaded rows are equivalent under the paired TOST (

p_{\mathrm{eq}}<0.05

Figure 15: Edge certainty under PreSTO-PowerMH across models and datasets.

Figure 16: Edge certainty under PreSTO-EntropyCut across models and datasets.

	Model calls		Transitions per call		Time per call (s)		Model-call time (s)
$N_{\mathrm{pf}}$	BFS	Predictive	BFS	Predictive	BFS	Predictive	BFS	Predictive	Speedup
8	27.85	23.40	3.59	4.27	16.15	15.35	$449.7_{\pm 19.1}$	$\mathbf{359.2}_{\pm 28.0}$	$1.25\times$
10	25.40	20.15	3.94	4.96	16.56	16.00	$420.7_{\pm 25.5}$	$\mathbf{322.3}_{\pm 28.5}$	$1.31\times$
12	23.40	18.90	4.27	5.29	17.58	17.13	$411.4_{\pm 38.7}$	$\mathbf{323.8}_{\pm 39.5}$	$1.27\times$
16	21.60	17.10	4.63	5.85	18.15	17.34	$392.1_{\pm 33.8}$	$\mathbf{296.6}_{\pm 45.8}$	$1.32\times$
20	20.35	15.30	4.91	6.54	18.43	17.95	$375.1_{\pm 35.1}$	$\mathbf{274.7}_{\pm 40.1}$	$1.37\times$

Appendix

Table 10: Predictive traversal rule against the default BFS traversal. Speedup is the total model-call time of BFS divided by that of the predictive rule.

Figure 17: First prefetched subtree under each traversal rule and prefetch budget. Prompts 1–3 of LCB V6 with Qwen3.5-9B under PreSTO-PowerMH, one row per prompt and rule; labels and colors follow Figure 7 d.

Engine	Sampler	Calls	Time per call (s)	Model-call time (s)	Speedup
vLLM	PowerMH	$100$	$5.39_{\pm 0.31}$	$539.24_{\pm 30.92}$	$1.40\times$
vLLM	PreSTO-PowerMH	$20.60_{\pm 1.31}$	$18.75_{\pm 0.43}$	$385.97_{\pm 21.25}$	$1.40\times$
SGLang	PowerMH	$100$	$3.57_{\pm 0.24}$	$357.40_{\pm 23.67}$	$1.97\times$
SGLang	PreSTO-PowerMH	$21.25_{\pm 1.86}$	$8.56_{\pm 0.10}$	$181.74_{\pm 15.25}$	$1.97\times$

Appendix

Table 11: PowerMH and PreSTO-PowerMH on vLLM and SGLang.

Figure 18: Prediction of the

\epsilon

-certainty class by feature set. Held-out accuracy against log loss of the logistic regression for PreSTO-PowerMH (a) and PreSTO-EntropyCut (b); upper left is better. Labels give the feature sets used: a single digit is one set, “1+3” combines Features 1 and 3, and “1–3” combines Features 1 through 3; comma-separated labels mark overlapping points. The cross marks the constant baseline, and the bold label is the configuration we use.

Figure 19: Effect of hidden layers on predicting the

\epsilon

-certainty class. Held-out accuracy against log loss with the cut fraction as the only feature; labels give the units per hidden layer, and “none” is the logistic regression. Markers as in Figure 18 .

Figure 20: Speed and memory against test-time baselines. Peak KV-pool use against time per prompt on the first 20 prompts of LCB V6 with Qwen3.5-9B; lower left is better.

Explore similar work

Sep 29, 2026

Explore Broadly, Reason Sharply: Push Small Models toward the Frontier via Sampling

Power-sharpened sampling is an inference-time alternative to reinforcement-learning (RL) post-training for enhancing reasoning in large language models (LLMs). High-probability sequences are amplified under the base model without parameter updates or external rewards, avoiding the costly optimization and jagged generalization of RL. However, this approach faces a fundamental exploration--exploitation trade-off, as % strong sharpening restricts exploration, trapping samplers in plausible but incorrect reasoning trajectories, whereas weak sharpening leaves the answer distribution diffuse. To resolve this trade-off, we introduce Parallel Power Tempering (PPT), instantiating power-sharpened LLM sampling via parallel tempering. Running multiple interacting replicas in parallel at different sharpening levels allows lower-power replicas to explore diverse reasoning trajectories and higher-power chains to further exploit higher-likelihood responses favored by the sharpened target. Specifically, we tailor \method{} to inference-time sampling by mitigating a truncation bias, identified in prior power samplers, and investigate effective swap strategies under finite memory and compute budgets. Extensive experimentation shows that \method{} substantially improves single-chain power-sharpened sampling and outperforms RL-post-trained models, producing higher-quality reasoning traces and even achieving performance comparable to frontier models.

Panagiotis Theodoropoulos, Nan Jiang, Xintong Duan +4

Georgia Institute of Technology · Morgan Stanley · University of Texas at El Paso +1

Efficient Language Model Reasoning Parallel Tempering Inference-Time Scaling

Ask PDF

Jan 29, 2026

Scalable Power Sampling: Unlocking Efficient, Training-Free Reasoning for LLMs via Distribution Sharpening

Reinforcement learning (RL) post-training is a dominant approach for improving the reasoning performance of large language models (LLMs), yet growing evidence suggests that its gains arise primarily from distribution sharpening rather than the acquisition of new capabilities. Recent work has shown that sampling from the power distribution of LLMs using Markov chain Monte Carlo (MCMC) can recover performance comparable to RL post-training without relying on external rewards; however, the high computational cost of MCMC makes such approaches impractical for widespread adoption. In this work, we propose a theoretically grounded alternative that eliminates the need for iterative MCMC. We derive a token-level decomposition of the autoregressive power distribution, showing that the global power distribution can be approximated by a token-level scaled low-temperature one, where the scaling factor captures future trajectory quality. Leveraging this insight, we introduce a training-free and verifier-free algorithm that sharpens the base model's generative distribution autoregressively. Empirically, we evaluate our method on math, QA, and code tasks, and show that our method matches or surpasses one-shot GRPO without relying on any external rewards, while reducing inference latency by over

10\times

compared to MCMC-based sampling.

Xiaotong Ji, Rasul Tutunov, Matthieu Zimmer +1

Huawei Noah’s Ark Lab · UCL Centre for Artificial Intelligence

Efficient Language Model Inference RL for Language Model Reasoning

Ask PDF

Jun 7, 2026

Sample Where You Struggle: Sharpening Base Model Reasoning via Entropy-Guided Power Sampling

Sampling from the sequence-level power distribution

p^α

elicits RL-level reasoning from base language models without any parameter updates, but the standard Metropolis--Hastings (MH), a Markov Chain Monte Carlo (MCMC) sampler, is both expensive and slow-mixing. We trace both to a structural mismatch:

p^α

mainly departs from

p

at a sparse, spatially clustered set of high-entropy decision points, yet MH proposes resampling positions uniformly along the prefix -- wasting compute on near-degenerate conditionals while under-mixing precisely where modes diverge. We propose Entropy-Guided Power Sampling (EGPS), a training-free and verifier-free sampler that re-derives its proposal from token-level entropy already in the forward pass. EGPS skips deterministic blocks, localizes each MCMC move to a high-entropy neighborhood, and applies Multiple-Try Metropolis at decision points -- making sampling cost scale with entropy mass rather than sequence length. On Qwen2.5-Math-7B, EGPS reaches best or tied-best accuracy on all three benchmarks (MATH500

75.8\%

, HumanEval

62.2\%

, GPQA

42.4\%

) at up to a

12.6\times

wall-clock speedup over the MH baseline.

Hong Guo, Nianhui Guo, Christoph Meinel +1

Hasso Plattner Institut · German University of Digital Science

Markov Chain Monte Carlo Inference-Time Scaling Efficient Language Model Reasoning LLM Inference Acceleration Adaptive Sampling

Ask PDF