FastH3's tile-128 route runs the ptx_blackwell CUDA forward on sm_100/sm_103 and, on every other GPU,
splits each 128-token tile into two 64-token halves on the Triton kernels. On the DGX Spark that split
route is what OmniRef's BF16 attention runs today. This adds a 128-token forward built for sm_121a
(mma.sync + TMA; the GB10 has no tcgen05 or tensor memory) and routes the tile-128 call to it on a
GB10, ahead of the split.
- csrc/attention/block_sparse_kernel_blk128_sm121.cuh and block_sparse_blk128_sm121.cu: the kernel and
its binding block_sparse_sm121_blk128_fwd(q, k, v, q2k_idx, q2k_num, variable_block_sizes) -> {out},
inference only, no LSE.
- CMakeLists.txt: built when TORCH_CUDA_ARCH_LIST has 12.1a or 12.1 and CUDA >= 13, so both the
aarch64 wheel and the GB10 Docker image carry it.
- fastvideo_kernel.block_sparse_attn_sm121: is_supported and an opaque custom op with a fake. It is a
separate module because block_sparse_attn_ptx_blackwell.is_supported also gates an LSE-requiring
autograd path.
- video_sparse_attn_h3: the tile-128 route tries it after ptx_blackwell and before split_tile128.
- Tests: kernel correctness against an FP32 reference (ragged blocks, empty selections, odd block
counts, many waves, a NaN-filled allocator, determinism), and the H3 route (fakes, plus a real-GB10
oracle test).
For Wan VSA inference with 64-token tiles, each attention layer scattered Q, K,
V and the gate into a tile buffer and then made four transposed copies. On
SM100, one Triton kernel (tile_to_bhsd) now does the scatter and the transposes
in one pass. Training, 256-token tiles, other GPUs and builds without the kernel
use the previous path. FASTVIDEO_DISABLE_VSA64_FUSED_LAYOUT=1 turns the fused
path off.
On GB200 the layout and attention step goes from 15.7 ms to 8.6 ms. FastWan 1.3B
frames are bit-exact with the previous path at sequence parallel 1 and 2.
On SM100, the Attn-QAT training kernel now uses a faster route for non-causal
BF16 attention with equal query and key lengths of 16384 tokens or more. The
forward pass processes full KV tiles without bounds masks, and a separate masked
step handles the partial tail tile. The backward pass uses 64x128 dQ and dK/dV
tiles. Outputs, softmax statistics and all gradients are bitwise identical to the
previous route. FASTVIDEO_ATTN_QAT_SM100_OPTIMIZED=0 turns off the SM100 route,
and FASTVIDEO_ATTN_QAT_SM100_WIDE_BWD=0 keeps the 64x64 backward tiles.
On GB200 with 6 heads and 31200 tokens, the forward kernel time goes from 9.14 ms
to 6.96 ms and the backward kernel time from 17.41 ms to 15.85 ms.
SolitaryThinker