fix: sparse BF16 LUT crash, lossless round-trip, and preprocess speed
- _block_map_lut_torch: replace the boolean-mask assignment
(filled_matrix[~block_map] = 10_000_000) with torch.where. On the XPU backend
that assignment silently misses entries, leaving 0s that become -1 LUT entries
and drive _lut_to_block_map's scatter_ out of bounds (ScatterGatherKernels
assert). Replace the same unsafe pattern (final_tile_map[~mask] = True) with
|= in sparse_attention and the triton preprocess.
- _lut_to_block_map: scatter True at the recovered block positions instead of
scattering the valid mask. Masked-out tail indices all point at slot 0, so
their trailing False writes were silently dropping a selected block at
position 0; the round-trip is now lossless (empty rows cleared explicitly).
- sage_sparse_bf16: route to the native BF16 sparse kernel when available,
keeping the dense-mask fallback for unsupported configs. The fallback
materializes a [B,Hq,Sq,Skv] mask and OOMs at long sequences.
- _fill_block_map_torch: vectorize the O(k_blocks) python loop (thousands of
sequential kernel launches at long sequences) into a gather/cumsum/scatter.
Preprocess at seq 75600 drops from ~1170 ms to ~450 ms.
auto-round
f9e86122 - fix: sparse BF16 LUT crash, lossless round-trip, and preprocess speed
