sparse_prefill_fwd (sm100): skip trailing fully-invalid top-k blocks when topk_length is not provided

[zkyue]

Problem

When topk_length is not passed to flash_mla_sparse_fwd, every row of the sm100 sparse
prefill forward kernel runs ceil(topk / B_TOPK) k-blocks, including trailing blocks whose
indices are all invalid (e.g. rows padded with -1 to a fixed width, which the docstring
explicitly allows). Those blocks already skip the TMA gather (should_skip_tma), but the MMA
and softmax warpgroups have no skip path: both UMMA GEMMs, the TMEM P retrieval, the
masking/exp2 pass and the whole barrier chain still execute. ncu shows tensor-pipe busy
cycles are a constant ~1024 cycles per k-block whether or not the block contains any valid
index (busy-cycle ratios match tile-count ratios exactly across padded widths), while DRAM
bytes stay identical — i.e. each fully-padded block costs a full compute iteration and buys
nothing.

On B200 (bf16, s_q = 8192, h_q = 64, d_qk = d_v = 512), with front-packed rows whose valid
counts ramp causally up to ~1027 and are padded with -1 to topk = 1152, this is a ~9% wall
time overhead vs. passing topk_length (and ~14% in noisier ambient conditions).

Fix

The pipeline's barrier phases are indexed by the k-block index, so blocks cannot simply be
skipped from the middle of the loop (that would deadlock the 3-warpgroup pipeline). Instead,
when the topk_length argument is absent, the kernel prologue derives an equivalent value:
each warp scans the row's index vector from the tail for the last valid entry
(0 <= idx < s_kv) and clamps the k-block loop exactly as an explicitly-passed topk_length
would.

To keep the scan off the critical path, its first-round loads (the trailing 256 indices) are
issued at the very beginning of the kernel — so they are in flight while the prologue
(barrier init, Q TMA, TMEM allocation) runs — and the result is consumed right after the
existing __syncthreads(). Only rows whose padding exceeds 256 indices pay additional scan
rounds, and those rows then save whole k-blocks. Every warp computes the same value from the
same read-only data, so the loop bound stays uniform across all warpgroups (and across the
2-CTA cluster of the head128 kernel) with no extra synchronization.

Correctness notes:

• Invalid indices before the last valid one keep the existing masking path, so arbitrary
  (not just front-packed) index layouts remain correct.
• Rows with no valid index at all derive topk_length = 0 and take exactly today's
  topk_length = 0 path (num_k_blocks = max(..., 1), one fully-masked block, output 0,
  lse = +inf, max_logits = -inf).
• The derived length never cuts off a valid index, so the documented NaN caveat for
  explicitly-passed topk_length (a valid index beyond topk_length[i] pointing at NaN
  payload) cannot be triggered by this path.
• When topk_length is passed, the scan is skipped and behavior is unchanged.

Scope: fwd/head64 and fwd/head128. The fwd_for_small_topk kernel (selected for
h_q = 128 with topk <= 1280) computes its per-job topk_length inside the CLC persistent
outer loop from divergent contexts (single-elected-thread and 16-lane callers), so the same
warp-cooperative prologue scan does not drop in; it is left unchanged and can be addressed
separately if desired.

Numbers

B200, CUDA 13.3, bf16, s_q = 8192, h_q = 64, h_kv = 1, d_qk = d_v = 512. "Ramp-padded" =
front-packed valid counts ramping causally up to ~1027, padded with -1 to width 1152;
"fully-valid" = width 1152 with no invalid index in any row. Medians of 100-iteration
interleaved runs, 3 alternating rounds per build, same GPU:

workload	before (9241ae3)	after	delta
ramp-padded, no topk_length	1.248 ms	1.144 ms	-8.4%
ramp-padded, topk_length passed	1.103 ms	1.101 ms	unchanged
fully-valid, no topk_length	1.280 ms	1.287 ms	+0.5%
fully-valid, topk_length passed	1.280 ms	1.279 ms	unchanged

• tests/test_flash_mla_sparse_prefill.py: all 617 cases pass on the patched build.
• Bitwise: out / lse / max_logits are bit-identical between patched and unpatched
  builds across ramp-padded, scattered-interior-invalid (-1 and >= s_kv mid-row),
  all-invalid, and fully-valid layouts, for head64 and head128 (regular kernel, topk 2048)
  at d_qk 512 and 576, with and without attn_sink, with and without topk_length
  (32 workload dumps x 3 tensors).
• ncu spot check (ramp-padded workload, default clock control, one launch after 5 warmups):
  tensor-pipe busy cycles per SM drop from 1.020M (no topk_length, before) to 0.844M
  (after), exactly matching the topk_length-passed run (0.844M); dram__bytes_read is
  unchanged (587.2 MB -> 587.5 MB); sass global_ld instructions drop 2.39M -> 2.27M
  (fewer per-block index loads outweigh the scan's tail reads). The residual wall-time gap
  vs. the topk_length path (~4%) is scan latency on the ramp rows, not tensor work.
• Register/spill status: head64 124 regs (same as before), head128 128 regs; no new ptxas
  spill warnings (the STACK:8 on fwd_for_small_topk decode pre-exists at 9241ae3).

Alternatives considered

• Scanning eagerly at the top of the kernel (before the prologue) costs a full memory
  round trip per CTA on the critical path: with 1 CTA/SM and ~55 sequential waves this
  showed up as +2.7% .. +19% on unpadded workloads. The issue-early/consume-late split
  removes that (fully-valid regression 0.2-0.9% across measurement windows).
• Wider scan rounds (512/1024 indices) did not reduce the residual gap and increased the
  fully-valid overhead (more registers held across the prologue).
• A per-k-block skip inside the pipelined loop is not possible without re-numbering
  barrier phases across the three warpgroups (phases are indexed by k), which this PR
  deliberately avoids.