AI Research

WIDE: Boosting Adaptive LLM Inference via Token-lev... | AI Research

Key Takeaways

WIDE is a framework designed to improve the efficiency of Large Language Models (LLMs) by enabling token-level dynamic width pruning.
Pruning is a promising approach for improving the efficiency of LLMs.
To address these challenges, we present WIDE, the first end-to-end differentiable token-level dynamic width pruning framework designed for both prefill and decode scenarios.
Through a two-stage training pipeline, WIDE learns effective token-wise sparse execution patterns and achieves substantially better quality retention than existing approaches.
At 50% sparsity, WIDE provides 55.1% performance boost when compared to the state-of-the-art dynamic depth pruning under calibration-only settings.

Paper AbstractExpand

Pruning is a promising approach for improving the efficiency of LLMs. Existing static structured pruning methods are hardware-friendly and can deliver practical throughput gains, but their input-agnostic computation allocation often causes substantial accuracy degradation under aggressive sparsity. Recent dynamic sparsity methods improve quality retention by adapting computation to individual inputs, yet they remain largely limited to coarse-grained structural decisions and their practical acceleration under real-world inference scenarios remains challenging. To address these challenges, we present WIDE, the first end-to-end differentiable token-level dynamic width pruning framework designed for both prefill and decode scenarios. WIDE enables fine-grained computation allocation by allowing each token to dynamically select attention-head groups and FFN-channel groups, extending dynamic pruning beyond layer-level decisions to neuron-block-level granularity. Through a two-stage training pipeline, WIDE learns effective token-wise sparse execution patterns and achieves substantially better quality retention than existing approaches. To make such fine-grained dynamic pruning practical, we further propose a pruning--kernel co-design framework that decomposes dynamic sparsity acceleration into mask reordering, hardware-agnostic block-level skipping, and hardware-dependent intra-block skipping, enabling efficient execution across different granularities. At 50% sparsity, WIDE provides 55.1% performance boost when compared to the state-of-the-art dynamic depth pruning under calibration-only settings. Under prefill and decoding inference workloads, WIDE achieves close-to-theoretical kernel-level speedups of up to 1.98x for prefill and 4.95x for decoding, as well as 1.68x and 1.55x end-to-end acceleration. Our code is available at this https URL .

WIDE is a framework designed to improve the efficiency of Large Language Models (LLMs) by enabling token-level dynamic width pruning. While traditional pruning methods remove entire layers or submodules permanently, WIDE allows models to dynamically select specific attention-head groups and FFN-channel groups for each individual token during inference. This approach aims to provide finer-grained computation allocation, allowing models to maintain higher quality while reducing computational costs.

Fine-Grained Dynamic Pruning

Existing dynamic pruning methods often operate at a coarse level, deciding whether to execute or skip entire layers. WIDE moves beyond this by introducing a bottleneck router that makes decisions at the neuron-block level. Each token independently selects which attention heads and FFN channels to activate. This is achieved through a two-stage training pipeline: first, a router-only calibration phase optimizes the routing components while keeping base model parameters frozen; second, an optional LoRA recovery phase is used to restore model performance.

Pruning-Kernel Co-Design

To ensure that fine-grained pruning translates into actual speedups, the authors introduce a pruning-kernel co-design framework. A common challenge with dynamic sparsity is that irregular execution patterns can lead to overheads from gather-scatter operations and device-to-host synchronization. WIDE addresses this by using a mask-reordering process that packs active tokens into structured layouts. This allows the system to skip inactive computation at three levels:

CTA-level: Hardware-agnostic skipping of fully inactive blocks.
Load-level: Skipping inactive memory-loading packets.
MMA-level: Skipping inactive tensor-core computation fragments.

Performance Results

The researchers evaluated WIDE using Llama 3.1-8B and Llama 3.2-3B models. At 50% sparsity, the framework achieved a 55.1% performance boost compared to state-of-the-art dynamic depth pruning in calibration-only settings. In terms of end-to-end inference, WIDE demonstrated speedups of 1.68x for prefill and 1.55x for decoding. The authors report that the kernel-level design provides up to 1.98x speedup for prefill and 4.95x for decoding compared to baseline operations.

Considerations for Implementation

The effectiveness of WIDE relies on the balance between pruning flexibility and kernel efficiency, which is controlled by the group size (G). Smaller group sizes offer more flexibility for pruning but may reduce the potential for hardware acceleration, while larger group sizes align better with GPU tiling strategies. The framework is designed to be compatible with existing optimizations like Grouped-Query Attention (GQA), ensuring that the dynamic pruning process does not conflict with standard inference optimizations.

Comments (0)

No comments yet

Be the first to share your thoughts!