[kernel] GB10 VSA tile-128: second review pass (device check, NaN rows, routing, linkage)

- Binding: require compute capability 12.1. Other archs in the same build get
  an empty kernel that returned uninitialized output without an error.
- Kernel: a NaN softmax sum now gives a NaN output row (l != 0 instead of
  l > 0), as on the ptx_blackwell and Triton routes; rows without keys stay 0.
- Kernel header: the kernel, encode_map and launch get internal linkage, so a
  second translation unit can include the header (it failed to link before).
  The launch limits move into the header (MaxBatchHeads, MaxRows, MaxTiles),
  used by both launch() and the binding and mirrored in Python.
- Header comments: describe the QK/PV overlap that PhaseLead=8 gives; drop
  the stale cp.async comment and the unused cp_async_mbar_arrive/prefetch_l2.
- H3 backend: an odd tile count without tile()'s partner tile falls back to
  split instead of reaching ptx_blackwell; tile() no longer resolves the
  tile-128 route (or warns) under an eager fine-attention override.
- Python module docstring states the replacement contract for the kernel.
This commit is contained in:
SolitaryThinker
2026-10-20 15:21:51 +00:00
parent 3f0d72f22f
commit 8503617df3
7 changed files with 93 additions and 38 deletions
@@ -10,6 +10,12 @@ backward and is never reached by grad-tracking calls.
It is a separate module, not another capability in ``block_sparse_attn_ptx_blackwell``, because
that module's ``is_supported`` also gates ``block_sparse_attn_from_indices``'s autograd op, which
always asks for the LSE.
Replacing the kernel: callers (the VSA-H3 backend) use only ``is_supported``,
``within_launch_limits`` and ``block_sparse_attn_sm121``. A new kernel keeps them, together with the
``block_sparse_sm121_blk128_fwd`` binding, by providing ``vsa_sm121::Vsa128Args``, the ``Max*``
limits and ``vsa_sm121::launch`` in the header; ``fastvideo-kernel/tests/test_block_sparse_sm121.py``
is its acceptance test.
"""
from typing import Tuple
@@ -28,8 +34,7 @@ _HAS_VSA_SM121 = _FWD is not None
_SUPPORTED_COMPUTE_CAPABILITIES = {(12, 1)}
BLOCK = 128
HEAD_DIM = 128
# Launch limits: the grid's y dimension is batch * heads; TMA addresses K and V by int32 row; the
# kernel packs a walk's start entry (< num_blocks) into 22 bits.
# Launch limits, mirroring vsa_sm121::MaxBatchHeads, MaxRows and MaxTiles in the kernel header.
_MAX_BATCH_HEADS = 65535
_MAX_ROWS = 2**31 - 1
_MAX_BLOCKS = 2**22 - 1