推論系統
WIDE Dynamically Prunes Attention and FFNs for Every Token, Accelerating End-to-End LLM Decoding by 1.55×
WIDE brings dynamic pruning down to token-level attention-head groups and FFN channel groups, supporting both prefill and decode. At 50% sparsity, the study reports up to a 4.95× decoding speedup for core layers, but the end-to-end gain falls to 1.55×.
Structured pruning for large language models typically removes entire layers, attention heads, or FFN channels according to a fixed configuration. This is easy to deploy, but it cannot adjust computation based on token difficulty. The newly released WIDE instead adds a lightweight router to each layer, allowing every token to independently select which query-head groups and contiguous FFN channel groups to execute. The K and V projections remain dense to prevent dynamic decisions from disrupting the KV cache.
Training proceeds in two stages. First, the backbone is frozen while binary routing is trained using Gumbel-Softmax. LoRA can then be applied selectively to recover quality lost through pruning. At runtime, WIDE rearranges masks and combines hardware-agnostic block skipping with hardware-tuned intra-block skipping, ensuring that fine-grained sparsity translates into real execution gains rather than merely lower FLOP counts. The paper reports that, at 50% sparsity, WIDE delivers 55.1% higher performance than existing dynamic depth-pruning methods. Core-layer prefill and decode are accelerated by up to 1.98× and 4.95×, respectively, while end-to-end speedups reach 1.68× and 1.55×.
The code has been released and includes the training pipeline for Llama 3.1 8B, FSDP2 support, evaluation configurations, and a standalone dynamic-pruning kernel library. Under the reference configuration using four A100 40GB GPUs, router calibration takes about nine hours, with LoRA recovery requiring another 13 hours. This indicates that WIDE is not a zero-cost, plug-and-play optimization.
Engineering teams should next examine end-to-end performance across different batch sizes, context lengths, and GPU architectures. The paper’s largest speedup figures primarily reflect core-layer results. Production serving remains constrained by routing overhead, memory transfers, scheduling, and operators that are not pruned. It has also not yet been demonstrated that WIDE can be applied directly to quantized models, MoE architectures, or mainstream inference servers.