Back Home

AI 推論系統

WIDE Dynamically Prunes Model Width for Each Token, Accelerating End-to-End Decoding by 1.55×

WIDE lets each token independently select attention heads and FFN channel groups, extending dynamic pruning from layer-level decisions to neuron blocks. The paper reports end-to-end speedups of 1.68× for prefill and 1.55× for decoding at 50% sparsity, but the method still relies on custom kernels and additional training.

Expedition 35 Crew Image courtesy of the NASA Johnson Space Center, Image Science & Analysis Laboratory Derivative work including grading, noise reduction, lens distortion and vignetting correction and dust spot removal: Julian Herzog · Public domain · Image source
zh-Hant

The “width” of a large language model is typically fixed before deployment: every token passes through the same number of attention heads and feed-forward network channels. WIDE, released on July 30, instead allocates compute on a per-token basis, allowing different tokens to dynamically select attention-head groups and FFN-channel groups. Easier tokens can skip more blocks, while harder ones retain more of the model’s full width. This is more granular than dynamic depth methods that only skip entire layers, and it covers both prefill and token-by-token decoding.

Training takes place in two stages. The system first learns sparse selections for each token, then adapts the model to the dynamic execution paths. The real challenge lies in hardware: scattered masks do not necessarily save GPU time. The team therefore introduced mask reordering, which rearranges selected heads and channels into more contiguous blocks, then combines hardware-agnostic block skipping with device-specific intra-block skipping. This enables the method to reduce not only theoretical FLOPs but also kernel execution time and end-to-end latency.

The paper states that, at 50% sparsity and using only calibration data, WIDE preserves more quality than existing dynamic depth pruning methods, improving the performance metric by 55.1%. Its custom kernels deliver speedups of up to 1.98× for prefill and 4.95× for decoding; after accounting for the rest of the inference overhead, the end-to-end figures fall to 1.68× and 1.55×, respectively. This gap is precisely what matters in deployment: pruning half the computation does not mean service throughput will double.

The code has been added to EIT-NLP’s LLM-Pruning repository, making it easier to inspect the training procedure and kernel design. However, the current results still come from models, GPUs, batch sizes, and sequence lengths selected by the authors. Dynamic routing also adds costs for masking, reordering, and kernel maintenance. The next question is whether WIDE can be integrated with continuous batching runtimes such as vLLM and SGLang while preserving its gains across different GPUs, quantized models, and high-concurrency serving workloads.

Sources

  1. WIDE: Boosting Adaptive LLM Inference via Token-level Dynamic Width Pruning
  2. WIDE source code in EIT-NLP/LLM-Pruning