Flash-MSA Steps In: Tackling Million-Token Training

Flash-MSA Makes Its Move: Sparse Attention Kernels Target Million-Token Training

Long-context training has always been bottlenecked by the same thing: attention. As sequence length scales linearly, compute and memory usage still grow quadratically — and that hasn’t changed since the birth of the Transformer. Over the past two years, most solutions have fallen into two camps: one squeezes performance at the IO level through the FlashAttention family, while the other cuts compute directly through sparsification. In July this year, a project called Flash-MSA (Flash Multi-Sparse Attention, also commonly referred to as Flash Sparse Attention / FSA) combined these two approaches and pointed directly at the million-token training scenario.

This is not just another paper prototype. The project released a complete Triton implementation on GitHub, along with striking benchmark results: compared with DeepSeek’s NSA (Native Sparse Attention), it achieves up to a 3.5× speedup on H20 and 2.9× on H200; compared with full attention, the gain reaches as high as 6.4×. In real training runs on Llama3-8B, Qwen3-14B, and Qwen2.5-32B, end-to-end acceleration averages 1.86×.

Filling NSA’s Gaps: What Flash-MSA Actually Fixes

To understand the significance of Flash-MSA, you first need to understand what problems it is addressing.

DeepSeek’s NSA (Native Sparse Attention), introduced earlier this year, was a milestone for sparse attention because it brought sparsity into the training phase instead of limiting it to inference-time approximations like earlier methods. NSA’s core idea is to split attention into three branches: compression, selection, and sliding window. Among them, “selected attention” chooses the most relevant KV blocks to compute and is the primary source of sparse-attention gains.

The issue lies at the Kernel level. NSA’s selected attention uses a query-grouping strategy: the outer loop iterates over query tokens, while the inner loop iterates over KV blocks. The efficiency of this implementation depends heavily on GQA group size. Large GQA groups fill Tensor Cores efficiently during matmul operations. Small GQA groups become problematic.

Modern LLMs typically use GQA group sizes of 4 or even smaller (Llama3 and Qwen3 are in this range). Under these conditions, NSA must apply padding to query heads in order to satisfy Hopper architecture matmul hardware requirements (minimum dimension size of 8). These padded computations are pure overhead — extra memory reads and wasted multiply-accumulate operations.

The result is an awkward situation: under many modern LLM configurations, unoptimized NSA can actually run slower than full attention. After all the effort of introducing sparsity, performance gets worse instead of better. That is precisely the gap the Flash-MSA team targeted.

Three Kernels, One Decoupling

Flash-MSA does not modify the NSA algorithm itself. Instead, it splits selected attention computation into three independent Kernels:

• Main Kernel: batches together all query tokens attending to the same KV block and computes them in one pass, writing intermediate results to a buffer. This is the core idea — instead of handling each query independently and suffering matrix-shape mismatches, Flash-MSA aggregates queries by KV block, naturally producing Tensor Core-friendly matmul shapes.
• Reduction Kernel: accumulates partial results from the buffer back into query tokens.
• Online Softmax Kernel: independently handles softmax online statistics to avoid entangling them with the main compute path.

The most interesting part of this structure is that it changes the semantics of loop ordering. Original NSA works as “queries search for KV.” Flash-MSA becomes “KV blocks collect queries.” It sounds like merely swapping inner and outer loops, but with small GQA groups, a KV block is often referenced by multiple query batches. Those batches together fill the Tensor Core, eliminating the need for padding.

There are tradeoffs, of course: non-contiguous memory access and additional buffer memory. The team introduced two optimizations — “early return” and “compact buffer management.” According to the paper, under a 64K token, T=16, d=128 configuration, the buffer consumes roughly 1GB of HBM, which is considered acceptable.

Notably, Flash-MSA uses a conditional fallback design: when GQA group size is ≥ 8, the original NSA implementation is already efficient enough, so Flash-MSA falls back automatically. This kind of “knowing when to step aside” behavior is uncommon in Kernel libraries and shows solid engineering judgment.

The Numbers: Small GQA Is Its Home Turf

Here are the key benchmark comparisons:

| Configuration | vs. NSA | vs. Full Attention | |------|---------|-------------| | H20 Peak | 3.5× | 6.4× | | H20 Average | 1.8× | 2.4× | | H200 Peak | 2.9× | 4.9× | | H200 Average | 1.4× | 2.3× |

Breaking it down by forward and backward passes:

• Forward: up to 2.36× faster than NSA, averaging 1.62×
• Backward: up to 4.32× faster, averaging 2.59×

The fact that backward propagation gains exceed forward gains is especially important. One major reason sparse attention has struggled to scale in training is that backward passes repeatedly recompute index tensors, which easily become bottlenecks. Flash-MSA skips much of this overhead during backpropagation, effectively moving sparse attention from “inference-friendly” toward “training-friendly.”

The team also ran end-to-end training tests on three model scales: Llama3-8B, Qwen3-14B, and Qwen2.5-32B.

• Compared with NSA: up to 1.25× faster, averaging 1.09×
• Compared with full attention: up to 2.47× faster, averaging 1.86×

What does 1.86× actually mean? If pretraining a long-context model previously took 30 days, it can now finish in about 16 days. In today’s compute-constrained environment, that matters far more than shaving off a fraction of a percent in parameter count.

Why Developers Should Pay Attention

Stepping back, two major trends in AI have become especially clear this year. First is the context-length arms race — Gemini 2’s million-token context, Claude’s memory systems, and GPT-5’s long-document handling are all pushing context windows outward. Second is the Kernel-level efficiency revolution — FlashAttention-3, FlashInfer, SGLang’s custom kernels, and vLLM’s PagedAttention all demonstrate that “the same algorithm can become twice as fast with a different implementation.”

Flash-MSA sits at the intersection of these two trends. It does not invent a new sparse algorithm or new hardware. Instead, it “translates” the existing NSA algorithm into GPU-friendly code. This kind of work may not look glamorous, but for teams actually spending large-scale compute budgets on training, its practical value far exceeds another SOTA benchmark.

Some concrete use cases:

• Million-token pretraining: full training at this sequence length was previously almost a luxury due to backward-pass memory and time costs. Flash-MSA’s average 2.59× backward acceleration significantly lowers that barrier.
• Long-context SFT / RLHF: domains like legal analysis, code repositories, and scientific literature inherently require models to process hundreds of thousands of tokens. These teams often cannot afford full attention but also hesitate to use poorly optimized sparse methods. Flash-MSA provides a practical middle ground.
• Inference acceleration spillover: training and inference often share the same Kernel stack. Although Flash-MSA primarily targets training, its selected-attention Kernel optimizations also benefit long-context inference.

Problems Still Unsolved

A few caveats are worth noting. Flash-MSA is not a silver bullet:

1. It only targets the NSA family. If you are using sparse methods like MoBA, Longformer, or BigBird, Flash-MSA offers no benefit. Sparse attention is already highly fragmented, and Kernel-level fragmentation will likely get worse.
2. It depends on Triton. The Triton ecosystem is maturing, but for engineers deeply experienced in CUDA optimization, it can still feel like a black box. In extreme optimization scenarios, hand-written CUDA may remain faster.
3. Buffer overhead grows with batch size and sequence length. 1GB is manageable at 64K context lengths, but what about 128K or 256K? The team has not released data for more extreme configurations.
4. Its relationship with FlashAttention-3 is not fully integrated yet. In theory, Flash-MSA’s approach could stack with FA3’s IO optimizations, but the current implementation does not share the same code path.

Judging from the project’s update cadence, development is active. GitHub issues already show users running long-context fine-tuning on Qwen3-32B, and the reported performance largely matches the paper.

The Longer-Term Signal

Viewed in a broader context, Flash-MSA reflects an ongoing shift: the center of innovation in attention mechanisms is moving from “algorithms” toward “Kernels.”

Before FlashAttention, researchers competed over smarter attention algorithms. After FlashAttention, the competition shifted toward who could extract more from hardware through better Kernels. NSA’s sparse algorithm was introduced only months before Flash-MSA began aggressively optimizing its Kernel implementation — feedback cycles have shortened by an order of magnitude compared with the past. The coupling between algorithms, hardware, and systems is tightening, and training frameworks are increasingly evolving toward Kernel-library architectures.

For individual developers, the implication is straightforward: if you are working on long-context systems, following algorithm papers alone is no longer enough. You need to learn how to read Triton code and analyze Kernel profiles. That is not exaggeration — it is quickly becoming a baseline engineering skill for the next wave of AI infrastructure jobs.

Among the models currently available through OpenAI Hub, the Qwen3 and DeepSeek series both use relatively small GQA group architectures. If you are experimenting with long-context workloads on these models, you can directly test large-window inputs through the API and observe the practical effects. Low-level optimizations like Flash-MSA ultimately surface as faster inference responses and support for longer contexts.

References

• Relaxed-System-Lab/Flash-Sparse-Attention (GitHub) — Official Flash-MSA / FSA repository, including Triton Kernel implementations and reproduction scripts
• In-Depth Analysis of DeepSeek NSA Native Sparse Attention Idea and Kernel (Zhihu) — Detailed Chinese analysis of NSA algorithms and Kernel internals, useful background reading before studying Flash-MSA