通过共享前缀压缩加速因果Transformer推理，提升1.5倍以上效率。
RadixMLP -- Intra-batch Deduplication for Causal Transformers
- 基于前缀字典树动态合并重复序列，仅计算一次共享部分。
- 在真实重排序任务中提速1.44-1.59倍，长前缀下最高达5倍。
- 无需修改模型，单次前向传播完成，适合大规模服务部署。
因果Transformer模型在批量推理时，常处理具有相同前缀的序列（如系统提示、少量示例或共享查询）。标准推理引擎对每个序列独立处理，重复计算相同的MLP激活值。我们提出RadixMLP，利用MLP、LayerNorm、线性投影和嵌入的位置无关特性，将批次动态映射到前缀字典树，将共享片段压缩为紧凑表示进行逐位置计算，并仅在注意力边界处还原结果。该方法无状态，仅需一次前向传播。在MS~MARCO v1.1数据集上使用Qwen3模型（0.6B至8B参数）的端到端服务测试中，真实重排序负载下获得1.44-1.59倍加速，合成基准中长共享前缀场景下最高达5倍加速。代码已开源：https://github.com/michaelfeil/radix-mlp。
原文摘要 · Abstract (English)
Batch inference workloads for causal transformer models frequently process sequences that share common prefixes, such as system prompts, few-shot examples, or shared queries. Standard inference engines treat each sequence independently, redundantly recomputing identical MLP activations for every copy of the shared prefix. We introduce RadixMLP, a technique that exploits the position-wise nature of MLPs, LayerNorms, linear projections, and embeddings to eliminate this redundancy. RadixMLP dynamically maps batches to a prefix trie, gathering shared segments into a compressed representation for position-wise computation and scattering results back only at attention boundaries. RadixMLP is stateless and operates within a single forward pass. In end-to-end serving benchmarks on MS~MARCO v1.1 with Qwen3 models (0.6B to 8B parameters), RadixMLP achieves 1.44-1.59$\times$ speedups in realistic reranking workloads, with up to $5\times$ speedups on synthetic benchmarks with longer shared prefixes. Our code is available at https://github.com/michaelfeil/radix-mlp.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。