HND vs NHD memory access for sparse KV tiles

This visualizes one fixed attention head reading selected sparse KV blocks. The selected sparse pattern is the same in both layouts, but the memory address pattern is very different.

Selected KV block for target headCurrently prefetched / loaded stepNot used by this head

1. Same sparse pattern, different physical memory order

Simplified setup: B=1, H=4, N=12 KV blocks. Each cell represents one KV block vector/tile for one head. In a real kernel, each KV block contains many tokens and the full head dimension.

HND = [B, H, N, D]

Memory is grouped by head first. For a fixed head, sequence blocks are compact.

NHD = [B, N, H, D]

Memory is grouped by token/block first. Heads are interleaved inside each sequence block.

2. Sparse prefetch stream

SpargeAttn uses a LUT-driven sparse stream. The logical selected blocks are identical, but the physical address stride is larger under raw NHD.

Logical selected blocks
LUT deltas by block id
Current animation step
Not started

HND physical address stream

Address model: addr = h * N + n. For one fixed head, moving from selected block n to n+1 is one block step.

NHD physical address stream

Address model: addr = n * H + h. For one fixed head, moving from block n to n+1 jumps over the other heads.

3. Why HND is better for sparse attention

HND mental model

For target head h, selected KV blocks are compact per-head chunks. Sparse jumps come only from the LUT pattern.

NHD mental model
For target head h, same-head sequence blocks are separated by all other heads. Sparse jumps combine with head interleaving.
Practical takeaway: optimized sparse-attention kernels prefer HND [B,H,N,D] because each selected K/V tile for one head is compact and easier to load with cp.async. If input is NHD [B,N,H,D], many implementations rearrange it to HND before the optimized sparse kernel.