RadixMLP -- Intra-batch Deduplication for Causal Transformers

TL;DR

RadixMLP addresses substantial redundant compute in batch inference for causal transformers by exploiting the per-token nature of MLPs, LayerNorms, and projections. It builds a dynamic prefix trie to compactly represent tokens with identical causal histories, enabling most computations to run in a reduced space and scattering results only at attention boundaries. The method is stateless, supports autograd, and integrates with ragged batching to preserve causal correctness while achieving significant speedups (1.4x–1.59x in real workloads and up to 5x on synthetic long-prefix cases). It demonstrates practical gains in MS MARCO reranking tasks with Qwen3 models and provides open-source kernels and integration guidance, positioning RadixMLP as a complementary, cache-free alternative to KV-based optimizations for batch inference. Limitations include reduced benefits for low-redundancy batches and longer-context workloads, with training and broader deployment considerations highlighted for future work.

Abstract

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.

Figures (7)

• Figure 1: RadixMLP integration into causal transformers. The hidden state remains in compact space ($N'$ tokens) between layers. Position-wise operations (pre-attention norm, projections, MLP) run in compact space. Only the attention operation requires the original space ($N$ tokens), with scatter before and gather after attention (see Algorithm \ref{['alg:radix_mlp_causal_lm']}).
• Figure 2: RadixMLP speedup over prefix lengths with fixed suffix length. Larger models benefit more as position-wise operations dominate the FLOP budget. Qwen3-8B achieves 5.0$\times$ speedup at 2048-token prefixes. Figure \ref{['fig:radixmlp_suffix_comparison']} shows additional configurations.
• Figure 3: RadixMLP speedup with Suffix=1024 tokens. With longer suffixes, the compression ratio decreases, resulting in reduced speedups. Qwen3-8B achieves 2.50$\times$ speedup at 2048-token prefixes. \ref{['tab:qwen3-4b-radixmlp-forward']} provides the raw numbers for Qwen3-4B.
• Figure 4: MS MARCO latency distribution: When running separate 1287 requests against TEI, response latency can vary based on the total number of tokens in a request. To show the Pareto improvement of RadixMLP, the three plots show two different settings (RadixMLP enabled/disabled and max batch tokens 16384 or 65536) and their impact on the perceived throughput and latency of a deployment.
• Figure 5: Comparison of TEI against vLLM against the dataset b (validation, regular order): With a RadixMLP compression ratio of around 0.61 and a reported KV Cache reuse ratio of 0.31, vLLM is noticeably faster than vanilla TEI across all configurations. Using RadixMLP shows on-par or improved performance over vLLM.