Polar Sparsity: High Throughput Batched LLM Inferencing with Scalable Contextual Sparsity
Authors: Susav Shrestha
Bradley Settlemyer
Nikoli Dryden
A. L. Narasimha Reddy
Venue: NeurIPS 2025 poster ()
Keywords: Large Language Models
Contextual Sparsity
Activation Sparsity
Head Sparsity
Select Head Attention
Batched Inference
Efficient Transformers
Machine Learning Systems
Primary area: deep_learning
TL;DR: Polar Sparsity scales contextual sparsity to large batches by exploiting stable attention head sparsity and using efficient GPU kernels, achieving up to 2.2× speedups with minimal accuracy loss.

Abstract: Accelerating large language model (LLM) inference is critical for real-world deployments requiring high throughput and low latency. Contextual sparsity, where each token dynamically activates only a small subset of the model parameters, shows promise but does not scale to large batch sizes due to union of active neurons quickly approaching dense computation. We introduce Polar Sparsity, highlighting a key shift in sparsity importance from MLP to Attention layers as we scale batch size and sequence length. While MLP layers become more compute-efficient under batching, their sparsity vanishes. In contrast, attention becomes increasingly more expensive at scale, while their head sparsity remains stable and batch-invariant. We develop Selective Head Attention with hardware-efficient, sparsity-aware GPU kernels, delivering up to \(2.2\times\) end-to-end speedups for models like OPT, LLaMA-2 \& 3, Qwen, Mistral across various batch sizes and sequence lengths without compromising accuracy. To our knowledge, this is the first work to demonstrate that contextual sparsity can scale effectively to large batch sizes, delivering substantial inference acceleration with minimal changes, making Polar Sparsity practical for large-scale, high-throughput LLM deployment systems.

Polar Sparsity: High Throughput Batched LLM Inferencing with Scalable Contextual Sparsity
Susav Shrestha Texas A&M University s1s7161@tamu, edu
Brad Settlemyer NVIDIA bssttlenyerrnniddi.. com
Nikoli Dryden
Narasimha Reddy
Lawrence Livermore National Laboratory dryden1 1aaann0nnnnoooooom gov
Texas A&M University reddyStamu edu
Abstract
Accelerating large language model (LLM) inference is critical for real-world deployments requiring high throughput and low latency. Contextual sparsity, where each token dynamically activates only a small subset of the model parameters, shows promise but does not scale large batch sizes due union of active neurons quickly approaching dense computation. We introduce Polar Sparsity, highlighting key shift in sparsity importance from MLP Attention layers as we scale batch size and sequence length. While MLP layers become more compute-efficient under batching, their sparsity vanishes. In contrast, attention becomes increasingly more expensive at scale, while their head sparsity remains stable and inviriiii We develop Selective Head Attention with hardware-efficien,, sparsity- aware GPU kernels, delivering up to 2.2x end-to-end speedups for models like OPT, LLaMA- 2 & 3, Qwen, Mistral across various batch sizes and sequence lengths without compromising accuracy. To our knowledge, this is the first work to demonstrate that contextual sparsity can scale effectively large batch sizes, delivering substantial inference acceleration with minimal changes, making Polar Sparsity practical for large-scale, high-throughput LLM deployment systems. Our code available at: https://github... com http://hhtbb.mmmmaaaaass/////aar Sparssity
Introduction
Modern LLMs have emerged as powerful tools capable of excelling at diverse tasks, leading to their widespread adoption in modern systems [56] 46 1S 52 However, their massive scale, typically involving billions of parameters, makes them computationally expensive and highlights the need for more efficient and cost-effective deployment solutions.
While sparsity and pruning have been extensively studied, their use in production LLMS remains limited due to poor workload generalization and inefficient hardware utilization caused by irregular 58 33]. A recent line of research has uncovered the phenomenon of contextual activation sparsity, where each input token dynamically activates only small subset of input-dependent neurons, enabling efficient acceleration without compromising model quality (39]. Activation sparsity emerges promising acceleration technique for LLMSs with structured sparsity which enables coalesced memory access and measurable wall-clock speedups.
memory access [26] [5] 59
37
Activation sparsity has been shown to be effective in reducing inference latency, but its application has been largely limited to single-query workloads 39 54 4 Modern inferencing systems depend
on batching maximize hardware efficiency and reduce serving costs [1
69 45
Neural
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

232
Ourput projection
attentics
call propect
323
Head
Head
Head
Head N
Head
Hred
Head
Head
Head
Head
Head
Head N
200
02300
19.1%
17.04
Input
batch
projection
Batch
(a) Transformer decnde latency hreakdown
(b) Selective Head Attention
Figure I: (a) Decode latency breakdown on A 100 GPUs (OPT-66B, sequence length 1920). As batch size increases, attention layers dominate end-to-end latency. (b) Selective Head Attention only activates the most critical heads for each request and accelerates high throughput inference.
scaling laws have driven the rise of massive models that rely multi-GPU inference, rendering single-query execution prohibitively expensive for real-world deployment 30 21 Recent work activation sparsity focus reducing latency by optimizing neuron activations in MLP and linear projection layers as these modules dominate runtime small batch sizes. However, as shown in Figure Ta this strategy breaks down under batched workloads. While batching improves the compute efficiency of linear layers, it has the opposite effect on attention, which scales with sequence length and batch size. Furthermore, union neuron sparsity, which varies dynamically with batches and layers, diminishes with batching (Figure 2). limiting the practical gains of existing approaches. This work explores techniques for integrating contextual sparsity into batched inference, enhancing both system throughput and computational efficiency.
Polar Sparsity refers the shift in sparsity importance from MLP layers to Attention layers batch size and sequence length increase. Current of-the-att sparsity methods primarily focus on model parameter sparsity, where only subset of model parameters is activated reduce computation and memory IO. We define head sparsity as the phenomenon where, for a given token, only subset of attention heads contribute significantly to the output while the remaining heads have negligible effect. In large-batch inference, the cost of accessing model parameters is largely amortized, since the entire batch utilizes the same model weights. In contrast, each batch has unique key-value (KV) cache, making attention layers memory I/O expensive. While contextual sparsity in model parameters diminishes as batch sizes increase, attention head sparsity remains stable and batch invariant. We introduce Selective Head Attention (Figure (b]. which activates only the most critical heads for each request, preserving overall sparsity and improving compute and memory efficiency.
We leverage these properties to build contextual sparsity-aware Selective GEMM and Selective FlashAttention kernels that reduce memory I/O and compute, enabling scalable and tigh-throughput inference. To the best of our knowledge, we are the first work to show that contextual sparsity is scalable with batch size, offering higher gains at larger batch sizes. Our main contributions as follows:
1. We show that activation sparsity in the MLP layers degrades with batch size due to union activations, while attention head sparsity remains stable and batch-invariant.
2. We design Selective GEMM kernels with a layer-wise top-k optimization strategy for dynamic MLP activations, achieving up 5.5x speedup.
3. We introduce Selective FlashAttention kernels that support Head/Group sparsity with per- query activation, reducing memory I/O and compute, achieving up 2.8x speedup.
4. Polar Sparsity delivers up 2.2x improvement in batched decoding throughput with negligible accuracy loss, and can be seamlessly integrated into wide range of L.Mss, including those without ReLU activations, unlock substantial inference acceleration.

---

2 Background and Related Works
Contextual sparsity can be leveraged at runtime through sparsity predictors or learned routers that dynamically select the neurons to activate based on the input vectors. Similar Mixture-of-Expent (MoE) architectures, it activates only a subset of the model conditioned on the input, but differs by operating at a finer granularity, selecting individual neurons rather than routing to large expert blocks.
Early models employing ReLU activations demonstrated high activation sparsity, with over 90% of the feed-forware network outputs being zero- valuud, allowing significant computational savings. This indicates that, for a given input, only a small fraction of neurons contribute meaningfully at each decoding step. However, recent LLMS adopt smoother activations like SwiGLU, which yield denser activations and limit the benefits of sparsity. To reintroduce sparsity, several recent studies have proposed returning ReLU or developing sparse variants [44 [53] B]. Several subsequent works have been performed to further increase the sparsity by introducing novel sparsity-enhancing techniques and activation functions
33
37
68
50
66
Additional discussions can be found in Appendix [A]
Several works explore token sparsity attention, leveraging the fact that many tokens contribute little to model output [14 47
62 40). This is orthogonal to head sparsity. MoA [65] and MoH [29] introduce MoE-style routing over attention heads, improving accuracy and theoretical efficiency but without wall-clock gains- showiing that reduced FLOPs do not always yield faster inference. Deja Vu 39 and TEAL 37 also exploit head sparsity but only within the projection layers: they compute QKV vectors for selected heads and copy the corresponding KV cache into smaller buffers, incurring costly memory I/O and limited scalability. Moreover, skipping heads during QKV projection can lead to incomplete context in later decoding. To address these issues, we retain dense QKV projections for KV-cache consistency and apply head sparsity within the attention kernel. We further design an V/Oefficient Selective FlashAttention kernel that scales efficiently to large batches while preserving sparsity benefits. Finally, while prior activation sparsity methods target single-query inference, we extend activation sparsity the large-batch regime via batch-invarian head sparsity and custom GPU kernels, enabling scalable,
47
I
LLM inference.
Motivation and Problem Formulation
The inference process in LLMs consists of two stages: prefill and decode. The prefill stage processes the input prompt in single forward pass to build and cache the KV representations at each transformer layer, and typically compute-bound. The decode stage then generates tokens autoregressively using the cached KV data; this stage is generally I/O-bound, as each token requires minimal computation but frequent memory access. While prefill has quadratic complexity in sequence length, is highly parallelizable across tokens and GPUs; by contrast, the decode stage is sequential, 1/O-bound, and dominates wall-clock latency in large batch, long-sequence generation. This work focuses on optimizing the autoregressive decode stage under batched inference, aiming to improve throughput and efficiency real-world serving scemarios.
Figure la shows the decode latency breakdown across key modules MoLP attention, QKV projection, output projection, and others (including communication, layer norm, and embedding projections) for the OPT-66b model. At small batch sizes, latency is dominated by linear layers, making prior work on activation sparsity highly effective in reducing decoding latency. However, as batch size increases, linear layers benefit from improved compute efficiency, while attention latency grows nearly linearly. quickly becoming the dominant bottleneck. This shift underscores critical insight: optimizing linear layers alone is insufficient at scale. Our work builds on this observation and focuses on improving decoding efficiency by addressing attention's growing cost in batched autoregressive inference.
3.1 Accelerating MLP Layers
In the decode stage, the input to the MLP block a 3D tensor x E RBxlxd, where B is the batch size and d is the model dimension. The two projection matrices in the MLP block are W1, W2 E RdxD. where is the hidden dimension of the feed-forward network.
With contextual sparsity, only subset of neurons in the MLP block are activated for a given input. Let S C (D) denote the set of active neurons for a single sequence or token in the decode stage. Under batching, compute the union of active neurons across all sequences in the batch. That is,
3

---

(a) OPT 66b batch neuron activations
(b) ReLU Llama-2-7b batch neuron activations
Figure 2: Batch neuron activations in MLP layers in sparse models. Union neuron sparsity MLP
layers decreases rapidly with batch size, illustrating why MLP sparsity is less effective at scale.
neuron is retained it is activated for any sequence in the current batch. Let SB S [D] denote this union set for batch size B. The sparsified MLP computation becomes:
MLP Sg (x) = (xW1,s, W2.5a²
where W1,s W2.Sg E [dx are the dynamically pruned weight matrices corresponding to the active neurons in the batch, and a is the activation function. As batch size increases, the union SB typically grows, reducing the overall sparsity and limiting computational savings. Figure 2a shows the average union activation across transformer layers for different batch sizes in the OPT 66b model during the decode stage. We observe strong layer-wise trend: early layers maintain high sparsity, while deeper layers exhibit progressively higher union activation. Even at large batch sizes, union activation in early layers remains low as the average per-token activation is under 1 1%. This trend is encouraging, as it reveals an opportunity to accelerate early layers through sparsity, even in lcage-scale batching. In contrast, sparsity patterns in sparsified LLaMA-2-7B models sensitive to batching. Figure shows that replacing SwiGLU with ReLU yields high initial sparsity that quickly diminishes with batching. See Appendix B for additional discussion MLP sparsity.
more
3.2 Accelerating Attention Layers
In the decode stage, the input to the attention layer is batch of inputs x E R Bxlxd, where the batch size and d is the model dimension. After applying the QKV projections and reshaping. the query tensor has shape Q E IR BxHxxxxd, while the key and value tensors are updated and retrieved from the KV cache with shape K, V BBxH.. Hx. xd., where H is the number of attention heads, Hgv is the number key-value heads, d, = d/ H is the per-head dimension, and N is the current sequence length of each batch! Each query attends to its own key-value tensor, with self-attention computed independently and in parallel across both heads and batch dimensions.
For each attention head i [H] in each batch E [B]. the attention is computed as:
(Qri.n (uii..Kppha
Attention (Qb,ine Ko VoSaa) = softmax
Voli..
ddp
Since each sequence in the batch attends over its own KV tensors, the attention computation and memory I/O scale linearly with both batch size and sequence length during decode. Importantly, exploiting head-level sparsity MLP sparsity, this form of sparsity is invariant to batching, as each sequence computes attention independently. Thus, attention head sparsity remains reliable target for acceleration even at large batch sizes.
significantly reduce both compute and memory 1/O at scale. Unlike
Some implementations represent the tensors Q Bx1xXxx1 and K, V E R BB

---

Lane-31/05
loe e.
<p--6ib
Lamm-2-13b-17
Pop (0g srae)
0.5
Importoce
315
206
0.
20
Head LctoaatortLayer
Layer T1cee
(a) Head Sparsity Perplexity
(b) Attention layer importance
Figure 3: (a) Perplexity increases gradually when only activating the most important attention heads/layer, showing that many heads can be skipped with minimal degradation. (b) Layer 0 has high importance score across a range of models, motivating the use dense attention in early layers.
To study head sparsity, we conduct experiment where, at each transformer layer, only the top-k attention heads, ranked by output L2 norm, are activated, while the outputs of non-activated heads are masked to zero. We then measure the model's perplexity Figure Ba illlusttates the relationship between perplexity and head sparsity, with the relative increase in perplexity annotated each sparsity level. We observe that by activating the most important heads at each layer, the perplexity does not increase drastically up to 50% head sparsity, across all the tested models. Interestingly, we observe that head sparsity increases with model size. For example, the OPT-66b model shows only a 5% increase perplexity at 30% head activation, while the smaller OPT-6.7b model experiences the expectation that larger models, which possess more attention heads, inherently activate more heads at comparable activation levels. As result, larger models present a greater opportunity for acceleration through increased sparsity.
subset of the Wikitext-2 dataset [42]
substantial 37.4% 4% increase. This trend aligns with
Recent work has also highlighted variation in attention importance across layers. We compute per-layer attention importance using the scoring method from (23). As shown in Figure layer consistently exhibits the highest importance across range of models. Based on this observation, we apply dense attention at layer 0 and enforce uniform head sparsity in all subsequent layers.
Polar Sparsity
We define Polar Sparsity as the emergence of distinct sparsity effectiveness regimes in large-scale LLM inference: MLP sparsity accelerates small-batch, low-latency inference, whereas Attention head sparsity unlocks high-throughhut inference at scale. In this section, we present system for accelerating high-throughput LLM inference via scalable contextual sparsity.
4.1 Dynamic Sparsity in MLP Blocks
Our fine-grained MLP router builds on recent advances in sparse inference 33 54 enabling Bx 1 the router is
selective activation at the neuron level. Given a hidden state for a batch, x E implemented lightweight two-layer feedforward network with bottleneck intermediate layer, and we trained using a binary cross-entron loss with the AdamW optimizer. Further details are provided in Appendix
Prior methods often use static top-k activation thresholds across all layers, which suboptimal for batched inference due to the dynamic and layer-specific nature of neuron activations that scale with batch size. address this, we propose dynamic top-k mechanism that adapts the number of active neurons per layer based target recall. A simple greedy algorithm (Algorithm 2 selects the minimal top-k neurons per layer to meet target recall, using the router's output logits and ground-truth activations. We optimize these thresholds offline to maintain high recall (99%) while minimizing
5

---

Denes CoMe
Dens fleviin Sdat PlohAly
CON
10
05
03-
16
Select GEMM kernel
(b) Select FlashAttention kernel
Figure 4: Polar Sparsity Kemells. Kernels demonstrate near linear speedup with respect to sparsity, A100 GPUs, OPT 66b, batch size 64, seqlen 1920
computation. At inference time, we aggregate the predicted activations across the batch to produce single neuron index tensor that identifies the active neurons for the entire batch.
To efficiently compute sparse MLP activations, we design custom GPU kernel that fuses indexing and matrix multiplication (Algorithm 3). avoiding the overhead of separate gather-scatter operations. Unlike prior work limited to sparse GEMV, our kernel supports batched GEMM workloads and achieves near-linear speedups with increasing sparsity, with improvements of up to 5.5 over the dense baseline (Figure 4a Implementation details are provided in Appendix D
4.2 Stable Sparsity with Selective Head Attention
We design our attention routers as single-layer fully connected networks that predict head activations based on attention output norms. The router is trained similarly to the MLP routers: for each input, the top-k attention heads with the highest output norms are considered active and used as supervision targets. The attention router produces scalar logits for each attention head, which are used to select the top-k heads for each batch. Since attention routing performed independently for each instance, different batches may activate entirely different sets of heads. This results in a batch head index tensor that records the active head indices for each batch. Selective Head Attention (SHA) then uses this tensor to perform attention exclusively on the activated heads within unified kernel.
Algorithm Selective Head FlashAttention (Decode)
Require: Q E JR Bxxxxx K, V I H batch heac index Bxupp MERAM 1/Va
Output: O E Bxxxxixd (written to HBM)
1: Determine target batch index b and top-k index & assigned this unit from the grid dimensions.
2: head _id__idx<4 + batch_head_______ hae__hha__inee deea_index((b_, k]
> Get the actual head index to compute
3: B. [MERAM/(C /(4d) (Oacey lace; Mace) (0,0, -00);
[Nse/ B.]
4: Load E I2xxx from Q[b, head idx, 0,
+ Get the activated query vector for the batch
5: for j to T: do
6: Aistant (j B,, kend min Be, Noo); Load Kj. V, from K, Vlb, head idx, histart Kend; :]
7: S, = <{q@K] K ); m, = max((S)):
exp(S, m,): 1):: =.P,
8:
Whew max my): a Maaen Manw, 3 d. -maes Inew = abace + Big
9: Once +- (alace Oace 3(P @V (V,)) lane &- Inew Thace &-
end for
11: Write Oace to Obb, haaahhaaddaaaae 0,:]
Given the inputs Q E BBxHxIxdn and K, V E BBBxHXNxd, in HBM, we aim to compute the attention output O Bxxxxxx for only the activated attention heads for each batch. Our goal is to reduce the amount of memory access and compute by the factor of induced head sparsity. To achieve this, we fuse the head activation logic into a single GPU kernel and perform selective head attention using a sparsity- aware variant of FlashAttention.. During batched decoding, each batch and each head is processed in parallel by a different CUDA thread-block or a Triton program. As outlined in Algorithm [ we modify the FlashAttenticn algorithm [1] to incorporate head sparsity by indexing into the relevant heads during kernel initialization using batch head index tensor. All memory access logic is updated accordingly to ensure only data from active heads are read from and

---

1.
00
3,5
02
Figure 5: Accuracy Attention Density. Left: OPT 66b model with dynamic sparse MLP + Select Head Attention. Middle: LLaMA 2-7b model with Select Head Attention Right: LLaMA 3.1 70b model with Select Group Attention. Dense attention in layer O used for all models.
written to, enabling efficient selective head computation within a unified kernel. Figure 4b shows the forward pass runtime of our SHA kernel compared to dense and standard sparse baselines. The SHA kernel exhibits near-linear speedup with respect to head sparsity, achieving 2.8x improvement at 30% sparsity over the dense bascline. For newer models with group query attention (GQA), we embrace group sparsity query heads within group share the KV cache (2).
5 Evaluation
Experiment Setting: We first evaluate Polar Sparsity on pretrained models, including OPT [64].
LLaMA-2 (57], LLaMA-3 211. Mistral [28] across multiple model sizes. The evaluation is conducted
on nine downstream tasks: COPA (20), OpenBookQA 43), PIQA (7]. RTE [19), Winogrande [48),
HellaSwag [63], MMLU (24), ARC-easy, ARC-challenging 9 We utilize the Im-eval-harness
framework to measure the accuracy of the models these benchmarks (I7). We further evaluate
instruction-ture models on MMLU-PRO 60 and LongBench (6). Specifically, we assess LLaMA-
3.1-8B-Instrue Mistral-7B -B-sttuct and Qwen-2.5-144 1.--144BInnttrctt [46].
To train the routers, we collected 400,000 token samples from random text sequences extracted from the Wikitext-2 training dataset. All experiments are performed on NVIDIA DGX A100 80GB GPU node servers. Sparsity applied to both the ML.P and attention layers only in OPT models. For all other models, we apply only attention head sparsity, since their MLP layers use Re activations, making MLP sparsification less effective (Appendix BB. We scale each model up batch size of 512 until it reaches the memory limit. We built on top of the FlashAttention triton kernel and utilized CUDA graphs to measure the decoding throughput with the included routers.
5.1 Benchmark Evaluation
Figure shows the zero-shot accuracy on downstream tasks vary the attention density. At each data point, we activate the attention heads/groups with the highest output logits as predicted by the routers. For OPT models, also apply dynamic layer-wise sparsity to MLP layers. Across all models, we observe that most tasks can be solved accurately even under high attention head sparsity with minimal degradation up critical threshold. The evaluation results show that this critical point varies with the architecture and size of the model. Stable average accuracy is maintained down to 30% attention density for OPT-66b, 50% for smaller OPT-6.7b. LLaMA-2 7b/13b models, and 62.5% 5 for GQA models like LLaMA 3.1 70b.
Table and Manndppreeem presents the accuracy of the sparsified models their respective critical attention density thresholds. Polar Sparsity achieves performance comparable to the dense baseline, with average accu- racy differences within 1%. Instruction tuned models also sustain strong performance LongBench under Polar Sparsity, underscoring its effectiveness on challenging generative benchmarks.
We observe that some challenging tasks like RTE, MMLU, MMLU PRO are more sensitive to head sparsity and require more active heads. This rask-dependent sensitivity is particularly encouraging, as it suggests that most tasks can be served with only the most critical heads, while harder tasks
7

---

Table 1: LLM zero-shot evaluation at critical thresholds, Polar Sparsity (PS) is competitive with the dense baseline with average accuracy within 1%.
Model
COPA
OBQA
PIQA
RTE
WG
HS
MMLU
ARC-E
ARC-C
Average
OPT 6.7B
0.81
0.276
0.763
0.552
0.653
0.499
0.265
0.657
0.305
0.531
OPT 6.7B PS-0.5
0.83
0.282
0.755
0.527
0.636
0.488
0.252
0.647
0.300
0.524
OPT 66B
0.85
0.304
0.787
0.603
0.690
0.557
0.263
0.711
0.369
0.570
OPT 66B PS-0.3
0.83
0.296
0.769
0.592
0.677
0.546
0.264
0.693
0.361
0.560
LLaMA 2
0.781
0.690
0.572
0.418
0.433
0.608
I.aMA PS-0.5
0.779
0.439
0.594
LLaMA 2 13B
0.91
0.350
0.791
0.653
0.722
0.600
0.521
0.794
0.485
0.647
LLaMA 2 13B PS-0.5
0.92
0.352
0.790
0.578
0.728
0.600
0.473
0.783
0.473
0.633
LLaMA 3.1 70B
0.92
0,370
0.831
0.697
0.799
0.665
0.753
0.872
0.606
0.724
LLaMA 3.1 70B PS-0.625
0.91
0,340
0.823
0.729
0.793
0.650
0.732
0.853
0.590
0.712
Mistral 7B
0.92
0.332
0.803
0.686
0.738
0.609
0.591
0.796
0.489
0.663
Mistra 7B PS-0.5
0.92
0.340
0.801
0.671
0.736
0.608
0.562
0.793
0.483
0.657
Table 2: Evaluation of instruction-tuned LLMs at critical sparsity thresholds. Polar Sparsity maintains strong performance on generative tasks.
Model
MMLU PRO
LongBench-e
Mistral-7B-inst
0.247
0.392
Mistral-7B-ing. PolarSparse-0.5
0.244
0.388
Llama-3 3. 3.1-8B-inst
0,409
0.443
Llama-3. 3.1-8B-inst Polar5parse-0.625
0.384
0.429
Qwen-2.5 Bren-2.5-14B-ins. Qwen-2.5-14Biimm PolarSparse-0.625
0.497
0.421
0.457
0.414
could potentially activate more for higher accuracy within the same batch. Given that head sparsity is batch-invariant, this opens the door to context-aware, per-sequence head activation, paving the way for lossless sparse inference, a direction we discuss in future work.
Although recent activation sparsity literature vary in their model and benchmark selection, found that several recent methods include LLaMA-2-70 for evaluation. This allows us to perform a fair comparison in Table Polar Sparsity outperforms is competitive with the state-of-the-ar activation sparsity methods across most benchmarks. Importantly, our approach maintains accuracy and scales efficiently with batch size key limitation of existing techniques.
Table 3: Benchmark results on LLaMA-2-7b using different sparsity approaches. Zero-shot evaluation unless noted. "_" indicates the metric was not reported in the original paper. Zero-shot OpenBook QA, RTE are unreported in prior work. Bolded numbers indicate the highest among sparse methods. Underlined values denote accuracy within 1% of the dense baseline, competitive performance.
Method
CОРA
PIQA
Winegrande
HellaSwag
MMLU (5-shot)
ARC-e
ARC-c
baseline
0.87
0.781
0.690
0.572
0.458
0.763
0.433
ProbePruning 40%
3..
0.749
0,575
0.617
0,355
ReLUfication
44
0.83
0.779
0.686
0.548
0.386
0.738
0.396
ProSparse
53
0.77
0.757
0,640
0.455
CATS-50%
(33
0.769
0.675
0.571
0.421
0.744
0,412
TEAL-50%
37
0.778
0.673
0.405
GRIFFIN-50%
16
0.86
0.778
0.571
0.745
0.428
R-Sparse-500s 67
0.773
0.674
0.543
0.746
0.408
PolarSparse-505
0.89
0.779
0.687
0.568
0.381
0.762
0.439

---

2333
Searite
Searsty
300
Botih
(a) OPT 6.7b
(b) OPT 66b
Figure 6: OPT models sparse decoding throughput, seq len 1920 using pipeline parallelism. (a) OPT 6.7b, critical threshold 50%. (b) OPT 66b, critical threshold 30%. Polar Sparsity delivers up to 2.2x higher throughput than dense and up 2x than standard activation sparsity at scale.
1200
Pola Sporsity
Fila Sprstyy
140
100
Dccodh Throgg
2
32 320 256
300
32
322
100
(a) LLaMA-2-7b
(b) LL 1-1bbb
Figure 7: LLaMA models sparse decode throughput results using pipeline parallelism. (a) LLaMA-2- 7b seq len 3968, critical threshold 50%.(b) LLaMA-3.1 seq len 8192, critical threshold 62.5%. Polar Sparsity delivers up to 1.85x higher throughput than dense baseline at scale.
5.2 Generation Throughput
In this section, we present the decoding throughput of various models at different batch sizes and sequence lengths their respective critical threshold. We used sequence lengths of 1920 for OPT, 3968 for LLaMA-2, and 8192 for LLaMA-3 to evaluate sparsity and system performance increasing workload scales. Additional results for varying sequence lengths are Appendix E.2
3968
shows the throughput results of OPT models using Deja Vu-style (39 activation sparsity and Polar Sparsity with the dense baseline different batch sizes. We substitute DejaVu's sparse GEMV with our Selective GEMM with dynamic top-k to support efficient batched execution. At batch size Polar Sparsity performance is similar that of conventional activation sparsity. As batch size increases, the effectiveness of conventional activation sparsity diminishes due to the reduced union sparsity across sequences. In contrast, Polar Sparsity leverages the stable, batch-invariant sparsity of attention heads and scales efficiently. Polar Sparsity delivers up to 2.2x higher decoding throughput than dense and more than DejaVu at large batch sizes. In LLaMA-2 models, where sparsity is applied only to attention layers, speedups become apparent at larger batch sizes, reaching up to 1.85 x. ( throughput results for I.LaMA 3.1 70B. Despite higher attention density (62.5%), Polar Sparsity achieves 1.51x speedup at large batch sizes. While primarily report

---

throughput, Polar Sparsity naturally reduces inter-token latency by similar margin and accelerates per-query generation. Appendix provides a detailed evaluation across models and workloads.
Limitations and Future Work
Polar Sparsity is most effective in large-scale, high-throughput settings, particularly during batched decoding. Its benefits diminish for small-batch inference due to limited GPU workload and reduced parallelism. Additionally, we fix the top-k activated heads/groups per layer, whereas a dynamic, input- or layer-adaptive strategy could yield better efficiency and accuracy. Head sparsity could also be combined with token sparsity for potential multiplicative gains. Our experiments use benchmarks with context lengths up 16k tokens. As newer LLMs reach million-token contexts, extending our analysis to these ultra-long settings is an important direction for future work. We expect our approach to provide efficiency gains in attention variants like Multi-Quers [49] and Multi-Latten Attention [2]: however, the critical threshold could be higher similar to GQA. GQA models have naturally less KV diversity in the attention layers and group sparsity also inherently weaker compared to head sparsity, as it may overlook important heads that reside in inactive groups. This limitation likely contributes to the faster aceuracy degradation observed in GQA models. Our current evaluation focuses on greedy decoding. Extending Polar Sparsity to beam search and speculative decoding, where sparsity patterns may differ, offers an exciting direction for future research.
Our approach maintains average accuracy within 1% of the original model at the critical threshold, and this minor drop could potentially be recovered through targeted fine-tuning. Further gains may be possible by training models to induce higher attention head or group sparsity. Since head/group sparsity is batch invariant, there is promising opportunity explore task-aware query-sensitive routing, allocating more heads to harder queries and fewer to easier ones within the same batch. Moreover, decoding difficulty can vary across tokens, even within a single sequence. Some tokens may be predicted with fewer active heads, allowing higher sparsity to be applied adaptively at each decoding step. fine-gu fine-grained router that dynamically selects heads based on context and difficulty could unlock lossless sparse inference with even greater throughput. We see this as promising direction toward more adaptive, efficient, and task-aware LLM inference systems.
7 Conclusion
Our work highlights the scalability and effectiveness of contextual sparsity for accelerating batched LLM inference. We introduce Polar Sparsity, a key insight showing that batch size and sequence length grow, the importance of sparsity transitions from MLP layers, where union activation increases, to Attention layers, where head-level sparsity remains stable and batch-invariant.1 To exploit this prop- erty, we develop Selective Head Attention with sparsity-awarr aware GPU kernels that execute computations only for activated heads and neurons. Together, these optimizations deliver consistent speedups across a wide range of models, batch sizes, and sequence lengths with minimal impact on accuracy. Our results are competitive with state-of-the-ard approaches and delivers up 2.2x end-to-end speedups in large-scale settings, affirming the practical viability of Polar Sparsity for efficient and scalable LLM serving. This method is step towards realizing scalable, high-performance, batched LLM inference that meets the growing demands of modern applications.
Acknowledgments
This work was performed under the auspices of the U.S. Department of Energy by Lawrence Livermore National Laboratory under Contract DE-AC52-07NA27344 (LLNL-CONT-2205579) This research was also supported in part by the National Science Foundation under Grant No. 2203033. A portion of this work was conducted during an internship at NVIDIA.
References
[1] Muhammad Adnan, Akhil Arunkumar, Gaurav Jain, Prashant J. Nair, Ilya Soloveychik, and Purushotham Kamath. Keyformer: Kv cache reduction through key tokens selection for efficient generative inference, 2024.
10

---

[2] Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints, 2023.
[3] Yash Akhauri, Ahmed F AbouElhamayed, Jordan Dotzel, Zhiru Zhang, Alexander M Rush, Safeen Huda, and Mohamed S Abdelfattah. Shadowllm: Predictor-based contextual sparsity for large language models, 2024.
[4] Keivan Alizadeh, Iman Mirzadeh, Dmitry Belenko, Karen Khatamifard, Minsik Cho, Carlo C Del Mundo, Mohammad Rastegari, and Mehrdad Farajtabar. Llm in flash: Efficient large language model inference with limited memory, 2024.
[5] Yuwei An, Zhuoming Chen, Chenyan Xiong, and Beidi Chen. Herd: Grouping before pruning for batch inference, 2025.
[6] Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. LongBench: A bilingual, multitask benchmark for long context understanding. In Lun-We Ku, Andre Martins, and Vivek Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 3119-3137, Bangkok, Thailand, August 2024. Association for Computational Linguistics.
[7] Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language, 2019.
[8] Mengzhao Chen, Wenqi Shao, Peng Xu, Jiahao Wang, Peng Gao, Kaipeng Zhang, and Ping Luo. Efficientqat: Efficient quantivation-wware training for large language models, 2024.
[9] 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, 2018
[10] Peizhuang Cong, Qizhi Chen, Haochen Zhao, and Tong Yang. Baton: Enhancing batch-wise inference efficiency for large language models via dynamic re-batching, 2024.
[11] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory efficient exact attention with io-awareness, 2022.
[12] DeepSeek-AI, Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, et al. Deepseek- v2: strong, economical, and efficient mixture-of-expprts language model, 2024.
[13] DeepSeek-AI, Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang, Deli Chen, Dongjie Ji, Erhang Li, et al. Deepseek-v3 technical report, 2025.
[14] Aditya Desai, Shuo Yang, Alejandro Cuadron, Ana Klimovic, Matei Zaharia, Joseph E. Gonza- lez, and lon Stoica Hashattention: Semantic sparsity for faster inference, 2024.
[15] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 4171 4176,, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
[16] Harry Dong, Beidi Chen, and Yuejie Chi. Prompt-prompted adaptive structured pruning for efficient Ilm generation, 2024.
[17] 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.

---

[18] Aayush Gautam, Susav Shrestha, and Narasimha Reddy, Token-driven gammatune: Adaptive calibration for enhanced speculative decoding. 2025.
[19] Danilo Giampiccolo, Bernardo Magnini, Ido Dagan, and Bill Dolan. The third PASCAL recognizing textual entailment challenge. In Satoshi Sekine, Kentaro Inui, Ido Dagan, Bill Dolan, Danilo Gimpiicolo, and Bernardo Magnini, editors, Proceedings of the ACL-PASCAL Workshop on Textual Entailment and Paraphrasing, pages 1-9, Prague, June 2007. Association for Computational Linguistics.
[20] Andrew Gordon, Zornitsa Kozareva, and Melissa Roemmele. SemEval-2001 task 7: Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In Eneko Agirre. Johan Bos, Mona Diab, Suresh Manandhar, Yuval Marton, and Deniz Yuret, editors, *SEM 2012: The First Joint Conference on Lexical and Computational Semantics - Volume 1: Proceedings of the main conference and the shared task, and Volume Proceedings of the Sixth International Workshop on Semuntic Evaluation (SemEval 2012), pages 394-398 Montréal, Canada, 7-8 June 2012. Association for Computational Linguistics.
[21] 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, et al. The llama herd of models, 2024.
[22] Zhiyu Guo, Hidetaka Kamigaito, and Taro Watanabe. Attention score is not all you need for token importance indicator in KV cache reduction: Value also matters. In Yaser Al- Onaizan, Mohit Bansal, and Yun-Nung Chen, editors, Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages
Miami, Florida,
USA, November 2024. Association for Computational Linguistics.
[23] Shwai He, Guoheng Sun, Zheyu Shen, and Ang Li. What matters in transformers? not all attention is needed, 2024.
[24] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding, 2021.
[25] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network, 2015.
[26] Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in deep learning: pruning and growth for efficient inference and training in neural networks. J. Mach. Learn. Res., 22(1), January 2021.
[27] Cheng- Yu Hsich, Chun-Liang Li, Chih-Kuan Yeh, Hootan Nakhost, Yasuhisa Fujii, Alexander Ratner, Ranjay Krishna, Chen- Yu Lee, and Tomas Pfister. Distilling step-by-step! outperforming larger language models with less training data and smaller model sizes, 2023.
[28] Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Banford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b, 2023.
[29] Peng Jin, Bo Zhu, Li Yuan, and Shuicheng Yan. Moh: Multi-head attention mixture-of-head attention, 2024.
[30] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, 2020.
[31] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, SOSP '23, page 611-626, New York, NY, USA, 2023. Association for Computing Machinery.
[32] Qi Le, Enmao Diao, Ziyan Wang, Xinran Wang, Jie Ding, Yang, and Ali Anwar. Probe pruning: Accelerating Ilms through dynamic pruning via model-probing, 2025.
12

---

[33] Donghyun Lee, Je-Yong Yong Lee, Genghan Zhang, Mo Tiwari, and Azalia Mirhoseini. Cats: -omtextually-aware thresholding for sparsity in large language models, 2024.
[34] Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding, 2023.
[35] Wei Li, Lujun Li, Mark Lee, and Shengjie Sun. Adaptive layer sparsity for large language models via activation correlation assessment. In A. Globerson, L.. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors, Advances in Neural Information Processing Systems, volume 37, pages 109350-109380. Curran Associates, Inc., 2024.
[36] Zonglin Li, Chong You, Srinadh Bhojanapalli, Daliang Li, Ankit Singh Rawat, Sashank J. Reddi, Ke Ye, Felix Chern, Felix Yu, Ruiqi Guo, and Sanjiv Kumar. The lazy neuron phenomenon: On emergence of activation sparsity in transformers, 2023.
[37] James Liu, Pragaash Ponnusamy, Tianle Cai, Han Guo, Yoon Kim, and Ben Athiwaratkun. Training-free activation sparsity in large language models, 2025.
[38] Yifei Liu, Jicheng Wen, Yang Wang, Shengyu Ye, Li Lyna Zhang, Ting Cao, Cheng Li, and Mao Yang. Vptq: Extreme low-bit vector post-training quantization for large language models, 2024.
[39] Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivas- tava, Ce Zhang, Yuandong Tian, Christopher Ré, and Beidi Chen. Deja contextual sparsity for efficient Ilms at inference time. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org. 2023.
[40] Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du, Tao Jiang, Chao Hong, Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang, Zhiqi Huang, Huan Yuan, Suting Xu, Xinran Xu, Guokun Lai, Yanru Chen, Huabin Zheng, Junjie Yan, Jianlin Su, Yuxin Wu, Neo Y. Zhang, Zhilin Yang, Xinyu Zhou, Mingxing Zhang, and Jiezhong Qiu. Moba: Mixture of block attention for long-context Ilms, 2025.
[41] Yuqi Luo, Chenyang Song, Xu Han, Yingfa Chen, Chaojun Xiao, Zhiyuan Liu, and Maosong Sun. Sparsing law: Towards large language models with greater activation sparsity, 2025.
[42] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
[43] Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? a new dataset for open book question answering. In Ellen Riloff, David Chiang, Julia Hockenmaier, and Jun't ichi Tsujii, editors, Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 2381-2399, Brussels, Belgium, October-November 2018. Association for Computational Linguisties.
[44] Iman Mirzadeh, Keivan Alizadeh, Sachin Mehta, Carlo C Del Mundo, Oncel Tuzel, Golnoosh Samei, Mohammad Rastegari, and Mehrdad Farajtabar. Relu strikes back: Exploiting activation sparsity in large language models, 2023.
[45] Bowen Pang, Kai Li, and Feifan Wang. Optimizing Ilm inference throughput via memory-aware and sla-constrained dynamic batching, 2025.
[46] Qwen, :, An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, Huan Lin, Jian Yang, Jianhong Tu, Jianwei Zhang. Jianxin Yang, Jiaxi Yang. Jingren Zhou, Junyang Lin, Kai Dang. Keming Lu, Keqin Bao, Kexin Yang, Le Yu, Mei Li, Mingfeng Xue, Pei Zhang, Qin Zhu, Rui Men, Runji Lin, Tianhao Li, Tianyi Tang, Tingyu Xia, Xingzhang Ren, Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang. Yu Wan, Yuqiong Liu, Zeyu Cui, Zhenru Zhang, and Zihan Qiu. Qwen2.5 technical report, 2025.
[47] Hyun rae Jo and Dongkun Shin. A2sf: Accumulative attention scoring with forgetting factor for token pruning in transformer decoder, 2024.
13

---

[48] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale, 2019.
[49] Noam Shazeer. Fast transformer decoding: One write-head is all you need, 2019.
[50] Jiho Shin, Hoeseok Yang, and Youngmin Yi. Sparseinfer: Training- frree prediction of activation sparsity for fast Ilm inference, 2025.
[51] Susav Shrestha, Aayush Gautam, and Narasimha Reddy. Storage access optimization for efficient gpu-centric information retrieval. The Journal of Supercomputing, (1(4):613, 2025.
[52] Susav Shrestha, Narasimha Reddy, and Zongwang Li. Espn: Memory-efficient multi-vector information retrieval. In Proceedings of the 2024 ACM SIGPLAN International Symposium on Memory Management, ISMM 2024, page 95-107, New York, NY, USA, 2024. Association for Computing Machinery.
[53] Chenyang Song, Xu Han, Zhengyan Zhang, Shengding Hu, Xiyu Shi, Kuai Li, Chen Chen, Zhiyuan Liu, Guangli Li, Tao Yang, and Maosong Sun. ProSparse: Introducing and enhancing intrinsic activation sparsity within large language models. In Owen Rambow, Leo Wanner, Marianna Apidianaki, Hend Al-Khalifa, Barbara Di Eugenio, and Steven Schockaert, editors, Proceedings of the 3 st International Conference on Computational Linguisties, pages 2626- 2644, Abu Dhabi, UAE, January 2025. Association for Computational Linguistics.
[54] Yixin Song, Zeyu Mi, Haotong Xie, and Haibo Chen. Powerinfer: Fast large language model serving with a consumer-grade gpu, 2024.
[55] Yupeng Su, Ziyi Guan, Xiaoqun Liu, Tianlai Jin, Dongkuan Wu, Graziano Chesi, Ngai Wong, and Hao Yu. Lim-barber: Block-aware rebuilder for sparsity mask in one-shot for large language models, 2024.
[56] Gemini Team, Petko Georgiev, Ving Ian Lei, Ryan Burnell, Libin Bai, Anmol Gulati, Garrett Tanzer, Damien Vincent, Zhufeng Pan, Shibo Wang, Soroosh Mariooryad, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context, 2024.
[57] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, et al. Llama 2: Open foundation and fine-tuned chat models, 2023.
[58] Tim Valicenti, Justice Vidal, and Ritik Patnaik. Mini-gpts: Efficient large language models through contextual pruning. 2023.
[59] Hanrui Wang. Zhekai Zhang, and Song Han. Spatten: Efficient sparse attention architecture with cascade token and head pruning. In 2021 IEEE International Symposium High--Prformance Computer Architecture (HPCA). IEEF, February 2021.
[60] Yubo Wang, Xueguang Ma, Ge Zhang, Yuansheng Ni, Abhranil Chandra, Shiguang Guo, Weim- ing Ren, Aaran Arulraj, Xuan He, Ziyan Jiang, Tianle Li, Max Ku, Kai Wang, Alex Zhuang, Rongqi Fan, Xiang Yue, and Wenhu Chen. Mmlu-pro: a more robust and challenging multi-task language understanding benchmark. In Proceedings of the 38th International Conference on Neural Information Processing Systems, NIPS '24, Red Hook, NY, USA, 2025. Curran Associates Inc.
[61] Haihang Wu. Llm-bip: Structured pruning for large language models with block-wise forward importance propagation, 2024.
[62] Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, Y. X. Wei, Lean Wang, Zhiping Xiao, Yuqing Wang, Chong Ruan, Ming Zhang, Wenfeng Liang, and Wangding Zeng. Native sparse attention: Hardware-aligned and natively trainable sparse attention, 2025.
[63] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can a machine really finish your sentence? In Anna Korhonen, David Traum, and Lluís Marquez, edi- tors, Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages Florence, Italy, July 2019. Association for Computational Linguistics.
14

---

[64] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke 7ettlemoyer. Opt: Open pre-trained transformer language models, 2022.
[65] Xiaofeng Zhang, Yikang Shen, Zeyu Huang, Jie Zhou, Wenge Rong, and Zhang Xiong. Mixture of attention heads: Selecting attention heads per token. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang, editors, Proceedings of the 2022 Conference on Empirical Methods Natural Language Processing, pages 4150-4162, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics.
[66] Zhengyan Zhang, Yankai Lin, Zhiyuan Liu, Peng Li, Maosong Sun, and Jie Zhou, Moefication: Transformer feed-forward layers are mixtures of experts, 2022.
[67] Zhenyu Zhang, Zechun Liu, Yuandong Tian, Harshit Khaitan, Zhangyang Wang, and Steven Li. R-sparse: Rank-aware activation sparsity for efficient IIm inference, 2025.
[68] Haizhong Zheng, Xiaoyan Bai, Xueshen Liu, Z. Morley Mao, Beidi Chen, Fan Lai, and Atul Prakash. Learn to be efficient: Build structured sparsity in large language models, 2024.
[69] Zhen Zheng, Xin Ji, Taosong Fang, Fanghao Zhou, Chuanjje Liu, and Gang Peng. Batchlim: Optimizing large batched Ilm inference with global prefix sharing and throughput-oriented token batching, 2025.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction clearly state the core contributions of the paper, including the design of Polar Sparsity inferencing system for large language models and the achieved throughput improvements. These claims are substantiated by experimental results and discussions in Sections[5 and limitations are acknowledged in section 6
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
The limitations are discussed in Section 6 where we note that our approach does not accelerate the prefill stage and may yield smaller gains for short-sequence and small batch decoding tasks. We also highlight that sparsity thresholds are statically defined per model, and future work could explore dynamic sparsity adaptation across decoding steps.
Justification:
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
15

---

The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: This paper focuses empirical and systems- contributions for acceler- ating large language model inference through contextual activation sparsity and does not include any formal theoretical results proofs.
Guidelines:
The answer NA means that the paper does include theoretical results. the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
The experimental setup. model configurations, sparsity strategies, and hardware details are described detail Sections We also include discussions and additional results in the appendix to facilitate reproducibility of our main findings.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
16

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that are appropriate to the research performed.
also be provided via detailed
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibilityy which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes)
Justification: We provide anonymized code and scripts in the supplemental material along with instructions to reproduce the main experimental results. The code includes setup details, model configurations, and commands for running key benchmarks.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g.. for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
7 7/nips. wwwpcccccccc/.....c....co..eetle
wwtppppgguidese//ppppppppopppaep guwdee/eee//eee CooeeSumisssii ssssspolicyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
17

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e-g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Section describes the model architectures and router designs, while Section 5 ooullines the inference setup across different models, including how OPT models apply sparsity to both MIPP and Attention layers, whereas Llama models sparsify only the Attention layers. The experiments cover varying batch sizes and sequence lengths. Additional router training details and results provided in the appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report error bars selected figures such as the neuron batch activation graphs, where activation patterns vary slightly across batches. For throughput results, use CUDA Graphs for inference, which show minimal variability across different runs. To ensure robustness, evaluate method across range of batch sizes, sequence lengths, and model scales, as detailed in Section and Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
18

---

Answer: [Yes]
Justification: We provide hardware details in Section including GPU models (NVIDIA A100 DGX Node). Experiments were conducted on internal compute nodes with 80 GB of GPU memory with pipeline parallelism (Section()) and tensor parallelism (Section E). While exact execution time per run is not provided. we deseribe throughput benchmarks across models and batch sizes. Preliminary experiments and router training consumed additional compute.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't It make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
The research adheres to the NeurlPS Code of Ethics. It involves no human subjects, private data, or potentially harmful applications, and all models and benchmarks used are publicly available and appropriately cited.
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations their jurisdiction..
Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We discuss broader impacts in Appendix F+ highlighting both the positive potential of enabling efficient and accessible LLM inference, well risks related the increased scalability of generative models. We also emphasize the importance of responsible deployment practices.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinfornation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
19

---

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: This work does not introduce release any new pretrained models, datasets, or generative systems with high risk for misuse. All experiments are conducted using publicly available models under standard licensing terms.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We use publicly released models such as OPT and LLaMA, which are cited appropriately in the main text. We respect their licenses, and do not modify redistribute model weights. All party libraries and benchmarks are used accordance with their respective licenses.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcodd com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
20

---

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We release new Triton kernels for Selective GEMM and Selective Head/Group Attention as part of our sparsity- -awar inference system. These assets are documented with usage instructions, kernel configurations, and integration guidance in the supplemental material and accompanying repository.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/code/modell as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This work does not involve any crowdsourcing, user studies, or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This work does not involve any research with human subjects and therefore does not require IRB approval or equivalent review.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
21

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the methods in this research?
Answer: [NA]
While the research focuses on improving inference efficiency for existing large language models, the models themselves are used in standard ways for evaluation purposes. No LLM was used as component in developing core methods or generating original model outputs novel or non-standard manner.
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standarr components.
Please refer to our LLM policy for what should or should be described.
22

---

A Extended Related Works
Inference efficiency has been a central focus in machine learning systems, with a diverse set of strategies such as sparsity, pruning, quantization, distillation, speculative decoding, SSD-based offloading, and advanced memory management techniques proposed to reduce computational cost, memory footprint, and latency 8 38 27 18
25 5T 31]. These techniques address distinct bottlenecks and are often complementary, enabling combined use for greater gains.
A.1 Contextual Activation Sparsity
Recent work on activation sparsity in large language models has explored a variety of mechanisms to reduce inference costs without sacrificing accuracy. Early studies revealed that only a small fraction of neurons in transformer MLP layers are active for any given input, establishing the foundation for sparsity-aware inference (3).. Building on this, ReLUfication and ProSparse showed that replacing GELU or SwiGLU with ReLU or progressively encouraging ReLU-like behavior exposes inherent sparsity and improves inference efficiency, even without retraining 44 [33]
Training-free methods such as TEAL and R-Sparse apply post-hoc magnitude-based pruning across
models like L-a and L-LMMA33 and preserve accuracy 337
67]. CATS improves on this by
making the threshold contextually adaptive, dynamically pruning low-activation neurons based on input features [33]. Deja Vu and ShadowILM train lightweight predictors to anticipate important neurons and attention heads token-by-token, enabling dynamic sparsity at inference time (39 3). Our work draws inspiration from previous predictive approaches. Complementary to predictor- based sparsity, LTE and GRIFFIN optimize structured sparsity patterns during pretraining or in a prompt-aware manner, enabling the model to internalize which parts of the architecture to use for different inputs 68 [6. Finally, hardware-awar methods like PowerInfer leverage the heavy- distribution of neuron activations to map "hot" neurons onto GPU memory while streaming "cold" ones from the CPU, delivering efficient offloading based inference
54
Existing methods for improving batch efficiency using contextual sparsity are limited in scope. PowerInfer considers batch sizes up to 32 but targets CPU-based offloading systems, where the baseline is already constrained by CPU and PCIe bandwidth [54]. Herd attempts to group sequences with similar sparsity patterns to increase overlap [5]. but operates batch sizes <4 and is difficult to apply in dynamic, real-world random queries, where identifying on-the-fly fly sequences that share similar activation patterns becomes prohibitively challenging. These approaches sparsify only the MLP block and speedup decreases with larger batch sizes. In contrast, our method handles arbitrary sequences and scales effectively to large batch sizes, aligning with realistic deployment scenarios.
A.2 Attention Head Sparsity and Token Sparsity
Several recent works explore token sparsity to reduce KV cache and attention computation during generative inference. HashAttention introduces semantic sparsity by mapping tokens a learned Hamming space, enabling efficient selection of pivotal tokens using hash collisions [4]. A2SF
proposes accumulative attention scoring with a forgetting factor to fairly rank token importance across time steps in autoregressive decoders [47] Keyformer and VATP both address token pruning by extending beyond attention scores: Keyformer selects a small set of "key tokens" based on attention concentration [1, while VATP incorporates the value vector norm into token importance estimation (22. Native Sparse Attention introduces a hardware-aligned sparse attention mechanism with trainable sparsity patterns suited for long-sequence modeling (62), and MoBA applies mixture- of-expert routing principles to attention tokens by dynamically selecting relevant blocks for each query 44). These approaches are largely orthogonal to method and can potentially be combined with Polar Sparsity for multiplicative gains.
Two recent methods, Mixture of Attention Heads (MoA) 65 and MoH: Mixture-ofe oe-d Attention 29 treat attention heads as experts and learn routers to activate a subset per token. MoA requires full model training from scratch, while MoH involves extensive fine-tuning of pre-trained models to learn dynamic head weighting. Although both demonstrate accuracy improvements, they offer limited efficiency gains due to reliance on dense attention kernels. In contrast, our method keeps the model backbone fixed, trains only lightweight routers for head activation, and leverages our custom Select Head Attention kernel to achieve real wall-clock speedups. We believe prior work was limited
23

---

by the absence of efficient sparse attention primitives, and with the integration of our kernel, these and future approaches have the potential achieve both high accuracy and meaningful efficiency gains. We hope this unlocks new opportunities and renews momentum in head sparsity research for accelerating L.M inference.
B
Extended Study on Batch Activation
B.1 MLP Activations
300
300
502
25
25
25
OPT 6.7b
(b) Prosparse Llama-22 2-bb
Figure 8: OPT and Prosparse Llama Neuron Batch Activations
To better understand the behavior activation sparsity under batching, we conduct an empirical study across multiple models and sparsification strategies. In this study, neuron is considered active if its output greater than zero. We collect true MLP activation statistics by running a forward pass over randomly selected samples from the WikiTexx-2 dataset. For each model, group the samples into batches of varying sizes and compute the average union neuron activation and its standard deviation for each layer.
OPT models, which use ReL.U activations, exhibit strong inherent sparsity. shown Figure Sa early layers are significantly sparser than deeper ones, trend consistent across all OPT models. While increasing batch size reduces overall sparsity, the early layers retain enough sparsity to benefit from selective execution. To test the limits of sparsity, we simulated a batch of 10,000 samples and observed that the activations the initial layers do approach dense and are not, in fact, dead neurons. However, for practical batch sizes, these layers are sufficiently sparse support accelerated inference. Recent work has shown that sparsity in ReLU-based models tends to increase during pretraining [41). The OPT models, having undergone extensive large-scale pretraining, likely benefit from this effect, resulting in more structured and persistent activation sparsity.
Figure 8b shows that Prosparse introduces deeper-layer sparsity, but with high variance and similarly degrading trends larger batch sizes. Unlike OPT, these Sparse-LLaae LLaMA variants are only lightly fine-tuned on smaller datasets, which likely limits the emergence of stable sparsity patterns observed in fully pretrained models. Given these observations, we focus our selective execution LLaMA models attention head sparsity, where the activation structure remains more consistent across batches.
B.2 Attention Head Activations
Figure presents heat maps of the head activation counts for the attention heads across all layers in the OPT 6.7B and Llama 7B models on 100,000 random token samples from the WikiText-2 dataset. The visualizations reveal that activation patterns vary substantially both from layer to layer and from head to head. In both architectures, some heads are activated significantly more often than others, leading to skewed distribution of usage. Most layers show relatively uniform spread of activations across all heads. These findings suggest that future work should focus on designing dynamic head

---

opt-6.75
Liama2-7B
40003
100993
22000
Fledd
Usead Inces
(a) OPT-6.7b
(b) Llama-2-7b
Figure 9: Head activation heat map
allocation strategies that adapt to each layer's unique activation and importance profile, rather than relying on a fixed global threshold.
C Sparsity Prediction
Each router is a two-layer feedforward network with hidden dimension of 1024, trained independently for each transformer layer. To partially hide the latency of the MLP router, we overlap its execution with the Attention layer. is trained using supervised learning with ground-truth neuron activations derived from dense MLP outputs. inference time, the router produces scalar logits for each neuron, which are used to rank and select the top-k neurons for activation. The attention head routers use a single layer feedforward network as described in Section 4.2 Since the attention routers are much smaller than the MLP routers, we simply run them synchronously before each attention layer. The routers are optimized as binary classifiers using a binary cross-entropy loss, with the AdamW optimizer. We use batch size of 64, a learning rate of le-4, and early stopping over maximum of 20 epochs. The LLM parameters are frozen during router training. Supervision data is collected from inference runs on the WikiText-2 dataset, as described in Section To determine minimal top-k values for the MLP layers, we apply a simple greedy algorithm (Algorithm 2 that incrementally adjusts threshold to meet the target recall of 99%. This calibration is performed offline for each OPT model variant.
Our MLP router follows the design used in prior contextual sparsity work [39 54
Algorithm 2 Greedy Top-k Selection to Meet Target Recall
Require: R (router), H (hidden states), 7' (true activations), ko (initial top-k), Trarget (target recall), 8 (step size)
1: k + ku
2: 7-0 0
3: while < Furget do
4:
Å R(H)
Predict activations
5:
< COMPUTERECAIL
6:
k<
7:
end while
8:
return
C.1 Router Ablation Study
Figure shows the decode latency of MLP and Attention blocks along with their respective routers, and includes the measured inference latency of sparse kernels across selected sparsity levels. The MI.P router introduces approximately four times higher latency compared to the attention router. At higher activation levels, the combined latency of the MLP router and sparse MLP approaches,
25

---

Latency
Labency
25-
233
L20010 (1)
Lateng
0.0
00
100660
Houter
spa'5o
Sparse 20%
sparse
Dense
Nouter
sparse Sime
Sparse 80hs
Figure 10: OPT 66b, MLP and Attention Decode Latency with Router Different Sparsity Levels. Batch size 64, Seq len 1920. MLP router 4 x more expensive than Attention router.
exceeds, that of the dense MLP. To address this, overlap the MLP router execution with the attention block, following a strategy similar to Deja Vu (39]. in order to hide its cost. However, under batched inference, it is not always possible fully mask the MLP router latency due to high GPU utilization. In practice, we observe that overlapping saves approximately 0.1 of router latency.
and in
In contrast, the attention router is significantly smaller and more efficient, enabling synchronous execution without incurring high overhead. This results in substantial performance gains: at 50% head activation, the attention block achieves latency reduction of approximately 1.4 ms, even after including the router cost. Under similar conditions, the MLP router yields only modest improvement of around 0.21 ms. This further highlights the scalability head sparsity for large batch workloads.
D Sparse Kernels
Standard implementations of selective matrix multiplication typically involve separate indexing of W1.sa W2.9p followed by dense GEMM operations, introducing unnecessary memory overhead. To avoid this, we fuse the indexing and GEMM into single kernel that dynamically processes only the activated neurons specified by the neuron index tensor as shown in Algorithm 3 To ensure coalesced memory access, we store the weight matrices with the neuron dimension, d, as the innermost (contiguous) dimension. Unlike prior work that often targets sparse matrix-vector (GEMV) operations, our kernel optimized matrix-matrix (GEMM) operations, enabling efficient execution for arbitrary batch sizes.
Algorithm 3 Sparse Fused GEMM Kernel
1: Input: A E IR MXK B E RRXN index vector I
2: Output: C' Mx
3: Initialize thread IDs, memory offsets
4: Gather indexed columns rows of using/
5: Compute C = A B Bxelected via block-wise multiplication
6: Apply activation function (e.g., ReLU)
7: Store result back to memory
Figure [] demonstrates that our selective MLP and attention kernels exhibit strong performance and deliver consistent speedups across NVIDIA GPU families (A5000, A A100, and H100). As MLP and Attention density increases, latency scales near linearly, indicating that the kernels maintain high efficiency and effectively exploit sparsity without incurring significant overheads.
26

---

4100 Dene
CAO 201
60
02
0.5
66
Acardian Deralyy
(a) Selective MLP Kernel
(b) Selective Attention Kernel
Figure 11: Selective MLP and Attention Kernel Performance across GPU Families. OPT 66b, batch size 64, sequence length 1920.
E Extended Throughput and Latency Evaluation
F.
Pipeline Parallel Execution
900
Activstio Spanityy (0ejo Md Folr Soursity lour)
6622
900
100
200
1.25
2007
(a) OPT 30b
(b) LLaMA 2-13b
Figure 12: Additional Sparse Decoding Throughput Results. (a) OPT 30b critical attention density 40%, seq len 1920 (b) LLaMA-2-13b attention density 50%, seq len 3968
In this section, we present our results using pipeline parallel setup without batching. Pipeline parallelism is technique that divides large model into sequential stages, distributing the com- putation across multiple GPUs to improve memory efficiency and throughput, and it is the most commonly used technique for inferencing.
Figure 12a shows decoding throughput for the OPT model family. For the OPT-6.7B model, despite around 90% sparsity in the MIP layer, the speedup is modest for small batch sizes due to insufficient workload to fully utilize the GPU's parallel capabilitics, highlighting the diminishing returns of sparsity in smaller models when the GPU is already underutilized. However, larger batch sizes, speedups increase to 1.83x as attention computation becomes more memory-intensive. In contrast, the larger OPT-66B model shows more significant improvements, with speedups ranging from 1.66x at batch size to 2.2x at batch size 64, indicating better scaling with workload size.
Figure 7 12b show the decoding throughput and speedup for the Llama model family. Since we sparsify only attention layers, speedups are limited smaller batch sizes due to overhead of the router and the dominance of ML.P layers in computational cost. As batch size increases, the computational burden shifts towards the attention layers, making sparsity more effective, For Llama 2 models, this results in speedups of up to 1.85x.
27

---

E.1.1 Tensor Parallel Execution
Done
165
1456
1,00
1000
300
800
500
222
400
200
16601
(a) OPT 66b TP world size
OPT 66b, TP world size
Figure 13: OPT 66b model with Tensor Parallel Sparse Decoding Throughput.
Figure [3 shows the decoding throughput and speedup for the OРT-66B model with tensor parallelism (TP) set and 4. Tensor parallelism splits individual model layers across multiple devices so each computes portion of the operation in parallel, reducing memory and computation per device, and is commonly used during training. As the TP degree increases, we observe that the speedup for smaller batch sizes decreases. This is because the workload is divided among multiple GPUs, and applying sparsity further reduces the workload per GPU, leading to diminishing returns when the workload already small. This is similar what we observed the OPT-6.7B model, where sparsity had limited impact at low batch sizes due to under utilization. However, at larger batch sizes, where the computational and memory I/O costs become more significant, Polar Sparsity achieves speedups of up to 1.8x, highlighting its effectiveness in reducing overhead at scale.
E.2 Decode Latency Analysis
The figure 14 illustrates the decode latency of OPT 6.7b and OPT 66b models using dense, standard activation sparsity (Deja Vu), and our proposed Polar Sparsity methods across varying sequence lengths, with fixed batch size of 16. Across both model scales, Polar Sparsity consistently achieves lower latency than both Dense and Deja Vu baselines, offering up to 2 speedup over Dense and up 1.52x Deja Vu in this workload. The figure 15 show the decode latency of LLaMA-2-7B and LLaMA-3.1 70B models across increasing sequence lengths using dense and Polar Sparsity, with batch size fixed to 16. Polar Sparsity consistently outperforms the dense baseline, achieving up
Denca
Dence Activation Sparsity (Deja Wi Sparisty
Activation Spansity (Deja W)
Polar Sperialyy
Dec neccatn Laterey
Diocaato Lattenyyy
43-
2>
126
755
1054
1950
326
355
112
1004
1250
Sequence Length
Sequence Lenggh
(a) OPT 6.7b
(b) OPT 66b
Figure OPT models decode, inter-token latency with fixed batch size 16. Polar Sparsity reduces latency up to 2x compared to dense baseline and up to 1.52x compared to Deja Vu in this workload.
28

---

Donse Spiratty
Polar Speriily
Detode Labercy (roo
Dctode: Laterayy (moo
:20
333
3024
2041
X01
3043
106
912
1234
20100
.007
8192
Sequence Length
Sequence Length
(a) LLaMA-2-7b
(b) LLaMA-3-70b
Figure 15: LLaMA models decode, inter-token latency with fixed batch size of 16. Polar Sparsity reduces latency up to 1.77 x compared to the dense baseline in this workload.
1.77x speedup. As with throughput, even greater gains can be expected larger batch sizes as the attention module dominates latency.
F
Broader Impacts
Our work improves the efficiency of LLM inference by leveraging contextual activation sparsity, enabling faster and more resource-efficient deployment. This has the potential to make LLMs more accessible in low-resource settings, and support wider adoption in academic and industrial applications. By lowering computational costs, our approach contributes toward more sustainable Al development.
At the same time, accelerating LLM inference can lower the barrier to mass deployment of generative models, which may increase the risk of misuse, such as generating disinformation or harmful content scale. While our method does not alter model behavior or training data, we recognize that efficiency gains can amplify existing ethical concerns. encourage responsible deployment practices and safeguards when integrating such techniques into real-world systems.
29


Reviews:
[{"review_id": "hieu7KpCmI", "reviewer": "Reviewer_nrZF", "summary": "This paper proposes Polar Sparsity, a batch-friendly inference framework that combines a dynamic, layer‐wise top-k MLP router with a custom Selective GEMM kernel and a lightweight attention router with sparsity-aware Selective FlashAttention kernels to exploit high-throughput contextual sparsity. Built on top of FlashAttention and CUDA graphs, Polar Sparsity achieves throughput improvements across various models, demonstrating that contextual sparsity can be practically scaled to high-throughput LLM deployment with minimal changes.", "questions": "See Weaknesses.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The paper is well written and easy to understand.\n2. The figures are well plotted to help understand the effectiveness of Polar Sparsity.\n\nWeaknesses:\n1. Under the single-query or small-batch setting—typical for interactive chat or RLHF—the speedup ratio is limited due to the extra routing overhead.\n2. More experiments should be conducted on harder reasoning tasks (e.g., GSM-Hard, MATH).\n3. Throughput performance could be tested on more devices, e.g., H100, consumer RTX, and AMD.\n4. No performance comparisons to the methods deploying the batch inference with structured sparsity.\n5. The training-based sparsification with customized sparse kernel design makes it hard to quickly deploy on a new model.", "was_revised": true, "final_justification": "The rebuttal has addressed most of my concerns. The proposed polar sparsity approach effectively resolves the batch inference issue associated with contextual sparsity in LLMs. While I acknowledge that polar sparsity leads to a relatively notable performance drop on math reasoning benchmarks, I recognize its contribution to enabling sparse batch inference. Therefore, I raise my score to weak acceptance.", "rebuttal": "We sincerely thank the reviewer for their time and for providing thoughtful and constructive feedback on our submission. We are encouraged that the reviewer found the paper well-written and the figures clear. We have carefully considered the weaknesses raised and have performed additional experiments and analysis to address them. We believe these additions significantly strengthen the paper and we hope the reviewer will reconsider their rating.\n\nBelow, we address each of the weaknesses point-by-point.\n\n## 1. Performance in Single-Query or Small-Batch Settings\n   \nWe appreciate the reviewer's observation regarding performance in low-batch settings. Our primary focus with Polar Sparsity is to maximize performance for high-throughput LLM serving, a dominant use-case in real-world cloud deployments. Most service providers leverage techniques like continuous batching to serve large numbers of concurrent requests (including chat models with RLHF), maximizing hardware utilization and cost-effectiveness [1]. Our work provides a novel solution to reduce latency and boost throughput in this critical, batch-heavy environment.\n\nThat said, we want to emphasize that the overhead of our method is minimal.\n\n+ Minimal Router Overhead: We have conducted extended router evaluation in Appendix C.1, where our ablation study shows that the attention routers are exceptionally lightweight and add negligible overhead.\n\n+ Meaningful Low-Batch Speedups: Even in low-batch settings, Polar Sparsity delivers tangible performance gains. As shown in Figures 5 and 6 of our original submission, we achieve meaningful speedups, demonstrating the efficiency of our approach even when not operating at maximum throughput.\n\nWhile the most dramatic speedups are realized in high-throughput scenarios, our highly optimized routers and kernels ensure that the gains from sparsity consistently outweigh the minimal overhead, making our framework beneficial even in smaller-batch contexts.\n\n## 2. Performance on Harder Reasoning Tasks\n\nThis is an excellent suggestion. To demonstrate the robustness of Polar Sparsity on more complex tasks, we have conducted new experiments on the MMLU-PRO, GSM8K and LongBench benchmarks. These results, which we will incorporate into the revised manuscript, show that our sparse models maintain high accuracy with only a minor, graceful trade-off.\n\nSummary of New Benchmark Results for instruction tuned chat models:\n\n| Subject                    | Llama-3.1-8B-Inst_dense | Llama-3.1-8B-Inst_sparse_0.625 | Mistral-7B-Inst_dense | Mistral-7B-Inst_sparse_0.5 |\n|---|---|---|---|---|\n| MMLU_PRO_Biology           | 0.62 | 0.60 | 0.47 | 0.50 | \n| MMLU_PRO_Business          | 0.44 | 0.43 | 0.21 | 0.20 | \n| MMLU_PRO_Chemistry         | 0.26 | 0.24 | 0.12 | 0.11 | \n| MMLU_PRO_Computer Science  | 0.42 | 0.41 | 0.28 | 0.29 | \n| MMLU_PRO_Economics         | 0.53 | 0.51 | 0.34 | 0.35 | \n| MMLU_PRO_Engineering       | 0.24 | 0.27 | 0.16 | 0.15 | \n| MMLU_PRO_Health            | 0.51 | 0.49 | 0.31 | 0.30 | \n| MMLU_PRO_History           | 0.41 | 0.37 | 0.25 | 0.29 | \n| MMLU_PRO_Law               | 0.28 | 0.28 | 0.17 | 0.17 | \n| MMLU_PRO_Math              | 0.39 | 0.35 | 0.17 | 0.17 | \n| MMLU_PRO_Other             | 0.46 | 0.44 | 0.32 | 0.29 | \n| MMLU_PRO_Philosophy        | 0.43 | 0.42 | 0.26 | 0.24 | \n| MMLU_PRO_Physics           | 0.34 | 0.34 | 0.19 | 0.18 | \n| MMLU_PRO_Psychology        | 0.59 | 0.59 | 0.44 | 0.43 | \n| MMLU_PRO Overall Accuracy  | 0.41 | 0.40 | 0.25 | 0.24 | \n| GSM8K                      | 0.75 | 0.71 | 0.34 | 0.32 |\n| LongBench                  | 0.443 | 0.429 | 0.392 | 0.388 |\n\nThese results confirm that Polar Sparsity preserves model capabilities on challenging reasoning and long-context tasks, with overall accuracy showing only a minor degradation (e.g ~1% on MMLU_PRO Overall Accuracy) relative to the dense models. \n\nCrucially, we want to highlight that this is the first work to demonstrate wall-clock speedups by enforcing sparsity on attention heads. The minor accuracy degradation opens an exciting avenue for future work toward completely lossless sparsity. As we explain in the paper, the batch-invariant nature of our head sparsity router means we can dynamically apply different sparsity levels to different queries within the same batch. This would allow \"harder\" queries to use more heads, preserving quality, while \"easier\" queries benefit from higher acceleration, a unique and powerful capability of our framework.\n\n## 3. Throughput Performance on More Devices\n\nWe agree that demonstrating performance across a range of hardware is vital. We have now benchmarked our sparsity-aware Selective FlashAttention kernel on a variety of modern GPUs, including the NVIDIA RTX A5000, A100, and H100.\n\nThe results below show that our kernel's speedup is consistent and significant across different hardware tiers.\n\nAttention Kernel Latency (ms) and Speedup (Batch=64, SeqLen=2k, Num Heads = 72, MHA)\n| Head Activation | A5000 Speedup | A100 Speedup | H100 Speedup |\n| :---: | :---: | :---: | :---: |\n| 1.0 (Dense) | 1.00x | 1.00x | 1.00x |\n| 0.8 | 1.17x | 1.22x | 1.24x |\n| 0.6 | 1.52x | 1.57x | 1.62x |\n| 0.5 | 1.80x | 1.82x | 1.91x |\n| 0.3 | 3.18x | 2.76x | 3.08x |\n| 0.1 | 7.93x | 5.59x | 8.23x |\n\nThe data confirms that our approach is not architecture-specific. At a typical 50% head activation, we achieve a ~1.8x-1.9x speedup across all tested GPUs. While the kernel was developed on A100, its excellent performance on both consumer and datacenter hardware is highly encouraging. We will include these kernel benchmarks in the paper and commit to providing full end-to-end throughput results for all models on these devices in the revised manuscript.\n\n## 4. Comparison to Structured Sparsity Methods\n\nWe thank the reviewer for this question and would like to clarify our positioning. We consider contextual sparsity to be an advanced, dynamic form of structured sparsity. We did compare Polar Sparsity with Deja Vu, a state-of-the-art method that applies structured sparsity to MLP and linear layers. Our results in Figure 5 show that **Polar Sparsity is up to 2x faster than Deja Vu**, marking a significant improvement.\n\nOur work is distinct from traditional structured pruning techniques [2] for two critical reasons:\n\n+ Addressing the Right Bottleneck: Most structured sparsity methods focus only on MLP/linear layers. In large-batch, high-throughput inference, the attention mechanism is often the primary bottleneck. Our work is the first to introduce custom kernels that bring the benefits of structured sparsity directly to the attention computation, which is a key novelty and contribution.\n\n+ Training Efficiency: Traditional methods often require full, computationally intensive sparse training to prune model weights. Our approach is far more practical, requiring only the training of a lightweight router.\n\n## 5. Deployment on New Models\n\nWe appreciate the reviewer’s concern about the practicality of our training-based method. We wish to clarify that our approach is designed for fast and easy deployment and does not require full model retraining or fine-tuning.\n\nThe MLP and attention routers are very lightweight and are trained completely separately from the frozen, pre-trained model weights. As detailed in Appendix C and the provided supplementary code, this process is highly efficient. For example, training the attention router for each model takes less than an hour on two A5000 GPUs.\n\nThis makes our framework significantly more agile and easier to apply to new models compared to methods that depend on extensive and costly retraining of the entire model.\n\n## Conclusion\n\nIn conclusion, we have demonstrated that Polar Sparsity is a practical and impactful framework that (1) excels in the critical high-throughput serving scenario, (2) maintains strong performance on hard reasoning tasks, (3) generalizes across different modern GPU architectures, (4) outperforms existing structured sparsity methods by tackling the attention bottleneck, and (5) is lightweight and easy to deploy.\n\nWe are grateful for the reviewer's feedback, which has helped us improve the paper. We believe these new results and clarifications address the concerns raised and demonstrate the significance of our work. We respectfully hope the reviewer will consider increasing their score.\n\n\nReferences\n\n[1] Orca: A Distributed Serving System for Transformer-Based Generative Models.\n\n[2] Accelerating Sparse Deep Neural Networks. arXiv:2104.08378."}, {"review_id": "elKvDeddtZ", "reviewer": "Reviewer_sVp2", "summary": "LLMs have been known to have very sparse activations. With specific choices of rectifying activations in MLP layers, such as ReLU, it has been observed that the activations are sparse. Similarly, attention heads have been shown to sparsely activate in the sense of their contributions to the output representation. The costs of these layers during inference depend on the batch size and length of the sequences involved. For short sequences in unbatched settings, the MLP layers tend to be the bottleneck, with runtime dominated by the time to read weights from memory. For longer sequences, the cost of reading the KV cache in the attention layers dominates the inference time. On the other hand, batching makes the MLP layers less sparse as the union of activations over the batch samples slowly approaches the dense computation. Furthermore, the MLP layers are memory-bound until the batch size hits a sufficient size, and increasing the batch size doesn't add to the cost. However, attention layers still maintain sparsity since the KV caches are independent for each batch item. The authors thoughtfully chose to optimize inference for the long-sequence and large-batch-size regime, where attention is the bottleneck. They do so by developing GPU kernels to exploit head sparsity effectively.", "questions": "- Is \"head sparsity\" a standard term? I was a bit confused initially about what it referred to. It might help define it early on in the introduction/abstract.\n- The head sparsity remains stable and batch-invariant only for independent batch samples. What about use cases where we wish to generate multiple samples for the same batch item? The verifier in speculative decoding, autocomplete usecases, beam search, etc. sample multiple tokens from the same sequence.\n- The paper claims decode stage dominates end-to-end latency, especially when generating long outputs spanning hundreds or thousands of tokens. I don't think this is necessarily true because the prefill stage has a quadratic computational complexity O(L^2) where L is the prefix length, and the decode stage has linear complexity O(L * D) where D is the number of decode steps and D << L.\n- The paper claims to be the *first* to demonstrate that contextual sparsity scales with batch size. I'd argue that this has been a well-known fact, and most works aimed towards long-context efficiency would continue to work in batched settings. It's a soft opinion and is very subjective, but I'm curious to know if the authors have looked into works focusing on long-context efficiency.\n- Are there any dead neurons in the initial layers? Fig 1 (b) seems to be suggestive of dead neurons and could be suitable for pruning.\n- What are the percentages in Fig 2 (a)? What layer(s) are being plotted? It can be guessed from the text in 3.2, but I think the figure caption should be self-standing.\n- What is the loss function used to train the routers?\n- What does accuracy refer to in Figure 4?\n- Have the naive implementations of the gather + attention operations been tested with different compilers? The operations sound like they can be fused easily by a compiler. I am not sure if the publicly available versions of the compilers support fusing arbitrary operations with custom calls yet (and I guess flash attention would require a custom call), but I don't think flash attention would help with performance until you hit a certain batch size.", "limitations": "Yes, mostly. I think this method wouldn't scale for multiple independent samples for the same sequence (like in beam search or the verifier in speculative decoding).", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "I'm leaning towards increasing the score if the authors make good-faith effort to address the weaknesses below and answer the questions raised.\n\n**Strengths:**\n1. The paper overall is a complete piece of work and is written well.\n2. The evaluations are thorough and cover a very broad set of models.\n3. Impactful work.\n\n**Weakness:**\n1. No description of the loss functions used to train the routers in the main text.\n2. No error estimates for results (Fig 4 and Table 1). The standard error for a lot of downstream evaluations are in several percentage points and is therefore extremely important to report error bars.\n3. In general, downstream evals tend to be noisy and statistically unreliable. On the other hand, training metrics such as perplexity/top1-accuracy tend to be statistically reliable, but often poorly correlate with downstream performance. However, reporting results on both sets of metrics may improve the fidelity of the results. It's important since even a 1% gap in training top1 accuracy at scales of 70B can be as costly as doubling the model size, given the diminishing returns of the scaling laws. \\*\n4. There should be evaluations on long-form text generation. Is it possible that the approximation holds good most of the time, but the model fails catastrophically occasionally? But such occasional 1 in 100 failures can completely break decoding in long-form generation. The averages in the evaluations, most of them short-form generation, could easily hide these small modes.\n\n\\* Since this work is primarily approximating a pretrained model, I suspect that good retention on training metrics could likely correlate with good retention on downstream performance, as approximations don't fundamentally alter the model to cause dramatic changes in downstream evaluations. I am not sure; this is just a guess, and I might be wrong.\n\n**Presentation Issues:**\n1. Please ensure that the figure captions can stand on their own. What are the percentages in Fig 2 (a)?\n\n**Minor nitpicks and soft opinions:**\n1. I found it quite hard to understand the list of contributions at the end of the introduction. There are too many new terms and pharses being introduced without context, and they are not likely to be understandable without reading the entire paper (which defeats its purpose). Example terms: \"Selective\" GEMMs, \"layer-wise top-k optimization strategy\", \"Group sparsity\".\n2. The list of contributions uses the phrase \"negligible accuracy loss\". It would be useful to report the exact top1 accuracy loss to make it clear; \"negligible\" is very subjective. I'd be concerned with even a 0.5% loss for very large models because that's the gap you see with doubling the model size at those scales.\n3. There could be very specific failure modes with this kind of approximation, like the one highlighted in weakness comment 4. Would it be possible to construct cases where this approximation fails? The fact that some downstream evals seem to be impacted more than others suggests that there could be specific instances where this approximation fails.\n4. I think the following excerpt from the paper is just one interpretation of the results. It could also be the case that most downstream don't catch issues. What is common in the evals that do worse than others?\n> This task-dependent sensitivity is particularly encouraging, as it suggests that most tasks can be served with only the most critical heads, while harder tasks could potentially activate more for higher accuracy within the same batch.\n5. Table 2 doesn't add value to the paper if the generation throughput or some other runtime performance quantity is not reported. It's possible to trade throughput for quality; hence, reporting just quality isn't useful.", "was_revised": true, "final_justification": "The authors have identified and executed the idea well. There are some concerns with loss in performance on certain downstream tasks, but their method might still be Pareto optimal (which is not established in the paper but probably is the case). The new kernel as a contribution doesn't feel sophisticated enough to warrant a publication, but it's possible that I missed a few details; the authors do not disclose the code nor did they document enough implementation details highlighting the complexity. I'm retaining my score at 4, but would give 4.5. I'd have increased to 5 if the implementation is technically complex. Overall, the concern is whether this should have been a bug report to the compiler's repository or a paper. It would have been interesting to see results on other compilers such as XLA.", "rebuttal": "We sincerely thank the reviewer for their constructive feedback. These points have helped us improve the paper's rigor and clarity. We have taken action on all suggestions as detailed below.\n\n---\n## **Responses to Weaknesses**\n\n## **1. Router Loss Function**  \nWe agree this information is better placed in the main text. Details on our router training, including the use of Binary Cross-Entropy (BCE) loss, were in Appendix C.\n\n## **2. Error Estimates for Results**  \nThis is an excellent suggestion. We will update all figures and tables with error bars in the revised version to ensure statistical reliability.\n\n## **3. Inclusion of Perplexity/Top-1 Accuracy Metrics**  \nWe agree that including training metrics would strengthen the evaluation. In Section 3 and Figure 2(a), we used the relative change in perplexity as a guiding signal to study head sparsity. We will report these training metrics for all evaluated models in the revised manuscript to ensure a more complete analysis.\n\n## **4. Evaluation on Long-Form Generation**  \nThis is a critical point. We have run extensive new experiments on LongBench-e with three instruction-tuned models. The results confirm your intuition: our high-fidelity approximation leads to robust downstream performance without catastrophic failures. Our sparse models show only a minor, graceful degradation. Accuracy reported out of 100. We have also included additional results in the response to reviewer nrZF for MMLU_PRO.\n\n| Metric                | Llama-3.1-8B-Inst_dense | Llama-3.1-8B-Inst_sparse_0.5 | Llama-3.1-8B-Inst_sparse_0.625 | Mistral-7B-Inst_dense | Mistral-7B-Inst_sparse_0.5 | Qwen2.5-14B-Inst_dense | Qwen2.5-14B-Inst_sparse_0.5 |\n|---|---|---|---|---|---|---|---|\n| 2wikimqa              | 15.28 | 14.28 | 14.50 | 32.30 | 31.99 | 13.76 | 12.52 |\n| gov_report            | 34.66 | 31.79 | 33.56 | 28.64 | 28.60 | 30.99 | 27.34 |\n| hotpotqa              | 17.50 | 15.17 | 15.46 | 38.71 | 37.11 | 14.08 | 14.35 |\n| lcc                   | 69.56 | 60.72 | 63.70 | 56.78 | 56.11 | 67.90 | 60.41 |\n| multi_news            | 25.45 | 25.08 | 25.56 | 22.60 | 22.18 | 21.53 | 21.70 |\n| multifieldqa_en       | 27.79 | 27.92 | 26.61 | 32.91 | 32.09 | 35.18 | 32.12 |\n| passage_count         | 15.75 | 18.22 | 18.68 | 4.04 | 4.84| 7.64 | 11.05 |\n| passage_retrieval_en  | 97.12 | 94.24 | 95.46 | 31.67 | 37.33 | 89.64 | 94.61 |\n| qasper                | 11.83 | 11.85 | 12.49 | 30.04 | 28.34 | 10.67 | 10.94 |\n| repobench-p           | 55.31 | 45.19 | 48.76 | 47.14 | 41.77 | 47.62 | 40.15 |\n| samsum                | 42.73 | 38.39 | 40.79 | 40.84 | 39.74 | 44.69 | 41.01 |\n| trec                  | 71.00 | 71.67 | 71.67 | 61.33 | 61.33 | 75.00 | 75.00 |\n| triviaqa              | 92.81 | 89.72 | 90.72 | 83.11 | 82.68 | 88.43 | 86.61 |\n| LongBench-e average   | 44.37 | 41.86 | 42.92 | 39.24 | 38.78 | 42.09 | 40.60 |\n\n| Model  / LongBench-e accuracy by context length | 0-4k | 4-8k | 8k+ |\n|---|---|---|---|\n| Llama-3.1-8B-Inst_dense              | 44.01 | 44.30 | 44.80 |\n| Llama-3.1-8B-Inst_sparse-0.5         | 42.45 | 42.13 | 41.02 |\n| Llama-3.1-8B-Inst_sparse-0.625       | 43.35 | 43.00 | 42.41 |\n| Mistral-7B-Inst_dense                | 44.26 | 39.09 | 34.37 |\n| Mistral-7B-Inst_sparse_0.5           | 44.39 | 38.34 | 33.60 |\n| Qwen2.5-14B-Inst_dense               | 42.45 | 41.51 | 42.31 |\n| Qwen2.5-14B-Inst_sparse              | 40.46 | 40.47 | 40.87 |\n\nAcross all these new, intensive evaluations, we observed no decoding collapses. This provides strong empirical evidence of stability for long-form generation. We will add a summary of these new results to the paper and are happy to include further analysis in the appendix.\n\n---\n\n## **Addressing Presentation Issues**\n\nWe will revise all the captions to be self-contained. The percentages in Figure 2(a) indicate the relative increase in validation perplexity over the dense baseline, as detailed in Section 3.2.\n\n## **Addressing Soft Opinions**\n\n- **Introduction Clarity, \"Negligible Accuracy Loss\":**  In the revision, we will rewrite the introduction to be more self-contained, briefly defining key terms as they are introduced. We will also report precise accuracy instead of using subjective terms to make our claims more accurate. \n\n- **Constructing Failure Cases:**  It is certainly possible to construct cases where our approximation is less effective, though this would require significant additional work. We see this as an exciting continuation of our research. Being able to identify such cases would allow for a dynamic, query-wise head activation strategy, paving the way for potentially lossless sparse inference.\n\n- **Analysis of Task Sensitivity:** While we have not observed catastrophic failures, we have noticed that harder reasoning, code, math benchmarks (e.g., MMLU-PRO Math, GSM8K, lcc, repobench) are more susceptible to accuracy degradation than other tasks. This aligns with the idea that more active heads perform better for harder queries.\n\n- **Value of Table 2:** This is a very helpful suggestion. We will include generation throughput in our revision, covering both single-query and large-batch scenarios to provide a more complete evaluation of the proposed approaches.\n---\n\n## **Answers to Questions**\n\n### **Q: Is \"head sparsity\" a standard term?**  \n**A:** Head sparsity isn't yet a standard term. We will define it early on to make the idea clear. \n\n\n### **Q: Does batch-invariance hold for use cases like beam search or speculative decoding?**  \n**A:** \nIn standard beam search, a beam width of k is handled by expanding the batch dimension, treating each beam as an independent sequence. Therefore, our batch-invariance property holds, as each parallel beam can activate a different set of heads. However, for a highly optimized implementation where beams can share the KV cache and IO for the shared prefix, the head sparsity can potentially decrease depending on the overlap of head activations. \n\nSpecDecode: The drafting stage can still take advantage of our approach and we should be able to make batched speculative drafting more efficient. The verification stage can potentially have decreased head sparsity as different verification tokens within the batch can activate different heads. Exploring these topics would be an exciting future line of work.\n\n### **Q: Is the decode stage always the bottleneck over prefill?**  \n**A:** You are correct about the prefill cost in long-prompt, short-generation scenarios (D≪L). Our work focuses on use cases where the decode stage is the practical bottleneck, such as long-form generation (D≥L) and multi-turn chat. In these applications, the decode stage's sequential, memory-bound nature dominates user-perceived latency, especially since the initial prefill cost is an amortized, one-time expense. From a systems perspective, while prefill has quadratic complexity, it is highly parallel and often achieves near-full GPU utilization. In contrast, the decode stage is inherently sequential and memory-bandwidth bound, making its latency the practical bottleneck in these use cases. We recognize our original claim lacked this precision and will revise the paper to clarify this specific context.\n\n### **Q: Is this the \"first\" work to show contextual sparsity scales with batch size?**  \n**A:** Our claim was made in the specific context of contextual sparsity; existing methods often report reduced speedup with larger batches. In Appendix A.1, we survey recent contextual sparsity approaches and document this behavior. In Appendix A.2 we discuss token level sparsity techniques that skip or prune tokens for improved long context efficiency. We believe combining those methods with Polar Sparsity could yield even greater throughput gains. We agree with your general point that most attention optimization should work well in batched settings. Our key contribution is being the **first to design, build, and demonstrate practical, high-performance GPU kernels (Selective FlashAttention) that convert theoretical *attention head sparsity* into measured, wall-clock speedups in high-throughput scenarios.**\n\n### **Q: Are there dead neurons in the initial layers?**  \n**A:** This is a sharp observation. As we discuss in Appendix B.1, these are not 'dead' neurons but rather a visualization of extremely high contextual sparsity, where very few neurons activate for any single query.\n\n### **Q: What does \"accuracy\" refer to in Figure 4?**  \n**A:** The 'accuracy' in Figure 4 refers to the zero-shot accuracy on downstream tasks. We will revise the caption and make sure all our graphs are self-sufficient. \n\n### **Q: Have you tested against advanced compilers that could fuse these operations?**  \n**A:** We tested the naive implementation with torch.compile, and while it helps, it still doesn't perform better than the standard dense baseline.\n\nbatch size 64, 72 Heads, seq len 2k, MHA, 50% sparsity, A5000 GPU\n| method | latency (ms) |\n|---|---|\n| naive gather + flashattn | 14.7 |\n| torch.compile(gather + flashattn) | 9.8 |\n| dense flashattn | 8.7 |\n| select flashattn (ours) | 4.8 |\n\n---\n\n## **Conclusion**\n\nWe'd like to conclude by expressing our genuine appreciation for your review. We are truly grateful for a review that was not only technically rigorous but also deeply engaged with the core value of our work. Your feedback has been pivotal in helping us clarify our contributions and better articulate the paper's impact. We believe the manuscript is significantly stronger now and hope we have fully addressed your points to merit an increased score."}, {"review_id": "J4CHE3qYub", "reviewer": "Reviewer_bAJa", "summary": "This paper proposes Polar Sparsity, which sparsifies the MLP and attention mechanism during the LLM decoding stage.", "questions": "- In Figure 1, I think the x label is wrong. Isn't it \"sequence length\" rather than \"batch size\"? The latency of attention is the only variable on that figure, so I suspect that the label is wrong.\n- Can you point out how S_B is constructed? I cannot well understand how an activated neuron is selected for each batched sequence. Is it just a simple top-k of input activation? If so, what does the input statistics look like? Can you make a heatmap of those tensors? And you mentioned layer coherences, but the visualization of that fact is insufficient, as I think. Can you add more statistics or visualization, such as a heatmap or histogram, to them? I feel like this paper is well written about efficiency and implementation, and not very friendly on methodological explanation and statistical visualizations. Just adding two or three figures about statistics that are needed for intuitions might improve the readability of this paper significantly.\n- I think the input sequence should not be random for neural activation statistics sampling (Figure 1b). Can you sample it from real data, such as PG19?", "limitations": "Please refer to the weakness section for more information. \n\nDespite their mention of MQA and MLA as their first limitation, I do not agree that this limitation significantly draws down their effectiveness. I suggest that authors think about another limitation of this work; therefore, it further helps future researchers.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 2, "strengths_and_weaknesses": "# Strength\n1. Very advanced and well-constructed implementation based on OpenAI Triton. \n1. Well scales up with larger batch size, thanks to well-constructed block-wise sparsification rather than scalar or fine-grained level sparsification methodologies.\n\n# Weakness\n1. Limited evaluation of context length. I think this paper should evaluate their model on a longer context, because serving a longer context is much more challenging and a more effective scenario when we optimize the Transformers (including the attention mechanism). I think this would improve the effectiveness of this method in longer sequence inference, such as RAG and Agentic AI. Moreover, the main assumption of activation sparisity is not also well studied with long-context. \n1. Limited evaluation of sparsity. In Figure 4, the plot only deals with attention sparsity, while the methodology can also dynamically sparsify the MLP. Furthermore, I do not agree that \"density\" is an \"efficiency metric\", because mostly FLOPs or mathematical asymptotic computational cost does not well match with real-world GPU kernel latency. But I believe authors are able to easily change this to actual latency on various conditions of GPUs (consumer-grade, server-grade, multi-node), while considering their excellent implementation.\n1. More performance drawdown happens on a better model. This might be trivial or not, but I think the ideal trend of this should be better model --> better efficiency. However, if we see the MMLU scores in Figure 4, the Llama 3.1 70B shows a more significant drawdown of the MMLU score compared to OPT. OPT maintains its original performance up to 20%, while Llama 3.1 starts to fall from 60% density. I am disappointed about this result, so I hope authors can test on other modes such as Qwen3 or Mistral so we can show that this is only a special phenomenon on Llama 3.1.\n1. Only applied to the decoding part. I think we can apply the sparsification method in the prefill stage also, so we can accelerate the speed of long-sequence prefill, which is usually bottlenecked on long-code prompting.", "was_revised": true, "final_justification": "I still have a concern about long context Evaluations, because the authors only did the subset of InfBench, and the results are almost same as random performance. \n\nHowever, it is clear that this method is effective in certain condition, so I will increase my voting to borderline accept.", "rebuttal": "We thank the reviewer for the thoughtful and constructive feedback. We are encouraged by the recognition of the implementation quality and scalability of our approach. Below, we provide detailed responses and updates:\n\n## Response to Weakness \n\n### 1. Context length evaluation\n\nThank you for this excellent suggestion. We agree that evaluating on long-context benchmarks is crucial for demonstrating the real-world utility of our method. To address this, we have conducted new experiments on the LongBench-e dataset with several state-of-the-art instruction-tuned models. The full set of results can be viewed in the response to reviewer sVp2.\n\n| Model  / LongBench-e accuracy by context length | 0-4k | 4-8k | 8k+ |\n|---|---|---|---|\n| Llama-3.1-8B-Instruct_dense              | 44.01 | 44.30 | 44.80 |\n| Llama-3.1-8B-Instruct_sparse-0.5         | 42.45 | 42.13 | 41.02 |\n| Llama-3.1-8B-Instruct_sparse-0.625       | 43.35 | 43.00 | 42.41 |\n| Mistral-7B-Instruct-v0.1_dense           | 44.26 | 39.09 | 34.37 |\n| Mistral-7B-Instruct-v0.1_sparse_attn_0.5 | 44.39 | 38.34 | 33.60 |\n| Qwen2.5-14B-Instruct_dense               | 42.45 | 41.51 | 42.31 |\n| Qwen2.5-14B-Instruct_sparse              | 40.46 | 40.47 | 40.87 |\n\nOn long-context benchmarks, model quality remains high. At 50% attention sparsity, the average performance drop is minimal across models (e.g., just -0.46 points for Mistral-7B and -1.49 for Qwen2.5-14B). While challenging coding tasks are more sensitive, many other tasks show negligible change. This modest accuracy trade-off is paired with a significant 1.77x speedup in long-context decoding (Appendix E.2). We will also include latency and throughput evaluations for the newly tested models in the revised manuscript.\n\n### 2. Evaluation of sparsity\n\nWe agree that density, theoretical FLOPs are insufficient and that real-world latency is the most meaningful metric. We have used \"density\" in Figure 4 primarily to illustrate the trade-off between model accuracy and the degree of attention head activation. The actual wall-clock speedups are presented in Figure 3, 5, 6, Appendix E. Using our optimized kernels, we can improve batch decoding up to 2.2x as reported in the original submission.\n\nRegarding MLP sparsity, you are correct that Figure 4 focuses on attention. In those experiments, MLP sparsity was determined dynamically by our algorithm as described in Section 4.1.\n\nTo directly address your concern about hardware-specific performance, we have benchmarked our optimized Triton kernels on both consumer (A5000) and server-grade (A100, H100) GPUs.\n\nAttention Kernel Latency (ms) and Speedup (Batch=64, SeqLen=2k, Num Heads = 72, MHA)\n| Attention Density | A5000 Dense | A5000 Sparse | A5000 Speedup | A100 Dense | A100 Sparse | A100 Speedup | H100 Dense | H100 Sparse | H100 Speedup |\n|---|---|---|---|---|---|---|---|---|---|\n| 1.0 | 8.67 | 8.85 | 0.98x | 3.26 | 3.30 | 0.98x | 2.12 | 2.14 | 0.99x |\n| 0.8 | 8.67 | 7.43 | 1.16x | 3.26 | 2.67 | 1.21x | 2.12 | 1.71 | 1.24x |\n| 0.6 | 8.67 | 5.71 | 1.52x | 3.26 | 2.08 | 1.56x | 2.12 | 1.31 | 1.62x |\n| 0.5 | 8.67 | 4.82 | 1.80x | 3.26 | 1.78 | 1.82x | 2.12 | 1.11 | 1.91x |\n| 0.3 | 8.67 | 2.73 | 3.18x | 3.26 | 1.18 | 2.76x | 2.12 | 0.69 | 3.07x |\n| 0.1 | 8.67 | 1.09 | 7.92x | 3.26 | 0.58 | 5.58x | 2.12 | 0.25 | 8.22x |\n\nThe results show consistent and significant speedups that scale linearly with sparsity across different GPU architectures. We achieve ~1.8-1.9x speedup at 50% sparsity.\n\n### 3. Performance drawdown\n\nThis is a very insightful point. We agree the ideal trend is for better models to show more resilience to sparsity. As shown in our original Figure 2, this trend holds within a model family (e.g., OPT-66B is more robust than OPT-6.7B). The variance between families (OPT vs. Llama 3.1) likely stems from different architectures and training recipes.\n\nAs requested, we evaluated Mistral-7B to see if the Llama 3.1 result was an outlier.\n\n| attn_density | arc_challenge | arc_easy | copa | hellaswag | mmlu | openbookqa | piqa | rte | winogrande | average |\n|---|---|---|---|---|---|---|---|---|---|---|\n| 1 | 0.489 | 0.796 | 0.920 | 0.609 | 0.591 | 0.332 | 0.803 | 0.686 | 0.738 | 0.663 |\n| 0.8 | 0.491 | 0.796 | 0.920 | 0.609 | 0.591 | 0.340 | 0.802 | 0.682 | 0.736 | 0.663 |\n| 0.6 | 0.488 | 0.792 | 0.920 | 0.607 | 0.577 | 0.342 | 0.804 | 0.668 | 0.734 | 0.659 |\n| 0.5 | 0.483 | 0.793 | 0.920 | 0.608 | 0.562 | 0.340 | 0.801 | 0.671 | 0.736 | 0.657 |\n| 0.4 | 0.474 | 0.789 | 0.920 | 0.606 | 0.535 | 0.330 | 0.801 | 0.646 | 0.733 | 0.648 |\n\nMistral-7B also maintains its performance robustly down to 50% density. The greater drop on harder tasks like MMLU suggests that while most queries benefit from sparsity, the most challenging ones may require higher head activations. This observation reinforces the value of our dynamic, batch-invariant design, which allows future work to build query-sensitive routers that activate more heads only when needed, making our efficient kernels a critical enabler for such techniques. We evaluated additional models on harder tasks like MMLU_PRO, GSM8K and observed similar trends (see response to reviewer nrZF). This does suggest that Llama-3.1 family could be a possible outlier in this data.\n\nWe want to emphasize that we only train lightweight routers leaving the base model frozen. Crucially, while any approximation technique can be lossy, our method's configurable and batch-invariant nature is a key strength. It allows different queries in a batch to activate different heads, paving the way for future work on query-sensitive routing to achieve nearly lossless performance. Our work provides the first efficient kernel implementation to make this direction practical. A recent line of work proposed to fine-tune the base model with head sparsity and has been shown to be able to completely recover accuracy [1]. While this work did not manage to show efficiency gains, we strongly believe that future research expanding this line of work will benefit from our optimized kernels and be able to extract wall-clock speedups with minimal to no accuracy loss. We include further discussions in Appendix A. \n\n[1]. MoH: Multi-Head Attention as Mixture-of-Head Attention. \n\n### 4. Applied only to decoding\n\nWe thank the reviewer for this insightful suggestion. The reviewer is correct that applying our sparsification method to the prefill stage is a valuable potential extension of our work.\n\nIn this paper, we chose to focus on the decoding stage as it presents a distinct, memory-bandwidth-bound bottleneck, especially during long-sequence generation, which contrasts with the compute-bound nature of the prefill stage. While sparsification could apply here, the performance trade-offs are different, and it may require a fundamentally different approach, kernels to be effective. \n\nWe agree that accelerating prefill is an important problem. We will add this to the future work section of our paper and clarify the scope of our current method in the introduction. We appreciate the reviewer for pointing out this promising research direction.\n\n## Response to Questions\n\n### 1. Figure 1 Label\n\nThe x-axis in Figure 1 is correctly labeled 'batch size'. The figure illustrates a key motivation for our work: during autoregressive decoding, attention latency scales linearly with batch size and context length, while the latency of other components like MLPs does not. The latency of other modules increases, but at a much slower rate.\nThis disproportionate scaling makes attention the primary bottleneck at larger batch sizes, which our method directly targets.\n\n### 2. How is S_B constructed? \n\nS_B is the set of unique active neurons for an entire batch, constructed as follows:\n\n+ A small router network predicts an importance score for each neuron for every sequence in the batch.\n\n+ For each sequence, we select the top-k neurons based on these scores.\n\n+ S_B is then the union of these selected neurons across all sequences in the batch. Therefore, a neuron is used in the computation if it is selected for at least one sequence.\n\nWe appreciate the feedback on improving the methodological intuition. We have included some activation statistics in Appendix B which shows the average activation per layer with standard deviation across a range of batch sizes. We will add the requested heatmaps and histograms of neuron activation to the revised manuscript to better illustrate these dynamics. Additionally, we already have included heat maps for head activations in Appendix B.2.\n\n### 3. Input sequences in Figure 1b\n\nThe term \"random sequences\" was meant to convey that the input queries were sampled from a diverse dataset without regard to semantic similarity, ensuring a varied activation pattern. We will clarify this in the manuscript and use samples from a standard dataset like PG19 for our new visualizations to improve realism and reproducibility.\n\n## Response to limitations\n\nWe use fixed attention head activation/layer to enable clean benchmarking of our architecture. However, results on challenging benchmarks like MMLU suggest that uniform activation is suboptimal. This motivates a dynamic attention head router, analogous to our MLP router, that can allocate more heads to complex queries. Our high-performance kernels make this direction feasible, and we will include this discussion in the revised manuscript.\n\n## Conclusion\n\nWe thank you once again for your time and insightful feedback. We have incorporated extensive new experiments and analyses directly addressing your concerns about evaluation on long contexts, different hardware, and newer models. We are confident that the manuscript has been substantially strengthened by these additions and hope you will view the revised submission favorably."}, {"review_id": "z6mVqxQ1tO", "reviewer": "Reviewer_cRa2", "summary": "This work proposes an approach for computation reduction at inference time, targeting batch inference via leveraging contextual sparsity in MLP and Attention layers. Given the diminishing gains of sparse MLP execution at large batch sizes, the proposed approach mostly focuses on the attention layers, which are shown to exhibit stable sparsity pattern with the increase of number of samples processed in a single pass. Activated neurons/attention heads are determined via lightweight MLP. And efficient GEMM and Attention kernel is developed and shown to achieve ~2x higher throughput at small or no performance drop. Polar Sparsity is validated on several language models families - OPT, Llama-2 and Llama-3.1.", "questions": "- Which framework is adopted for inference? Is the inference benchmark conducted within HuggingFace transformers of some efficient inference engine, such as vLLM or SGLang? If using transformers - which versions of torch/transformers/flash_attention are adopted? The reported speed-ups may vary noticeably depending on a specific version. Are the models compiled prior to execution?\nOPT models are very outdated as of 2025. Would be interesting to see, how would the Figure 2 (a) look for newer generations of models, such as Llama-3 or Qwen-2.5/3.\n- Wikitext-2 may be not the best source of data for router training. A higher quality data source (such as FineWeb-Edu) may yield better results on benchmarks.\n- I think it would be interesting as well to apply Polar Sparsity on a long-context tasks, such as LongBench suite.\n\n[References]\n- [1] Penedo, Guilherme, et al. \"The fineweb datasets: Decanting the web for the finest text data at scale.\" Advances in Neural Information Processing Systems 37 (2024): 30811-30849.\n- [2] Bai, Yushi, et al. \"Longbench: A bilingual, multitask benchmark for long context understanding.\" arXiv preprint arXiv:2308.14508 (2023).", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n- The training of sparse routers is quite cheap - relatively small amount of data and finetuning budget is sufficient to get good results.\n- The provided speed-up is quite decent, given that the accuracy drop is quite small for 50/62.5% sparsity. \n\n**Weaknesses**\n* Conceptually, the idea doesn’t appear to be novel. Use of routers for contextual sparsity was previously explored in prior art (referenced DejaVu and Power Infer works). The main contribution of the work is the specific focus on the Attention heads and design of an efficient kernel. I believe that some analysis of the learned attention gates, which heads are activated given some input or some study, that could explain higher drops on MMLU compared to other tasks would strengthen the work.\n* The evaluation involves only log-likelihood tasks. The impact of skipping some neurons/heads is likely to be more pronounced in generative tasks. I would propose trying one (or more) of {GSM8K [1], IFEval [2], MMLU-Pro [3]} benchmark. Generative evals are closer to real-world tasks.\n\n[References]\n- [1] Cobbe, Karl, et al. \"Training verifiers to solve math word problems.\" arXiv preprint arXiv:2110.14168 (2021).\n- [2] Zhou, Jeffrey, et al. \"Instruction-following evaluation for large language models.\" arXiv preprint arXiv:2311.07911 (2023).\n- [3] Wang, Yubo, et al. \"Mmlu-pro: A more robust and challenging multi-task language understanding benchmark.\" The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track. 2024.", "was_revised": true, "final_justification": "After reading the rebuttal and the responses addressed to me and the other reviewers, I have decided to raise my score to `Accept`.\n\nThe evaluation results demonstrate the robustness of the proposed method across diverse settings. The increase in decoding throughput is quite significant, highlighting the great promise of the method in high‑throughput scenarios.", "rebuttal": "We sincerely thank reviewer for their time and for providing insightful, constructive feedback. We are encouraged that the reviewer recognized the efficiency of our router training and the decent speedups our method provides. We have carefully considered the weaknesses and questions raised and have conducted extensive new experiments that we believe substantially strengthen the paper.\n\nBelow, we address each point in detail.\n\n## Response to Weakness\n\n### 1. Novelty and Analysis of Attention Gates\nWe appreciate the reviewer's perspective on the conceptual novelty. While we build on the established idea of using routers for contextual sparsity (as seen in DejaVu and PowerInfer), we wish to clarify our key, novel contribution:\n\n**Polar Sparsity is the first work to design, build, and demonstrate practical, high-performance GPU kernels (Selective FlashAttention) that convert theoretical attention head sparsity into measured, wall-clock speedups for large-batch, long-sequence inference with up to 2.2x end-to-end throughput improvements.**\n\nPrior work has either focused exclusively on MLP layers, where sparsity benefits diminish with batch size, or has failed to demonstrate tangible end-to-end latency improvements from attention sparsity. Our primary innovation lies in tackling the attention bottleneck directly and providing the efficient implementation necessary to make this approach practical. We have included an extended discussion of related work in Appendix A to further position our contributions.\n\nImportantly, our work provides a foundation for future research to extend head sparsity into a lossless acceleration framework. Our approach has two key properties that enable this: the batch-invariant nature of head sparsity and the use of lightweight routers with a frozen base model. As we propose in our future work section, these properties would allow for a powerful extension where head sparsity is selected dynamically based on query complexity. This would permit different queries within the same batch to activate different numbers of heads, paving the way for a potentially lossless acceleration framework for modern LLMs. Our work provides the foundational kernels that make dynamic, query-aware routing a practical reality.\n\nRegarding the analysis of learned gates, this is an excellent suggestion. We have observed that more challenging reasoning and coding tasks experience a slightly larger accuracy drop, which supports the intuition that difficult queries do require the activation of more attention heads. A deep analysis of which specific heads activate for certain tasks is a fascinating direction for future research. We have provided a preliminary study on head activation in Appendix B.2 and commit to adding further analysis to strengthen the paper.\n\n### 2. Evaluation on Generative and Long-Context Tasks\n\nThis is a critical point, and we agree that evaluating on generative tasks is essential. In response to this and similar feedback from other reviewers, we have conducted extensive new evaluations on several challenging generative and long-context benchmarks on new instruction tuned models.\n\n**New Results on MMLU-PRO, GSM8K, and LongBench**\n\nWe evaluated our sparse models on MMLU-PRO and GSM8K. The results show that our method maintains strong performance with only a minor, graceful degradation in accuracy.\n\n| Subject                    | Llama-3.1-8B-Inst_dense | Llama-3.1-8B-Inst_sparse_0.625 | Mistral-7B-Inst_dense | Mistral-7B-Inst_sparse_0.5 |\n|---|---|---|---|---|\n| MMLU_PRO_Biology           | 0.62 | 0.60 | 0.47 | 0.50 | \n| MMLU_PRO_Business          | 0.44 | 0.43 | 0.21 | 0.20 | \n| MMLU_PRO_Chemistry         | 0.26 | 0.24 | 0.12 | 0.11 | \n| MMLU_PRO_Computer Science  | 0.42 | 0.41 | 0.28 | 0.29 | \n| MMLU_PRO_Economics         | 0.53 | 0.51 | 0.34 | 0.35 | \n| MMLU_PRO_Engineering       | 0.24 | 0.27 | 0.16 | 0.15 | \n| MMLU_PRO_Health            | 0.51 | 0.49 | 0.31 | 0.30 | \n| MMLU_PRO_History           | 0.41 | 0.37 | 0.25 | 0.29 | \n| MMLU_PRO_Law               | 0.28 | 0.28 | 0.17 | 0.17 | \n| MMLU_PRO_Math              | 0.39 | 0.35 | 0.17 | 0.17 | \n| MMLU_PRO_Other             | 0.46 | 0.44 | 0.32 | 0.29 | \n| MMLU_PRO_Philosophy        | 0.43 | 0.42 | 0.26 | 0.24 | \n| MMLU_PRO_Physics           | 0.34 | 0.34 | 0.19 | 0.18 | \n| MMLU_PRO_Psychology        | 0.59 | 0.59 | 0.44 | 0.43 | \n| MMLU_PRO Overall Accuracy  | 0.41 | 0.40 | 0.25 | 0.24 | \n| GSM8K                      | 0.75 | 0.71 | 0.34 | 0.32 |\n\nWe also tested our models on the LongBench-e benchmark, which specifically evaluates long-form generation. The results demonstrate the stability and robustness of our approach, with no catastrophic failures observed.\n\n| Metric                | Llama-3.1-8B-Inst_dense | Llama-3.1-8B-Inst_sparse_0.5 | Llama-3.1-8B-Inst_sparse_0.625 | Mistral-7B-Inst_dense | Mistral-7B-Inst_sparse_0.5 | Qwen2.5-14B-Inst_dense | Qwen2.5-14B-Inst_sparse_0.5 |\n|---|---|---|---|---|---|---|---|\n| 2wikimqa              | 15.28 | 14.28 | 14.50 | 32.30 | 31.99 | 13.76 | 12.52 |\n| gov_report            | 34.66 | 31.79 | 33.56 | 28.64 | 28.60 | 30.99 | 27.34 |\n| hotpotqa              | 17.50 | 15.17 | 15.46 | 38.71 | 37.11 | 14.08 | 14.35 |\n| lcc                   | 69.56 | 60.72 | 63.70 | 56.78 | 56.11 | 67.90 | 60.41 |\n| multi_news            | 25.45 | 25.08 | 25.56 | 22.60 | 22.18 | 21.53 | 21.70 |\n| multifieldqa_en       | 27.79 | 27.92 | 26.61 | 32.91 | 32.09 | 35.18 | 32.12 |\n| passage_count         | 15.75 | 18.22 | 18.68 | 4.04 | 4.84| 7.64 | 11.05 |\n| passage_retrieval_en  | 97.12 | 94.24 | 95.46 | 31.67 | 37.33 | 89.64 | 94.61 |\n| qasper                | 11.83 | 11.85 | 12.49 | 30.04 | 28.34 | 10.67 | 10.94 |\n| repobench-p           | 55.31 | 45.19 | 48.76 | 47.14 | 41.77 | 47.62 | 40.15 |\n| samsum                | 42.73 | 38.39 | 40.79 | 40.84 | 39.74 | 44.69 | 41.01 |\n| trec                  | 71.00 | 71.67 | 71.67 | 61.33 | 61.33 | 75.00 | 75.00 |\n| triviaqa              | 92.81 | 89.72 | 90.72 | 83.11 | 82.68 | 88.43 | 86.61 |\n| LongBench-e average   | 44.37 | 41.86 | 42.92 | 39.24 | 38.78 | 42.09 | 40.60 |\n\nThese new results confirm that Polar Sparsity preserves model capabilities on the very tasks where high-throughput inference is most critical, adding significant weight to our evaluation. We will integrate these findings into the revised manuscript.\n\n## Response to Questions \n\n### 1. Inference Framework, Installed version of libraries\n\nWe evaluate our kernels and end-to-end generation performance on top of the FlashAttention library (FlashAttention library includes inference and training scripts) with CUDA graphs enabled. We selected FlashAttention library to test our inference system for its ease of integration with the new Selective FlashAttention kernels and its close performance with vLLM. We are actively working on integrating our kernels with future versions of vLLM. \nWe have also included the entire code base in the supplemental materials which includes our triton kernels and inference scripts. \nWe used fused kernels for efficient layernorms, MLPs but did not use torch.compile in our latency, throughput measurements. This would be a valuable addition which can further reduce the latency and we are happy to update our graphs with the compilation enabled. \n\nPyTorch version: 2.6.0+cu124\nFlashAttention version: 2.6.3\nTransformers version: 4.50.3\n\n### 2. Figure 2 (a) with newer models\n\nThis is a great point. While Figure 2(a) focused on the OPT family, the majority of our evaluations (throughput in Fig 5, 6 and new accuracy benchmarks) feature modern models like Llama2, Llama-3.1, Mistral-7B, and Qwen2.5-14B. These newer models exhibit similar or even more pronounced attention dominance as these models support larger context lengths. We will gladly add a new version of Figure 2(a) with these modern architectures in the revised version of the paper.\n\n### 3. Wikitext-2 for router training.\nThank you for this excellent suggestion. We agree that using higher-quality data like FineWeb-Edu could potentially improve router performance and downstream task accuracy. We used Wikitext-2 as a standard, well-established dataset to ensure a reproducible and comparable baseline. Exploring the impact of different training data sources is a valuable direction for future work.\n\n### 4. Polar Sparsity on a long-context task.\nWe completely agree. As shown in our response to your second weakness, we have now benchmarked our method on the LongBench suite. The results demonstrate that our approach is highly effective for long-context scenarios, maintaining robust accuracy while providing significant speedups.\n\n## Conclusion\n\nWe are grateful for the reviewer's sharp questions and valuable suggestions, which have prompted us to significantly strengthen our paper's evaluation. We have provided extensive new benchmarks on challenging generative and long-context tasks, clarified the novelty of our high-performance kernels, and detailed our experimental setup. We believe these additions directly address the concerns raised and more clearly demonstrate the practical impact and significance of Polar Sparsity. We sincerely hope the reviewer will consider these updates favorably and reconsider their rating."}]

Author rebuttal:

Meta-review: The paper proposes Polar Sparsity, a batch-friendly framework that scales contextual sparsity to large-batch LLM inference, achieving up to 2.2× throughput gains across multiple models and hardware with minimal accuracy loss. While remaining concerns persist (e.g., downstream performance on reasoning and math tasks), the proposed method is well validated and delivers promising inference speedups, which is also integrable with existing solutions. Overall, this is a practical work that leverages the beneficial conflicts of contextual sparsity between Attention and MLP. I also suggest the authors to incorporate the valuable results from the rebuttal period into the revised manuscript.
Decision: Accept (poster)