[wave] NSA: selection attention forward kernel

[harsh-nod]

Parent

Part of #1243 — DeepSeek NSA kernels for MI350

Description

Implement the forward pass of the selection attention branch — the core sparse attention kernel in NSA. Given top-k block indices, this kernel gathers the selected KV blocks and computes attention only over those blocks.

Operation

Input:  Q [B, M, H, D], K [B, N, G, D], V [B, N, G, D], block_indices [B, M, G, T]
Output: O_slc [B, M, H, D], LSE_slc [B, H, M]

Where T = block_count (number of selected blocks per query).

Algorithm

For each query position m:

1. Load query vector q[m]
2. For each selected block index t in block_indices[m]:
   • Gather K, V for tokens [t*block_size : (t+1)*block_size]
   • Compute QK^T, apply causal mask, softmax
   • Accumulate weighted V using online softmax (log-sum-exp trick)
3. Store output and LSE

Requirements

• GQA support: heads within a group share block_indices
• Online softmax across all T blocks (numerically stable via running max + LSE)
• Causal masking within gathered blocks
• FP16 matmuls, FP32 accumulation
• Return LSE (needed for backward pass)
• Support both block_size=64 and block_size=32

MI350 considerations

• This is the most performance-critical kernel — it determines the speedup over dense attention
• Block-gather pattern is irregular — plan for LDS staging of gathered KV blocks
• Each program instance handles one (batch, query_pos, gqa_group) — parallelize across these
• BLOCK_H dimension processes all heads in a GQA group together (register tiling)
• Inner loop over T blocks — unroll for small T (16) if register pressure allows
• The tl.multiple_of(col, SELECTION_BLOCK_SIZE) hint in the reference is critical for coalesced loads on MI350

Performance target

• For B=1, M=64k, T=16, block_size=64, G=8, D=128: should achieve >8x speedup over dense FA at same seqlen

Depends on

• [wave] NSA: architecture mapping & design doc for MI350 #1244 (design doc)
• [wave] NSA: top-k block selection kernel #1247 (top-k selection kernel, for block_indices)

References

• _sel_attn_fwd_kernel in tilde-research/nsa-impl/nsa/selection.py
• NSA paper Section 3.2