CUDA kernel learning note Block size: 128 in repo API Focus: K/V tile prefetch

Dense FlashAttention vs Block-Sparse Attention

This page compares the prefetch pattern in the dense FlashAttention-style path and the block-sparse path in mit-han-lab/Block-Sparse-Attention. The important code-level point: sparse K tile selection is resolved at runtime from a compressed block mask, then the K pointer jumps by leap. The shown V staging, however, still advances one block at a time in the loop.

1. High-level difference

Dense FlashAttention-style path

For one Q block, the kernel walks K/V blocks in reverse contiguous order.

K7→ K6→ K5→ K4→ K3...

next_K = current_K - 1
K_ptr += -kBlockN * k_row_stride
V_ptr += -kBlockN * v_row_stride

The next K/V tile address is predictable from the loop index n_block.

Block-sparse path

For one Q block, valid K blocks come from a runtime blockmask list.

K7→ K6 K5→ K4 K3 K2→ K1

next_K = blockmask.mask_val(mask_idx)
leap   = current_K - next_K
K_ptr += -kBlockN * leap * k_row_stride

The next selected K tile is known only after reading sparse metadata at runtime.

2. Where the sparse pattern is resolved

The repo first converts the boolean block mask into a reversed integer list per Q-block row:

boolean mask row:   [1, 0, 0, 1, 0, 1]
converted row list: [5, 3, 0, -1, -1, -1]

In the CUDA kernel, fwdBlockmask computes a pointer to the row list using batch_idx, sparse head id, and the current Q-block row. Then mask_val(i) returns the selected K-block id, and max_no_larger(target) binary-searches the first selected K-block that is legal under the current causal/local range.

3. Prefetch timeline: dense path

Dense path is a regular software pipeline. K is initially staged for the last legal block. Each loop iteration stages V for the current block and stages the next K block by subtracting exactly one block.

Loop step

Prologuecopy Q, copy K[n]

Iter nwait K[n]

Iter ncopy V[n]

Iter nQK GEMM

Iter ncopy K[n-1]

Iter nsoftmax + PV

Address rule

K startn_block_max - 1

Currentn_block

V-1 block

Computealways

K next-1 block

Outputaccumulate

for n_block = n_max-1 ... n_min:
    stage V[n_block]
    compute Q @ K[n_block]^T
    stage K[n_block - 1]
    compute softmax(scores) @ V[n_block]

4. Prefetch timeline: block-sparse path

Sparse path keeps the FlashAttention-style pipeline but inserts runtime sparse control. The current dense loop position n_block is compared with next_block_col_idx from the mask. If they differ, the code treats this as is_skip.

Loop step

Metadatamask_val(i)

Comparen_block != next

If skipzero / masked scores

If hitQK GEMM

K prefetchjump by leap

V stagingshown as one-block step

Address rule

Next Kfrom blockmask

Skip?runtime branch

Computeskip QK/PV

Computenormal QK/PV

K ptr-leap

V ptr-1 in shown code

next = blockmask.mask_val(mask_idx)

if n_block != next:
    # skip compute for this logical block
    leap = n_block - next
    K_ptr += -kBlockN * leap * k_row_stride
else:
    # real selected block
    mask_idx += 1
    next = blockmask.mask_val(mask_idx)
    compute QK and PV

Code-level nuance: the K prefetch path is sparse-aware through leap. In the shown forward code, the V staging statements still advance tVgV by -kBlockN * v_row_stride, not by leap. So the safest interpretation is: sparse K tile selection is runtime-indirect; V staging is not visibly the same sparse jump in these snippets.

5. Side-by-side comparison

Aspect	Dense FlashAttention-style path	Block-sparse path in this repo
Pattern knowledge	Compile-time / loop-structure predictable. Every legal K block is visited.	Runtime. The block mask is passed from Python and read by CUDA.
Next K tile	`n_block - 1`	`blockmask.mask_val(mask_idx)`, then pointer jumps by `leap`.
K prefetch address	Contiguous: subtract one K block.	Sparse-aware: subtract `leap` K blocks.
V prefetch / staging address	Contiguous: subtract one V block.	Nuanced: shown statements subtract one V block, not `leap`.
Skipped block compute	No skip; all legal blocks compute QK/PV.	Skip branch exists; when `n_block != next_block_col_idx`, it avoids the normal QK/PV path for that logical block.
Metadata overhead	Almost none for block selection.	Extra blockmask loads, binary search / iterator logic, branch, and pointer leap calculation.
Best mental model	Dense tiled attention with predictable one-tile lookahead.	FlashAttention-like pipeline plus runtime sparse iterator; K prefetch follows selected sparse blocks.

6. Tiny example

Dense sequence

Legal blocks: 7, 6, 5, 4, 3, 2, 1, 0

7654 3210

K prefetch: 7 → 6 → 5 → 4 → 3 → 2 → 1 → 0
V staging:  7 → 6 → 5 → 4 → 3 → 2 → 1 → 0

Sparse sequence

Selected blocks from mask: 7, 4, 1

7654 3210

K prefetch: 7 → leap 3 → 4 → leap 3 → 1
V staging:  code snippets show step-by-step V pointer movement
Compute:     only selected logical blocks take normal QK/PV

7. Takeaways

What to remember

Dense FA: next K/V addresses are statically predictable from the loop counter.
Block-sparse: next selected K address depends on runtime blockmask metadata.
The code resolves this by: compressed Python blockmask → CUDA fwdBlockmask → mask_val / max_no_larger → leap → sparse K pointer update.
Do not overstate the V side: the visible code path uses leap for K pointer advancement, while V staging appears as a one-block decrement in the shown snippets.

Source mapping

Main source files: flash_fwd_kernel.h, flash_blockmask.h, block_sparse_attn_interface.py, and README.md.

Generated as a learning visualization. It uses pseudocode and code variable names, not a full copy of the source.