S2-Attention: Hardware-Aware Context Sharding Among Attention Heads

Weakness 1: Besides the 1.3B model, we pre-trained two 7B models on 3T tokens, as well as four models for the continual pre-training setting. We first report the pre-training results here. We can see the S2-L8V16 and S2-L8V16-Inter both achieve strong performance.

We also resent the downstream task performance of continual pre-training Llama 3.1 in addition to the 128k needle retrieval performance. We can observe all the downstream remains on par with the dense version.

Weakness 2:

We proof-read the paper and clear all the typos.

Q1:

A1: We eval 1.3B with context length extension to 32k, which has perfect needle score as shown in Figure 2(b). As we don’t do post training, it’s hard to do further on more complex tasks. For the 7B models, we trained on 128K context and got perfect scores as shown in Figure 6.

Q2:

We thank the reviewer for pointing out the missing information that we lost at refactoring. We’ll add them to the paper. But here are the details. For studying the speed-up of MergeQ and SplitD to attention-op , the experiments are conducted on a 1.3B model scale. (batch size, heads, head dim)=(4, 16, 128). In sparse block size is 16 for merge-Q experiments and 64 otherwise. Vertical stride is 16 for merged-Q while varied otherwises. Local windows are all kept to 512 tokens. As for the end-to-end training and inference speed-up, for 1.3B models, the speed-up benchmark is done with L1V15 and 64 block size. For 7B models, the e2e speed-up benchmark is done with L4V24+dense10-22 with 64 block size .

Q3:

As we observe on-par converging speed for sparse models and dense models, the token throughput can actually be translated to end-to-end training speed-up. The sample training curve is shown in the Appendix A.

Q4:

Yes. At decoding, we roughly save (1-(1-local_window_size)/stride_size) of memory.

Weakness 1:

We thank the reviewer for this suggestion. One major point of this paper is stating that plug-in-and-play KV eviction methods do not work well with the underlying working mechanism of the hardware, primarily due to suboptimal GPU memory access and allocation. Therefore, they typically cannot bring wall-clock speed-up when compared to real serving systems like vLLM, SGLang, FlashInfer, and TensorRT, nor can they be easily implemented with these systems due to incompatibility with foundational techniques such as continuous batching and PagedAttention.

More specifically, from a memory management perspective, (1) frequent manipulating small chunks of memory at arbitrary location, (2) different workload for each head/layer will cause significant memory fragmentation, instead of reducing the memory footprints. However, (1) and (2) are exactly what many of the KV eviction methods are doing. From a computational perspective, (2) will also hurt the benefit of parallelism and can hardly beat vanilla GEMM in reality.

Because of these, the claimed flexibility of KV eviction methods, e.g., H2O and Quest, is only theoretical and can hardly be used in production with today’s accelerators. In fact, they are not adopted in any major LLM providers, nor the open-source community like vLLM.

We propose S2-Attention aiming to bridge the gap between innovations in the research context and their real-world deployment. Meanwhile, S2-Attention is designed as a general attention alternative, instead of solely targeting for inference optimization. We actually experiment it for pre-training and context extension scenarios and observe on par results with dense attention and much faster training speed.

Weakness 2:

We thank the reviewer for pointing out this. We optimize the performance of finer-granularity because it’s an obvious feature that can help the community to explore sparsity more freely. Putting it straight, it’s similar to how improving magnifying glass to microscope would help science discovery.

For example, the KV eviction methods usually need token-level sparsity. If now the system restricted users can only mask/keep every 64 tokens, which is the case to be compatible with FlashAttention, the exploration and study would be largely limited.

Thank you for your explanation. However, your claim that prior sparse attention work, such as Quest, is limited to theoretical speedups and unsuitable for production is incorrect. FlashInfer [1] has already implemented block-wise sparse attention. In fact, prior works, including Quest and MInference, have demonstrated measured speedups using either the CUDA API provided by FlashInfer [2] or customized CUDA kernels. Ignoring these implementations misrepresents the state of sparse attention research. Therefore, I strongly agree with Reviewer 8jB7 that this work must include comparisons with other optimized sparse attention kernels tailored for prefilling (e.g., MInference) and decoding (e.g., Quest).

Moreover, since S2-Attention requires training from scratch, the absence of performance benchmarks against dense models on diverse tasks and datasets, such as GSM8k and LongBench, is a critical omission that weakens the evaluation of the proposed method.

[1] https://github.com/flashinfer-ai/flashinfer/issues/367 [2] https://github.com/mit-han-lab/Quest

Thanks for bringing up Quest. I think there is some misunderstanding on the paper.

1. This paper is not just about inference or KV-cache saving and inference speed. It is also about new pretraining architecture. Our method and kernel work for both training and inference, while the FlashInfer can only do inference. FlashInfer benefits inference, while ours benefits both at training and inference. Therefore, the scope of our work is broader. We enjoy the optimization done by FlashInfer, but in my understanding, it currently doesn’t support head-heterogenous sparsity patterns, nor usage in training.

2. From an industry perspective, training from scratch/requiring continual training is not a disadvantage. On one hand, plug-in-and-play is nice, but it probably not the top priority in real-world LLM deployment. What the real-world deployment will avoid the most is performance degradation and unpredicted model behavior, which KV eviction or Quest like methods usually suffer from due to the inconsistence between training and inference. The cost of regression in production is a lot more severe than the extra training before deployment.

3. For comparison with Quest:

   a.) Quest does not save memory. Its implementation still needs to store the entire KV cache, plus the metadata for topk computation. It only saves memory bandwidth and computation (the so-called memory movement, i.e., loading from HBM to SRAM). Quest is actually motivated by the IO overhead and memory fragmentation issues in H2O, but trade-off the memory saving to avoid these overhead. However, S2-Attention’s design can achieve both targets. Like H2O, our method also saves memory for storing the KV-cache. Please note that for a single sequence of 128k input in a small 7B model, takes up 64GB of memory, meaning that you can only handle 1 user input at a time with a 80GB GPU (or 4 users input of length 32k), which is quite expensive. This is also part of the reason why Quest limits the batch size on each device to 1.

   b.) In terms of computation, with a kv-cache budget of 2k in a 32k input, Quest achieves a speed up of 4.86x on attention, with 30% is for finding the top-k (without that, it is 7.03x ) as from Figure 9 of their paper. However, for S2-Attention, with a 2k budget out of 32k input (sparsity of 1/16), our speed up 2.3X over Quest. Most of the difference comes from Quest’s large overhead and the memory fragmentation. GPUs work better if loading a consecutive chunk, but top-k does not guarantee that even with paged KV. That’s what we mean hardware-aware.

   c.) Another problem of Quest, is that the K in their “lossless” top-k budget in their benchmark changes through tasks. However, in a real application, we typically don’t know what task the user’s input is. S2-Attention don’t have this issue as inference and training are the same.

Quest’s downstream benchmark is mostly on Longbench, but Longbench for a lightly pre-training model will not be great for benchmark (see point 3 below). So it is hard to compare the latency by controlling the accuracy (also due to point c above) or compare accuracy by controlling latency(our method could attend to almost 2x more tokens, with pre-trained backs the attention pattern, as well as memory saving) . However, from all the points above plus the results of downstream tasks in our paper, it is almost obvious that our method will likely balance latency and accuracy better, beside that fact that it also serves as a much faster pre-training method.

4. Regarding the GSM8k, LongBench, or even MMLU. For pretrained models using open-source web data(e.g., fineweb-edu), it is challenging to achieve a sensible number for these tasks within hundreds of billions of tokens (already very expensive, 256 A100, a bit more than one day). We’d probably need to train trillions of tokens to get sensible scores for a single model variation, which makes paper on pre-training prohibitive, let alone doing lots of ablations. Another way is to use non-open source synthetic data, e.g., Phi-3, which hurts reproducibility and is way too expensive for most researchers. Indeed, cross-entropy or PPL is a very good metric and correlates well with downstream performance in realistic pre-training research, as long as the tokenizer and data are the same. That’s why many pre-training papers[1,2] only include loss and the tasks we presented on the paper.

I hope the reviewer understands that this is not “inference/kv-cache” centric paper and the cost & limitation on publishing paper related to pre-training. The amount of exploration, insights, and invested resources are way beyond.

[1] Yoco Neurips 2024: https://openreview.net/forum?id=25Ioxw576r&referrer=%5Bthe%20profile%20of%20Shaohan%20Huang%5D(%2Fprofile%3Fid%3D~Shaohan_Huang1) [2] Mamba: https://openreview.net/forum?id=AL1fq05o7H

Overall response to Weakness 1: Despite their conceptual elegance, previous KV cache reduction techniques often fail to translate the efficiency gains on paper to real-world applications. This is because, similar to FlashAttention vs Approximate Attention, the main bottleneck for sparse inference is memory management overhead and lack of computation kernel. S2-Attention aims to bridge the gap between innovations in the research context and their real-world deployment. Weakness 1.a: We thank the reviewer for correctly pointing out that our work draws a conclusion that differs from some previous works [1]-[5]. A key insight of our work is that many previous KV cache reduction methods, including those cited by the review, are often incompatible with the underlying working mechanism of the hardware used for serving these models. As a result, their promising savings on paper rarely translate speedup in practice and are not adopted in industry providers like OpenAI/Microsoft and Anthropic. This is the precise reason that, to the best of our knowledge, none of them are implemented in, or compared to, actual inference optimization systems widely used in practice such as, vLLM, SGLang, FlashInfer. Therefore, we consider our findings on the limitations of these prior works to be novel, as is the S2-Attention architecture, which is designed based on the principles learned from these limitations. If the reviewer can provide evidence contradicting our claims, we are happy to revise them. Otherwise, we respectfully disagree with the review’s concerns about the rigor of our claims and kindly ask the reviewer to reconsider their evaluation.

About weakness 1.b: Those are great questions. Maybe we didn’t convey the message very well, due to length constraints.

The key principle here is that efficient sparsity should avoid memory fragmentation, which [1]-[5] and most of the KV eviction works overlook. For example, it's not feasible to release/allocate small chunks frequently, e.g., 2 bytes, of memory at arbitrary, non-deterministic locations decided on GPU, as many of these works do. It also won't save memory, but only cause more memory fragmentation, e.g., by having different sparse attentions across instances in the same batch. "Absolute positions rather than relative ones" ensures that the deterministic pattern among instances and avoids memory fragmentation. "Vertical line" means a KV block should be kept for a few decoding steps for more efficient memory management, compared to non-consecutive usages (fragmentation and memory waste) or immediate/frequent release/allocation (hurt performance).

Take dilated attention as an illustrative example, as shown in Figure 5a. Every token attends to 1 token every 2 tokens depending on the relative distance. Now token n attends to tokens n-2, n-4, …. So we need to cache (n-2, n-4, …) for it. Let’s move on to token n+1, the next token. It requires the KV at tokens n-1, n-3,... Therefore, it requires caching all of them for decoding starting from token n, which ends up keeping all post tokens in memory despite not being used in many steps. This is typical sparse attention that causes internal fragmentation. Moreover, runtime eviction like [1]-[5]introduces even more fragmentation in addition to what’s mentioned above. When the KV is different for each head or each layer, it will cause similar fragmentation to that caused by various sequence lengths in a batch as pointed in paged attention. Moreover, computation-wise, the workstream could be highly imbalanced and unpredictable, which hurt the benefit of foundational tensor/pipeline/sequence parallelism.

This observation inspires us to study a KV-cache efficient pattern. For sliding-window-like attention, KV-cache is efficient despite that they are relative to token position. But for tokens outside of the sliding window, i.e., the “global” tokens attended by many tokens far away, attention patterns based on relative positions will cause fragmentation, as in the example above.

That’s how we come to the definition of efficient kv-cache pattern, If (k_i , v_i) is attended by q_{j+l}, (k_i , v_i) must also be attended by q_j. j >=i, l >0 . If you pause and think about it from the argument above, you can see that otherwise, we’ll have (k_i, v_i) at hand (as we need it for token j+l) when we are decoding token j, yet we don’t use it, which isn’t “the best” for decoding token j as we have access to more information.

This definition very naturally shows that token i above, has to be attended by a continuous segment of tokens after i before it is dropped completely. Obviously this is based on “absolute position” (position i) instead of relative position. Hence the example patterns shown on Figure 5(b), where you see the “vertical lines”

I hope this explanation of “kv-cache efficient pattern” makes things clearer. Due to the length of the paper, we didn’t include all these detailed arguments/steps.

For [4], MInference, in many cases, is significantly slower than flash_attn backend of vLLM. The advantage is shown when the prompt is sufficiently long to cover the overhead, e.g., >100k. For numerical comparison, at 262015 prompt length, we're 2.28X faster than MInference. At 128k prompt length, we're 1.91X faster than MInference.

For [5], Quest, with a kv-cache budget of 2k in a 32k input, Quest achieves a speed up of 4.86x speed-up on attention, with 30% is for finding the top-k (without that, it is 7.03x ). However, for S2-Attention, with a 2k budget out of 32k input (sparsity of 1/16), our speed up 2.3X over Quest. Most of the difference comes from Quest’s large overhead and the memory fragmentation. GPUs work better if loading a consecutive chunk, but top-k does not guarantee that even with paged KV. That’s what we mean hardware-aware. Meanwhile, Quest's bonus become significantly thinner when moving to A100/H100, we suspect the original speed-up partially came from the fact that FlashAttention/FlashInfer algorithms become weaker on commercial GPUs like RTX4090, which Quest benchmarked on. Also, [5] actually supports some of our main argument for those sparse KV, as [5] is motivated by the IO and fragmentation issues in H2O, but [5] trade-off memory saving for this purpose. S2-Attention's design can achieve both targets without trading-off memory saving for reducing IO and fragmentation overhead.

Both [4] and [5] have issues for batch size > 1 on single device, and for correctly handling tensor parallel and pipeline parallel, which further weaken the practical of those systems. It also makes our comparison infeasible, as we can easily use TP=2 for S2-Attention to larger the throughput advantages over the [4][5]. As discussed in previous comments, dynamic KV eviction (though Quest doesn't evict), will have too many impractical overheads to be implemented in real serving systems.

Another point is, both [4][5] needs to specify different hyperparameters for the framework to work and get the claimed speed-up. MInference get the minimum FLOPs needed for different tasks to maintain model quality. Similarly, Quest needs different K for their topK on different tasks. However, for realistic deployment where user query is unpredictable, it's not feasible to do so.

We also need to stress that S2-Attention is not a "KV eviction/inference optimization only" paper. S2-Attention and the open-sourced kernel, are designed for both training and inference. We bridge the long-lasting absence of an optimized sparse self-attention framework, which can facilitate model architectures greatly. However, [4][5] and most of the inference optimization framework cannot do this. On the other hand, "plug-in-and-play" is not the top priority of actual LLM deployment and efficiency optimization in industry. We ourselves deployed and optimized several most widely used LLM endpoints in the world. The lessons we learned are (1) model quality regression, e.g., unpredicted behavior, is the worst thing we want to avoid, while any plug-in-and-play sparse module can cause degradation and change model behavior, (2) to the best of our knowledge, most of our counterparts and ourselves have the time and computation (continual)train a more efficient model variants, and are willing to do so as long as it's beneficial to long-term production benefits.

[a] MInfernce degradation and imcompatibilities: https://github.com/vllm-project/vllm/issues/5751#issuecomment-2225698913 [b] MInference TP failure: https://github.com/microsoft/MInference/issues/63 [c] MInference slower than vLLM: https://github.com/microsoft/MInference/issues/18 [d] MInference performance degradation: https://github.com/microsoft/MInference/issues/83, https://github.com/microsoft/MInference/issues/43 [e] QUEST TP: https://github.com/mit-han-lab/Quest/issues/8

Q1:

Thanks for the great question. It is true that the FLOPS among heads are imbalanced. But it is not true that the computation is decided by the most intensive head. The parallelization is not just on head level, but also on token chunk level, i.e., the computation on one-head is not sequential. The q will be chunked and parallel (indeed, even kv can also be chunked and parallel, but mostly we typically don’t find it necessary in training but only at inference). So you see that for a head, the most intensive one is the last chunk (attend to the most kv). Therefore, over all heads, the most intensive threadblock is the one handling the last chunk of the least-sparse head. But chunks close-by have similar computations as well. Remember that GPU’s assignment of threadblock to SMs is random and based on availability, so as long as the number of thread blocks is high (significantly higher than SMs), it will even out. This is the case in length like 4k (32 heads x 4k/64 block size=2k thread block >> 106 SMs in A100).

In addition, the patterns (local-stride) used in the paper (Figure 1), all mostly balanced (but not required even from computation perspective as mentioned above). We did try mixing sparse heads with dense heads as well, and we don’t observe computation downside due to head imbalance. We didn’t include it in the paper due to length.

Q2:

On an NVDA GPU, a warp consists of 32 threads (fixed) , there are threads that executes together. Warp and thread blocks are separate concepts. In our code, a block of 64 tokens is computed in a threadblock, which can contain 1, 2 or more warps (it is a hyperparameter to tune). The mapping of data and warps is complex. A warp (32 threads) can handle 16 tokens, 32 tokens, 64 tokens or even more. There is no such a restriction of 32 threads handling 32 tokens, but rather based efficiency consideration.

Q3:

For the 32k length, we continual pretrain the 1.3B model whose max sequence length was 8192 to get context extension. We follow this training protocol as previous studies suggest. “depth perfect” means the extended model can perfectly retrieve the information, namely “needle”, wherever the needle is put in a 32k context.

Q4:

Yes, we follow a fixed, rule-based pattern for MergeQ. As long as the vertical stride pattern is used, we can easily decide the SRAM access logic. We merged q-blocks when sparse block size are smaller (e.g., 16 tokens) than efficient kernel block size (64). Softmax is fused in an online way similar to FlashAttention.

Q5:

For the forward/backward figures, it is 100 times for benchmark, with 25 warm-up in advance using “triton.testing.do_bench”. For inference in vLLM, we use the default, 30 times with 10 times of warmup in advance. The warm-up here is mostly for optimization like tuning(Triton) or cudagraph(vLLM) building.

In all those tests, we found that p50, p95 and average are almost identical (<1%). Therefore, we don’t bother to report the standard deviation. For GPU (A100, H100, not those for displaying), as long as it is solely occupied by the benchmarking process, it seems to be very stable across runs, unlike CPU which is never completely free. That’s why in FlashAttention 1 and 2 papers, no standard deviation was mentioned either.