From 8503617df396e0e33e6c1507bda78684b027d894 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Tue, 20 Oct 2026 15:21:51 +0000 Subject: [PATCH] [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. --- .../attention/block_sparse_blk128_sm121.cu | 19 ++++---- .../block_sparse_kernel_blk128_sm121.cuh | 47 ++++++++----------- .../block_sparse_attn_sm121.py | 9 +++- .../tests/test_block_sparse_sm121.py | 13 +++++ .../backends/video_sparse_attn_h3.py | 8 +++- .../attention/test_vsa_h3_heads_first_tile.py | 1 + .../attention/test_vsa_h3_ref2va_regions.py | 34 ++++++++++++++ 7 files changed, 93 insertions(+), 38 deletions(-) diff --git a/fastvideo-kernel/csrc/attention/block_sparse_blk128_sm121.cu b/fastvideo-kernel/csrc/attention/block_sparse_blk128_sm121.cu index aefba6cd8..189f5d9cd 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_blk128_sm121.cu +++ b/fastvideo-kernel/csrc/attention/block_sparse_blk128_sm121.cu @@ -10,7 +10,6 @@ #include #include -#include #include #include "block_sparse_kernel_blk128_sm121.cuh" @@ -38,6 +37,11 @@ std::vector block_sparse_sm121_blk128_fwd(torch::Tensor q, torch: check_tensor(q, "q", at::kBFloat16, device); check_tensor(k, "k", at::kBFloat16, device); check_tensor(v, "v", at::kBFloat16, device); + // Other archs in the same build get an empty kernel (it is compiled for sm_121a only), which would return + // empty_like garbage without an error. + const cudaDeviceProp* prop = at::cuda::getDeviceProperties(device.index()); + TORCH_CHECK(prop->major == 12 && prop->minor == 1, "block_sparse_sm121_blk128_fwd needs a GB10 (sm_121), got sm_", + prop->major, prop->minor); TORCH_CHECK(q.dim() == 4, "q must be [batch, heads, seq, 128], got ", q.sizes()); TORCH_CHECK(k.sizes() == q.sizes() && v.sizes() == q.sizes(), "q, k and v must have one shape, got ", q.sizes(), ", ", k.sizes(), " and ", v.sizes()); @@ -56,13 +60,12 @@ std::vector block_sparse_sm121_blk128_fwd(torch::Tensor q, torch: "q2k_idx must be [", batch, ", ", heads, ", ", blocks, ", max_kv], got ", q2k_idx.sizes()); TORCH_CHECK(q2k_num.dim() == 3 && q2k_num.size(0) == batch && q2k_num.size(1) == heads && q2k_num.size(2) == blocks, "q2k_num must be [", batch, ", ", heads, ", ", blocks, "], got ", q2k_num.sizes()); - // The grid is (num_blocks, batch * heads). - TORCH_CHECK(batch * heads <= 65535, "batch * heads must be at most 65535, got ", batch * heads); - // TMA addresses K and V by int32 row, and the walker packs its start entry into 22 bits. - TORCH_CHECK(batch * heads * seq <= std::numeric_limits::max(), - "batch * heads * seqlen must be below 2^31, got ", batch * heads * seq); - TORCH_CHECK(blocks < (int64_t{1} << 22) && q2k_idx.size(3) < (int64_t{1} << 22), "too many blocks: ", blocks, - " (max_kv ", q2k_idx.size(3), "); both must be below 2^22"); + TORCH_CHECK(batch * heads <= vsa_sm121::MaxBatchHeads, "batch * heads must be at most ", vsa_sm121::MaxBatchHeads, + ", got ", batch * heads); + TORCH_CHECK(batch * heads * seq <= vsa_sm121::MaxRows, "batch * heads * seqlen must be at most ", + vsa_sm121::MaxRows, ", got ", batch * heads * seq); + TORCH_CHECK(blocks <= vsa_sm121::MaxTiles && q2k_idx.size(3) <= vsa_sm121::MaxTiles, "too many blocks: ", blocks, + " (max_kv ", q2k_idx.size(3), "); both must be at most ", vsa_sm121::MaxTiles); auto out = torch::empty_like(q); if (q.numel() == 0) return {out}; diff --git a/fastvideo-kernel/csrc/attention/block_sparse_kernel_blk128_sm121.cuh b/fastvideo-kernel/csrc/attention/block_sparse_kernel_blk128_sm121.cuh index e401e5a83..2f76a9808 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_kernel_blk128_sm121.cuh +++ b/fastvideo-kernel/csrc/attention/block_sparse_kernel_blk128_sm121.cuh @@ -11,9 +11,9 @@ // How, on mma.sync (the GB10 has no wgmma, tcgen05 or tensor memory): one CTA of 512 threads per (query tile, // batch * head), one CTA per SM (99 KB of shared memory). // - Two groups of eight warps, 16 query rows each: the QK group computes S = Q K^T and the online softmax and -// publishes P (bf16) and the rescale factors through shared memory; the PV group accumulates O += P V. Each -// PV warp waits for its QK partner half-way through the partner's next QK burst, so the two groups' mma -// bursts overlap instead of finishing together. +// publishes P (bf16) and the rescale factors through shared memory; the PV group accumulates O += P V. A +// QK warp computes the next chunk's scores while its PV partner multiplies the previous P by V; the pair +// meets (pair_sync) only before P is overwritten, so the two groups' mma bursts overlap. // - K and V stream in 64-key chunks through TMA (two 64x64 boxes, 128-byte swizzle) into separate rings; // the last warp to finish reading a slot refills it. // - Each CTA walks its selection starting (launch index mod SMs) / SMs of the way through, so the CTAs that @@ -49,6 +49,12 @@ struct Vsa128Args { int max_kv; // the last dimension of q2k_idx }; +// The largest call one launch covers. The binding checks them before launching, and +// fastvideo_kernel/block_sparse_attn_sm121.py's _MAX_* mirror them for the callers that route around a launch. +constexpr int64_t MaxBatchHeads = 65535; // batch_heads: the grid's y dimension +constexpr int64_t MaxRows = INT32_MAX; // batch_heads * tiles * 128: TMA addresses K and V by int32 row +constexpr int64_t MaxTiles = (1 << 22) - 1; // tiles and max_kv: the walker packs its start entry into 22 bits + namespace ptx { @@ -59,12 +65,7 @@ __device__ __forceinline__ void cp_async_16(uint32_t dst, const void* src, uint3 } // Closes the group of the cp.async copies this thread issued since the last commit (an empty group if none). -// UNUSED by the current kernel, deliberately. The main loop publishes a stage's copies through that stage's -// mbarrier (`cp.async.mbarrier.arrive.noinc`), which waits for every thread's copies into the stage and makes -// them visible -- something a commit group cannot do, because a group is private to the thread that issued it -// and `cp.async.wait_group` below retires only this thread's own. Both wrappers stay because they are the right -// tool for a thread-private wait; re-introducing cp_async_wait() beside the loop's per-stage waits would not -// restore anything, it would only add a wait that cannot see the other warps' copies. +// Only the Q load uses cp.async; K and V stream through TMA. __device__ __forceinline__ void cp_async_commit() { asm volatile("cp.async.commit_group;\n" ::: "memory"); } @@ -92,12 +93,6 @@ __device__ __forceinline__ void mbar_arrive(uint32_t addr) { asm volatile("mbarrier.arrive.shared::cta.b64 %0, [%1];\n" : "=l"(state) : "r"(addr) : "memory"); } -// The mbarrier `addr` is signalled when this thread's prior cp.async copies have completed and are visible to -// the threads that observe the barrier's phase flip. One arrival per thread per pass. -__device__ __forceinline__ void cp_async_mbar_arrive(uint32_t addr) { - asm volatile("cp.async.mbarrier.arrive.noinc.shared::cta.b64 [%0];\n" ::"r"(addr) : "memory"); -} - // True once phase `phase` (0 or 1, alternating) of the mbarrier `addr` has completed; does not block. __device__ __forceinline__ bool mbar_try_wait(uint32_t addr, uint32_t phase) { uint32_t done; @@ -118,11 +113,6 @@ __device__ __forceinline__ void mbar_wait(uint32_t addr, uint32_t phase) { } } -// Brings the 128-byte line holding `p` into L2. -__device__ __forceinline__ void prefetch_l2(const void* p) { - asm volatile("prefetch.global.L2 [%0];\n" ::"l"(p)); -} - // Loads four 8x8 matrices of 16-bit elements: lanes 8i..8i+7 give the addresses of the 8 rows (16 bytes each) // of matrix i, and r[i] receives, in every lane, elements (lane / 4, 2 * (lane % 4) + {0, 1}) of matrix i. __device__ __forceinline__ void ldmatrix_x4(uint32_t (&r)[4], uint32_t addr) { @@ -169,6 +159,7 @@ __device__ __forceinline__ float exp2_approx(float x) { namespace split128 { using bf16 = __nv_bfloat16; constexpr int Tile = 128, Dim = 128, Chunk = 64; +// QK k-steps (of 8) a warp issues before waiting for its PV partner to release P; 8 waits after the whole burst. constexpr int PhaseLead = 8; constexpr int GroupThreads = 256, Threads = 512; constexpr int OperandBytes = 16384, BoxBytes = 8192; @@ -399,7 +390,7 @@ __device__ __forceinline__ void output(bf16* out, uint8_t* staging, const float #pragma unroll for (int r=0; r<2; ++r) { const float l=sums[row0+8*r]; - keep[r]=row0+8*r0.f; + keep[r]=row0+8*r(rows)}; const cuuint64_t strides[1]={Dim*2}; const cuuint32_t box[2]={64,64}, elem_strides[2]={1,1}; @@ -613,12 +604,14 @@ CUresult encode_map(CUtensorMap* map, const void* ptr, int64_t rows) { } } // namespace split128 -inline cudaError_t launch(const Vsa128Args& args,cudaStream_t stream) { +// The kernel's whole host interface. Internal linkage, like the kernel and its helpers: a translation unit that +// includes this header gets its own copy, so a second includer (a benchmark, a newer kernel built beside this +// one) links without duplicate symbols. +static cudaError_t launch(const Vsa128Args& args,cudaStream_t stream) { if (args.batch_heads<=0 || args.tiles<=0) return cudaSuccess; const int64_t rows=int64_t(args.batch_heads)*args.tiles*split128::Tile; - if (rows>INT32_MAX) return cudaErrorInvalidValue; - // The walker keeps its start entry (< count <= tiles) in bits 9..30 of an int. - if (args.tiles>=(1<<22) || args.max_kv>=(1<<22)) return cudaErrorInvalidValue; + if (args.batch_heads>MaxBatchHeads || rows>MaxRows || args.tiles>MaxTiles || args.max_kv>MaxTiles) + return cudaErrorInvalidValue; CUtensorMap kmap,vmap; if (split128::encode_map(&kmap,args.k,rows)!=CUDA_SUCCESS || split128::encode_map(&vmap,args.v,rows)!=CUDA_SUCCESS) return cudaErrorInvalidValue; diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm121.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm121.py index 04cdf0cc0..e1ceb8f14 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm121.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm121.py @@ -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 diff --git a/fastvideo-kernel/tests/test_block_sparse_sm121.py b/fastvideo-kernel/tests/test_block_sparse_sm121.py index 0d9d28c36..b6d2d3a04 100644 --- a/fastvideo-kernel/tests/test_block_sparse_sm121.py +++ b/fastvideo-kernel/tests/test_block_sparse_sm121.py @@ -150,6 +150,19 @@ def test_padding_keys_holding_nan_or_inf_do_not_leak(fill): torch.testing.assert_close(out.float()[has], expected[has], atol=0.02, rtol=0.02) +def test_nan_in_a_valid_query_row_reaches_its_output_row(): + """Like the other routes, a NaN input row gives a NaN output row instead of a silent zero.""" + q, k, v, idx, num, vbs = make_case(num_blocks=4, topk=2, seed=29) + expected, has = reference(q, k, v, idx, num, vbs) + q[0, 0, 5] = float("nan") + out = run_with_nan_pool(q, k, v, idx, num, vbs) + assert torch.isnan(out[0, 0, 5]).all() + rest = has.clone() + rest[0, 0, 5] = False + assert torch.isfinite(out[rest]).all() and (out[~has] == 0).all() + torch.testing.assert_close(out.float()[rest], expected[rest], atol=0.02, rtol=0.02) + + def test_launch_limits(monkeypatch): """Sizes one launch cannot cover: is_supported says so instead of the binding raising.""" def meta(*shape): diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index 5c332646d..b3adf671c 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -774,7 +774,10 @@ class MiniMaxH3VSAImpl(AttentionImpl): else: ptx_blackwell_requested = envs.FASTVIDEO_VSA_PTX_BLACKWELL.get() # Tile 128 picks its kernel by device; only ptx_blackwell needs the partner tile. - if attn_metadata.tile_elems == 128: + # An eager fine-attention override reads the logical rows and runs no kernel. + if attn_metadata.tile_elems == 128 and self.fine_attention_override is not None and not compiling: + ptx_blackwell_route = False + elif attn_metadata.tile_elems == 128: ptx_blackwell_route = (not grad_mode and not regional_compiling and self._tile128_route( x.device, x.dtype, int(x.shape[-1]), compiling)[0] == "ptx_blackwell") else: @@ -1002,6 +1005,9 @@ class MiniMaxH3VSAImpl(AttentionImpl): logger.warning_once("VSA-H3 tile 128: a call exceeds the GB10 kernel's launch limits; " "it runs on the Triton tile-64 kernels") kernel, reason = "split", f"{reason}; sm_121a: the call exceeds its launch limits" + if kernel == "ptx_blackwell" and (n_tiles + pair_tiles) % 2 != 0: + # Only tile() adds the partner tile; a call that skipped it cannot run there. + kernel, reason = "split", "ptx_blackwell: an odd tile count without the partner tile" if kernel == "ptx_blackwell": use_ptx_blackwell = True elif kernel == "sm121": diff --git a/fastvideo/tests/attention/test_vsa_h3_heads_first_tile.py b/fastvideo/tests/attention/test_vsa_h3_heads_first_tile.py index 51dce050d..d8a40de32 100644 --- a/fastvideo/tests/attention/test_vsa_h3_heads_first_tile.py +++ b/fastvideo/tests/attention/test_vsa_h3_heads_first_tile.py @@ -40,6 +40,7 @@ def _impl() -> vsa.MiniMaxH3VSAImpl: impl._regional_compile_ptx_blackwell_enabled = None # tile() resolves the tile-128 route to decide on the ptx_blackwell partner tile. impl._tile128_routes = {} + impl.fine_attention_override = None return impl diff --git a/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py b/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py index dfcfae709..ee81ff131 100644 --- a/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py +++ b/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py @@ -558,6 +558,40 @@ def test_forward_tile128_sm121_call_past_launch_limits_runs_split(tile128, monke assert out.shape == (1, n_tiles * 128, _HEADS, _DIM) +def test_forward_tile128_odd_call_without_partner_tile_runs_split(tile128, monkeypatch): + """ptx_blackwell needs an even tile count; a call that skipped tile()'s partner tile falls back to Triton.""" + meta = _build_r2v(sparsity=0.5, tile_size=128, ref_keep_rate=0.25) + n_tiles = meta.variable_block_sizes.numel() + assert n_tiles % 2 == 1 + calls = [] + monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", + lambda q, k, v, mask, vbs: (calls.append(q.shape), (torch.zeros_like(q), None))[1]) + impl = _impl(_DIM) + raw = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, dtype=torch.bfloat16) + query, key, value = impl.preprocess_qkv(raw, meta)[:, :n_tiles * 128].chunk(3, dim=0) + out = impl.forward(query, key, value, None, meta) + assert tile128.calls == [] and len(calls) == 1 and calls[0][2] == n_tiles * 128 + assert out.shape == (1, n_tiles * 128, _HEADS, _DIM) + + +def test_tile128_fine_attention_override_skips_kernel_routing(tile128): + """An eager override runs no kernel: tile() neither probes the route nor adds the partner tile.""" + meta = _build_r2v(sparsity=0.5, tile_size=128, ref_keep_rate=0.25) + n_tiles = meta.variable_block_sizes.numel() + assert n_tiles % 2 == 1 + probes, calls = [], [] + tile128.is_supported = lambda q, variable_block_sizes: probes.append(q.device) or True + impl = _impl(_DIM) + impl.fine_attention_override = lambda q, k, v, mask, metadata: calls.append(q.shape) or torch.zeros_like(q) + raw = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, dtype=torch.bfloat16) + tiled = impl.preprocess_qkv(raw, meta) + assert tiled.shape[1] == n_tiles * 128 and probes == [] + query, key, value = tiled.chunk(3, dim=0) + out = impl.forward(query, key, value, None, meta) + assert len(calls) == 1 and tile128.calls == [] and probes == [] + assert out.shape == (1, n_tiles * 128, _HEADS, _DIM) + + def test_forward_tile128_prefers_ptx_blackwell_over_sm121(tile128, monkeypatch): meta = _build_r2v(sparsity=0.5, tile_size=128) sm121 = _FakeSm121()