[CuTe, SM100] Deadlock in varlen + block-sparse + SplitKV forward: packed `split_idx` is not unpacked on the block-sparse paths (regression from #2559)

[JiaxuanBai]

Summary

On SM100, the CuTe forward kernel deadlocks (GPU spins at 100%, kernel never returns) whenever varlen + block sparsity + SplitKV are combined. This is a regression introduced by #2559 (c46b8144, "Varlen Dynamic Persistent scheduler and metadata"): the scheduler now packs the per-batch dynamic num_splits into the top 16 bits of split_idx, but the SM100 block-sparse code paths don't consistently unpack it, so different warps in the same CTA disagree on the tile's KV block range and the softmax/correction/MMA mbarrier handshake never completes.

Note: #2705 (first-tile flag preservation) does not fix this — it addresses a different scheduler issue and the hang still reproduces with it applied.

Reproduction

pytest -x "tests/cute/test_mask_mod_varlen.py::test_varlen_block_sparse_coarse_kv_requires_offsets[packed-3-14]"

The test hangs in its first (valid) forward call. All parametrizations of test_varlen_block_sparse_splitkv_matches_unsplit hang the same way. With pytest -n, this shows up as workers crashing with "node down: Not properly terminated" once the per-test timeout kills them.

• GPU: B200 (SM100)
• Bisect: c75d019d passes → c46b8144 ([CuTe,Sm100] Varlen Dynamic Persistent scheduler and metadata #2559) hangs; still hangs at current main.
• test_varlen_block_sparse_coarse_kv_requires_offsets[seqused-1-15] (num_splits=1) does not hang, consistent with the analysis below (no split packing when there is effectively no split).

Root cause

Since #2559, when scheduler metadata provides num_splits_dynamic_ptr, the varlen schedulers pack the per-batch dynamic split count into the work tile's split_idx:

• tile_scheduler.py, SingleTileVarlenScheduler._decode_work_tile / DynamicPersistentVarlenScheduler.get_current_work:
  split_idx = split_idx | (num_splits << 16)

The dense path handles this correctly — BlockInfo.get_n_block_min_max unpacks internally when pack_split_idx is set:

if const_expr(self.pack_split_idx):
    num_splits = split_idx >> 16
    split_idx = split_idx & 0xFFFF

But in flash_fwd_sm100.py the block-sparse paths bypass get_n_block_min_max and call the block-sparse helpers (produce_block_sparse_loads_sm100, get_total_block_count, softmax_block_sparse_sm100) directly — and the unpacking is inconsistent across warps:

warp	what it passes to the block-sparse helpers
load	packed split_idx (the existing & 0xFFFF is inside the not use_block_sparsity branch)
MMA	packed split_idx (no unpack at all)
softmax	unpacked split_idx, but static num_splits (mO.shape[0])
correction	unpacked split_idx, but static num_splits

All of these helpers funnel into split_block_range:

blocks_per_split = cute.ceil_div(block_count, num_splits)
block_begin = cutlass.min(split_idx * blocks_per_split, block_count)   # packed split_idx >= 65536
block_end = cutlass.min(block_begin + blocks_per_split, block_count)

With a packed split_idx (≥ 1 << 16), block_begin clamps to block_count, i.e. an empty range. So for every varlen + block-sparse + SplitKV tile:

• load/MMA warps see 0 KV blocks → skip the tile, never produce S, never arrive mbar_P_full_*;
• softmax/correction warps see a non-empty range → follow the non-empty-tile path of the mbarrier contract (see NOTE [SM100 block-sparse empty tiles: mbarrier contract] in block_sparse_utils.py) and wait for S from MMA forever.

Result: all persistent CTAs spin at 100% GPU utilization and the kernel never terminates.

Suggested fix

Mirror the dense-path unpacking on every SM100 block-sparse call site, and pass the dynamic split count to the helpers. E.g. in each affected warp:

num_splits_dyn = num_splits
if const_expr(self.is_split_kv and block_info.pack_split_idx):
    num_splits_dyn = split_idx >> 16
    split_idx = split_idx & 0xFFFF

and use (split_idx, num_splits_dyn) for produce_block_sparse_loads_sm100, get_total_block_count (MMA/softmax/correction) and softmax_block_sparse_sm100. Six small edits in flash_fwd_sm100.py in total. (Alternatively the unpacking could live inside the block-sparse helpers themselves, gated on a pack_split_idx constexpr, matching get_n_block_min_max.)