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()