Junxian Ma / Writings

TL;DR

• More sparsity does not necessarily mean more speed. Reducing computation is only half the problem. End-to-end performance also depends on whether the selected data can be moved and reused regularly, and on how much work it takes to find that data.
• NSA creates regularity before loading the data. It selects contiguous 64-token blocks and shares each selection across 16 query heads, allowing the sparse workload to map directly to a regular GPU tile.
• DSA reconstructs regularity after loading the data. It permits arbitrary token selection but shares token IDs across all 128 main heads. FlashMLA then moves scattered rows into SMEM and repacks them into regular tiles. The cost is that Discovery (Indexer plus exact Top-K) must still scan every candidate.
• DSA’s bottleneck is not strongly tied to address randomness; it lies in finding the selected tokens. In the tested H100 sparse-prefill configuration, latency point estimates for eight random row permutations stayed within 0.47% of sorted access. In the public DSA pipeline, however, Discovery (Indexer plus Top-K) was already slower than FlashMLA at the 64K checkpoint, and the gap widened rapidly with context length.
• The system must therefore optimize Discovery, not only the final sparse kernel. At the systems level, fusing the Indexer with Top-K can avoid writing and rereading the complete FP32 score matrix. At the algorithmic level, each query must avoid scanning the entire KV history. The former reduces constant-factor overhead in the current implementation; the latter can fundamentally change long-context scaling.

1. Opening

Sparse Attention is easy to write down. The hard part begins when the selected KV entries no longer form a regular matrix: how should a GPU load them, compute with them, and reduce their gradients? NSA and DSA answer that question through two different co-design paths.

1. NSA creates regularity before the load. Contiguous 64-token blocks and 16-head shared selection let the algorithm map directly to a regular tile.
2. DSA reconstructs regularity after the load. All 128 main heads share the selected token IDs, and FlashMLA moves scattered rows into SMEM before computing on them. The cost is that candidate discovery—Indexer plus Top-K—can still grow quadratically.

Together they lead to the central claim of this article:

A sparsity ratio tells us only how much computation was removed. Actual speed depends on whether the selected data can be moved and reused regularly—and on how much work it takes to discover that data in the first place.

We first establish a common way to read tensors and tiles, then derive the algorithmic dataflow, kernel ownership, and memory movement of NSA and DSA. Two H100 experiments appear only where they answer a specific question: one tests how sensitive the DSA sparse kernel is to selected-row address patterns; the other tests how the quadratic Discovery term hidden behind the sparse main kernel grows with context length.

This article does not re-explain dense Attention or FlashAttention. It begins where block IDs and token IDs are produced and asks how a kernel consumes them. [PAPER], [CODE], [DERIVATION], and [MEASURED] mark claims from papers, public code, our derivations, and formal measurements.

2. Three Concepts Are Enough to Read the Rest

Many shapes appear below, but the reading method has only three parts.

A program is one independent Triton kernel instance; in CUDA terminology it is roughly one CTA (thread block). SMEM is on-chip memory shared within a CTA. HBM/global memory is larger and slower device memory.

First, a logical tensor is not necessarily an HBM tensor. One Attention row has

and logically produces $S,P:[1,T_K]$ and $O:[1,d_v]$. This does not mean that $S$ or $P$ must be written to HBM. We will distinguish the logical shape defined by the algorithm, the materialized shape preserved across kernels, and the tile shape processed by one CTA/program.

Second, every output needs an explicit owner. Here the owner is the program that ultimately writes that output region. Suppose one program owns $B_R$ softmax rows and streams over $B_{\mathrm{kv}}$ KV entries at a time:

With score scale $\alpha$, the local work is always two regular matrix products:

The program keeps only the online-softmax state for the rows it owns:

As each KV tile arrives, it updates the row-wise maximum $m$, rescales the old $z$ and $O_{\mathrm{acc}}$ to the new maximum, and accumulates the current probabilities and value numerator. Selected blocks or rows can therefore stream in without materializing the complete probability matrix.

Third, the key question is not only “how much was skipped?” but “where does irregularity stop?” For each path we track four things:

1. how selection is produced and whether it materializes an intermediate tensor;
2. what output one program owns;
3. how many rows or heads reuse one KV load;
4. where a dynamic index becomes a regular tile again.

We count one FMA as 2 FLOPs. Unless noted otherwise, byte counts use BF16 useful payload; cache-line overfetch, page/TLB behavior, and allocator effects are discussed where they matter.

3. NSA: Make the Algorithm Produce an Executable Block Workload

Takeaway. NSA is not efficient merely because it “selects 16 blocks.” Each ID names 64 contiguous tokens, and all 16 heads in one KV group share those IDs. The algorithm therefore already produces something close to a regular $[16,64]$ interaction tile.

NSA decomposes long-context Attention into three access patterns:

The key derivation is not merely “sum three branches,” but

The algorithm comes from the NSA paper. The kernels used in the paper’s experiments are not fully public, so our code-level discussion of online Top-K, the forward grid, and backward uses the FLA community implementation as one concrete path. Choices specific to that implementation are identified explicitly. [PAPER] [CODE/community] [GAP]

3.1 Algorithmic Contract: What Each Sparse Branch Contributes

Omitting the batch axis and using a sequence-major layout, a sequence of length $T$ has

NSA targets GQA. The paper’s efficiency experiments use

The query-head group corresponding to KV head $g$ is

The algorithmic parameters of the three branches are:

Here $L_t=t+1$ is the length of the causal prefix at position $t$. It contains

complete compression windows, and

causally visible selection blocks.

The three branches are not three masks over one shared KV representation. The paper gives each branch its own K/V projection:

Therefore, even if multiple branches access the same token, we cannot first take a union of token IDs and treat the result as one attention operation. The branches use different representations, normalize independently, and only then combine their outputs.

Compression answers “what is roughly present far away?” Each KV head maps an overlapping window of length $\ell=32$ through a learned compressor into one compressed K/V row. Adjacent windows have stride $d=16$, so a full sequence produces approximately $C_T\approx T/d$ compressed K/V rows.

The local shapes for position $t$ and group $g$ are

With a fixed stride $d$, the total compression workload across all queries remains

It is smaller than dense attention by a constant factor of roughly $d$, but it is not linear in sequence length.

The sliding window answers “what is happening nearby?” The local branch accesses only

providing fixed, contiguous local context that tiled attention can handle naturally.

Selection answers “which distant details should be expanded again?” Each branch performs its own softmax and produces

which are then combined by independent sigmoid gates:

These gates need not sum to 1. In the execution DAG, compression attention and sliding attention can run in parallel; selection attention must wait for the routing IDs.

3.2 Routing: From Compression Probabilities to Top-16 Block IDs

Selection does not train another router. It reuses compression-Attention probabilities to convert “which compressed windows matter?” into “which 64-token blocks should be expanded?” The process has three steps: spatial remapping, a 16-head reduction, and Top-16.

Spatial remapping. Let

be the compression probability for query head $h$. Equation 9 of the paper converts it into a selection-block score:

An out-of-range compressed-window index is treated as zero. The paper’s configuration satisfies

This double sum handles the fact that overlapping compression windows and selection blocks do not correspond one-to-one.

GQA group reduction and Top-$n$. The $G=16$ query heads sharing a KV head must use the same selection IDs, so

The concrete configuration in Section 4 of the paper further specifies that, among the 16 selection slots, one initial block and the two most recent local blocks are always active; only the remaining slots are filled by the dynamic ranking above. [PAPER]

The fixed initial block deterministically creates a hot row with $\Theta(T)$ fan-in. The two local slots move forward with the query: under the current-plus-previous rule alone, any particular block serves at most about $2\ell'=128$ query positions. Dynamic ranking can create additional hot spots. [DERIVATION]

At a mature position,

This is an index tensor with an explicit slot order, not a mathematical set. Near the beginning of a sequence, fewer than 16 causally valid blocks may exist. The kernel may still maintain 16 selection slots, but padding must be marked by $-1$ or by a block count. “A fixed 16 slots” does not mean “16 valid blocks always exist.”

The complete shape chain is

compression probabilities       [T,Hq,≈T/16]
  → Eq. 9 spatial remap          [T,Hq,≈T/64]
  → reduce G=16 query heads      [T,Hkv,≈T/64]
  → Top-16                       [T,Hkv,16]

But where are the probabilities? Standard FlashAttention does not write the complete probability matrix back to HBM; it usually returns only the output and LSE. Therefore, “reuse the compression probability” does not by itself explain how that probability reaches Top-k.

At $T=65536$, the actual number of complete windows is $C_T=4095$. To illustrate the upper bound for a fixed-shape allocation, let the padded capacity be

Materializing per-head FP32 compressed scores at that capacity,

requires

per layer.

If we first sum over each GQA group before the Equation 9 remap, materializing

still requires

If instead we are discussing group-reduced block scores after the remap, let $M_T=\lceil T/\ell'\rceil=1024$. Their shape should be

corresponding to

These three numbers describe different locations in the dataflow and must not be conflated.

The current FLA implementation demonstrates one feasible mechanism: save the compression LSE, recompute probabilities in blocks from Q and compressed K, and maintain Top-$n$ online using bitonic merges, thereby avoiding materialization of the quadratic score tensors above. This path trades recomputation for memory capacity: the Top-k kernel must scan compressed K with Q once more, so routing still contains a $\Theta(T^2/d)$ QK scan and candidate merging. It eliminates HBM materialization of the quadratic score tensor, not quadratic pair scoring itself.

The accompanying path also simplifies compression to non-overlapping mean pooling, shares K/V, and imposes implementation constraints that the GQA group size be at least 16 and a power of two. None of these is part of the paper algorithm’s general definition. FLA therefore proves only that one infrastructure path is feasible; it does not prove that the paper authors used the same implementation. [CODE/community] [GAP]

3.3 Selection Forward: Logical Gather, Physical Block Tile

From 16 IDs to 1,024 Logical Tokens

Routing produces discrete block IDs:

Each $j_s$ points to a contiguous 64-token block in selection K/V:

For a mature query with all 16 slots valid, the logical gather is

The logical attention shapes are

Here 1024 is the padded capacity, not the number of valid tokens at every position. Some slots may be invalid near the beginning of the sequence, and a causal block containing the current position may expose only part of its tokens. An implementation must apply the ID mask, range mask, and causal mask together.

The $[1024,d]$ above is only the logical shape after concatenation. In the original KV storage, it corresponds to 16 mutually disjoint 64-token blocks, each internally contiguous in token index; it is not one contiguous 1,024-token source range. If the physical layout is token-major $[T,H_{kv},D]$, adjacent token rows for a fixed head have a fixed stride and need not form a byte-contiguous $64D$ interval. What the algorithm truly improves is the addressing structure: one block requires only one irregular base address, from which 64 rows are expanded by regular affine offsets, rather than 64 separate indirect token lookups.

If we actually constructed these two tensors for every query/group, the external gather would need both to handle 16 irregular block bases and to assemble them into a compact selected-KV tensor written back to HBM. The K/V payload alone would be

in BF16. This creates an extra write, after which attention reads the same data again. The correct approach is to retain only the IDs and load the original K/V by block inside the attention kernel.

What Does One Program Ultimately Write?

The most natural owner for selection forward is

We first expand all input shapes and grid symbols used by this source path. For fixed-length inputs, FLA uses

Here:

• $B$ is the batch size;
• $T_Q,T_K$ are the numbers of query and KV tokens per sequence in the current invocation;
• $H_q,H_{kv}$ are the numbers of query and KV heads;
• $G=H_q/H_{kv}$ is the number of query heads served by one KV head;
• $d_k,d_v$ are the Q/K feature and V/output feature dimensions per head;
• $N_{\mathrm{slc}}$ is the slot capacity along the final dimension of the selection-ID tensor, written as S in the source; the configuration here uses 16;
• $I^{\mathrm{slc}}$ stores block IDs. When block_counts is a tensor, it records the number of valid slots for each query/group. When it is a scalar capacity, the kernel scans all $N_{\mathrm{slc}}$ slots;
• $B_V$ is the tile width that one program computes along the value-feature axis;
• $N_V=\lceil d_v/B_V\rceil$ is the number of tiles required to cover the complete value dimension.

The NSA numerical configuration used here is

FLA further assumes that $Q$ corresponds to the last $T_Q$ tokens of each KV sequence. Thus, on the fixed-length path, the absolute position of query slot $t$ within the KV sequence is not simply $t$, but

The subsequent causal mask and selected-block start positions are both evaluated relative to $p_t$.

The launch grid in the source is a three-dimensional program-count tuple:

The source names the three program IDs i_t, i_v, and i_bh. To avoid confusing i_v with the value tensor $V$, we rename the second ID $\nu$:

They represent, in order, the query position, the value-feature tile, and a flattened batch/KV-head position. The third ID is unflattened as

to recover batch index $a$ and KV-head/group index $g$. We intentionally use $a$ for the batch here to avoid confusion with the KV block ID $b$ below. The query heads served by group $g$ remain

Next define the value-feature interval owned by program $\nu$:

Program $(t,\nu,i_{bh})$, equivalently $(a,t,g,\nu)$, is then the unique writer of the following output slice:

The physical accumulator is allocated as $[G,B_V]$; a final tile narrower than $B_V$ is handled by a feature mask. This also makes the axis order explicit: after removing batch flattening, the grid is $(T_Q,N_V,H_{kv})$, not $(T_Q,H_{kv},N_V)$.

If $N_V>1$, the same $(a,t,g)$ expands into $N_V$ programs. Each program recomputes the same QK/softmax for its own value-feature slice, but the $\mathcal D_\nu$ intervals are disjoint, so there is no output write conflict. LSE is not tiled along the value dimension; the source permits only the $\nu=0$ program to write $\mathrm{LSE}^{\mathrm{slc}}[a,t,\mathcal H_g]$.

The paper uses $d_v=128$, and the current FLA path also takes $B_V=128$, so

Each $(a,t,g)$ then corresponds to exactly one program, which exclusively owns the complete selection output $[G,d_v]=[16,128]$. The program subsequently loads and keeps resident the logical query tile

Optional implementation note: variable length does not change ownership. The public variable-length path uses a packed $B=1$ representation. If sequence $i$ contains $T_Q^{(i)}$ query tokens and $T_K^{(i)}$ KV tokens, respectively, then

The grid then becomes

The first program ID $u\in[0,T_Q^{\mathrm{total}})$ is a packed query slot; token_indices_q[u] maps it to $(i_{\mathrm{seq}},t_{\mathrm{local}})$. Let

denote the starting offset of sequence $i$ in packed Q. The physical output owner is then

The corresponding causal absolute position of this query within its sequence’s KV storage is

Thus, varlen changes only how fixed-length $(a,t)$ coordinates are located. The $\nu$ and $g$ grid axes, and the ownership rule that “one program uniquely writes one output slice,” remain unchanged.

How a 64-Token Block Becomes a Regular Tile

The program then loops over the 16 selection slots. We use $B_{\mathrm{kv}}$ for the number of KV-token rows loaded and computed in one inner-loop iteration. In the current FLA path, it corresponds to the source parameter BS and exactly equals NSA’s selection block length $\ell'$:

For a valid block ID $j_s$, the program computes $\mathrm{base}=j_s\ell'$ and loads

This is the logical tile at the algorithmic level. In the current Hopper branch of FLA,

so matching Triton’s dot-product feature tile requires the physical construction

where the $192\rightarrow256$ tail of the feature dimension is handled by masking/padding. The logical contractions are therefore

The physical padded execution shape for QK is

This $192\rightarrow256$ execution shape is a conclusion about the current FLA Hopper path. The non-Hopper branch caps BK at 128; with $d_k=192$, it produces two feature tiles and triggers the source’s NK == 1 assertion. The $[16,256][256,64]$ shape therefore cannot be extrapolated directly to that path.

The only dynamic component is the block base computed before each load. Once the data reaches the on-chip tile, computation returns to a fixed shape.

We deliberately do not use $B_K$ for the token tile here. The NSA paper’s token-tile notation $B_K$ corresponds to BS in the current FLA source, whereas FLA’s BK denotes the padded feature width—for example,

Mixing the two notation systems would mistake “64 tokens” for “256 feature elements.”

Online normalization. Across the 16 blocks, the program uses the exact online-softmax merge introduced above. It maintains only the running maximum, normalizer, and output accumulator for each query head; it does not write $P^{\mathrm{slc}}\in\mathbb{R}^{16\times1024}$. After all selected blocks have been processed, it directly writes

The output therefore has a unique writer, and the forward pass needs no atomic reduction on the output.

Why head sharing is an execution contract. For one 64-token block, the QK and PV multiply-add count is approximately

With BF16 K/V, the primary K/V payload for this iteration is

Considering only the K/V payload shared by 16 heads, the local arithmetic intensity is

This is not a roofline figure for the entire kernel: it excludes Q, output, indices, LSE, masks, and launch overhead. It only shows that the more query heads reuse the same K/V block, the more easily the cost of one irregular address lookup is amortized by regular matrix multiplication.

For a mature query, the main selection loop has 16 slots. Near the beginning of a sequence, the same control structure still runs, but invalid slots are masked. NSA provides a “fixed-capacity, padding-compatible workload,” not an absolutely fixed amount of work without boundary conditions.

3.4 Forward and Backward Use Different Sparse Graphs

Takeaway. The forward relation query → selected blocks finds an owner only for the output. If we retain query-centric ownership, $dK^{\mathrm{slc}}/dV^{\mathrm{slc}}$ requires floating-point atomics or a partial-gradient reduction. The FLA community implementation audited here instead transposes the relation into KV block → selecting queries, then lets each KV-block program perform its own many-to-one reduction.

The natural owner in the forward pass is the query, because each program can independently complete one $O^{\mathrm{slc}}_{t,\mathcal H_g,:}$. The backward pass, however, has two reduction directions.

The discussion below covers only the backward pass of the selection branch. Complete NSA also includes gates, a compression branch, a sliding branch, and the merging of gradients from all three branches into shared Q/input tensors.

Recompute Probabilities, Preserve the Sparse Graph

Selection backward does not need to save the complete probability matrix. It recomputes score/probability tiles from $Q,K^{\mathrm{slc}},I^{\mathrm{slc}}$, and $\mathrm{LSE}^{\mathrm{slc}}$, then combines them with $dO^{\mathrm{slc}}$ to recover local gradients. For batch/sequence $a$, query $t$, group $g$, and selected block $b$, the local shapes are

Without activation checkpointing, the FLA selection autograd state includes at least

as well as block counts and sequence metadata required in variable-length cases. The current FLA autograd path actually saves q, k, v, o, lse and retains metadata such as block indices in the context. The accurate conclusion is “$P^{\mathrm{slc}}$ is not saved,” not “the forward pass only needs to save $O^{\mathrm{slc}},\mathrm{LSE}^{\mathrm{slc}},I^{\mathrm{slc}}$.”

$dQ$: the Query Remains the Owner

For query $t$, only the blocks it selected contribute:

The local shape for each block is

The same program can loop over all slots and accumulate

in registers before writing it back exactly once. Thus, dQ remains well suited to query-centric ownership.

Why $dK/dV$ Cannot Remain Query-Owned

One KV block $b$ may be selected by many queries. Define

Its gradients are many-to-one reductions across queries:

Transpose the Sparse Graph and Make Each KV Block the Owner

If we keep the query as owner, different programs update the same $dK^{\mathrm{slc}}_b,dV^{\mathrm{slc}}_b$ concurrently. The implementation must then use many floating-point atomics or write partial buffers and reduce them later. A more natural approach is to transpose the sparse relation first:

forward adjacency:   query → selected KV blocks
backward adjacency:  KV block → all selecting queries

and then assign one program to each $(a,g,b)$ KV block.

CSR representation. Let the batch size be $B$, and let the number of blocks in a complete sequence be

The padded shape of the input IDs is $[B,T,H_{kv},n]$. The upper bound on edge capacity allocated for csr_indices by the current FLA implementation is

The actual number of valid edges is determined by the ID range, block count, and causal condition:

Flatten $(a,g,b)$ into a one-dimensional row:

The CSR metadata shapes are

where only $\texttt{csr\_indices}[0:E_{\mathrm{valid}}]$ is a valid prefix. Row $r=r(a,g,b)$ corresponds to

This is one-dimensional flattened CSR. It must not be written as a nonexistent two-dimensional index $\texttt{offsets}[g,b]$.

CSR is constructed with two counting/scatter passes:

traverse valid selection edges
  → atomically count the fan-in of each (batch,group,block)
  → prefix sum to obtain csr_offsets
  → traverse the edges again and scatter query IDs using an atomic cursor

Atomics operate only on compact integer metadata; the expensive floating-point reductions for $dK^{\mathrm{slc}}/dV^{\mathrm{slc}}$ obtain a unique owner.

Take

Then

so the int32 csr_indices buffer allocated by the current FLA implementation is 16 MiB.

For a complete causal self-attention sequence, suppose early positions use all distinct visible blocks. The first 15 query blocks of 64 tokens then have $1,\ldots,15$ valid slots respectively, after which the count reaches 16. In this case,

This number is smaller than the allocation upper bound; the two should not both be denoted by the same $E$.

Meanwhile,

int32 offsets require only about 16 KiB. For comparison, materializing a uint8 inverse mask of shape $[T,H_{kv},M]$ requires

CSR reduces inverse metadata from $\Theta(T\,H_{kv}\,M)$ to $O(E_{\mathrm{alloc}}+B\,H_{kv}\,M)$, at the cost of two edge traversals, integer atomics, a prefix sum, and subsequent irregular query gathers.

KV-centric dKV tile. A dKV program owns

and gradient accumulators of the same shapes. It reads $B_{\mathrm{qry}}$ query positions at a time from the corresponding CSR row and expands the $G$ heads of each query:

The batched query and output-gradient shapes are

Local recomputation and reduction once again form regular matrix multiplications:

For example, when $B_{\mathrm{kv}}=64,B_{\mathrm{qry}}=4,G=16$, we have $R=64$, and the central score-recomputation tile is

CSR does more than compress metadata. It rewrites the conflict “many queries write the same KV-gradient address” into “one KV owner sequentially reads its own queries,” and then restores a regular Tensor Core tile within that owner.

Write Conflicts Disappear; the Fan-in Tail Remains

Define

When block_counts=16 is a fixed scalar capacity, each $dQ^{\mathrm{slc}}$ program scans 16 padded slots. FLA also accepts a per-query tensor block_counts, in which case the runtime $N_S$ may be shorter. The loop length of a dKV program, by contrast, is determined by $NQ_{a,g,b}$. A fixed initial/sink block, as well as popular blocks created by dynamic ranking, may have fan-in far above the mean and form long-tail CTAs; the moving local slots themselves create only bounded local fan-in. Splitting a hot row can improve load balance but reintroduces a partial reduction or atomic operations. The backward problem therefore changes from “how many FLOPs can be skipped?” to “how should a many-to-one reduction be scheduled?”

We do not turn this long tail in fan-in into a performance experiment here. On the one hand, the backward kernel used by the paper’s authors is not public. On the other hand, changing the routing topology in the community implementation changes fan-in skew, query-gather locality, and CSR row order together, making it difficult to attribute latency differences to a single variable. FLA backward therefore moves irregularity from floating-point write conflicts into CSR row length, query-gather locality, and CTA load balance. We report no performance measurement for these factors; we treat them only as execution consequences confirmed by the source and do not claim that they explain the backward latency of the paper authors’ kernel. [DERIVATION] [CODE/community] [GAP]

3.5 NSA’s Regularity Comes from the Algorithmic Contract, Not Kernel Magic

Under the paper’s default configuration, NSA selection maps efficiently to a kernel because the algorithm deliberately provides three execution contracts:

The forward pass brings dynamic block bases into the program’s inner loop and restores regular $[16,64]$ computation after the load. The FLA community backward path analyzed here then transposes the sparse graph so that a KV block becomes the unique owner of $dK^{\mathrm{slc}}/dV^{\mathrm{slc}}$.

This also reveals three boundaries that must not be crossed.

First, NSA is not evidence that “arbitrary block sparsity automatically accelerates.” Hardware efficiency arises only when contiguous blocks, head sharing, a fixed slot capacity, and explicit ownership all hold simultaneously.

Second, the token budgets of the selection and sliding branches are approximately fixed at mature positions, but compression still accesses roughly $L_t/d$ compressed K/V rows. For fixed $d$, full-layer compression remains $\Theta(T^2/d)$.

Third, FLA demonstrates a complete feasible mechanism comprising online Top-k, group-centric forward, and CSR backward. It cannot prove that the paper authors’ unreleased implementation uses the same routing materialization, tile parameters, or inverse-adjacency construction.

The paper reports that, with its $G=16,\ell'=64,n=16$ configuration and on its A100 test environment, the authors’ Triton NSA reaches 9.0× forward and 6.0× backward speedup over Triton FlashAttention-2 at 64K. This result can be attributed to the complete system reported in the paper; it cannot be attributed directly to any single kernel choice isolated from the community code in this article. [PAPER] [GAP]

In other words, the most reusable part of NSA is not one particular sparsity formula, but a co-design method:

First make the algorithm emit sparse units that a GPU can consume, then choose a unique owner separately for the forward output and the backward gradient.

4. From Blocks to Tokens: Why DSA Can Use Finer Selection

NSA and DSA are not simply an “old design” and a “new design.” NSA selects blocks so that training and inference can consume contiguous tiles directly. DSA selects tokens because MLA exposes a different and sufficiently strong reuse axis. The hardware constraint did not disappear; the algorithm changed how it satisfies that constraint.

4.1 Infra: Contiguous Blocks Are No Longer the Only Source of Regularity

Consider the problem NSA faced. Arbitrary token selection creates many scattered KV rows, and during training those rows also participate in the backward reductions for $dK/dV$. NSA therefore fixes the selection unit to a contiguous 64-token block. The cost is that selection can say only “this region matters,” not select just the relevant tokens inside it. NSA trades selection granularity for regularity before the load. [PAPER]

DSA starts from a different base. DeepSeek-V3.2 already uses MLA, so each token stores one shared latent-KV entry. DSA further makes all 128 main heads of a query share one set of token IDs. A scattered row may be difficult to load, but after loading it can serve many heads. FlashMLA can therefore leave random addressing to the producer while Tensor Core consumers operate on regular SMEM tiles. Contiguous blocks are no longer the only way to obtain hardware efficiency. [PAPER] [CODE]

From the infrastructure perspective, the design changes along three axes:

This finer selection is not a free upgrade. NSA routes on a compressed token axis. DSA preserves full token resolution in the Indexer and then performs exact Top-K. It bounds the main Attention by $O(L\kappa)$ without removing the $\Theta(L^2)$ Discovery term. A more accurate description is: DSA trades stronger cross-head reuse and a specialized gather kernel for token-level selection, then moves the system bottleneck from sparse main Attention toward the Indexer and Top-K. [DERIVATION]

4.2 Training: Attach a New Router to an Existing MLA Checkpoint

DeepSeek did not train DSA from scratch. It continued from DeepSeek-V3.1-Terminus after extending that model to a 128K context:

1. Dense warm-up: keep dense Attention, freeze the main model, and train the Indexer with a KL objective derived from the dense-Attention distribution.
2. Sparse training: enable Top-2048 and let the main model adapt through the language-modeling loss; the Indexer continues to use a separate KL loss. Its input is detached from the main-model graph so that the two gradient paths remain isolated.

The point of this procedure is to retain the existing MLA parameterization while adding a router and a sparse kernel, rather than training the entire Attention system from scratch. [PAPER]

Evidence boundary. The DeepSeek report does not provide a direct NSA-versus-DSA ablation or enumerate reasons for abandoning NSA. What is public is each paper’s algorithmic contract, DSA’s use of MLA/MQA sharing, and its continued-training procedure. The explanation above—why the design can move from blocks to tokens—is a systems inference from those facts.

We can now follow the actual DSA dataflow: Indexer scores, Top-K token IDs, the MLA execution representation, and finally the FlashMLA kernel.

5. DSA: Turn Arbitrary Token Selection Back into Regular Matrices

Takeaway. DSA permits arbitrary token selection but makes all 128 main heads share one set of IDs. Only the FlashMLA producer handles scattered addresses; once rows enter SMEM, the consumer still sees regular matrices.

The complete DSA forward path has three stages:

The table below keeps only the notation needed to follow those shapes.

$D_V^{\mathrm{MHA}}=128$ must not be confused with FlashMLA’s latent-value dimension $D_C=512$. The main Attention also has two algebraically equivalent representations with different execution shapes:

5.1 Indexer: Low-Dimensional, but Still Scanning the Full History

Step 1: Score Every Query–Token Pair

The Indexer assigns a cheap routing score to every visible token and selects Top-$\kappa$:

where

Across all queries and visible tokens, the final routing-score matrix is $I:[T_Q,T_K]$. Projection, RoPE, Hadamard, FP8, and DeepGEMM fusion all serve the same goal: compute this matrix without writing the much larger $[T_Q,64,T_K]$ intermediate to HBM.

One DSA layer also contains two distinct kinds of scores:

and

$I$ is the Lightning Indexer’s routing score, used only to produce token addresses. $A^{\mathrm{slc}}$ is the attention logit recomputed by the main MLA over the selected tokens, and only this score enters the main Attention softmax.

The Indexer therefore decides where to look, not how to weight what is found.

Step 2: Produce Low-Dimensional Queries, a Shared Key, and Head Weights

The normalized hidden states

first produce the low-rank query latent shared by the main MLA and the Indexer:

Rather than performing another large projection directly from the 7168-D hidden state, the Indexer derives its query from $C^Q$:

After reshaping:

The released DeepGEMM interface takes a materialized FP8 Q as input. The public code does not reveal whether the production system further fuses this projection, rotation, and quantization.

Each historical token, by contrast, has only one shared Indexer key vector:

The persistent cache has logical shape

Thus, the 64 Indexer query heads share one historical key rather than maintaining a separate K cache for every head.

The Indexer also projects query-dependent head weights directly from the current hidden state:

$u^I_{t,j}$ is the base weight of Indexer head $j$ for query $t$; this branch does not pass through a softmax. Absorbing the two fixed normalization factors used by the official reference, define

Only Indexer Q reuses $C^Q$. Indexer K and these weights are both projected directly from $X$.

Step 3: RoPE, Hadamard, and FP8

Each 128-D Indexer Q/K vector is split as

The first 64 dimensions use non-interleaved RoPE, unlike the interleaved RoPE used by the main MLA. This is not a notational detail: if the two sides use different pair layouts, their dot product at a given position changes.

After rotation, a normalized Hadamard transform $\mathcal H_{128}$ is applied:

Before quantization, therefore,

The Hadamard transform does not change the exact real-valued inner product; it spreads outliers so that subsequent E4M3 quantization is more stable. “Inner-product preserving” must not be extrapolated to mean “bitwise identical after quantization.”

A note on the two symbols named $H$. $H_I=64$ is a scalar—the number of Indexer heads. $\mathcal H_{128}\in\mathbb R^{128\times128}$ is the Hadamard transform matrix. Its entries are $\pm1/\sqrt{128}$; it is fixed, orthogonal, and not learned. The implementation need not materialize this matrix. A fast Hadamard transform evaluates it through hierarchical additions and subtractions in $O(128\log128)$ time. Q and K use the same $\mathcal H_{128}$ after partial RoPE and before FP8 quantization, so the real-valued inner product is preserved exactly before quantization and only approximately afterward.

Let $q^{(8)},k^{(8)}$ denote the E4M3 data, and let $s^Q_{t,j},s^K_s$ denote the corresponding scales. The official reference folds the query scale and both normalization factors into the query-dependent weight:

The scorer’s implementation semantics can then be written as

Here, $q^{(8)},k^{(8)}$ explicitly refer to vectors after partial RoPE, Hadamard transformation, and FP8 quantization; they are no longer interchangeable with the pre-transform $q^I,k^I$.

The logical payload of one Indexer K-cache entry is

Compared with a full main-KV row, this is what makes the bandwidth of the full-history scan substantially smaller.

Step 4: Fuse 64 Heads On Chip and Write Only the Final Score

Mathematically, for a fixed query $t$, we first have

The global logical shape is

After ReLU, weighting, and reduction across heads:

Writing $[T_Q,T_K,64]$ to HBM would make the low-dimensional Indexer create an intermediate 64 times larger than its final score tensor. The released DeepGEMM scorer therefore fuses the matrix multiplication, ReLU, query-dependent weighting, and 64-head reduction, writing only the final FP32 $I$.

One of its core tiles, expressed in the GEMM direction used by the code, can be summarized as

The final step applies ReLU, query-dependent weighting, and reduction over the 64 Indexer heads, then maps candidate-major accumulator fragments directly to the logical $[2,256]$ output addresses. There is no separate transpose stage. The intermediate 64-head result $Z$ remains on chip; HBM sees only the final score tile for two queries and 256 candidates.

The released scorer, however, does not include Top-K. The public DeepGEMM pipeline still allocates and materializes the complete rectangular carrier

The main scorer writes the legal candidate range of each row. When clean_logits is enabled, a separate cleanup kernel fills invalid and future regions. Exact Top-K then scans the complete rectangular result again.

DeepGEMM has therefore implemented fusion within the scorer: FP8 matrix multiplication, ReLU, query-dependent weighting, and 64-head reduction occur in one kernel, and the per-head scores $[T_Q,64,T_K]$ are never written to HBM. But this is not Indexer–TopK fusion. The released scorer still emits the full $I:[T_Q,T_K]_{\mathrm{FP32}}$, which a separate exact Top-K implementation must scan again.

Step 5: Top-K Turns Scores into Token Addresses

For query position $t$, let $p_t$ be the final causally visible position. Mathematically,

$J_{t,r}$ is a token address, not a routing score. All 128 main heads of the same query share $J_t$, while the selected sets of different queries can be entirely different.

The public interface also requires two shape/dtype adaptations.

First, torch.topk in the official inference reference returns INT64 indices and uses min(index_topk, end_pos) for the fixed trailing dimension of one call; the causal mask keeps future positions invalid.

Second, under the DSA configuration studied here, FlashMLA sparse-prefill accepts a fixed width $\kappa=2048$:

If a row has only $\kappa_t<\kappa$ legal tokens, all remaining slots must be replaced by an invalid sentinel, such as -1 or an address no smaller than $T_K$. Repeating a legal token to fill the row is incorrect, because the softmax would count that token more than once.

Converting the reference’s INT64 Top-K result into a fixed-width INT32 kernel tensor is therefore real data-preparation work, not a dtype annotation that can be omitted from a shape diagram.

Top-K IDs are discrete control flow. Except at sorting boundaries, the language-modeling loss cannot provide an ordinary continuous gradient through the decision of whether a token is selected. During dense warm-up, the main model is frozen and the Indexer is trained with a KL loss against a target obtained by summing dense attention over the main heads and L1-normalizing along the sequence axis. During the sparse stage, the main model trains on the language-modeling loss while the Indexer continues to train on a KL loss restricted to the current selected set. The paper explicitly detaches the Indexer input, so the Indexer receives only $\mathcal L_I$ while the main model receives only the language-modeling loss. The main model therefore adapts to sparse selections produced dynamically by the current Indexer, not to a fixed set of token IDs. This explains how routing is learned, but it does not change the execution cost of exact Top-K at inference time.

5.2 MLA Bridge: Why the Main Kernel Is 576-D MQA

The Original MHA Form Re-expands the KV Cache

The Indexer produces addresses; the main MLA still computes the actual attention weights. FlashMLA’s input shapes become clear only after rewriting MLA from its original MHA form into an equivalent MQA execution form.

Start from the original score. The shared query latent is projected by

Each head is split into

where $q^R$ receives the main MLA’s interleaved RoPE. On the KV side,