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

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

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

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

Preprint

WIDE: B OOSTING A DAPTIVE LLM I NFERENCE VIA T OKEN - LEVEL DYNAMIC W IDTH P RUNING Haozhe Hu1 Hao Wu1 Peiran Yin1 Chao Han1 Yunpu Ma2 Xiaoyu Shen1∗ Ningbo Institute of Digital Twin, Eastern Institute of Technology, Ningbo 2 Munich Center for Machine Learning, LMU Munich 1

[email protected]; [email protected]

arXiv:2607.28418v1 [cs.AI] 30 Jul 2026

A BSTRACT Pruning is a promising approach for improving the efficiency of large language models (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 attentionhead groups and FFN-channel groups, extending dynamic pruning beyond layerlevel 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 codesign framework that decomposes dynamic sparsity acceleration into mask reordering, hardware-agnostic block-level skipping, and hardware-dependent intrablock 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. These results establish WIDE as an effective fine-grained dynamic width pruning framework that pushes the frontier of token-wise dynamic structured pruning. Our code is available at https://github.com/EIT-NLP/LLM-Pruning/tree/main/WIDE.

I NTRODUCTION

The rapid advancement of LLMs has enabled their widespread deployment across a broad range of real-world applications (Team et al., 2026; DeepSeek-AI, 2026; Meng et al., 2026; Dai et al., 2026; Chen et al., 2026). However, as model continue to scale and increasingly complex agent harness systems emerge, the efficient LLM serving has become increasingly challenging. Beyond system-level and hardware-specific optimizations (Kwon et al., 2023; Zheng et al., 2024; Lin et al., 2024; Zadouri et al., 2026), model pruning has also evolved rapidly as a model-level approach to improving inference efficiency. By removing redundant structures along the depth and width dimensions, pruning reduces model size and computational cost, thereby offering a complementary path toward efficient LLM/VLM serving (Cheng et al., 2024; Frantar & Alistarh, 2023; Ma et al., 2023; Ashkboos et al., 2023; Men et al., 2025; Wu et al., 2025; 2026b;a). ∗

Corresponding Author

Preprint

Currently, LLM pruning is dominated by static schemes, which remove predefined structures such as entire layers (Kim et al., 2024), attention or feed-forward network (FFN) sublayers (Zhong et al., 2025), rows or columns of weight matrices (Li et al., 2025), and individual neurons (Fang et al., 2024). Once calibrated, these units are permanently removed, making static pruning simple to deploy and compatible with existing hardware backends. However, applying the same pruning decisions to every input often sacrifices model quality, especially under aggressive sparsity. Dynamic pruning addresses this limitation by introducing lightweight routers that allocate computation according to individual tokens (Raposo et al., 2024; Jiang et al., 2024; Zhao et al., 2025; Han et al., 2025). Despite their improved flexibility, existing dynamic pruning methods primarily operate at the depth level, deciding whether each token should execute or skip entire layers or submodules. Such coarse-grained decisions limit the achievable quality–efficiency trade-off, as tokens that require only partial computation within a layer may still lose useful capacity when the entire module is removed (Shrestha et al., 2025; Gautam et al., 2026). Meanwhile, pushing dynamic pruning toward finer-grained token-wise allocation introduces additional system challenges: irregular execution patterns can prevent the reduced computation from translating into practical inference acceleration (He et al., 2026; Hu et al., 2026). Therefore, two fundamental questions remain: (1) Can dynamic pruning move beyond coarsegrained depth decisions to achieve finer computation allocation and better quality retention? (2) Can such fine-grained dynamic decisions be efficiently executed to deliver real end-to-end speedups? Motivated by these challenges, we present WIDE, the first end-to-end token-wise dynamic width pruning framework that jointly optimizes pruning granularity and GPU execution. WIDE pushes dynamic pruning beyond layer- and sublayer-level routing by enabling each token to independently select fine-grained groups of attention heads and FFN channels. Specifically, lightweight bottleneck routers select group-query-attention (GQA)-aligned attention-head groups and configurable FFNchannel groups independently for every token. The differentiable router-only calibration learns these decisions with end-to-end training, with an optional LoRA recovery (Hu et al., 2021) further restores quality. The group size explicitly controls the trade-off between allocation flexibility and kernel efficiency. To translate these fine-grained decisions into practical acceleration, WIDE further introduces a unified GPU acceleration framework through pruning–kernel co-design. Instead of explicitly materializing token-specific sparse tensors, WIDE first performs mask reordering to transform arbitrary token-wise routing patterns into structured execution layouts. Active tokens from each routing group are packed into tile-aligned prefixes, converting irregular sparsity into CTAlevel regularity. Based on this representation, WIDE kernels progressively eliminate unnecessary computation at multiple granularities: hardware-agnostic predicates remove fully inactive CTAs, while architecture-aware predicates further skip inactive memory-loading packets and tensor-core computation fragments. This design enables fine-grained dynamic sparsity to preserve efficient GPU execution patterns and deliver real acceleration for both prefill and decoding workloads. Our main contributions are summarized as follows: (1) We introduce WIDE, the first end-to-end token-wise dynamic pruning framework that achieves fine-grained width-level computation allocation. Unlike previous dynamic pruning approaches that perform token-wise decisions at the layer or sublayer level, WIDE enables each token to dynamically select fine-grained attention-head groups and FFN-channel groups within each Transformer block. This pushes dynamic pruning beyond coarse structural skipping toward neuron-block-level computation allocation, providing substantially finer-grained control over model capacity while maintaining structured GPU execution. (2) We present the first unified framework that unlocks practical acceleration for fine-grained dynamic sparsity. We identify synchronized mask-to-index conversion and irregular gather–scatter execution as two key obstacles preventing dynamic sparsity from yielding practical speedups. We thus introduce standalone mask-reordering preprocessing together with multi-stage intra-block predication, allowing the resulting kernels to approach the ideal speedup of the prunable operations in both prefill and decoding. (3) We advance the pruning Pareto frontier through pruning–kernel co-design. To the best of our knowledge, WIDE is the first framework to combine token-wise dynamic width pruning with GPU kernel co-design for both prefill and decoding. At a 50% target sparsity, WIDE 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.68× and 1.55× for prefill and decoding, respectively. The kernellevel analysis also demonstrate up to 200x speedup and 1,000x peak memory reduction compared to naive gather–scatter baselines.

Preprint

FFN

A en on Head Group Pruning

Router (lightweight bo leneck)

ℝD×r

Wq ℝB×T×Hk M

ℝD×r

W1 Group Query A en on

ℝ

r×2Hk

🔥

Group Pruning (G=32, 64, 128…)

Router (lightweight bo leneck)

🔥

WoT

ℝB×T×(D′/G)

Wup

ℝ

r×2(D′/G)

❄ / LoRA

Wgate

WdownT

Figure 1: The pruning dimension of WIDE in the attention and FFN layer, with Wq , Wup , and Wgate pruned over out features (N-axis), and Wo , Wdown over in features (K-axis). The output feature size of routers are set to twice of the group number NG to align with the Gumbel Softmax. tt

R ELATED W ORKS

Static Pruning Among existing pruning methods, static pruning remains the dominant. Given a small set of calibration samples, these methods estimate the importance of predefined units, such as Transformer layers, attention or FFN sublayers, attention heads, rows or columns of weight matrices, and individual neurons, and permanently remove low-importance units until the target sparsity is reached. According to the pruning unit, existing work can be broadly categorized into depth and width pruning. Depth pruning removes computation along the layer dimension: Shortened LLaMA prunes complete Transformer blocks using perplexity-based or Taylor and heuristic-based importance followed by retraining (Kim et al., 2024); BlockPruner further evaluates attention and FFN residual blocks separately with calibration perplexity (Zhong et al., 2025); recent CoopPruner models layers as interacting players and estimates their marginal contributions with surrogate-assisted Shapley values (Ding et al., 2025). Width pruning instead removes channels, neurons, or attention heads within layers. Representative methods include SliceGPT, which exploits orthogonal invariance and PCA on calibration activations (Ashkboos et al., 2023); Týr-the-Pruner, which searches for a global non-uniform sparsity allocation over FFN channels and attention heads (Li et al., 2025); and recent Deterministic Differentiable Pruning (DDP) (Huang et al., 2026), which learns deterministic structured masks under an ℓ0 sparsity constraint via augmented Lagrangian optimization, avoiding stochastic mask sampling and its train–test mismatch. Beyond structural pruning, methods like SparseGPT Frantar & Alistarh (2023), Wanda Sun et al. (2023), and MaskLLM Fang et al. (2024) further incorporate semi-structured or unstructured masks to individually pruning each neuron. These static methods are deployment-friendly, but their input-agnostic pruning decisions limit their ability to allocate computation adaptively across tokens. Dynamic Pruning Dynamic pruning allocates computation adaptively during inference instead of permanently removing a fixed set of parameters. Existing methods mainly focus on token-wise depth or sublayer routing. Mixture-of-Depths (MoD) routes only the top-k tokens through each Transformer layer under a fixed compute budget (Raposo et al., 2024). D-LLM equips layers with dynamic decision modules to determine whether each token should execute or skip a network unit, together with a KV-cache eviction policy for skipped tokens (Jiang et al., 2024). SkipGPT further decouple attention and FFN modules, using a differentiable routing mechanism to construct a two-stage training Zhao et al. (2025). Informed Routing relaxes binary execute-or-skip decisions by predicting unit outputs with a lightweight feature forecaster (LFF) and thus construct a three-stage recovering (Han et al., 2025). PolarSparse and FastForward instead introduce the dynamic width pruning for batch decode / prefill only (Gautam et al., 2026; Shrestha et al., 2025). While these methods improve adaptive computation, they mostly operate at coarse granularity or specific scenario; WIDE instead studies fine-grained dynamic width pruning and couples its routing design with GPU-oriented block skipping.

WIDE: P RELIMINARY, T RAINING , AND I NFERENCE D ESIGN

As previously introduced, WIDE is a dynamic width pruning framework that selectively activates attention head groups and weight matrix blocks under token-wise routing. In this section, we first 3

Preprint

elaborate on the minimal atomic pruning unit of WIDE, then followed by its two-stage training pipeline, and the customized acceleration framework for dynamic pruning. 3.1

P RELIMINARY

Consider a standard Transformer layer with input token embeddings X ∈ RB×T ×D with B for batch size, T for token length, and D for hidden size, WIDE follows the granularity of SkipGPT (Zhao et al., 2025), which decomposes the entire layer into two computation parts: attention and FFN, with Tq , Tk denotes the token length of query and key value, Hq , Hk denotes the number of query / key value attention heads, d = HDq denotes the head size, and D′ denotes the intermediate size of FFN’s up scaling. The overall pruning granularity of WIDE is shown in Figure 1, where each token allocates their computation budget based on the atomic pruning group size G. Attention Pruning Unit For the attention layer, considering the compatibility with the standard FlashAttention tiling scheme that launches over B, Hq , Tq dimensions and consumes one head at each step (Dao, 2023), we assume the head as the minimal pruning atom to skip d or multiple-d under each token routing. Such pruning strategy can be equivalently represented with a structured mask of the form like: [Q, K, V] = [XWq⊤ , XWk⊤ , XWv⊤ ] (1)     X′ = Attn.(Q, K, V) ⊙G Mattn  Wo⊤ | {z }

where Mattn ∈ {0, 1}B×Tq ×Hk , G =

Hq D = ×d Hk Hk

RB×Tq ×Hq ×d

(2) Here ⊙G represents to apply gating based on group size. The outer mask Mattn also prevent to pruning over Wk and Wv , which can lead to cache eviction. Additionally, since that recent LLMs predominantly adopt GQA (Ainslie et al., 2023), WIDE sets attention’s G to the head group size Hq Hk × d for preserving compatibility with existing GQA decoding optimizations (e.g., XQA (NVIDIA, 2026; Ye et al., 2025)) while minimizing the degradation of pruning flexibility. FFN Pruning Unit Unlike attention layers, FFN layers comprise only consecutive general matrix multiplication (GEMM) and element-wise operations. For a standard GEMM operation with its left-hand-side operand (LHS) A ∈ RM ×K and right-hand-side operand (RHS) B ∈ RN ×K , the tiling strategy typically slices all three logical dimensions M , N , and K, with tile sizes (BM, BN, BK) ∈ T 3 , T = {2k | k ∈ Z⩾4 }. Based on this, WIDE adopts commonly used tiling sizes as the G for FFN to prune the intermediate dimension D′ into contiguous group. This dimension corresponds to the out-feature in Up/Gate projections and the in-feature in Down projections, ensuring a single mask matrix to represent the entire FFN layer pruning:    ⊤  ⊤ ⊤ X′ = (ϕ(XWgate ) ⊙ (XWup )) ⊙G Mffn  Wdown | {z }

D′

where Mffn ∈ {0, 1}B×T × G , G ∈ T (3)

RB×T ×(D′ /G,G)

where ϕ denotes element-wise activation function. The G in FFN unit provides a more flexible choice than attention unit. As GEMM instructions typically support minimum block widths of 16, the available choices for G lie in the same tiling-size set T , with a trade-off between pruning performance (G ↓) and potential acceleration gain (G ↑). Routing Components Based on the pruning units and their corresponding mask representations described above, we can design the router component for WIDE, which maps the input activation to NG binary decisions. To flexibly handle varying dimensions and make compatible with Gumbel Softmax sampling during training, WIDE employs a bottleneck architecture, with the router implementations for attention/FFN layers defined as follows: R ∈ RB×T ×NG ×2 = reshape (XW1 W2 ) where W1 ∈ RD×r , W2 ∈ Rr×2NG i ( h 1 argmaxc∈{0,1} R...,c = 0 , when inference, M= GumbelSoftmaxhard when training, τ (R)...,0 , 4

(4) (5)

Preprint

A Preprocessing: Mask reordering GEMM RHS B [K, N], N = (NG, G)

GEMM LHS A [M, K]

g0 M

1 (ac ve)

sort

Sorted Indices II [M, NG]

0 (skip)

M / BM Tiles

…

1 (ac ve)

0 (skip)

gNG−1

! Sorted Mask M M

Route Mask M M

NG Groups

0 2 4 5 7

9 1 3 6 8

GEMM Results C [M, (NG, G)] gNG−1 g1

…

NG Tiles

B Kernel: Dense pipeline with mul -granularity skip CTA le ini alize (prepare mask)

prologue (gather A & B from HBM)

Before prologue: CTA early exit

skip (i, j) CTA le if all tokens in mi skip group gj

total CTA les ((M / BM) x NG) processed

… M / BM Tiles

… …

…

par al skip

…

NG Tiles Granularity: CTA

MMA

Prologue & mainloop: row-wise loading skip

! M

…

K / BK Tiles

9 1 3 6

— skip — skip

… sm80+ cp.async (1 row per thread)

load A & B

Mainloop: MMA skip ! M

8 16

128

…

G skip —

K / BK Tiles sm100+ TMA gather4 (4 row per warp)

Granularity: thread / warp

epilogue (sca er C to HBM)

…

load B le (dense)

load A from row indices

early exit

…

mainloop (repeat over K les on chip)

load A & B

…

sm80+ wmma (16 row per group)

…

sm90 wgmma (64 row per group)

…

sm100 tcgen05.mma (128 row per group)

Granularity: warp

Figure 2: Workflow of WIDE for GEMM N-axis pruning. (a) Given group size G, WIDE generates a mask matrix of shape [M, NG ] and sorts it along the M dimension, so that active rows of each group are clustered within BM -token tiles. This preserves tile-based GEMM without explicit mask-to-index conversion. (b) Inside the kernel, WIDE first skips fully inactive [BM, G] blocks. For blocks with active entries, it follows the standard GEMM pipeline and applies architecture-dependent intra-block skipping during A loading and MMA execution. titt ti

titi ti

where τ is the temperature, NG is set to D/G for attention, and D′ /G for FFN. Here r ≪ D is typically chosen from {16, 32}. In each output pair, class 0 denotes execution and class 1 denotes skipping. 3.2

T WO - STAGE T RAINING

To fully unlock the potential of WIDE, a two-stage training pipeline is introduced to restore model’s original performance. During the first stage that called Router training, all parameters except the router remain frozen, with calibration data employed only to optimize the routing components. In the next optional LoRA tuning stage, LoRA are incorporated into all base models’ linear modules at each layer to further support performance recovery. Both stages utilize not only the standard language modeling loss but also an auxiliary sparsity loss to ensure the pruning meets the target sparsity. For a model with N layers, the overall objective is similar to SkipGPT: N 1 X 1 X i 1 X i L = LLM + α S − 1− M + 1 − M (6) attn ffn 2N i=1 |Miattn | |Miffn |

where α is a hyperparameter controlling the weight of the sparsity loss, |M| denotes the total element number of the mask, S denotes the target sparsity budget, e.g., 50%. 3.3

I NFERENCE - TIME G ROUP S KIPPING WITH U NIFIED M ASK R EORDERING

Unlike MoE-style routing, dynamic pruning does not assume a fixed compute budget for each token. A naive implementation therefore typically needs to explicitly convert routing masks into active indices and then use gather-scatter operations to feed only active tokens into GEMM or attention kernels. Let C = AB⊤ , where A ∈ RM ×K , B ∈ RN ×K , and C ∈ RM ×N . This strategy introduces two major sources of overhead: (1). Device-to-host synchronization. The active token layout must be derived from input-dependent masks before constructing the compact gathered matrix, which can trigger device-to-host synchronization and break CUDA Graph execution.

Preprint

GEMM

Kernel overall wall-clock time

(2). Extra memory movement and storage. For (a) M=16 -- N = 14,336, K = 4,096 @ 50% sparsity dynamic depth pruning, ignoring mask and index57.15 Dense 27.55 processing overhead, explicit gather-scatter increases the memory traffic of a dense GEMM by roughly Pruned 25.24 150.56 3pM (N +K)+N K+M N , where p denotes the active ra0 40 80 120 160 M K+N K+M N Latency ( s) tio. For finer-grained dynamic width pruning, each (b) M=4096 -- N = 14,336, K = 4,096 @ 50% sparsity token may select a different subset of B. A naive 1604.87 1628.35 Dense N-axis implementation therefore incurs traffic of 734.20 1024.63 Pruned M K + 3pM N K + 3pM N + M N , while its K-axis 0 400 800 1200 1600 counterpart incurs 3pM K + 3pM N K + M N , both Latency ( s) require dominant intermediate storage of pM N K for token-specific weight subsets Bsub (see Appendix A). Figure 3: GEMM latency breakdown of dense As shown in Fig. 3, these non-GEMM overheads are and naive dynamic depth pruning. difficult to hide in both prefilling and decoding, while such implementations also struggle with more complex tiled operators such as flash attention (Dao, 2023). WIDE therefore introduces a unified mask-reordering preprocessing step with fused gather-scatter kernels that remain close to the original tile-based pipeline, as illustrated in Fig. 2. Under this notation, WIDE represents dynamic routing decisions with a binary mask M ∈ {0, 1}M ×NG , where NG = 1 for M-axis routing (i.e., dynamic depth) and NG = Dp /G for width pruning over the pruned feature dimension Dp . For each group g, the mask is sorted along the M f:,g and the corresponding row indices I:,g : dimension, producing a reordered mask M f:,g , I:,g ) = SortDesc(M:,g ). (M

(7)

Sorting each routing column packs active rows into a contiguous prefix, while I:,g gathers the corresponding activation rows and scatters their outputs back. This turns irregular token-wise routing into CTA-level regularity: except for at most one boundary tile per group, CTA row tiles are fully active / inactive. The active row in each admitted GEMM CTA consequently share a group-aligned weight tile Btile ∈ RBN ×BK , enabling multi-granularity predicated skipping in dense-style tiling and mainloop while avoiding token-specific weight materialization. Based on this layout, WIDE applies a unified skipping predicate over progressively finer row groups. Let S0 = BM,

S1 = Sld ,

S2 = Smma ,

(8)

where S0 is the CTA-level block size, while Sld and Smma are architecture-dependent granularities for A loading and MMA execution. For level ℓ ∈ {0, 1, 2}, group g, and row block m, WIDE computes min((m+1)Sℓ ,M )−1

s(ℓ) m,g =

i=mSℓ (ℓ)

fi,g M

(9)

If sm,g = 0, the corresponding unit at level ℓ is skipped; otherwise it is executed. This single predicate instantiates three stages: ℓ = 0 performs hardware-agnostic CTA early exit, ℓ = 1 skips inactive A load packets, and ℓ = 2 skips inactive MMA fragments1 . Under this abstraction, the proposed GEMM-MN, GEMM-K, and Attention kernels cover Q/Up/Gate projections, O/Down projections, and flash attention operations with customized routing group, ensuring WIDE to maximize the acceleration potential brought by sparsity while preserving its vanilla mainloops. The detailed implementation for each kernel is available at Appendix B.

E XPERIMENTS

4.1

S ETTINGS

Training The Llama3.1-8B and Llama3.2-3B (Grattafiori et al., 2024) models are introduced as the backbone model. Unless otherwise specified, all experiments use a subset of RedPajama-1T (Weber 1 ℓ ∈ {1, 2} are depended on the GPU architecture, e.g., cp.async vs Tensor Memory Accelerator (TMA) gather4, and warp-level wmma vs CTA-level wgmma and tcgen05.mma

Preprint

et al., 2024)2 as the calibration and LoRA recovery corpus. Both training stages are run for 10k steps with batch size 16 and maximum sequence length 4,096, using 4 NVIDIA A100-SXM4-40G GPUs with PyTorch FSDP2. For router training, the Gumbel-Softmax temperature is linearly annealed from 5 to 0.5, and the weight α is set to 20, following SkipGPT. For LoRA recovery, the LoRA r and α is set to 16 and 32, respectively, with additional 0.1 dropout ratio. Evaluation We evaluate all models with lm-evaluation-harness (Gao et al., 2024), using a maximum context length of 4,096. We report perplexity on WikiText2 (Merity et al., 2016) and zero-shot accuracy on ARC-Easy, ARC-Challenge (Clark et al., 2018), BoolQ (Clark et al., 2019), WinoGrande (Sakaguchi et al., 2021), PIQA (Bisk et al., 2020), OpenBookQA (Mihaylov et al., 2018), and HellaSwag (Zellers et al., 2019). The average accuracy and performance retention ratio is computed over the seven zero-shot classification tasks. Inference Implementation For kernel-level throughput evaluation, we use the Triton benchmark interface (Tillet et al., 2019) and report TFLOPs under CUDA Graph replay. For end-to-end acceleration, an ELANA-style profiling tool (Chiang et al., 2025) is introduced to measure single-step Time-To-First-Token (TTFT) and Time-Per-Output-Token (TPOT) under CUDA Graph execution. We implement WIDE kernels with both tile-level Domain Specific Languages (DSLs), e.g., Triton (Tillet et al., 2019) and Tilelang (Wang et al., 2025), following a lower-level implementation based on the CuTe template library in CUTLASS together with the TVM-FFI Just-In-Time (JIT) interface used by SGLang (Zheng et al., 2024). The CuTe implementation provides the most fine-grained control and enables intra-block skipping during both A loading and MMA execution, while the tile-level DSL implementations only support CTA-level skipping due to their coarser control granularity. All inference experiments are mainly conducted on GPUs with sm120 architecture (NVIDIA RTX 5090). Baselines We compare WIDE with representative static and dynamic pruning methods. For static depth pruning, we include Shortened LLaMA (Kim et al., 2024) and CoopPruner (Ding et al., 2025). For static width pruning, SliceGPT (Ashkboos et al., 2023), Týr-the-Pruner (Li et al., 2025), and DDP (Huang et al., 2026) are included. For dynamic pruning, D-LLM (Jiang et al., 2024) and SkipGPT (Zhao et al., 2025) is selected as the baselines. To ensure fair comparison, all methods use the same RedPajama-1T subset for calibration and LoRA recovery. The calibration stage follows each method’s original training budget, while the LoRA stage uses the same settings as WIDE. For SkipGPT, we additionally increase the router rank r to match the router capacity of WIDE. Table 1: Downstream performance on static depth, static width, and dynamic pruning methods under 25%/50% target pruning ratio and Llama3.1-8B/Llama3.2-3B backbone, with the 1st result in bold and 2nd in underlined. The numbers in parentheses denote the group size G of WIDE. The averaged proportion of performance retained is also reported. Methods Dense Shortened-ppl Shortened-taylor CoopPruner SliceGPT Týr-the-Pruner DDP D-LLM SkipGPT WIDE (32) WIDE (64) WIDE (128)

Llama3.1-8B 25% sparsity 50% sparsity WikiText2 ppl↓ Avg. Acc.↑ WikiText2 ppl↓ Avg. Acc.↑ 7.71 71.55 (100.00%) 7.71 71.55 (100.00%) 25.84 47.64 (66.59%) 473.15 39.61 (55.36%) 30.07 49.30 (68.90%) 4.85e+7 35.98 (50.29%) 22.96 53.33 (74.53%) 503.86 39.27 (54.89%) 22.57 53.37 (74.59%) 85.81 38.70 (54.10%) 12.89 61.43 (85.85%) 59.88 45.47 (63.55%) 12.41 64.34 (89.93%) 21.83 53.04 (74.14%) 69.45 38.88 (54.34%) 504.48 36.32 (50.76%) 15.63 51.67 (72.22%) 96.04 42.51 (59.42%) 8.49 70.06 (97.92%) 14.15 61.84 (86.42%) 10.37 64.88 (90.68%) 14.96 61.17 (85.50%) 10.04 63.65 (88.95%) 12.57 61.48 (85.93%)

https://huggingface.co/datasets/ZengXiangyu/RedPajama-Data-1T-Sample

Llama3.2-3B 50% sparsity WikiText2 ppl↓ Avg. Acc.↑ 9.77 64.91 (100.00%) 676.05 38.22 (58.89%) 2.32e+5 39.35 (60.63%) 1.79e+5 36.62 (56.41%) 96.69 37.62 (57.96%) 87.37 42.60 (65.63%) 31.20 48.03 (74.00%) 3877.53 34.99 (53.91%) 156.58 36.74 (56.60%) 17.28 57.00 (87.81%) 19.02 56.40 (86.89%) 18.11 55.42 (85.38%)

Preprint

Table 2: Downstream performance under LoRA configuration. Methods Dense Shortened-ppl Shortened-taylor CoopPruner SliceGPT Týr-the-Pruner DDP D-LLM SkipGPT WIDE (32) WIDE (64) WIDE (128)

4.2

Llama3.1-8B 25% sparsity 50% sparsity WikiText2 ppl↓ Avg. Acc.↑ WikiText2 ppl↓ Avg. Acc.↑ 7.91 71.36 (99.73%) 7.91 71.36 (99.73%) 12.33 57.24 (80.01%) 22.47 47.25 (66.04%) 12.12 59.83 (83.62%) 23.79 49.27 (68.86%) 12.03 61.25 (85.60%) 22.66 49.62 (69.36%) 14.88 60.72 (84.86%) 28.76 46.01 (64.30%) 11.72 64.05 (89.52%) 41.53 52.97 (74.03%) 12.10 64.24 (89.78%) 19.52 54.63 (76.36%) 10.29 66.83 (93.41%) 19.31 56.06 (78.35%) 9.87 69.48 (97.10%) 13.90 61.60 (86.09%) 8.61 70.18 (98.09%) 11.99 64.82 (90.59%) 10.32 69.82 (97.59%) 12.05 64.65 (90.36%) 12.51 69.53 (97.17%) 11.43 64.33 (89.92%)

Llama3.2-3B 50% sparsity WikiText2 ppl↓ Avg. Acc.↑ 9.86 65.00 (100.14%) 29.75 45.83 (70.61%) 31.50 45.84 (70.61%) 42.69 41.82 (64.43%) 31.13 42.76 (65.87%) 53.11 45.69 (70.38%) 25.66 48.88 (75.30%) 193.97 38.66 (59.56%) 18.91 54.33 (83.70%) 15.28 58.77 (90.53%) 15.17 58.36 (89.90%) 15.26 58.11 (89.52%)

M AIN R ESULTS

Calibration-only Results Table 1 compares all methods in the calibration-only setting. Two observations stand out. First, WIDE is already reliable at moderate sparsity. On Llama3.1-8B with 25% sparsity, the G = 32 variant retains 97.92% of the dense accuracy, while G = 64 still slightly exceeds DDP, the state-of-the-art static width pruning methods (90.68% vs. 89.93%). In contrast, the layer-wise dynamic baseline SkipGPT retains only 72.22%, suggesting that most of the gain comes from moving dynamic routing from layer selection to width allocation. Second, the advantage becomes much clearer in the more aggressive 50% setting. 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 similar at 8.97 points (57.00 vs. 48.03). Notably, these numbers use only router calibration; several WIDE calibration-only results are already beating LoRA baselines in Table 2.

LoRA Results Table 2 repeats the comparison after applying the same LoRA recovery recipe to every pruned model. The main trend remains unchanged, but the gaps are more informative after recovery. At 25% sparsity on Llama3.1-8B, the choice of group size has little effect: all three WIDE variants retain lossless performance of dense baseline. At 50% sparsity, LoRA narrows the gap for several baselines, especially SkipGPT, yet WIDE still stays around the 90% retention regime on both backbones. Its 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. The stricter uniform-sparsity and real-sparsity-aligned studies in the Appendix C also points in the same direction. GEMM-MN (Q + FFN up/gate) GEMM-K (O + FFN down)

(a) Prefill

Attention K/V projections (fixed)

B = 1, context length = 16,384

Dense

(b) Decode

WIDE

B = 1, context length = 16,384 361 s

WIDE

27.13 ms (1.66×) 15

CPU graph launch

Dense

45.08 ms

Element-wise ops (fixed) Extra ops (router / sort / zero-init)

231 s (1.57×) 0

Latency per transformer layer (ms)

100

200

300

Latency per transformer layer ( s)

400

Figure 4: Layer-wise latency breakdown for Llama3.1-8B at 50% sparsity, including three types of accelerable operations and the remaining fixed overhead.

Preprint

4.3

A NALYSIS ON P RUNING G ROUP, S PARSITY, AND ACCELERATION

For WIDE, it exposes two knobs that jointly determine the quality–speed trade-off: the group size G controls routing granularity, while the target sparsity set the computing budget. Their effects, along with the overall acceleration landscape will be analyzed in this section.

The Upper Bound of Speedup Before analyzing Figure 5: WIDE performance on Llama3.1the effects of the other design variables, we quantify 8B 50% sparsity with different group size. the maximum speedup that WIDE can achieve. We record a forward-pass trace on Llama3.1-8B at 50% Non-LoRA LoRA G WT2 Avg. Acc. WT2 Avg. Acc. sparsity with B = 1, T = 16,384, and G = 128 Vary FFN, Gattn = 512 (Fig. 4). The operations on the accelerated path de16 12.02 62.89 (87.89%) 11.37 64.92 (90.74%) liver average speedups of 1.82× for prefill and 1.92× 32 14.15 61.84 (86.42%) 11.99 64.82 (90.59%) for decoding, which close to the ideal upper bound. 64 14.96 61.17 (85.50%) 12.05 64.65 (90.36%) The remaining latency comes from KV projections, 128 12.57 61.48 (85.93%) 11.43 64.33 (89.92%) element-wise operations, CUDA Graph launches, and 256 14.51 60.27 (84.24%) 12.37 64.81 (90.58%) 512 22.12 53.84 (75.25%) 13.86 63.79 (89.15%) WIDE’s own overhead, including router execution Vary Attention, Gffn = 32 and the kernel initialization. These components ac128 17.59 62.68 (87.61%) 16.34 65.84 (92.02%) count for about 16.7% of total prefill latency and 256 12.72 61.91 (86.52%) 11.53 64.92 (90.74%) 29.1% of total decoding latency, which explains why 512 14.15 61.84 (86.42%) 11.99 64.82 (90.59%) the layer-wise speedup remains below the theoretical 1024 13.39 57.90 (80.92%) 11.71 64.53 (90.19%) speedup. 2048 14.38 58.51 (81.78%) 12.14 63.17 (88.28%) 4096

15.58

58.91 (82.34%)

12.62

63.34 (88.53%)

Effects of Pruning Group After the overall acceleration landscape, we then isolate the group size, which trades routing flexibility for kernel efficiency. Table 5 shows that quality is robust across G, where increasing Gffn from 16 to 256 at 50% sparsity decreases average accuracy by 2.62 points, and changing Gffn from 16 to 512 after LoRA tuning only lowers accuracy by 1.13 points, with similar trends at Table 2 and attention side. Figure 7 explains why larger groups are preferable in practice. GEMM-K approaches the ideal speedup only when G > 128, as small K-axis groups leave too little mainloop work to amortize pipeline overhead; GEMM-MN also needs G ⩾ 128 to preserve a tunable tiling space. Attention kernels are less sensitive, as head-wise pruning largely preserves the original pipeline and remains fast in both long-prefill and decoding regimes, but a GQA-aligned settings are still essential for GQA-packing decoding (Appendix C).

Avg. acc.

Performance over Sparsity To examine sensitivity to the pruning budget, we sweep target sparWIDE SkipGPT DDP LoRA w/o LoRA Dense sity from 10% to 70% on Llama3.2-3B and comLlama-3.2-3B pare WIDE with DDP and SkipGPT, representative static-width and dynamic-depth pruning baselines 60 (Figure 6). WIDE’s average zero-shot accuracy declines gradually across the full range, both before and 50 after LoRA recovery. In contrast, calibration-only 40 SkipGPT drops sharply at 20% sparsity. From 20% to 70%, calibration-only WIDE also remains more 0 10 20 30 40 50 60 70 accurate than LoRA-recovered SkipGPT. Figure 8 Sparsity (%) further reports end-to-end throughput across sparsity levels and group sizes. At 0% sparsity, the best setting Figure 6: Average zero-shot accuracy of with G ⩽ 128 retains 98.60% of dense throughput, WIDE(G = 32), DDP, and SkipGPT on indicating that routing and optimized kernels add lit- Llama3.2-3B across target sparsity. tle overhead. At 50% sparsity, the best configuration accelerates prefill by 1.68× and decode by 1.55×. Together, the two curves give a flexible quality–throughput trade-off, where WIDE requires minimal extra cost at low sparsity, and higher sparsity yields more measurable speedups without an abrupt loss of accuracy. 9

Preprint

DSL (Triton/Tilelang) CuTe

GEMM-MN (M), [N, K]=14336x4096 6

G=32

G=128

Prefill NC Prefill C

1x torch 2x torch 6

G=256

G=512

2 1 0

256

4k 16k

256

4k 16k

256

4k 16k

256

4k 16k

G=512 [N,K]=4096x4096

256

4k 16k

GEMM-K (M), [N, K]=4096x14336 4

G=32

G=128

G=256

G=512

2 1 0

G=512 [N,K]=4096x4096

Speedup over torch

2 1 0

256

4k 16k

256

4k 16k

256

4k 16k

256

4k 16k

256

4k 16k

Attention-Prefill (Tq), [Hq, Hk, D]=32x8x128 2.5 2

2.5

B=1

B=2

2.5 2

16k

32k

2.5

B=4

16k

32k

B=8

2.5 2

16k

32k

B=16

16k

32k

16k

32k

16k

32k

Attention-Decode (Tk), [Hq, Hk, D]=32x8x128 3

B=1

B=2

B=4

B=8

2 1 0

16k

32k

16k

32k

16k

32k

16k

32k

B=16

Figure 7: Kernel-level speedups of the four types of WIDE kernels on Llama3.1-8B shapes under varying input sizes and group sizes, with a random 50% sparsity mask. GEMM-MN covers the attention Q and FFN Up/Gate projections, while GEMM-K covers attention O and FFN Down projections. For GEMM, M denotes the number of input tokens. For attention, NC and C denote non-causal and causal kernels. The dense baseline is set to PyTorch cuBLAS for GEMM, and SDPA (FlashAttention/FlashDecoding backend) for Attention, with throughput calculated under full wall-clock time of each operation’s lifetime. G=32

G=64

G=128

G=256

(a) Prefill -- Batch Size = 1, Context Length = 16,384 2.50

2.50

1.22

Zoom: 0-20%

2.25

2.00

1.01

1.75 1.50

Dense

1.19

Zoom: 0-20%

2.25

0.79

Speedup

2.00

G=512

(b) Decode -- Batch Size = 1, Context Length = 16,384

1.25

1.50

0.86

1.25

1.00

0.75

0.50

1.02

1.75

0.50 0%

10%

20%

30%

40%

50%

60%

70%

Sparsity

10%

20%

30%

40%

50%

60%

70%

Sparsity

Figure 8: End-to-end speedup of WIDE over the dense Llama3.1-8B baseline across sparsity and group sizes.

4.4

ROUTING B EHAVIOR

To understand how WIDE allocate its compute budget, we inspect the learned routing patterns at both the layer and token levels. Figure 9 summarizes the routing distribution of the Llama3.1-8B checkpoint with 50% target sparsity and group size 128. Although the overall sparsity is close to the target (47.3%), the learned allocation is highly non-uniform. Attention carries most of the skipping, reaching 66.2% sparsity on WikiText2, while the FFN is more conservative at 28.5%. Across layers, both curves peak around layers 4–12 and 24–28, and dip near the model boundaries and layers 15–18. 10

Preprint

WIDE Attention

DDP Attention

WIDE avg.: 47.3%

Real sparsity (%)

WIDE FFN DDP FFN DDP avg.: 44.8% These profiles echo the layer sensitivity observed in Layer-wise real sparsity (Attn avg.: 66.2%, FFN avg.: 28.5%) static pruning studies (He et al., 2024; Huang et al., 100 2026), but are produced here by token-conditioned 80 group routing. For the token-wise routing behav60 ior, Figure 10 shows a representative case study on 40 HellaSwag. In here, he FFN routers keep semantic 20 content tokens such as boy, running, and track, 0 4 8 12 16 20 24 28 32 especially in layers 9–11 where the average FFN Layer sparsity is high, while articles like A, a, and the are often skipped. Similar behavior appears in the appendix cases across other tasks (Fig 14 to Fig 18). Figure 9: WIDE and DDP’s layer-wise sparThis analysis highlight the ability that WIDE can dy- sity distribution of Llama3.1-8B on Wikinamically allocate its budget across layer and tokens Text2, with target sparsity set to 50%. with semantic-awareness.

hellaswag case: hellaswag_3

Router layer

attn L24

attn L25 60

attn L26 ffn L1

ffn L9 20

ffn L10 ffn L11

Skip probability (%)

100

attn L2

y bo

e th

k tr

n w do

g ru

y bo

Question/context token

Figure 10: Token-wise routing probability for Llama3.1-8B in selected Hellaswag sample.

C ONCLUSION

In this work, we identify the potential limitation for current naive token-wise dynamic pruning pipeline under modern high-throughput inference system, and present a finer-grained group-based dynamic width pruning baseline with unified acceleration design for routing-based dynamic structural pruning. The presented WIDE uses lightweight bottleneck routers and differentiable masks to select GQA head groups and FFN parameter groups for each token, pushing structured pruning below the granularity of layers and sublayers. In the system side, we formulate mask-based dynamic pruning as two steps: mask reordering and progressive intra-kernel skipping. This abstraction lets the same routing decisions be consumed by customized attention and GEMM kernels with limited disruption to dense execution. Across calibration-only and LoRA-recovered settings, WIDE improves the quality-sparsity trade-off over existing structured static and dynamic pruning methods, while its kernels translate the learned sparsity into practical speedups under CUDA Graph execution for both batched prefill and decode workloads. These results suggest that fine-grained dynamic pruning needs to be designed together with the hardware execution path, rather than treated as a model-only compression problem.

R EFERENCES Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, and Sumit Sanghai. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. In The 2023 Conference on Empirical Methods in Natural Language Processing, December 2023. Saleh Ashkboos, Maximilian L. Croci, Marcelo Gennari do Nascimento, Torsten Hoefler, and James Hensman. SliceGPT: Compress Large Language Models by Deleting Rows and Columns. In The Twelfth International Conference on Learning Representations, October 2023. 11

Preprint

Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432–7439, 2020. Terry Chen, Zhifan Ye, Bing Xu, Zihao Ye, Timmy Liu, Ali Hassani, Tianqi Chen, Andrew Kerr, Haicheng Wu, Yang Xu, Yu-Jung Chen, Hanfeng Chen, Aditya Kane, Ronny Krashinsky, MingYu Liu, Vinod Grover, Luis Ceze, Roger Bringmann, John Tran, Wei Liu, Fung Xie, Michael Lightstone, and Humphrey Shi. AVO: Agentic Variation Operators for Autonomous Evolutionary Search. http://arxiv.org/abs/2603.24517, March 2026. Hongrong Cheng, Miao Zhang, and Javen Qinfeng Shi. A survey on deep neural network pruningtaxonomy, comparison, analysis, and recommendations. http://arxiv.org/abs/2308.06767, August 2024. Hung-Yueh Chiang, Bokun Wang, and Diana Marculescu. ELANA: A Simple Energy and Latency Analyzer for LLMs. http://arxiv.org/abs/2512.09946, December 2025. Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint arXiv:1905.10044, 2019. Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457, 2018. Weinan Dai, Hanlin Wu, Qiying Yu, Huan-ang Gao, Jiahao Li, Chengquan Jiang, Weiqiang Lou, Yufan Song, Hongli Yu, Jiaze Chen, Wei-Ying Ma, Ya-Qin Zhang, Jingjing Liu, Mingxuan Wang, Xin Liu, and Hao Zhou. CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation. http://arxiv.org/abs/2602.24286, February 2026. Tri Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. In The Twelfth International Conference on Learning Representations, October 2023. DeepSeek-AI. DeepSeek-V4: Towards highly efficient million-token context intelligence. https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro, 2026. Xuan Ding, Pengyu Tong, Ranjie Duan, Yunjian Zhang, Rui Sun, and Yao Zhu. Pruning as a Cooperative Game: Surrogate-Assisted Layer Contribution Estimation for Large Language Models. In The Fourteenth International Conference on Learning Representations, October 2025. Gongfan Fang, Hongxu Yin, Saurav Muralidharan, Greg Heinrich, Jeff Pool, Jan Kautz, Pavlo Molchanov, and Xinchao Wang. MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, November 2024. Elias Frantar and Dan Alistarh. SparseGPT: Massive Language Models Can be Accurately Pruned in One-Shot. In Proceedings of the 40th International Conference on Machine Learning, pp. 10323–10337. PMLR, July 2023. Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. The language model evaluation harness, 07 2024. URL https://zenodo.org/records/12608602. Aayush Gautam, Mukul Gagrani, Junyoung Park, Mingu Lee, Chiris Lott, and Narasimha Reddy. Fast Forward: Accelerating LLM Prefill with Predictive FFN Sparsity. https://arxiv.org/abs/2602.00397v1, January 2026. Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava 12

Preprint

Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, Danny Wyatt, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Francisco Guzmán, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Govind Thattai, Graeme Nail, Gregoire Mialon, Guan Pang, Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Evtimov, Jack Zhang, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Alwala, Karthik Prasad, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Kushal Lakhotia, Lauren Rantala-Yeary, Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin, Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira, Madeline Muzzi, Mahesh Pasupuleti, Mannat Singh, Manohar Paluri, Marcin Kardas, Maria Tsimpoukelli, Mathew Oldham, Mathieu Rita, Maya Pavlova, Melanie Kambadur, Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogoychev, Niladri Chatterji, Ning Zhang, Olivier Duchenne, Onur Çelebi, Patrick Alrassy, Pengchuan Zhang, Pengwei Li, Petar Vasic, Peter Weng, Prajjwal Bhargava, Pratik Dubal, Praveen Krishnan, Punit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong, Ragavan Srinivasan, Raj Ganapathy, Ramon Calderer, Ricardo Silveira Cabral, Robert Stojnic, Roberta Raileanu, Rohan Maheswari, Rohit Girdhar, Rohit Patel, Romain Sauvestre, Ronnie Polidoro, Roshan Sumbaly, Ross Taylor, Ruan Silva, Rui Hou, Rui Wang, Saghar Hosseini, Sahana Chennabasappa, Sanjay Singh, Sean Bell, Seohyun Sonia Kim, Sergey Edunov, Shaoliang Nie, Sharan Narang, Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhosale, Shun Zhang, Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarek Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher, Todor Mihaylov, Tong Xiao, Ujjwal Karn, Vedanuj Goswami, Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do, Vish Vogeti, Vítor Albiero, Vladan Petrovic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whitney Meers, Xavier Martinet, Xiaodong Wang, Xiaofang Wang, Xiaoqing Ellen Tan, Xide Xia, Xinfeng Xie, Xuchao Jia, Xuewei Wang, Yaelle Goldschlag, Yashesh Gaur, Yasmine Babaei, Yi Wen, Yiwen Song, Yuchen Zhang, Yue Li, Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan, Zhengxing Chen, Zoe Papakipos, Aaditya Singh, Aayushi Srivastava, Abha Jain, Adam Kelsey, Adam Shajnfeld, Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alexei Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Amos Teo, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Dong, Annie Franco, Anuj Goyal, Aparajita Saraf, Arkabandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisenman, Azadeh Yazdan, Beau James, Ben Maurer, Benjamin Leonhardi, Bernie Huang, Beth Loyd, Beto De Paola, Bhargavi Paranjape, Bing Liu, Bo Wu, Boyu Ni, Braden Hancock, Bram Wasti, Brandon Spence, Brani Stojkovic, Brian Gamido, Britt Montalvo, Carl Parker, Carly Burton, Catalina Mejia, Ce Liu, Changhan Wang, Changkyu Kim, Chao Zhou, Chester Hu, Ching-Hsiang Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer, Cynthia Gao, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang, Duc Le, Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery, Eleonora Presani, Emily Hahn, Emily Wood, Eric-Tuan Le, Erik Brinkman, Esteban Arcaute, Evan Dunbar, Evan Smothers, Fei Sun, Felix Kreuk, Feng Tian, Filippos Kokkinos, Firat Ozgenel, Francesco Caggioni, Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpern, Grant Herman, Grigory Sizov, Guangyi, Zhang, Guna Lakshminarayanan, Hakan Inan, Hamid Shojanazeri, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harrison Rudolph, Helen Suk, Henry Aspegren, Hunter Goldman, Hongyuan Zhan, Ibrahim Damlaj, Igor Molybog, Igor Tufanov, Ilias Leontiadis, Irina-Elena Veliche, Itai Gat, Jake Weissman, James Geboski, James Kohli, Janice Lam, Japhet Asher, Jean-Baptiste Gaya, Jeff Marcus, Jeff Tang, Jennifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Teboul, Jessica Zhong, Jian Jin, Jingyi Yang, Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie, Jonathan Torres, Josh Ginsburg, Junjie Wang, Kai Wu, Kam Hou U, Karan Saxena, Kartikay Khandelwal, 13

Preprint

Katayoun Zand, Kathy Matosich, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li, Kiran Jagadeesh, Kun Huang, Kunal Chawla, Kyle Huang, Lailin Chen, Lakshya Garg, Lavender A, Leandro Silva, Lee Bell, Lei Zhang, Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manav Avalani, Manish Bhatt, Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Keneally, Miao Liu, Michael L. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vyatskov, Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang, Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari, Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White, Navyata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikhil Mehta, Nikolay Pavlovich Laptev, Ning Dong, Norman Cheng, Oleg Chernoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli, Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji, Pedro Rittner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Pritish Yuvraj, Qian Liang, Rachad Alao, Rachel Rodriguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Rangaprabhu Parthasarathy, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang, Russ Howes, Ruty Rinott, Sachin Mehta, Sachin Siby, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Satadru Pan, Saurabh Mahajan, Saurabh Verma, Seiji Yamamoto, Sharadh Ramaswamy, Shaun Lindsay, Shaun Lindsay, Sheng Feng, Shenghao Lin, Shengxin Cindy Zha, Shishir Patil, Shiva Shankar, Shuqiang Zhang, Shuqiang Zhang, Sinong Wang, Sneha Agarwal, Soji Sajuyigbe, Soumith Chintala, Stephanie Max, Stephen Chen, Steve Kehoe, Steve Satterfield, Sudarshan Govindaprasad, Sumit Gupta, Summer Deng, Sungmin Cho, Sunny Virk, Suraj Subramanian, Sy Choudhury, Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo Koehler, Thomas Robinson, Tianhe Li, Tianjun Zhang, Tim Matthews, Timothy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar, Vishal Mangla, Vlad Ionescu, Vlad Poenaru, Vlad Tiberiu Mihailescu, Vladimir Ivanov, Wei Li, Wenchen Wang, Wenwen Jiang, Wes Bouaziz, Will Constable, Xiaocheng Tang, Xiaojian Wu, Xiaolan Wang, Xilun Wu, Xinbo Gao, Yaniv Kleinman, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi, Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu, Wang, Yu Zhao, Yuchen Hao, Yundi Qian, Yunlu Li, Yuzi He, Zach Rait, Zachary DeVito, Zef Rosnbrick, Zhaoduo Wen, Zhenyu Yang, Zhiwei Zhao, and Zhiyu Ma. The Llama 3 Herd of Models. http://arxiv.org/abs/2407.21783, November 2024. Chao Han, Yijuan Liang, Zihao Xuan, Daokuan Wu, Wei Zhang, and Xiaoyu Shen. Informed Routing in LLMs: Smarter Token-Level Computation for Faster Inference. http://arxiv.org/abs/2510.13831, October 2025. Shwai He, Guoheng Sun, Zheyu Shen, and Ang Li. What Matters in Transformers? Not All Attention is Needed. http://arxiv.org/abs/2406.15786, October 2024. Zicheng He, Anhao Zhao, Xiaoyu Shen, Chen Wu, and Lei He. SkipOPU: An FPGA-based Overlay Processor for Large Language Models with Dynamically Allocated Computation. https://arxiv.org/abs/2603.14785v1, March 2026. Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-Rank Adaptation of Large Language Models. In International Conference on Learning Representations, October 2021. Haozhe Hu, Hao Wu, Anhao Zhao, Longwei Ding, Peiran Yin, Yunpu Ma, and Xiaoyu Shen. Beyond FLOPs: Benchmarking Real Inference Acceleration of LLM Pruning under a GEMM-Centric Taxonomy. http://arxiv.org/abs/2606.09080, June 2026. Weiyu Huang, Pengle Zhang, Xiaolu Zhang, Jun Zhou, Jun Zhu, and Jianfei Chen. Deterministic Differentiable Structured Pruning for Large Language Models. http://arxiv.org/abs/2603.08065, May 2026. Yikun Jiang, Huanyu Wang, Lei Xie, Hanbin Zhao, Chao Zhang, Hui Qian, and John C. S. Lui. D-LLM: A Token Adaptive Computing Resource Allocation Strategy for Large Language Models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, November 2024. Bo-Kyeong Kim, Geonmin Kim, Tae-Ho Kim, Thibault Castells, Shinkook Choi, Junho Shin, and Hyoung-Kyu Song. Shortened LLaMA: Depth Pruning for Large Language Models with Comparison of Retraining Methods. http://arxiv.org/abs/2402.02834, June 2024. 14

Preprint

Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. Efficient Memory Management for Large Language Model Serving with PagedAttention. http://arxiv.org/abs/2309.06180, September 2023. Guanchen Li, Yixing Xu, Zeping Li, Ji Liu, Xuanwu Yin, Dong Li, and Emad Barsoum. Týrthe-Pruner: Structural Pruning LLMs via Global Sparsity Distribution Optimization. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, October 2025. Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Wei-Ming Chen, Wei-Chen Wang, Guangxuan Xiao, Xingyu Dang, Chuang Gan, and Song Han. AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration. Proceedings of Machine Learning and Systems, 6:87–100, May 2024. Xinyin Ma, Gongfan Fang, and Xinchao Wang. LLM-Pruner: On the Structural Pruning of Large Language Models. In Thirty-Seventh Conference on Neural Information Processing Systems, November 2023. Xin Men, Mingyu Xu, Qingyu Zhang, Qianhao Yuan, Bingning Wang, Hongyu Lin, Yaojie Lu, Xianpei Han, and Weipeng Chen. ShortGPT: Layers in Large Language Models are More Redundant Than You Expect. In Findings of the Association for Computational Linguistics: ACL 2025, pp. 20192–20204, Vienna, Austria, July 2025. Association for Computational Linguistics. ISBN 979-8-89176-256-5. doi: 10.18653/v1/2025.findings-acl.1035. Qianyu Meng, Yanan Wang, Liyi Chen, Wei Wu, Yihang Li, Wenyuan Jiang, Qimeng Wang, Chengqiang Lu, Yan Gao, Yi Wu, and Yao Hu. Agent harness for large language model agents: A survey. 2026. doi: 10.20944/preprints202604.0428.v3. Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv:1609.07843, 2016. Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv:1809.02789, 2018. NVIDIA. Tensorrt-llm. https://github.com/NVIDIA/TensorRT-LLM, 2026. David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, and Adam Santoro. Mixture-of-Depths: Dynamically allocating compute in transformer-based language models. http://arxiv.org/abs/2404.02258, April 2024. Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99–106, 2021. Susav Shrestha, Bradley Settlemyer, Nikoli Dryden, and A. L. Narasimha Reddy. Polar Sparsity: High Throughput Batched LLM Inferencing with Scalable Contextual Sparsity. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, October 2025. Mingjie Sun, Zhuang Liu, Anna Bair, and J. Zico Kolter. A Simple and Effective Pruning Approach for Large Language Models. In The Twelfth International Conference on Learning Representations, October 2023. Kimi Team, Tongtong Bai, Yifan Bai, Yiping Bao, S. H. Cai, Yuan Cao, Y. Charles, H. S. Che, Cheng Chen, Guanduo Chen, Huarong Chen, Jia Chen, Jiahao Chen, Jianlong Chen, Jun Chen, Kefan Chen, Liang Chen, Ruijue Chen, Xinhao Chen, Yanru Chen, Yanxu Chen, Yicun Chen, Yimin Chen, Yingjiang Chen, Yuankun Chen, Yujie Chen, Yutian Chen, Zhirong Chen, Ziwei Chen, Dazhi Cheng, Minghan Chu, Jialei Cui, Jiaqi Deng, Muxi Diao, Hao Ding, Mengfan Dong, Mengnan Dong, Yuxin Dong, Yuhao Dong, Angang Du, Chenzhuang Du, Dikang Du, Lingxiao Du, Yulun Du, Yu Fan, Shengjun Fang, Qiulin Feng, Yichen Feng, Garimugai Fu, Kelin Fu, Hongcheng Gao, Tong Gao, Yuyao Ge, Shangyi Geng, Chengyang Gong, Xiaochen Gong, Zhuoma Gongque, Qizheng Gu, Xinran Gu, Yicheng Gu, Longyu Guan, Yuanying Guo, Xiaoru Hao, Weiran He, Wenyang He, Yunjia He, Chao Hong, Hao Hu, Jiaxi Hu, Yangyang Hu, Zhenxing Hu, Ke Huang, Ruiyuan Huang, Weixiao Huang, Zhiqi Huang, Tao Jiang, Zhejun Jiang, Xinyi 15

Preprint

Jin, Yu Jing, Guokun Lai, Aidi Li, C. Li, Cheng Li, Fang Li, Guanghe Li, Guanyu Li, Haitao Li, Haoyang Li, Jia Li, Jingwei Li, Junxiong Li, Lincan Li, Mo Li, Weihong Li, Wentao Li, Xinhang Li, Xinhao Li, Yang Li, Yanhao Li, Yiwei Li, Yuxiao Li, Zhaowei Li, Zheming Li, Weilong Liao, Jiawei Lin, Xiaohan Lin, Zhishan Lin, Zichao Lin, Cheng Liu, Chenyu Liu, Hongzhang Liu, Liang Liu, Shaowei Liu, Shudong Liu, Shuran Liu, Tianwei Liu, Tianyu Liu, Weizhou Liu, Xiangyan Liu, Yangyang Liu, Yanming Liu, Yibo Liu, Yuanxin Liu, Yue Liu, Zhengying Liu, Zhongnuo Liu, Enzhe Lu, Haoyu Lu, Zhiyuan Lu, Junyu Luo, Tongxu Luo, Yashuo Luo, Long Ma, Yingwei Ma, Shaoguang Mao, Yuan Mei, Xin Men, Fanqing Meng, Zhiyong Meng, Yibo Miao, Minqing Ni, Kun Ouyang, Siyuan Pan, Bo Pang, Yuchao Qian, Ruoyu Qin, Zeyu Qin, Jiezhong Qiu, Bowen Qu, Zeyu Shang, Youbo Shao, Tianxiao Shen, Zhennan Shen, Juanfeng Shi, Lidong Shi, Shengyuan Shi, Feifan Song, Pengwei Song, Tianhui Song, Xiaoxi Song, Hongjin Su, Jianlin Su, Zhaochen Su, Lin Sui, Jinsong Sun, Junyao Sun, Tongyu Sun, Flood Sung, Yunpeng Tai, Chuning Tang, Heyi Tang, Xiaojuan Tang, Zhengyang Tang, Jiawen Tao, Shiyuan Teng, Chaoran Tian, Pengfei Tian, Ao Wang, Bowen Wang, Chensi Wang, Chuang Wang, Congcong Wang, Dingkun Wang, Dinglu Wang, Dongliang Wang, Feng Wang, Hailong Wang, Haiming Wang, Hengzhi Wang, Huaqing Wang, Hui Wang, Jiahao Wang, Jinhong Wang, Jiuzheng Wang, Kaixin Wang, Linian Wang, Qibin Wang, Shengjie Wang, Shuyi Wang, Si Wang, Wei Wang, Xiaochen Wang, Xinyuan Wang, Yao Wang, Yejie Wang, Yipu Wang, Yiqin Wang, Yucheng Wang, Yuzhi Wang, Zhaoji Wang, Zhaowei Wang, Zhengtao Wang, Zhexu Wang, Zihan Wang, Zizhe Wang, Chu Wei, Ming Wei, Chuan Wen, Zichen Wen, Chengjie Wu, Haoning Wu, Junyan Wu, Rucong Wu, Wenhao Wu, Yuefeng Wu, Yuhao Wu, Yuxin Wu, Zijian Wu, Chenjun Xiao, Jin Xie, Xiaotong Xie, Yuchong Xie, Yifei Xin, Bowei Xing, Boyu Xu, Jianfan Xu, Jing Xu, Jinjing Xu, L. H. Xu, Lin Xu, Suting Xu, Weixin Xu, Xinbo Xu, Xinran Xu, Yangchuan Xu, Yichang Xu, Yuemeng Xu, Zelai Xu, Ziyao Xu, Junjie Yan, Yuzi Yan, Guangyao Yang, Hao Yang, Junwei Yang, Kai Yang, Ningyuan Yang, Ruihan Yang, Xiaofei Yang, Xinlong Yang, Ying Yang, Yi Yang, Yi Yang, Zhen Yang, Zhilin Yang, Zonghan Yang, Haotian Yao, Dan Ye, Wenjie Ye, Zhuorui Ye, Bohong Yin, Chengzhen Yu, Longhui Yu, Tao Yu, Tianxiang Yu, Enming Yuan, Mengjie Yuan, Xiaokun Yuan, Yang Yue, Weihao Zeng, Dunyuan Zha, Haobing Zhan, Dehao Zhang, Hao Zhang, Jin Zhang, Puqi Zhang, Qiao Zhang, Rui Zhang, Xiaobin Zhang, Y. Zhang, Yadong Zhang, Yangkun Zhang, Yichi Zhang, Yizhi Zhang, Yongting Zhang, Yu Zhang, Yushun Zhang, Yutao Zhang, Yutong Zhang, Zheng Zhang, Chenguang Zhao, Feifan Zhao, Jinxiang Zhao, Shuai Zhao, Xiangyu Zhao, Yikai Zhao, Zijia Zhao, Huabin Zheng, Ruihan Zheng, Shaojie Zheng, Tengyang Zheng, Junfeng Zhong, Longguang Zhong, Weiming Zhong, M. Zhou, Runjie Zhou, Xinyu Zhou, Zaida Zhou, Jinguo Zhu, Liya Zhu, Xinhao Zhu, Yuxuan Zhu, Zhen Zhu, Jingze Zhuang, Weiyu Zhuang, Ying Zou, and Xinxing Zu. Kimi K2.5: Visual Agentic Intelligence. http://arxiv.org/abs/2602.02276, February 2026. Philippe Tillet, H. T. Kung, and David Cox. Triton: An intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, MAPL 2019, pp. 10–19, New York, NY, USA, June 2019. Association for Computing Machinery. ISBN 978-1-4503-6719-6. doi: 10.1145/3315508.3329973. Lei Wang, Yu Cheng, Yining Shi, Zhiwen Mo, Zhengju Tang, Wenhao Xie, Tong Wu, Lingxiao Ma, Yuqing Xia, Jilong Xue, Fan Yang, and Zhi Yang. TileLang: Bridge Programmability and Performance in Modern Neural Kernels. In The Fourteenth International Conference on Learning Representations, October 2025. Maurice Weber, Daniel Y. Fu, Quentin Anthony, Yonatan Oren, Shane Adams, Anton Alexandrov, Xiaozhong Lyu, Huu Nguyen, Xiaozhe Yao, Virginia Adams, Ben Athiwaratkun, Rahul Chalamala, Kezhen Chen, Max Ryabinin, Tri Dao, Percy Liang, Christopher Ré, Irina Rish, and Ce Zhang. Redpajama: an open dataset for training large language models. NeurIPS Datasets and Benchmarks Track, 2024. Hao Wu, Yingqi Fan, Dai Jinyang, Junlong Tong, Yunpu Ma, and Xiaoyu Shen. HiDrop: Hierarchical Vision Token Reduction in MLLMs via Late Injection, Concave Pyramid Pruning, and Early Exit. In The Fourteenth International Conference on Learning Representations, October 2025. Hao Wu, Junlong Tong, Xudong Wang, Yang Tan, Changyu Zeng, Anastasia Antsiferova, and Xiaoyu Shen. From data to model: A survey of the compression lifecycle in mllms. February 2026a. doi: 10.36227/techrxiv.177220375.55495124/v1. URL http://dx.doi.org/10.36227/ techrxiv.177220375.55495124/v1. 16

Preprint

Hao Wu, Xudong Wang, Jialiang Zhang, Junlong Tong, Xinghao Chen, Junyan Lin, Yunpu Ma, and Xiaoyu Shen. UTPTrack: Towards Simple and Unified Token Pruning for Visual Tracking. http://arxiv.org/abs/2602.23734, February 2026b. Zihao Ye, Lequn Chen, Ruihang Lai, Wuwei Lin, Yineng Zhang, Stephanie Wang, Tianqi Chen, Baris Kasikci, Vinod Grover, Arvind Krishnamurthy, and Luis Ceze. FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. In Eighth Conference on Machine Learning and Systems, May 2025. Ted Zadouri, Markus Hoehnerbach, Jay Shah, Timmy Liu, Vijay Thakkar, and Tri Dao. FlashAttention-4: Algorithm and kernel pipelining co-design for asymmetric hardware scaling, 2026. Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv:1905.07830, 2019. Anhao Zhao, Fanghua Ye, Yingqi Fan, Junlong Tong, Jing Xiong, Zhiwei Fei, Hui Su, and Xiaoyu Shen. SkipGPT: Each Token is One of a Kind. In Forty-Second International Conference on Machine Learning, June 2025. Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Sun, Jeff Huang, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E. Gonzalez, Clark Barrett, and Ying Sheng. SGLang: Efficient Execution of Structured Language Model Programs. http://arxiv.org/abs/2312.07104, June 2024. Longguang Zhong, Fanqi Wan, Ruijun Chen, Xiaojun Quan, and Liangzhi Li. BlockPruner: Finegrained Pruning for Large Language Models. In Findings of the Association for Computational Linguistics: ACL 2025, pp. 5065–5080, Vienna, Austria, July 2025. Association for Computational Linguistics. ISBN 979-8-89176-256-5. doi: 10.18653/v1/2025.findings-acl.262.

Preprint

T HE C OST FOR NAIVE G ATHER -S CATTER I MPLEMENTATIONS

In this section, we first provide a narrow, theoretical memory accounting for naive gather-scatter implementations of M-axis, N-axis, and K-axis dynamic pruning. We count logical data-element transfers caused by materializing and consuming the explicit data tensors in the PyTorch-style implementations below. Routing masks, indices, offsets, and their construction, storage, and traffic are excluded, as are hardware-specific cache effects, and framework workspace. We assume that A, B, and C use the same data type, so byte traffic is obtained by multiplying TM , TN , or TK by the data-type size. Consider a dense GEMM C = AB⊤ ,

A ∈ RM ×K ,

B ∈ RN ×K ,

C ∈ RM ×N

(10)

The ideal dense traffic is Tdense ≈ M K + N K + M N

(11)

which corresponding to reading A, reading B, and writing C. Naive M-axis pruning. M-axis pruning (i.e., dynamic depth pruning) selects a subset of token rows for the current operator. Let M ′ = pM be the number of active rows, where p is the active ratio, and ′ let IM ∈ ZM denote the active row indices. A naive PyTorch-style implementation gathers active rows, computes the compact GEMM, and scatters the result back like: # A: [M, K] # B: [N, K] # I_M: [M’] A_sub = A[I_M, :] C_sub = A_sub @ B.T

# [M’, K] # [M’, N]

C = torch.zeros(M, N, device=A.device, dtype=A.dtype) C[I_M, :] = C_sub # scatter back to dense output

When ignoring index traffic, the total memory traffic is TM ≈ 3M ′ K + N K + 3M ′ N + M N = 3pM K + N K + 3pM N + M N

(12)

The terms 3M ′ K and 3M ′ N come from materializing and rereading Asub and Csub , respectively, while M N comes from initializing the dense output tensor. The additional storage, excluding the final dense output C, is SM ≈ M ′ K + M ′ N = pM K + pM N

(13)

Naive N-axis pruning. N-axis pruning selects different N -axis groups for each token. Divide the N -axis into groups of size G. If each token selects R groups, the active N -axis extent is N ′ = RG = pN

(14)

where p is the active ratio along the N -axis. Let IN ∈ ZM ×R be the per-token selected group indices. The fully vectorized PyTorch-style implementation below first expands group indices and materializes token-specific subsets of B: # A: [M, K] # B: [N, K] # I_N: [M, R], group indices # G: group size # N_prime = R * G offset = torch.arange(G, device=A.device)

# [G]

J_N = I_N[..., None] * G + offset

# [M, R, G]

Preprint

J_N = J_N.reshape(M, R * G)

# [M, N’]

B_sub = B[J_N, :]

# [M, N’, K]

C_sub = torch.einsum("mk,mnk->mn", A, B_sub)

# [M, N’]

C = torch.zeros(M, N, device=A.device, dtype=A.dtype) C.scatter_(dim=1, index=J_N, src=C_sub)

# [M, N]

Ignoring index traffic and the small offset tensor, the overall memory traffic is TN ≈ M K + 3M N ′ K + 3M N ′ + M N = M K + 3pM N K + 3pM N + M N

(15)

The dominant term 3M N ′ K comes from reading selected rows of B, writing the gathered tensor Bsub , and reading Bsub again during token-wise dot products. Unlike dense GEMM, where B is shared across all M rows and read once as N K, the naive N-axis implementation replicates selected weight groups for each token. The additional storage is SN ≈ M N ′ K + M N ′ = pM N K + pM N

(16) ′

which is dominated by the per-token gathered weight tensor Bsub ∈ RM ×N ×K , where the remaining M N ′ entries belong to Csub . Naive K-axis pruning. For attention output and FFN down projections, K-axis pruning is applied over the reduction dimension K. Divide the K-axis into groups of size G. If each token selects R groups, the active K-axis extent is K ′ = RG = pK

(17)

where p is the active ratio along the K-axis. Let IK ∈ ZM ×R denote the per-token selected K-group indices. A token-specific PyTorch-style implementation can be written as: # A: [M, K] # B: [N, K] # I_K: [M, R], group indices # G: group size # K_prime = R * G offset = torch.arange(G, device=A.device)

# [G]

J_K = I_K[..., None] * G + offset J_K = J_K.reshape(M, R * G)

# [M, R, G] # [M, K’]

A_sub = torch.gather(A, dim=1, index=J_K) B_sub = B[:, J_K].permute(1, 0, 2)

# [M, K’] # [M, N, K’]

C = torch.einsum("mk,mnk->mn", A_sub, B_sub)

# [M, N]

Ignoring index traffic and the small offset tensor, the total memory traffic is TK ≈ 3M K ′ + 3M N K ′ + M N = 3pM K + 3pM N K + M N

(18)

The term 3M K ′ comes from reading selected entries of A, writing Asub , and reading Asub again during the token-wise dot product. The dominant term 3M N K ′ comes from reading selected entries of B, writing the token-specific tensor Bsub , and reading Bsub again during the einsum. The additional materialized data-tensor storage is SK ≈ M N K ′ + M K ′ = pM N K + pM K 19

(19)

Preprint

′

which is dominated by the token-specific gathered weight tensor Bsub ∈ RM ×N ×K , where the remaining M K ′ entries belong to Asub . These estimates highlight why the fully vectorized N-axis and K-axis implementations above are especially inefficient: they materialize token-specific subsets of B, incurring O(pM N K) intermediate storage and destroying the regular weight reuse of dense GEMM, which highlight the role of customized fuse kernel. The results are summarized in Table 3. Pattern Dense GEMM Dynamic pruning (M-axis) Dynamic pruning (N-axis) Dynamic pruning (K-axis)

Naive data-tensor traffic

Extra intermediate storage

MK + NK + MN 3pM K + N K + 3pM N + M N M K + 3pM N K + 3pM N + M N 3pM K + 3pM N K + M N

– pM K + pM N pM N K + pM N pM N K + pM K

Table 3: Theoretical data-tensor traffic and additional materialized intermediate storage of naive gather-scatter implementations. Routing masks, indices, and offsets are excluded, the storage calculation also excludes the final dense output C. Building on the theoretical analysis, we further benchmark the throughput and peak memory usage of the dense GEMM baseline, naive gather-scatter kernel, and optimized CuTe GEMM kernel for pruning along the M, N, and K axes. Since the naive implementation can be compatible with groupedGEMM by mask-reordering preprocessing, we also introduce it as a stronger baseline. The results are shown in Figure 11. For the coarsest-grained M-axis pruning setting, the explicit gather-scatter kernel still delivers a meaningful speedup but incurs an additional memory overhead of roughly 2x. For the finer-grained GEMM-N and GEMM-K settings, however, the naive implementation is constrained by the massive additional memory traffic, achieving only about 1.0 TFLOPS and requiring roughly 1,000x the memory footprint, where the optimized grouped-GEMM kernel still requires 20x more memory footprint and only achieves 50∼100 TFLOPS, failing to reach the dense GEMM baseline. In contrast, the fused CuTe kernel incurs virtually no additional memory-storage overhead and can already deliver speedups for matrix of size 20483 .

C U T E K ERNEL D ETAILS

At the abstraction level, all three CuTe kernels follow the same execution path as the framework in the main text: (1) standalone mask-reordering preprocessing, (2) hardware-agnostic block-level skipping at stage ℓ = 0, (3) a dense-style tiled pipeline with optional hardware-dependent intra-block skipping during operand gathering (ℓ = 1, corresponding to A loading in GEMM and Q loading in attention) and MMA execution (ℓ = 2), and finally, (4) a scatter epilogue. The latter two predicates are enabled only when selected by the architecture-specific kernel configuration. Herein, we first describe the metadata shared by all three kernels. A routing column is M:,g after flattening the batch and token dimensions for GEMM, and Mb,:,g for flash attention. As summarized in Algorithm 1, each routing column is sorted independently on the GPU. In this way, the resulting fj,g = MI ,g . Activation tensors index Ij,g maps logical row j back to its physical row, i.e., M j,g therefore remain in their original layout and are gathered/scattered directly through the reordered indices. Consecutive operators governed by the same routing decision can reuse this metadata, although zero initialization remains necessary for skipped blocks and atomic-add epilogues.

Preprint

400.8

400

TFLOPS

314.4 316.1

300

285.1

212.7

200

168.3 121.5

100 29.0

156.1 154.5

99.5 29.3

10243

20483

40963

TFLOPS

400

372.6

300

260.2 212.7

200

168.3 121.5

101.9

92.0

100

32.1 1.5

6.4

OOM

1.6

10243

20483

40963

(e) GEMM-K — Throughput 400

TFLOPS

317.9

300 246.3 212.7

200

168.3 121.5 90.6

100

65.1 33.3

1.1

6.4

10243

1.1

20483

OOM

40963

Incremental memory (GiB)

Grouped GEMM

Incremental memory (GiB)

Naive

Incremental memory (GiB)

Dense

(a) GEMM-M — Throughput

Matrix size (M = N = K)

CuTe kernel

OOM (32 GiB cap)

(b) GEMM-M — Peak memory 32 GiB test limit

10 1 1e-1

6.25e-2 6.25e-2 3.13e-2

3.12e-2 1.56e-2 1.56e-2

1e-2

7.84e-3

7.81e-3 3.91e-3 3.91e-3 1.97e-3

1.95e-3

1e-3

10243

20483

40963

(d) GEMM-N — Peak memory OOM 8.03

10 1.01

5.48e-1

1e-1

7.46e-2 3.24e-2

3.12e-2 1.08e-2

1e-2

2.03e-3

1.95e-3

1e-3

8.10e-3

7.81e-3

10243

20483

40963

(f) GEMM-K — Peak memory OOM 8.03

1.05

1.01

1.37e-1

1e-1

3.24e-2

3.12e-2 1.86e-2

1e-2 2.03e-3

1.95e-3

1e-3

8.10e-3

7.81e-3

10243

20483

40963

Matrix size (M = N = K)

Figure 11: Throughput and memory overhead benchmarks for M,N,K axis GEMM pruning based on dense GEMM baseline, naive gather-scatter kernel, grouped-GEMM kernel, and fused CuTe kernel at G = 128 and 50% sparsity.

Algorithm 1 Unified mask-reordering preprocessing Require: Binary routing mask M with NG groups f and gather/scatter index I Ensure: Reordered mask M 1: Form routing columns: flatten (B, T ) into M for GEMM; retain one length-Tq column per (b, g) for attention 2: for all routing columns x in parallel do 3: (e x, Ix ) ← S ORT D ESC(x) ▷ active entries form a prefix 4: end for f I) 5: return (M,

Preprint

For a sorted CTA row tile R0 = [m0 , min(m0 + BM, M )), the block-level predicate is simply fm ,g . Because the column is monotone, s(0) = 0 proves that the whole CTA and every s(0) = M 0 later CTA for this group are inactive. Conversely, if a full tile’s last entry is active, the kernel takes an all-active fast path and avoids loading per-row masks. For a load packet u spanning Sld rows and an MMA row fragment v spanning Smma rows, the finer predicates are _ _ fj,g , fj,g . s(1) M s(2) M (20) u = v = j∈Rld (u)

j∈Rmma (v)

The ℓ = 1 predicate is instantiated per thread or per warp according to the CuTe copy atom. The ℓ = 2 predicate is reduced with a warp vote instruction and therefore keeps uniform across all lanes issuing the same tensor-core instruction. In the pseudocode below, P REDICATES returns (s(0) , s(1) , s(2) , J ), where J contains the original row indices needed by the active packets. We omit CuTe layout construction, register-fragment permutations, and architecture-specific barrier phases that do not change the execution schedule. P REFETCH MN and P REFETCH K denote semantic producer operations: they conditionally gather active A packets according to s(1) , stage the dense B tile, and commit one stage. WAIT S TAGE and R ELEASE S TAGE denote the corresponding consumer operations.

Implementation of GEMM-MN GEMM-MN is the shared kernel for M-axis and N-axis pruning. Its N-axis path applies routing to the attention Q and FFN up/gate projections, while M-axis dynamic depth pruning is the NG = 1 special case of the same execution path. Let C = AB⊤ , with A ∈ RM ×K and B ∈ RN ×K , and partition N into groups of width G. Each CTA owns one (BM, BN ) output tile, so its N coordinate uniquely selects a routing group g that is selected by BM tokens. Algorithm 2 makes the CTA admission, prologue, pipelined mainloop, and epilogue path explicit. Algorithm 2 GEMM-MN multi-stage pipeline for an output tile (m0 , n0 ) f I), output tile (m0 , n0 ), group-aligned tile sizes Require: A, B, zero-initialized C, (M, (BM, BN, BK) with BN | G, and P pipeline stages Block-level admission (ℓ = 0) 1: g ← ⌊n0 /G⌋; R0 ← [m0 , min(m0 + BM, M )) 2: (s(0) , s(1) , s(2) , J ) ← P REDICATES(R0 , g) 3: if ¬s(0) then 4: return ▷ ℓ = 0: CTA early exit 5: end if Prologue 6: Cacc ← 0; NK ← ⌈K/BK⌉ 7: for q = 0, . . . , min(P − 1, NK ) − 1 do 8: P REFETCH MN(q, q mod P, s(1) , J ) ▷ ℓ = 1: predicated A loading 9: end for Pipelined mainloop 10: for kc = 0, . . . , NK − 1 do 11: s ← kc mod P ; WAIT S TAGE(s) 12: q ← kc + P − 1 13: if q < NK then ▷ refill one free stage 14: P REFETCH MN(q, q mod P, s(1) , J ) 15: end if (2) 16: for all MMA row fragments v with sv = 1 do 17: Cacc [v] ← MMA(Akc [v], Bkc , Cacc [v]) ▷ ℓ = 2: predicated MMA 18: end for 19: R ELEASE S TAGE(s) 20: end for Epilogue 21: Scatter active rows of Cacc to C[J , n0 : min(n0 + BN, N )] ▷ atomic add under split-K The important pipeline invariant is that each logical K tile still commits exactly one asynchronouscopy group: sparse predicates suppress only A load packets, while the dense B copy keeps the 22

Preprint

producer stage well-defined. Split-K variants use the same mainloop and atomically reduce the scattered partial outputs. For the GEMV (General Matrix-Vector Multiplication) variant, it can be dispatched in M ⩽ 4. Implementation of GEMM-K GEMM-K implements the K-axis path for the attention output and FFN down projections. Here the routed dimension is the reduction dimension: Cm,n =

NX G −1 g=0

Mm,g

Am,k Bn,k ,

k∈Kg

|Kg | = G.

(21)

Because different groups have different row permutations, each active group computes a partial output and scatters it with an atomic reduction, as shown in Algorithm 3. Algorithm 3 GEMM-K with multi-stage group-local pipelines f I), output tile (m0 , n0 ), group-aligned tile sizes Require: A, B, zero-initialized C, (M, (BM, BN, BK) with BK | G, and P pipeline stages 1: for K-groups g assigned to this split do Block-level admission (ℓ = 0) 2: (s(0) , s(1) , s(2) , J ) ← P REDICATES([m0 , min(m0 + BM, M )), g) 3: if s(0) then Group-local prologue 4: Cacc ← 0; NK,g ← ⌈|Kg |/BK⌉ 5: for q = 0, . . . , min(P − 1, NK,g ) − 1 do 6: P REFETCH K(g, q, q mod P, s(1) , J ) ▷ ℓ = 1: predicated A loading 7: end for Group-local pipelined mainloop 8: for kc = 0, . . . , NK,g − 1 do 9: s ← kc mod P ; WAIT S TAGE(s) 10: q ← kc + P − 1 11: if q < NK,g then 12: P REFETCH K(g, q, q mod P, s(1) , J ) ▷ refill one free stage 13: end if (2) 14: for all MMA row fragments v with sv = 1 do 15: Cacc [v] ← MMA(Ag,kc [v], Bg,kc , Cacc [v]) ▷ ℓ = 2: predicated MMA 16: end for 17: R ELEASE S TAGE(s) 18: end for Epilogue 19: Atomically scatter-add Cacc into C[J , n0 : min(n0 + BN, N )] 20: else 21: Skip the complete [BM, G] block ▷ ℓ = 0: block-level skipping 22: end if 23: end for Unlike GEMM-MN, GEMM-K scopes its pipeline to one active K-group. The producer–consumer ring and accumulator reset at each group boundary, while an inactive group skips its copies, synchronizations, MMA, and epilogue work entirely. This group-local design sacrifices prefetching across adjacent K-groups, but allows ℓ = 0 block-level skipping to bypass the entire inner pipeline. Advancing a shared ring over skipped groups would otherwise leave unmatched wait/commit operations and increase barrier-management overhead. It also creates substantial atomic-reduction contention because around ⌈K/G⌉ groups may scatter partial outputs to the same entries of C. We mitigate this contention with an outer loop that packs multiple K-groups into one CTA and executes them sequentially, thus reducing concurrent writebacks. For the GEMV variant in GEMM-K, it can be dispatched in M = 1, and M ∈ {2, 3, 4} for G ⩾ 256. On sm90+, TMA also supports asynchronous tensor reduction from shared to global memory through cp.reduce.async.bulk.tensor. This enables atomic writeback from shared memory, allowing a warp-specialized pipeline to overlap the atomic epilogue with the mainloop and potentially 23

Preprint

improve GEMM-K performance for small G. This optimization is less suitable for sm120: persistently storing C in shared memory consumes too much of the limited shared-memory capacity and substantially reduces occupancy. We therefore leave it as future work. Implementation of Attention For prefilling, a CTA owns a partial tiles query at batch b and head hq . The query head determines both its KV head for GQA and routing group g. Algorithm 4 follows the tiled online-softmax organization of FlashAttention (Dao, 2023), while fusing the same three predicates into the query-side work. Its path is query-tile admission, Q prologue, dense K/V pipelined mainloop, and scatter epilogue. Here O NLINE S OFTMAX updates (m, z, Oacc ) only for active rows and returns P, where a fully masked row returns P[r, :] = 0 and leaves its running state unchanged. Algorithm 4 Attention multi-stage pipelines f I), query tile (b, hq , m0 ), tile sizes (BM, BN ), and Require: Q, K, V, zero-initialized O, (M, PK , PV pipeline stages Block-level admission (ℓ = 0) 1: Map hq to KV head hkv and routing group g 2: (s(0) , s(1) , s(2) , J ) ← P REDICATES(b, [m0 , min(m0 + BM, Tq )), g) 3: if ¬s(0) then 4: return ▷ ℓ = 0: CTA early exit 5: end if fb,r,g = 1} 6: A ← {r ∈ [m0 , min(m0 + BM, Tq )) : M Q prologue (ℓ = 1: Q loading) 7: Async-gather the active Q[b, J , hq , :] tile according to s(1) ; wait until ready 8: (m, z, Oacc ) ← (−∞, 0, 0) 9: δ ← Tk − Tq ▷ query–key position offset 10: jmin ← leftpad(b) (or 0); jmax ← Tk 11: if causal then fb,r,g = 1}) 12: jmax ← min(Tk , δ + 1 + max{Jr : M 13: end if 14: tmin ← ⌊jmin /BN ⌋; tmax ← ⌈jmax /BN ⌉; NT ← tmax − tmin K/V pipelined mainloop 15: Prime independent K and V pipelines with their first min(PK − 1, NT ) and min(PV − 1, NT ) tiles 16: for t = tmin , . . . , tmax − 1 do 17: j ← t · BN ▷ aligned key-tile origin 18: Issue the next dense K and V tiles, if any, to free stages 19: Wait for the current K stage 20: S ← −∞; P ← 0 (2) 21: for all MMA row fragments v with √sv = 1 do ⊤ ▷ ℓ = 2: predicated QK MMA 22: S[v] ← MMA(Q[v], Kt , 0)/ d 23: end for 24: Mask invalid keys using aligned query positions J + δ and reset inactive rows to −∞ 25: (m, z, Oacc , P) ← O NLINE S OFTMAX(S, m, z, Oacc , A) 26: Wait for the current V stage (2) 27: for all MMA row fragments v with sv = 1 do 28: Oacc [v] ← MMA(P[v], Vt , Oacc [v]) ▷ ℓ = 2: predicated PV MMA 29: end for 30: Advance both consumer stages 31: end for Epilogue 32: Normalize active rows of Oacc by z and scatter them to O[b, J , hq , :] Once a CTA is admitted, K and V remain dense and their two producer–consumer pipelines advance once per logical key tile. Row-level predicates never suppress a barrier arrival or stage transition. They only suppress the gathered Q packets and the QK/PV MMA fragments, preserving barrier phase agreement across the CTA. For causal attention, the streamed KV range is shortened using the 24

Preprint

maximum aligned active query position J + δ. A streamed tile wholly inside the valid range and ending no later than the minimum aligned query position takes a fully-valid fast path; only boundary tiles evaluate element-wise causal predicates. For decoding (Tq = 1), mask reordering is bypassed and routing is checked directly at the block level, while the split online-softmax pipeline follows the same abstraction, similar to the GEMV variants of GEMM-MN and GEMM-K.

A DDITIONAL E XPERIMENTAL R ESULTS

Benchmark over Uniform Sparsity We further evaluate WIDE with a uniform sparsity constraint that allocates the pruning budget equally between the attention and FFN modules. Specifically, we decompose the sparsity loss into separate attention and FFN terms: P i P i ! N N 1 X 1 X Mattn Mffn α S− + S− (22) 1− 1− L = LLM + i 2 N i=1 N i=1 |Mattn | |Miffn | This design constrains the sparsity of the two modules independently rather than matching only their average, keeping their module-wise sparsity levels close to a 1:1 ratio. As shown in Table 4, the uniform constraint causes a larger accuracy drop at larger group sizes, but LoRA recovery restores most of the lost accuracy. After recovery, the uniform variant is within roughly 4% of the unconstrained WIDE variant and still significantly outperforms static structured pruning at the same budget. The layer-wise sparsity profiles in Fig. 12 support the same interpretation: enforcing uniformity between attention and FFN does not alter how sparsity is distributed across layers, it only shifts the profile without changing its shape. Thus, the allocation learned by WIDE appears to capture a stable layer-importance pattern induced by the router, suggesting that the router learns a locally optimal allocation for the provided data distribution. Table 4: Downstream performance and the real averaged sparsity on Llama3.1-8B with DDP, SkipGPT, WIDE and 50% target sparsity, where uniform denotes to force the attention and FFN allocating the same 50% pruning budget. Methods DDP – LoRA SkipGPT – uniform – LoRA – uniform + LoRA WIDE(32) – uniform – LoRA – uniform + LoRA WIDE(128) – uniform – LoRA – uniform + LoRA

Real Sparsity 41.6%/47.9% 41.6%/47.9% 56.7%/38.7% 49.1%/47.9% 62.7%/31.6% 47.6%/46.5% 63.1%/30.7% 47.8%/47.1% 66.1%/27.9% 47.6%/46.3% 63.0%/30.0% 48.8%/48.5% 66.2%/28.5% 48.5%/46.8%

WikiText2 21.83 19.52 96.04 104.91 13.90 16.17 14.15 14.97 11.99 13.62 12.57 21.58 11.43 13.03

ARC-E 55.13 55.77 35.69 33.16 72.01 66.20 69.74 61.87 72.90 68.39 66.37 51.64 71.04 68.73

ARC-C 33.36 34.56 24.57 24.91 42.24 37.54 41.55 37.29 44.20 40.78 39.33 30.89 41.64 39.08

BoolQ 65.69 68.26 57.65 48.72 65.81 64.40 70.31 63.82 77.22 73.43 69.51 57.13 75.38 73.06

Winogrande 60.77 61.80 52.49 50.59 64.80 63.85 63.69 64.96 65.98 66.85 67.32 55.01 67.72 66.61

PIQA 69.31 71.00 58.71 60.23 75.68 74.32 76.22 73.56 78.40 75.90 76.55 71.55 78.94 76.22

OpenbookQA 30.80 33.60 26.60 27.60 42.80 38.00 42.20 39.60 43.40 42.40 42.20 32.60 43.80 40.40

Hellaswag 56.28 57.46 41.92 31.90 67.87 64.45 69.17 64.76 71.66 68.83 69.14 58.17 71.85 67.30

Avg. Acc. 53.05 (74.14%) 54.64 (76.36%) 42.51 (59.42%) 39.58 (55.32%) 61.60 (86.09%) 58.39 (81.61%) 61.84 (86.42%) 57.98 (81.03%) 64.82 (90.59%) 62.36 (87.16%) 61.48 (85.93%) 50.99 (71.27%) 64.33 (89.92%) 61.62 (86.13%)

Full Downstream Benchmark Results Here are the full results for Llama3.1-8B and Llama3.2-3B (Table 5, 6, and 7).

Preprint

WIDE G=32 WIDE G=128

Real sparsity (%)

100

WIDE G=32 uniform WIDE G=128 uniform

(a) Attention router

(b) FFN router

80 60 40 20 0 4

Layer

Figure 12: WIDE’s layer-wise sparsity with different G and sparsity allocation.

Table 5: Llama3.1-8B performance with target sparsity 25%. WikiText2 (ppl↓)

ARC-E (Acc. Norm.↑)

ARC-C (Acc. Norm.↑)

Dense Shortened-ppl Shortened-taylor CoopPruner SliceGPT Týr-the-Pruner DDP D-LLM SkipGPT WIDE (32) WIDE (64) WIDE (128)

7.71 25.84 30.07 22.96 22.57 12.89 12.41 69.45 15.63 8.49 10.37 10.04

82.70 52.19 54.42 59.81 46.42 68.22 70.01 29.63 50.38 79.84 71.34 68.77

55.03 29.95 32.42 33.79 29.52 38.73 43.26 22.53 30.89 52.82 44.97 44.97

Dense Shortened-ppl Shortened-taylor CoopPruner SliceGPT Týr-the-Pruner DDP D-LLM SkipGPT WIDE (32) WIDE (64) WIDE (128)

7.91 12.33 12.12 12.03 14.88 11.49 12.10 10.29 9.87 8.61 10.32 12.51

81.52 67.63 70.58 73.53 62.92 71.25 70.16 77.36 80.68 80.30 79.50 79.29

54.01 39.08 40.78 42.41 36.86 43.17 43.69 48.38 51.96 53.50 52.56 51.45

Methods

BoolQ Winogrande (Acc. ↑) (Acc. ↑) Calibration-only 83.06 74.27 41.41 54.62 40.28 55.49 61.56 53.99 74.92 67.48 78.59 66.53 77.25 69.85 58.78 49.96 69.94 51.30 82.17 71.67 80.46 62.15 77.95 60.77 LoRA Tuning 83.30 73.32 54.25 58.48 64.04 59.83 63.52 62.51 76.17 68.82 78.50 68.74 78.78 68.35 76.39 68.27 81.47 70.88 82.66 70.88 81.47 71.19 80.58 70.88

PIQA (Acc. Norm.↑)

OpenbookQA (Acc. Norm.↑)

Hellaswag (Acc. Norm.↑)

Avg. Acc.↑

81.07 69.80 70.62 72.52 66.53 74.21 75.63 54.19 68.77 79.76 79.11 79.16

45.40 33.00 35.20 33.60 31.80 37.40 42.00 25.80 31.40 45.60 40.40 38.20

79.32 52.55 56.68 58.06 56.93 66.33 72.42 31.32 59.06 78.58 75.78 75.73

71.55 (100.00%) 47.65 (66.59%) 49.30 (68.90%) 53.33 (74.54%) 53.37 (74.59%) 61.43 (85.86%) 64.35 (89.93%) 38.89 (54.35%) 51.68 (72.23%) 70.06 (97.92%) 64.89 (90.69%) 63.65 (88.96%)

81.23 76.61 76.88 77.04 73.77 76.27 76.01 78.29 79.65 80.30 80.52 80.36

46.20 39.00 38.20 40.00 36.80 39.80 40.60 43.80 44.00 45.20 45.20 45.80

79.94 65.69 68.53 69.76 69.71 71.68 72.12 75.38 77.73 78.46 78.34 78.35

71.36 (99.73%) 57.25 (80.01%) 59.83 (83.63%) 61.25 (85.61%) 60.72 (84.87%) 64.20 (89.73%) 64.24 (89.79%) 66.84 (93.42%) 69.48 (97.11%) 70.19 (98.09%) 69.83 (97.59%) 69.53 (97.18%)

Preprint

Table 6: Llama3.1-8B performance with target sparsity 50%. Methods

WikiText2 (ppl↓)

ARC-E (Acc. Norm.↑)

ARC-C (Acc. Norm.↑)

Dense Shortened-ppl Shortened-taylor CoopPruner SliceGPT Týr-the-Pruner DDP D-LLM SkipGPT WIDE (32) WIDE (64) WIDE (128)

7.71 473.15 4.86e+7 503.86 85.81 59.88 21.83 504.48 96.04 14.15 14.96 12.57

82.70 32.20 29.17 32.87 30.30 44.65 55.13 27.02 35.69 69.74 69.28 66.37

55.03 24.40 25.94 23.21 21.84 24.91 33.36 23.12 24.57 41.55 38.99 39.33

Dense Shortened-ppl Shortened-taylor CoopPruner SliceGPT Týr-the-Pruner DDP D-LLM SkipGPT WIDE (32) WIDE (64) WIDE (128)

7.91 22.47 22.30 22.66 28.76 37.99 19.52 19.31 13.90 11.99 12.05 11.43

81.52 51.56 52.57 52.82 39.64 54.75 55.77 63.05 72.01 72.90 73.91 71.04

54.01 27.56 29.69 29.18 24.06 31.74 34.56 35.75 42.24 44.20 41.72 41.64

BoolQ Winogrande (Acc. ↑) (Acc. ↑) Calibration-only 83.06 74.27 53.00 49.49 38.20 49.88 56.09 49.33 55.47 51.53 58.44 54.61 65.69 60.77 50.52 47.20 57.65 52.49 70.31 63.69 67.43 66.69 69.51 67.32 LoRA Tuning 83.30 73.32 50.37 53.75 61.01 54.62 63.15 54.30 64.12 58.24 68.56 58.72 68.26 61.80 62.45 60.14 65.81 64.80 77.22 65.98 77.74 66.30 75.38 67.72

PIQA (Acc. Norm.↑)

OpenbookQA (Acc. Norm.↑)

Hellaswag (Acc. Norm.↑)

Avg. Acc.↑

81.07 57.40 53.70 56.80 53.97 64.68 69.31 52.72 58.71 76.22 75.95 76.55

45.40 25.80 27.40 26.20 25.80 29.40 30.80 25.20 26.60 42.20 41.40 42.20

79.32 35.03 27.63 30.45 32.05 41.62 56.28 28.48 41.92 69.17 68.50 69.14

71.55 (100.00%) 39.62 (55.37%) 35.99 (50.30%) 39.28 (54.90%) 38.71 (54.10%) 45.47 (63.55%) 53.05 (74.14%) 36.32 (50.77%) 42.52 (59.42%) 61.84 (86.43%) 61.18 (85.50%) 61.49 (85.94%)

81.23 68.12 69.75 67.95 60.60 70.56 71.00 72.36 75.68 78.40 78.07 78.94

46.20 32.00 33.20 30.20 28.00 33.80 33.60 37.40 42.80 43.40 42.80 43.80

79.94 47.44 48.92 49.80 47.43 56.09 57.46 61.31 67.87 71.66 72.03 71.85

71.36 (99.73%) 47.26 (66.05%) 49.97 (69.83%) 49.63 (69.36%) 46.01 (64.31%) 53.46 (74.72%) 54.64 (76.36%) 56.07 (78.36%) 61.60 (86.10%) 64.82 (90.60%) 64.65 (90.36%) 64.34 (89.92%)

Table 7: Llama3.2-3B performance with target sparsity 50%. Methods

WikiText2 (ppl↓)

ARC-E (Acc. Norm.↑)

ARC-C (Acc. Norm.↑)

Dense Shortened-ppl Shortened-taylor CoopPruner SliceGPT Týr-the-Pruner DDP D-LLM SkipGPT WIDE (32) WIDE (64) WIDE (128)

9.77 676.05 2.32e+5 1.79e+5 96.69 87.37 31.20 3877.53 156.58 17.28 19.02 18.11

72.01 33.88 34.01 25.55 29.12 38.38 47.14 27.19 30.89 62.84 61.74 57.66

46.67 23.55 25.60 24.57 21.67 23.20 29.78 24.32 20.99 35.41 34.64 35.32

Dense Shortened-ppl Shortened-taylor CoopPruner SliceGPT Týr-the-Pruner DDP D-LLM SkipGPT WIDE (32) WIDE (64) WIDE (128)

9.86 29.75 31.50 42.69 31.13 53.11 25.66 193.97 18.91 15.28 15.17 15.26

72.64 46.51 46.13 39.23 36.48 45.62 48.36 31.40 58.92 67.72 65.78 65.74

46.42 27.47 25.94 23.12 22.44 25.25 30.20 21.25 33.45 37.97 37.29 36.43

BoolQ Winogrande (Acc. ↑) (Acc. ↑) Calibration-only 73.88 69.22 44.19 52.75 50.92 51.93 50.67 50.04 54.58 48.61 62.17 53.35 62.42 56.51 40.86 49.57 43.30 50.04 65.20 62.51 66.67 59.98 62.02 60.77 LoRA Tuning 72.87 68.98 61.28 51.22 61.88 51.70 49.24 52.96 59.32 54.61 57.58 52.56 63.30 56.12 57.09 51.22 62.81 58.64 66.94 59.27 65.23 61.01 66.12 61.88

PIQA (Acc. Norm.↑)

OpenbookQA (Acc. Norm.↑)

Hellaswag (Acc. Norm.↑)

Avg. Acc.↑

78.02 56.58 55.71 52.23 53.21 59.63 65.34 50.65 55.77 73.01 72.36 73.23

40.60 25.20 25.60 27.00 25.80 27.40 30.00 26.00 25.80 38.40 38.60 39.20

73.98 31.44 31.70 26.25 30.38 34.06 45.04 26.35 30.38 61.60 60.84 59.76

64.91 (100.00%) 38.23 (58.89%) 39.35 (60.63%) 36.62 (56.41%) 37.62 (57.96%) 42.60 (65.63%) 48.03 (74.00%) 34.99 (53.91%) 36.74 (56.60%) 57.00 (87.81%) 56.40 (86.89%) 55.42 (85.38%)

77.91 64.85 64.85 63.11 59.46 65.34 67.14 57.18 72.31 74.54 75.14 74.92

42.20 29.60 28.20 29.20 26.60 30.00 30.20 24.40 36.60 40.60 39.80 37.00

74.00 39.91 42.16 35.89 40.40 43.46 46.84 28.07 57.59 64.32 64.26 64.69

65.00 (100.14%) 45.83 (70.61%) 45.84 (70.61%) 41.82 (64.43%) 42.76 (65.87%) 45.69 (70.38%) 48.88 (75.30%) 38.66 (59.56%) 54.33 (83.70%) 58.77 (90.53%) 58.36 (89.90%) 58.11 (89.52%)

Preprint

Decoding Speedup for Different Attention Group Size Figure 13 shows how the attention group H size affects decoding. In Gattn ⩽ Hkq × d, the kernel progressively transitions from vanilla multi-head H

attention to GQA-packing attention, where a single CTA partition can read and compute at most Hkq heads in parallel. Increasing Gattn beyond this threshold does not affect the tiling strategy used for attention decoding. DSL (TileLang)

CuTe

1x torch

2x torch

Speedup over torch

Attention-Decode, Tk=32,768, [Hq, Hk, D]=32x8x128 3

B=1

2.5 2

B=2

2.5 2

2.5

B=4

B=8

2.5 2

B=16

2 1

1 0

128 256 512 1024 2048 4096

Group size

Figure 13: Attention decoding speedup of different group size on Llama3.1-8B shapes, with a random 50% sparsity mask.

Preprint

Full Visualization Results on Routing Decision Figures 14 to 18 provide additional visualizations of WIDE’s token-wise routing decisions and reveal the same semantic-aware behavior in Figure 10. Figure 19 further compares the probability of retaining each routing group across low and high sparsity layers in the attention and FFN modules, where the patterns are differing across layers. arc_easy case: arc_easy_229

Router layer

attn L24

attn L25 60

attn L26 ffn L1

ffn L9 20

ffn L10 ffn L11

Skip probability (%)

100

attn L2

cu r oc

is ys

y gl

do es

ll ce

Question/context token

Figure 14: Token-wise routing probability for Llama3.1-8B in selected ARC-Easy sample. arc_challenge case: arc_challenge_28

Router layer

attn L24

attn L25 60

attn L26 ffn L1

ffn L9 20

ffn L10 ffn L11

Skip probability (%)

100

attn L2

gr ow

an t

lp s

at th

e uc pr od

is yn th es

os ph ot

es do

ha t

Question/context token

Figure 15: Token-wise routing probability for Llama3.1-8B in selected ARC-Challenge sample. openbookqa case: openbookqa_284

Router layer

attn L24

attn L25 60

attn L26 ffn L1

ffn L9 20

ffn L10 ffn L11

Skip probability (%)

100

attn L2

t ea

y lik el

t os m

it bb

ou w

Question/context token

Figure 16: Token-wise routing probability for Llama3.1-8B in selected OpenbookQA sample.

Preprint

openbookqa case: openbookqa_308

Router layer

attn L24

attn L25 60

attn L26 ffn L1

ffn L9 20

ffn L10 ffn L11

Skip probability (%)

100

attn L2

ak e

ed ne

Question/context token

Figure 17: Token-wise routing probability for Llama3.1-8B in selected OpenbookQA sample.

hellaswag case: hellaswag_1012

Router layer

attn L24

attn L25 60

attn L26 ffn L1

ffn L9 20

ffn L10 ffn L11

Skip probability (%)

100

attn L2

gs do

al l sm

as s

t ou

ks al w

Question/context token

Figure 18: Token-wise routing probability for Llama3.1-8B in selected Hellaswag sample.

(a) Attention selected-layer group sparsity

Layer

L2 low (29.5%)

100

L24 high (88.3%) L25 high (84.7%) 80

Route group

(b) FFN selected-layer group sparsity L1 low (0.1%)

Layer

40 L9 high (63.3%) 20

L10 high (60.1%)

L11 high (61.3%)

0 1

112

Route group

Figure 19: Group activation patterns of different layer in Llama3.1-8B on WikiText2.

Skip probability (%)

L26 high (85.7%)

Record · ID 414103 · SHA-256 6601faf3be237bcb
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.