🤖 AI Summary
This study addresses the limitation of existing LLM length prediction methods that rely solely on entropy signals while neglecting semantic distinctions. We propose ESTP, a novel framework that innovatively integrates self-attention weights as semantic importance indicators with entropy metrics. By employing joint token pooling and reusing prefill activations, ESTP achieves length-aware scheduling optimization with minimal overhead. Experiments on the ForeLen benchmark demonstrate significantly improved prediction accuracy. Furthermore, end-to-end evaluations confirm that ESTP effectively reduces padding ratios and enhances system throughput. These results validate its practical utility in overcoming the semantic deficiencies of pure entropy-based approaches and improving inference efficiency for large language models.
📝 Abstract
Efficient LLM serving is often bottlenecked by the need to pad sequences to a fixed maximum length, and this wastes compute and degrades throughput. Predicting output lengths in advance makes it possible to adopt length-aware scheduling, and this reduces the overhead. This advantage is especially pronounced in long-context reasoning and reinforcement learning applications. Existing approaches, such as entropy-guided token pooling, use token-wise entropy as their primary signal, but they tend to ignore differences in semantic content across tokens. So, important tokens are often underweighted, and tokens carrying little information receive disproportionate emphasis. This hurts the reliability of length prediction. We introduce ESTP (Entropy-and-Semantic Token Pooling), a lightweight framework that addresses this issue by combining entropy with attention-based importance scores. These scores are derived directly from the self-attention weights computed during the LLM prefill phase, and this allows ESTP to capture both uncertainty and semantic importance with minimal additional computation. Since the framework reuses prefill activations, it adds almost no extra memory overhead and introduces only minimal latency. On the ForeLen benchmark, ESTP outperforms baseline methods, achieves better prediction accuracy and lower error rates in most scenarios. When integrated with a length-aware scheduler in end-to-end system tests, it further helps improve overall throughput and reduce the padding ratio. Our results offer a practical and effective building block for length-aware LLM serving systems.