Title: Polar Sparsity: High Throughput Batched LLM Inferencing with Scalable Contextual Sparsity
Venue: NeurIPS 2025
Accepted: yes
Keywords: Large Language Models

Contextual Sparsity

Activation Sparsity

Head Sparsity

Select Head Attention

Batched Inference

Efficient Transformers

Machine Learning Systems
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n- 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).\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes Polar Sparsity, which sparsifies the MLP and attention mechanism during the LLM decoding stage.\n\nSTRENGTHS AND WEAKNESSES:\n# 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nPlease 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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nLLMs 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.\n\nSTRENGTHS AND WEAKNESSES:\nI'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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nYes, 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).", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}

comments: 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.

STRENGTHS AND WEAKNESSES:
Strengths:
1. The paper is well written and easy to understand.
2. The figures are well plotted to help understand the effectiveness of Polar Sparsity.

Weaknesses:
1. 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.
2. More experiments should be conducted on harder reasoning tasks (e.g., GSM-Hard, MATH).
3. Throughput performance could be tested on more devices, e.g., H100, consumer RTX, and AMD.
4. No performance comparisons to the methods deploying the batch inference with structured sparsity.
5. The training-based sparsification with customized sparse kernel design makes it hard to quickly deploy on a new model.

QUESTIONS:
See Weaknesses.

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-03