WIDE: Boosting Adaptive LLM Inference via Token-level Dynamic Width Pruning

AI summary

Overview

Research area: Efficient large language model (LLM) inference — specifically dynamic structured pruning of Transformer models combined with customized GPU kernel design.

Technical level: Advanced. The paper assumes familiarity with Transformer internals (attention heads, FFN channels, GQA, FlashAttention), pruning terminology (structured/unstructured, sparsity ratios, calibration), and GPU kernel concepts (GEMM tiling, tensor cores, CTA-level scheduling, CUDA Graph, Triton/Tilelang/CUTLASS).

Scope: The paper introduces WIDE, a framework for token-level dynamic width pruning that lets every token independently select groups of attention heads and FFN channels, together with a pruning–kernel co-design that converts these irregular decisions into practical prefill and decoding speedups on GPUs.

What This Paper Is About

Static pruning removes the same layers, heads, or neurons for every input, which is hardware-friendly but hurts accuracy at aggressive sparsity levels. Dynamic pruning adapts computation per token, but existing methods mostly decide only whether to skip entire layers or sublayers — a coarse choice that limits the quality–efficiency trade-off — and their irregular execution patterns rarely translate into real GPU speedups. WIDE addresses both sides of this problem at once: it pushes token-wise routing down to the level of attention-head groups and FFN-channel groups (the paper calls this "neuron-block-level"), and it pairs that routing with GPU kernels that reorder masks and skip inactive work at several granularities so the sparsity actually pays off in wall-clock time.

Key Contributions

1. Token-wise dynamic width pruning. WIDE is described as the first end-to-end token-wise dynamic pruning framework achieving fine-grained width-level computation allocation. Instead of layer- or sublayer-level execute-or-skip decisions, each token dynamically selects fine-grained attention-head groups and FFN-channel groups inside every Transformer block, using lightweight bottleneck routers and a differentiable mask formulation.

2. A unified framework for accelerating fine-grained dynamic sparsity. The authors identify two obstacles — synchronized mask-to-index conversion and irregular gather–scatter execution — and address them with standalone mask-reordering preprocessing plus multi-stage intra-block predication, allowing the kernels to approach the ideal speedup of the prunable operations in both prefill and decoding.

3. Pruning–kernel co-design that advances the Pareto frontier. WIDE is presented as the first framework to combine token-wise dynamic width pruning with GPU kernel co-design for both prefill and decoding. At 50% target sparsity it improves average zero-shot accuracy by up to 20.26 points over the strongest evaluated dynamic-depth baseline and achieves end-to-end speedups of 1.68x (prefill) and 1.55x (decoding), with kernel-level analysis showing up to 200x speedup and 1,000x peak memory reduction versus naive gather–scatter baselines.

4. A two-stage training pipeline. Router training keeps all non-router parameters frozen and optimizes only the routing components on calibration data; an optional LoRA recovery stage then restores quality. Both stages combine the standard language modeling loss with an auxiliary sparsity loss targeting a specified budget (for example 50%).

Main Findings

• Calibration-only quality at 50% sparsity: On Llama3.1-8B, WIDE improves over the strongest non-WIDE baseline by 8.80 average-accuracy points (61.84 vs. 53.04); on Llama3.2-3B the gap is 8.97 points (57.00 vs. 48.03). The paper also states a 55.1% performance boost over state-of-the-art dynamic depth pruning at 50% sparsity under calibration-only settings.

• Moderate sparsity is nearly lossless. On Llama3.1-8B at 25% sparsity, the G=32 variant retains 97.92% of dense accuracy, while G=64 slightly exceeds DDP, a state-of-the-art static width pruning method (90.68% vs. 89.93%). The layer-wise dynamic baseline SkipGPT retains only 72.22%.

• Most of the gain comes from granularity, not just dynamism. The authors attribute the improvement over SkipGPT to moving dynamic routing from layer selection to width allocation.

• LoRA results: At 25% sparsity on Llama3.1-8B, all three WIDE variants retain lossless dense performance. At 50% sparsity, the best configurations reach 64.82 average accuracy on Llama3.1-8B and 58.77 on Llama3.2-3B, improving over SkipGPT by 3.22 and 4.44 points and over DDP by 10.19 and 9.89 points.

• Speedup upper bound. A forward-pass trace on Llama3.1-8B at 50% sparsity with B=1, T=16,384, G=128 shows the accelerated-path operations deliver average speedups of 1.82x for prefill and 1.92x for decoding, close to the ideal bound. Non-accelerated components (KV projections, element-wise ops, CUDA Graph launches, router execution, kernel initialization) account for about 16.7% of total prefill latency and 29.1% of total decoding latency.

• Group size trades flexibility for kernel efficiency. Increasing G_ffn from 16 to 256 at 50% sparsity decreases average accuracy by 2.62 points; changing G_ffn from 16 to 512 after LoRA tuning lowers accuracy by only 1.13 points. On the kernel side, GEMM-K approaches ideal speedup only when G > 128, and GEMM-MN needs G ≥ 128 to preserve a tunable tiling space; attention kernels are less sensitive but require GQA-aligned settings.

• Throughput across sparsity. At 0% sparsity, the best setting with G ≤ 128 retains 98.60% of dense throughput, indicating routing and optimized kernels add little overhead. At 50% sparsity, the best configuration accelerates prefill by 1.68x and decode by 1.55x. Kernel-level speedups reach up to 1.98x for prefill and 4.95x for decoding.

• Sparsity sweep. From 10% to 70% target sparsity on Llama3.2-3B, WIDE's accuracy declines gradually; calibration-only SkipGPT drops sharply at 20% sparsity, and from 20% to 70% calibration-only WIDE remains more accurate than LoRA-recovered SkipGPT.

• Learned routing is highly non-uniform and semantic-aware. On the Llama3.1-8B checkpoint with 50% target sparsity and group size 128, overall sparsity is 47.3%. Attention carries most of the skipping at 66.2% sparsity on WikiText2, while the FFN is more conservative at 28.5%. Layer-wise, both curves peak around layers 4–12 and 24–28 and dip near the model boundaries and layers 15–18. In a HellaSwag case study, FFN routers keep content words such as "boy", "running", and "track" — especially in layers 9–11 where average FFN sparsity is high — while often skipping articles like "A", "a", and "the".

Methodology in Plain English

What gets pruned. WIDE decomposes each Transformer layer into attention and FFN parts. For attention, the smallest unit is an attention head, and the group size is set to the head group size H_q/H_k × d so the decisions stay aligned with grouped-query attention and with existing GQA decoding optimizations. For the FFN, the intermediate dimension D' is cut into contiguous groups whose size G is drawn from the same tiling-size set used by GEMM kernels (powers of two of at least 16), because GEMM instructions typically support minimum block widths of 16. A single mask matrix represents the pruning of the whole FFN layer.

How decisions are made. A lightweight bottleneck router maps input activations to N_G binary decisions per token (N_G equals D/G for attention and D'/G for FFN). During training the router uses hard Gumbel-Softmax sampling with temperature annealed from 5 to 0.5; at inference it takes an argmax. Class 0 means execute, class 1 means skip. The router's inner dimension r is typically 16 or 32.

How the model is trained. Stage one, router training, freezes everything except the routers and uses calibration data (a RedPajama-1T subset) to learn routing only. Stage two, optional LoRA tuning, adds LoRA adapters to every linear module to recover quality. Both stages optimize the language modeling loss plus an auxiliary sparsity loss weighted by alpha (set to 20, following SkipGPT) that pulls the average mask density toward the target budget S. Both stages run for 10k steps with batch size 16 and max sequence length 4,096 on 4 NVIDIA A100-SXM4-40G GPUs with PyTorch FSDP2; LoRA uses r=16, alpha=32, and 0.1 dropout.

How the speedups are obtained. Rather than materializing token-specific sparse tensors and using gather–scatter (which triggers device-to-host synchronization, breaks CUDA Graph execution, and inflates memory traffic), WIDE first sorts each routing column so that active rows form a contiguous prefix — turning irregular token-wise routing into CTA-level regularity, with at most one boundary tile per group. It then applies a single skipping predicate at three levels: CTA-level early exit (hardware-agnostic), skipping of inactive A-load packets, and skipping of inactive MMA fragments (both architecture-dependent). Three customized kernel families — GEMM-MN, GEMM-K, and Attention — cover the Q/Up/Gate projections, the O/Down projections, and flash attention. Kernels are implemented both in tile-level DSLs (Triton, Tilelang), which support only CTA-level skipping, and in a lower-level CuTe/CUTLASS implementation using the TVM-FFI JIT interface used by SGLang, which enables intra-block skipping during A loading and MMA execution.

Evaluation setup. Models are evaluated with lm-evaluation-harness at maximum context length 4,096: WikiText2 perplexity plus zero-shot accuracy on ARC-Easy, ARC-Challenge, BoolQ, WinoGrande, PIQA, OpenBookQA, and HellaSwag, with average accuracy and retention ratio computed over the seven classification tasks. Kernel throughput is measured through the Triton benchmark interface under CUDA Graph replay, and end-to-end acceleration through an ELANA-style profiling tool measuring TTFT and TPOT; inference experiments mainly use sm120-architecture GPUs (NVIDIA RTX 5090). Baselines include static depth pruning (Shortened LLaMA, CoopPruner), static width pruning (SliceGPT, Týr-the-Pruner, DDP), and dynamic pruning (D-LLM, SkipGPT), all calibrated on the same RedPajama-1T subset with the same LoRA recipe.

Why This Matters

Impact on research. The paper argues that fine-grained dynamic pruning should be designed jointly with the hardware execution path rather than treated as a model-only compression problem. It provides a concrete demonstration that token-level width routing can beat both static width pruning and dynamic depth pruning on the quality–sparsity trade-off while still producing measurable kernel and end-to-end speedups in both prefill and decoding regimes — a combination the authors state has not previously been shown together.

Real-world applications:

• High-throughput LLM serving: the prefill and decode speedups (1.68x and 1.55x end-to-end, up to 1.98x and 4.95x at the kernel level) target exactly the two phases that dominate serving cost, under CUDA Graph execution which is standard in production stacks.

• Agentic and long-context workloads: the paper motivates its work by increasingly complex agent harness systems and long contexts; the traced configuration uses T=16,384 tokens, and efficient prefill is directly relevant to long prompt processing.

• Memory-constrained deployment: the reported 1,000x peak memory reduction versus naive gather–scatter baselines addresses the storage overhead of materializing token-specific weight subsets, which matters when memory bandwidth, not FLOPs, is the bottleneck.

• Accuracy-sensitive deployment at aggressive compression: the calibration-only results (retaining 97.92% of dense accuracy at 25% sparsity, and around the 90% retention regime after LoRA at 50%) suggest a practically usable operating point for teams that cannot afford a full retraining pipeline.

Industry relevance. The design choices are deliberately compatibility-oriented: GQA-aligned attention groups, GEMM-tile-aligned FFN groups, and kernels that preserve the vanilla mainloop rather than replacing it. WIDE also preserves dense-style tiling so masks can be consumed without explicit index conversion, which keeps the approach compatible with CUDA Graph replay and existing inference backends — the practical conditions for any pruning method to be adopted in a serving stack.

Future Directions

• Closing the gap between theoretical and realized speedup. The non-accelerated components (KV projections, element-wise ops, CUDA Graph launches, router execution, kernel initialization) account for about 16.7% of prefill latency and 29.1% of decoding latency, which the paper identifies as the reason layer-wise speedup stays below the theoretical bound. Reducing that fixed overhead is an open systems problem.

• Extending fine-grained skipping to tile-level DSL implementations. Triton and Tilelang implementations currently support only CTA-level skipping because of their coarser control granularity, while the CuTe/CUTLASS path enables intra-block skipping during A loading and MMA execution. Broadening fine-grained control to higher-level DSLs would make the approach easier to adopt.

• Generalizing beyond the evaluated architectures and backbones. Evaluation covers Llama3.1-8B and Llama3.2-3B, with inference experiments mainly on sm120 GPUs. Whether the design transfers to other model families, other GPU architectures (the authors note the level-1 and level-2 predicates are architecture-dependent, citing cp.async vs TMA gather4 and wmma vs wgmma/tcgen05.mma), and other serving scenarios is left open.

• Smarter allocation of the sparsity budget. The learned routing is highly non-uniform (47.3% overall sparsity, but 66.2% in attention versus 28.5% in the FFN, with concentration around layers 4–12 and 24–28), which raises the question of whether explicit control over the per-layer or per-token budget could improve the quality–speed trade-off further. The appendices reference stricter uniform-sparsity and real-sparsity-aligned studies, but those results are not in the provided content.

Target Audience

This paper is most useful to researchers and engineers working on efficient LLM inference and model compression — particularly those interested in the intersection of dynamic sparsity and GPU kernel implementation. It will appeal to:

• Pruning and compression researchers comparing static versus dynamic, depth versus width, and coarse versus fine-grained routing strategies.
• Inference systems engineers who need pruning decisions to survive contact with CUDA Graphs, tiled GEMM mainloops, and flash attention kernels.
• Kernel developers interested in mask reordering, predicated skipping at CTA / load / MMA granularity, and co-design across Triton, Tilelang, and CUTLASS/CuTe.
• Practitioners deploying LLMs who need a concrete quality–throughput curve: near-lossless at low sparsity (98.60% of dense throughput at 0% sparsity, 97.92% dense accuracy retained at 25% sparsity) and meaningful acceleration at 50% sparsity.

Readers without a background in GPU kernel programming or Transformer internals will find the algorithmic idea accessible, but the acceleration framework will require substantial background knowledge.

Authors’ abstract

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 https://github.com/EIT-NLP/LLM-Pruning/tree/main/WIDE.