Probe and Skip: Self-Predictive Token Skipping for Efficient Long-Context LLM Inference

Zimeng Wu; Donghao Wang; Chaozhe Jin; Jiaxin Chen; Yunhong Wang

cs.CLJan 19, 2026

Probe and Skip: Self-Predictive Token Skipping for Efficient Long-Context LLM Inference

Authors: Zimeng Wu, Donghao Wang, Chaozhe Jin, Jiaxin Chen, Yunhong Wang

Organizations: School of Computer Science and Engineering, Beihang University, Beijing, China · State Key Laboratory of Virtual Reality Technology and Systems, Beihang University, Beijing, China

Abstract

Long-context inference enhances the reasoning capability of Large Language Models (LLMs), but incurs significant computational overhead. Token-oriented methods, such as pruning and skipping, have shown great promise in reducing inference latency, yet still suffer from inherently insufficient structure optimization, outdated selection criteria, and redundancy interference, resulting in suboptimal speed-accuracy trade-off. To address these issues, we propose a novel training-free framework dubbed Self-Predictive Token Skipping (SPTS), for efficient long-context LLM inference. Specifically, motivated by probing the influence of target layers prior to skipping, we design two selective token skipping strategies for typical structures, including Partial Attention Probing (PAP) for multi-head attention and Low-rank Transformation Probing (LTP) for feed forward network. The former selects informative tokens via partial forward attention computation, while the latter constructs a low-rank proxy network to predict token transformations. In addition, a Multi-Stage Delayed Pruning (MSDP) strategy reallocates skipping budgets and progressively removes redundant tokens across layers. Extensive experiments display the effectiveness of our method, achieving up to 2.46 $\times$ and 2.29 $\times$ speedups for prefilling and end-to-end generation, respectively, while maintaining state-of-the-art accuracy. The source code is available at https://github.com/Wuzimeng/SPTS.

Figures & tables

Figure 1. Efficiency-accuracy comparison on LLaMA-3.1-8B-Instruct ( Dubey et al., 2024 ) . (a) SPTS accelerates both prefilling and decoding, and achieves consistent speedups in first token (TTFT) and end-to-end (E2E) generation. (b) SPTS maintains higher accuracy at a more aggressive TTFT speedup ratio, with metrics averaged over the LongBench ( Bai et al., 2024 ) dataset.

Figure 2. Framework overview of SPTS. (Top) The overall data flow with SPTS. SPTS accelerates the LLM inference over both long textual and visual inputs by token skipping, reducing computation cost while preserving contextual information. (Bottom) Illustration of the main modules of SPTS. (a) Partial Attention Probing (PAP) performs token skipping in MHA by pre-computing a small fraction of attention to identify informative tokens, which then participate in reduced attention computation. (b) Low-rank Transformation Probing (LTP) enables token skipping in FFN by leveraging an offline-derived low-rank proxy network to probe token transformations, which are combined with attention-based information to guide token selection for reduced feed-forward. (c) Multi-Stage Delayed Pruning (MSDP) partitions the network into multiple stages, where the set of candidate tokens remains fixed within each stage and is reduced through pruning at stage boundaries.

Figure 3. Empirical observations on LLaMA-3.1-8B-Instruct. (a) Layer-wise average cosine similarity of tokens before and after MHA, FFN, and the full Transformer block (MHA+FFN). (b) Accuracy fluctuations induced by disabling MHA updates for all input tokens, or for all of them except the last, starting from a given layer. (c) Number of tokens contributing to 90% and 95% of cumulative attention across deep layers for input sequences of varying lengths. (d) Jaccard index between consecutive-layer token sets selected via the attention criterion under a progressively reduced token budget, across distinct input lengths.

Method	Single-Doc. QA		Multi-Doc. QA			Summarization				Few-shot Learning				Synthetic Task		Code Completion		Avg. (%)	TTFT ratio
Method	Qasper	MQA	HPQA	2WiKi	MuSiQue	GovRep	QMSum	MNews	VCSum	TREC	TQA	SAMSum	LSHT	Count	PassR	LCC	RepB-p
LLaMA-3.1-8B-Instruct
Full Model	45.16	54.71	55.64	44.73	30.66	34.99	25.26	27.17	17.16	72.50	91.66	43.98	46.50	7.00	99.50	63.02	56.01			47.98	1.00 $\times$
LazyLLM	45.56	53.17	53.47	45.50	30.44	32.76	24.84	26.19	16.99	71.00	91.25	43.28	46.00	7.50	99.50	58.66	55.17	47.13	1.63 $\times$
MInference	44.59	52.64	53.51	47.30	28.08	35.25	25.51	27.29	17.43	72.50	91.42	43.58	46.50	6.55	95.50	63.30	57.41	47.55	0.47 $\times$
FlexPrefill	45.18	55.52	56.95	42.75	32.14	34.55	25.16	26.86	17.39	71.00	90.74	43.50	42.00	3.16	81.50	62.93	60.89	46.60	0.96 $\times$

Table 1. Comparison of performance scores and TTFT acceleration ratios (TTFT ratio) by various methods for three representative LLMs on the LongBench dataset. The best results are highlighted in bold and the second-best ones are underlined.

Figure 4. Comparison of (a) TTFT acceleration ratios and (b) E2E acceleration ratios with different methods across distinct context lengths on LLaMA3.1-8B-Instruct.

Method	MME	SQA-I	MMB	GQA	TextVQA	Avg.
Upper Bound: 576 Tokens ( 100% )
Full Model	1860.8	69.5	64.6	61.9	58.3	64.1
Retain 192 Tokens ( $\downarrow$ 66.7% )
FastV	1612.0	67.3	61.2	52.7	52.5	58.3
LLaVA-PruMerge	1632.0	67.9	59.6	54.3	54.3	58.9
SparseVLM	1721.0	69.1	62.5	57.6	56.1	61.4

Table 2. Comparison of performance by various methods for LLaVA-1.5-7B on 5 benchmarks. The best results are highlighted in bold and the second-best ones are underlined.

Method	MME	SQA-I	MMB	GQA	TextVQA	Avg.
Upper Bound: 2880 Tokens ( 100% )
Full Model	1846.6	70.2	67.9	64.3	61.2	65.9
Retain 320 Tokens ( $\downarrow$ 88.9% )
FastV	1661.0	62.8	61.6	55.9	55.7	59.1
LLaVA-PruMerge	1534.0	66.4	61.3	53.6	50.6	57.3
SparseVLM	1533.0	66.1	60.6	56.1	58.4	59.2

Table 3. Comparison of performance by various methods for LLaVA-Next-7B on 5 benchmarks. The best results are highlighted in bold and the second-best ones are underlined.

Method	$8$ K	$16$ K	$24$ K	$32$ K
Full Model	1.00 $\times$	1.00 $\times$	1.00 $\times$	1.00 $\times$
+ PAP	1.08 $\times$	1.24 $\times$	1.35 $\times$	1.44 $\times$
+ PAP + LTP	1.26 $\times$	1.70 $\times$	1.98 $\times$	2.15 $\times$
+ PAP + LTP + MSDP	1.35 $\times$	1.91 $\times$	2.26 $\times$	2.46 $\times$

Table 4. Ablation of the PAP, LTP and MSDP modules on TTFT acceleration ratios across context lengths with LLaMA.

Method	$8$ K	$16$ K	$24$ K	$32$ K
Ours w/o PAP (Full Cache)	4.00	8.00	12.00	16.00
Ours with PAP	2.81	4.13	5.38	6.63
Memory Saving (%)	29.7	48.4	55.2	58.6

Table 5. Effect of the PAP module on KV Cache memory (GB) savings across varying context lengths with LLaMA.

Skipping Criterion	Dataset					Avg. (%)
	MQA	HPQA	SAMSum	Count	RepB-p	Avg. (%)
$S_{n}^{\mathrm{MHA}}$	52.96	53.39	44.23	5.59	52.79	46.31
$C_{n}^{\mathrm{FFN}}\cdot S_{n}^{\mathrm{MHA}}$	53.71	54.11	44.64	7.47	53.26	46.85

Table 6. Effect of low-rank proxy network on LLaMA. Results are reported on five representative LongBench subsets and the average score (Avg.) on the full dataset.

Method	Dataset					Avg. (%)
	MQA	HPQA	SAMSum	Count	RepB-p	Avg. (%)
Ours w/o DP	54.99	52.51	43.72	4.32	52.39	46.11
Ours with DP	55.37	53.66	44.23	7.24	55.20	47.80

Table 7. Effect of Delayed Pruning (DP) on LLaMA across five LongBench subsets and average (Avg.) on the full dataset.

Figure A. Illustration of different schemes in KV Cache management. (a) Decoding-phase pruning computes all tokens, offering no speedup in prefilling. (b) Prefilling-phase pruning discard tokens to reduce computation, but cached tokens exist only in the early consecutive layers, leading to noticeable information loss. (c) SPTS combines skipping with delayed pruning, allowing skipped tokens to be recomputed and cached later. It achieves the same computation costs as (b) while preserving richer contextual information.

$D_{\mathrm{low}}$	$r$	FLOPs	Qasper	HPQA	TQA	LSHT	LCC
512	128	590K	43.67	53.60	91.17	45.00	57.49
512	256	1180K	44.00	54.36	92.04	45.50	57.76
512	/	2097K	43.85	55.18	92.03	46.00	58.36
512	192	855K	43.85	54.11	91.59	45.00	57.65
256	192	836K	43.17	53.38	91.33	45.50	57.55
1536	192	1081K	43.87	53.93	91.64	46.00	57.79

Table A. Effect of

D_{\mathrm{low}}

and

r

on per-token per-projection FLOPs and accuracy for LLaMA across five representative LongBench subsets. ‘/’ denotes no decomposition on the corresponding dimension. The default setting is underlined.

Model	Total Layers	Skipping Layers		Stage Ending Layers	Active Token Budgets	Pruned Tokens
LLaMA-3.1-8B-Instruct	32		10 $\sim$ 28	(13, 18, 23, 28)	(9K, 7K, 4K, 2K)	1K
Qwen-2.5-7B-Instruct	28		9 $\sim$ 24	(12, 16, 24, 20)	(13K, 10K, 7K, 4K)	2K
openPangu-Embedded-1B-V1.1	26		11 $\sim$ 22	(13, 16, 19, 22)	(13K, 10K, 7K, 4K)	2K
LLaVA-1.5-7B (192 Tokens)	32		2 $\sim$ 24	(6, 12, 18, 24)	(316, 243, 169, 96)	32
LLaVA-1.5-7B (128 Tokens)	32		2 $\sim$ 24	(6, 12, 18, 24)	(184, 144, 104, 64)	32
LLaVA-1.5-7B (64 Tokens)	32		2 $\sim$ 16	(6, 11, 16)	(88, 56, 24)	16

Table B. Token skipping and pruning configurations for various models. ‘K’ denotes 1024 tokens.

Query Length	MQA	HPQA	2WiKi	MuSiQue	QMSum	TREC	LSHT	Count	LCC	Avg. (%)
64	54.33	52.67	43.95	31.31	24.29	63.00	45.00	7.38	63.32	47.00
4	54.53	54.83	45.39	30.58	25.23	70.50	46.50	4.77	62.85	47.53
1	55.37	53.66	46.69	30.43	24.73	72.50	46.50	7.24	63.14	47.80

Table C. Effect of query token length for PAP on LLaMA. Results are reported on representative LongBench subsets and the average score (Avg.) on the full dataset. The default setting is underlined.

Method	MQA	HPQA	2WiKi	MuSiQue	QMSum	TREC	LSHT	Count	LCC	Avg. (%)
MaxNorm-Add	55.27	53.82	45.46	29.13	24.60	72.00	45.50	6.84	62.45	47.48
Std-Add	55.00	53.63	44.72	28.99	24.63	71.50	45.50	6.88	62.57	47.40
Sequential Select	55.15	53.37	45.52	28.37	24.34	72.50	45.50	5.47	62.88	47.43
Multiply	55.37	53.66	46.69	30.43	24.73	72.50	46.50	7.24	63.14	47.80

Table D. Effect of various fusion strategies for the conditional transformation score on LLaMA. Results are reported on representative LongBench subsets and the overall average (Avg.) across the full dataset. The default setting is underlined.

Method	MQA	HPQA	2WiKi	MuSiQue	QMSum	TREC	LSHT	Count	LCC	Avg. (%)
$\rho=0.1$	55.03	54.12	46.50	28.75	24.80	72.50	46.00	6.32	62.97	47.67
$\rho=0.2$	55.37	53.66	46.69	30.43	24.73	72.50	46.50	7.24	63.14	47.80
$\rho=0.3$	55.27	54.41	45.51	29.95	24.77	72.50	46.00	8.24	63.15	47.81
w/o top- $\rho$	54.64	52.68	45.71	30.24	24.65	72.00	46.00	6.21	63.19	47.56

Table E. Effect of selecting top-

\rho

saliency values for channel importance on LLaMA. Results are reported on representative LongBench subsets and the average score (Avg.) on the full dataset. The default setting is underlined.

Starting Layer	Avg. (%)	TTFT Ratio
8	47.43	1.78 $\times$
10	47.80	1.68 $\times$
12	48.01	1.60 $\times$

Table F. Effects of starting layer (Left), stage number (Middle) and token budget (Right) on LLaMA, measured by LongBench average score (Avg.) and TTFT speedup ratio. The default setting is underlined.

Method	MQA	HPQA	2WiKi	MuSiQue	QMSum	TREC	LSHT	Count	LCC	Avg. (%)
SlimInfer	53.72	54.34	43.68	30.69	24.60	71.00	45.50	6.40	63.38	47.56
SPTS with C4	55.13	53.19	45.46	30.96	24.59	72.50	46.00	7.84	62.78	47.77
SPTS with Qasper	55.37	53.66	46.69	30.43	24.73	72.50	46.50	7.24	63.14	47.80

Table G. Effect of the calibration data source for LTP on LLaMA. Results are reported on representative LongBench subsets and the overall average (Avg.) across the full dataset, compared with the second-best SlimInfer ( Long et al., 2025 ) . The default setting is underlined.

Size	1	50	100	150	200	250	300
Avg.(%)	47.52	47.68	47.73	47.70	47.80	47.82	47.80

Table H. Effect of the calibration set size on LLaMA, measured by LongBench Avg. score. The default setting is underlined.

LLM Inference Acceleration

LLM inference acceleration develops algorithms and systems that reduce language-model latency, memory use, and serving cost while preserving output quality.

See all 1,048 papers

Training-Inference Consistent Segmented Execution for Long-Context LLMs
May 12, 2026Xianpeng Shang, Jiang Li, Zehua Duo +2Self-Attention Long-Context Language Modeling
College of Computer Science, Inner Mongolia University, Hohhot 010021, China · National & Local Joint Engineering Research Center of Intelligent Information Processing Technology for Mongolian, Hohhot 010021, China · Inner Mongolia Key Laboratory of Multilingual Artificial Intelligence Technology, Hohhot 010021, China +1
From Rigid to Dynamic: Entropy-Guided Adaptive Inference for Long-Context LLMs
Jun 8, 2026Zhanchao Xu, Haoyang Li, Qingfa Xiao +4Long-Context Language Model Inference LLM Inference Acceleration
Department of Computing, PolyU · DSA, HKUST(GZ) · CSE, HKUST
Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps
May 16, 2026Yanke Zhou, Yiduo Li, Hanlin Tang +6Self-Attention Long-Context Language Model Inference
Nanjing University · Alibaba Group

May 12, 2026J/K move · Enter open · S save

Training-Inference Consistent Segmented Execution for Long-Context LLMs

Xianpeng Shang, Jiang Li, Zehua Duo, Qianyi Cai+1

College of Computer Science, Inner Mongolia University, Hohhot 010021, China · National & Local Joint Engineering Research Center of Intelligent Information Processing Technology for Mongolian, Hohhot 010021, China · Inner Mongolia Key Laboratory of Multilingual Artificial Intelligence Technology, Hohhot 010021, China+1

Self-Attention Long-Context Language Modeling Long-Context Inference Efficient Language Model Inference

Open paper AskPDF arXiv

Transformer-based large language models face severe scalability challenges in long-context generation due to the computational and memory costs of full-context attention. Under practical computation and memory constraints, many inference-efficient long-context methods improve efficiency by adopting bounded-context or segment-level execution only during inference, while continuing to train models under full-context attention, resulting in a mismatch between training and inference execution and state-transition semantics. Based on this insight, we propose a training-inference consistent segment-level generation framework, in which training and inference follow the same segment-level forward execution semantics. During training, consistency with inference is enforced by restricting gradient propagation to KV states carried over from the immediately preceding segment, while permitting head-specific access to past KV states during the forward pass without involving them in gradient propagation. Across long-context benchmarks, our approach achieves performance comparable to full-context attention, while achieving competitive latency-memory trade-offs against strong inference-efficient baselines, and substantially improving scalability at very long context lengths (e.g., approximately 6x lower peak prefill memory at 128K compared to full-context attention with FlashAttention).

Abstract

LLM Inference Acceleration

Explore similar work

Training-Inference Consistent Segmented Execution for Long-Context LLMs

From Rigid to Dynamic: Entropy-Guided Adaptive Inference for Long-Context LLMs

Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps