From 3f0d72f22f1c4cfe7d5a8007ca650b1173071da5 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Tue, 20 Oct 2026 07:33:13 +0000 Subject: [PATCH] [kernel] GB10 VSA tile-128: review fixes for routing, limits and padding - Kernel: zero V padding rows on partial 64-key chunks (and skip 16-key steps wholly past the valid keys), so NaN/Inf padding in K or V cannot reach the output; cache the SM count and the shared-memory attribute per device in atomics instead of an unsynchronized function-static slot. - is_supported now checks the launch limits (batch * heads <= 65535, rows < 2^31, tiles < 2^22) via within_launch_limits; the binding checks the row limit too. The H3 route checks the limits per call and falls back to split-Triton instead of failing in the binding. - H3 route: one tile-128 route cache (ptx_blackwell / sm121 / split), resolved on a two-tile probe before tile(), so the sm121 and split routes no longer pad odd tile counts with the ptx_blackwell partner tile. - Warn once when a GB10 falls back to split-Triton; the fail-closed error names the sm_121a build. - CMake: accept 12.1f and +PTX spellings; warn when CUDA < 13 skips it. - Docs: bf16 fallback leaves FASTVIDEO_H3_VSA_FINE_ATTENTION unset on Spark. - Tests: the real-kernel oracle tests assert the route and spy on the kernel entry; new tests for NaN/Inf padding, launch limits, the GB10 warning and the over-limit fallback. --- docs/inference/fasth3_omniref.md | 7 +- fastvideo-kernel/CMakeLists.txt | 12 +- fastvideo-kernel/README.md | 2 +- .../attention/block_sparse_blk128_sm121.cu | 7 +- .../block_sparse_kernel_blk128_sm121.cuh | 54 +++++-- .../block_sparse_attn_sm121.py | 24 +++- .../tests/test_block_sparse_sm121.py | 31 ++++ .../backends/video_sparse_attn_h3.py | 136 ++++++++++++------ .../attention/test_vsa_h3_heads_first_tile.py | 2 + .../attention/test_vsa_h3_ref2va_regions.py | 88 +++++++++--- 10 files changed, 274 insertions(+), 89 deletions(-) diff --git a/docs/inference/fasth3_omniref.md b/docs/inference/fasth3_omniref.md index eae6ac143..f3c0a66d8 100644 --- a/docs/inference/fasth3_omniref.md +++ b/docs/inference/fasth3_omniref.md @@ -159,10 +159,11 @@ missing native kernel raises rather than silently selecting another mode. query tiles, per-token FP8 Q/K for prefix tiles, bf16 P.V and K smoothing. Needs no consumer-Blackwell extension but is slower, and its numerics differ from the native kernel (which also quantizes P.V). -3. **bf16 VSA**: set `FASTVIDEO_H3_VSA_FP4=0` and +3. **bf16 VSA**: set `FASTVIDEO_H3_VSA_FP4=0`. On the DGX Spark, leave + `FASTVIDEO_H3_VSA_FINE_ATTENTION` unset so the GB10 CUDA kernel runs the + 128-token tiles. On the other consumer CUDA GPUs, including sm_89, add `FASTVIDEO_H3_VSA_FINE_ATTENTION=emulator-bf16`. This runs unquantized - attention on the same tiles with Triton on consumer CUDA GPUs, including - sm_89. + attention on the same tiles with Triton. A nonempty fine-attention override takes priority over `FASTVIDEO_H3_VSA_FP4=1`. These overrides use the eager path; disable inference diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index 81089b152..825a737ca 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -487,15 +487,17 @@ if(BUILD_CXX_KERNELS) COMPILE_OPTIONS "${_VSA_FWD_GENCODE};-DVSA_BHSD=true") endif() # VSA block-sparse attention FORWARD for the GB10 (DGX Spark), 128-token blocks: an mma.sync + TMA kernel - # for sm_121a, the device the tcgen05 kernel above cannot run on. Inference only (no LSE). "12.1" without the - # suffix counts too (the GB10 Docker image builds that), since this one source carries its own sm_121a image. + # for sm_121a, the device the tcgen05 kernel above cannot run on. Inference only (no LSE). Every spelling of + # the GB10 counts: "12.1" without the suffix (the GB10 Docker image builds that), "12.1f" and "12.1+PTX", + # since this one source carries its own sm_121a image. set(ENABLE_VSA_SM121A OFF) - if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(12\\.1a?|121a?|sm_121a?)([; ,]|$)") + if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(12\\.1[af]?|121[af]?|sm_121[af]?)(\\+PTX)?([; ,]|$)") if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0) set(ENABLE_VSA_SM121A ON) else() - message(STATUS "fastvideo-kernel: skipping block_sparse_sm121 (sm_121a needs CUDA >= 13.0, have " - "${CMAKE_CUDA_COMPILER_VERSION})") + # A warning: the GB10 then runs tile-128 VSA on the slower split-Triton route. + message(WARNING "fastvideo-kernel: skipping block_sparse_sm121 (sm_121a needs CUDA >= 13.0, have " + "${CMAKE_CUDA_COMPILER_VERSION}); tile-128 VSA on a GB10 falls back to Triton") endif() endif() if(ENABLE_VSA_SM121A) diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index b0aadef80..4b1af2c1e 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -11,7 +11,7 @@ Compiled CUDA extensions (CMake, see the build summary printed at the end of eve | `fastvideo_kernel._C.fastvideo_kernel_ops` | TurboDiffusion INT8 GEMM, quant, RMSNorm, LayerNorm | `csrc/turbodiffusion/` | every arch in `TORCH_CUDA_ARCH_LIST` | always built | | same extension, optional part | ThunderKittens sliding-tile attention (`sta_fwd`) and VSA block-sparse (`block_sparse_fwd/bwd`) | `csrc/attention/*_h100.cu` | Hopper `sm_90a` only | `FASTVIDEO_KERNEL_BUILD_TK` (AUTO = ON iff `9.0a` is in the arch list; always OFF on aarch64 hosts — TK headers don't compile there) | | same extension, optional part | MiniMax-H3 block-sparse VSA forward (64- and 128-token blocks) | `csrc/attention/block_sparse*_blackwell.cu` | Data-center Blackwell `sm_100a`/`sm_103a` | ON iff `10.0a` or `10.3a` is in `TORCH_CUDA_ARCH_LIST` | -| same extension, optional part | MiniMax-H3 block-sparse VSA forward, 128-token blocks, inference (`block_sparse_sm121_blk128_fwd`, no LSE) | `csrc/attention/block_sparse*_sm121.cu*` | GB10 / DGX Spark `sm_121a` | ON iff `12.1a` or `12.1` is in `TORCH_CUDA_ARCH_LIST` and CUDA ≥ 13 | +| same extension, optional part | MiniMax-H3 block-sparse VSA forward, 128-token blocks, inference (`block_sparse_sm121_blk128_fwd`, no LSE) | `csrc/attention/block_sparse*_sm121.cu*` | GB10 / DGX Spark `sm_121a` | ON iff `TORCH_CUDA_ARCH_LIST` names the GB10 (`12.1a`, `12.1`, `12.1f`, with or without `+PTX`) and CUDA ≥ 13 | | same extension, optional part | fused NVLink Ulysses all-to-all | `csrc/comm/ulysses_all_to_all.cu` | CUDA | `FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A` (AUTO = ON with NCCL 2.29+ device headers and library; always OFF on ROCm) | | `fp4attn_cuda`, `fp4quant_cuda` | FP4 attention + quantization ("attn_qat_infer", modified SageAttention3) | `attn_qat_infer/` | consumer/workstation Blackwell `sm_120a` (CUDA ≥ 12.8) + `sm_121a` (DGX Spark GB10; CUDA ≥ 13) | `FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER` (AUTO = ON iff `12.0a` or `12.1a` is in the arch list) | diff --git a/fastvideo-kernel/csrc/attention/block_sparse_blk128_sm121.cu b/fastvideo-kernel/csrc/attention/block_sparse_blk128_sm121.cu index fdcd870ce..aefba6cd8 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_blk128_sm121.cu +++ b/fastvideo-kernel/csrc/attention/block_sparse_blk128_sm121.cu @@ -58,8 +58,11 @@ std::vector block_sparse_sm121_blk128_fwd(torch::Tensor q, torch: "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); - TORCH_CHECK(blocks <= std::numeric_limits::max() && q2k_idx.size(3) <= std::numeric_limits::max(), - "too many blocks: ", blocks, " (max_kv ", q2k_idx.size(3), ")"); + // 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"); 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 ad84c6310..e401e5a83 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_kernel_blk128_sm121.cuh +++ b/fastvideo-kernel/csrc/attention/block_sparse_kernel_blk128_sm121.cuh @@ -27,6 +27,7 @@ #include #include +#include #include #include #include @@ -38,7 +39,7 @@ namespace vsa_sm121 { struct Vsa128Args { const void* q; // [BH, S, 128] bf16; a tile's padding rows may hold anything, they are not read const void* k; // [BH, S, 128] bf16 - const void* v; // [BH, S, 128] bf16 + const void* v; // [BH, S, 128] bf16; a tile's padding keys may hold anything, K and V both void* out; // [BH, S, 128] bf16; every element is written const int32_t* q2k_idx; // [BH, T, max_kv]: each query tile's selected key tiles; entries past its count are not read const int32_t* q2k_num; // [BH, T]: how many key tiles each query tile selects @@ -359,13 +360,33 @@ __device__ __forceinline__ void store_p(uint32_t dst, const float (&p)[8][4], in :: "r"(addr), "r"(a0), "r"(a1), "r"(a2), "r"(a3) : "memory"); } } -__device__ __forceinline__ void pv(float (&o)[16][4], const uint32_t (&p)[4][4], const uint32_t (&v_addr)[4]) { +// A B fragment register's bf16 pair of keys (the lower key in the lower half), given how many keys from the +// pair's first one on are valid: both kept, the lower one only, or neither. +__device__ __forceinline__ uint32_t valid_keys(uint32_t b, int valid) { + return valid>=2 ? b : (valid==1 ? (b&0x0000ffffu) : 0u); +} +// Masked: the chunk holds fewer than 64 valid keys. P is zero past them, but the V rows there are whatever the +// caller's padding holds, and 0 * NaN or 0 * Inf would still reach O: skip the 16-key steps that lie wholly past +// the valid keys, and zero the invalid rows of the one that straddles them. +template +__device__ __forceinline__ void pv(float (&o)[16][4], const uint32_t (&p)[4][4], const uint32_t (&v_addr)[4], + int keys, int lane) { #pragma unroll for (int kk=0; kk<4; ++kk) { + if constexpr (Masked) { + if (16*kk>=keys) break; + } + // b[0], b[2] hold keys 16kk + c + {0, 1}; b[1], b[3] those 8 further on (c = 2 * (lane % 4)). `valid` + // counts the valid keys from 16kk + c on. + const int valid=keys-16*kk-2*(lane&3); #pragma unroll for (int dd=0; dd<8; ++dd) { uint32_t b[4]; ptx::ldmatrix_x4_trans(b,v_addr[dd&3]+2048*kk+BoxBytes*(dd>>2)); + if constexpr (Masked) { + b[0]=valid_keys(b[0],valid); b[1]=valid_keys(b[1],valid-8); + b[2]=valid_keys(b[2],valid); b[3]=valid_keys(b[3],valid-8); + } ptx::mma_bf16_16816(o[2*dd],p[kk],b[0],b[1]); ptx::mma_bf16_16816(o[2*dd+1],p[kk],b[2],b[3]); } @@ -516,11 +537,13 @@ __global__ __launch_bounds__(Threads,1) void forward(const Vsa128Args args, // prevents PV output staging from overwriting QK's K alias too early. asm volatile("bar.sync 0, 512;" ::: "memory"); } else { - int pending=0; + // Each V slot's valid key count, eight bits per slot. + int pending=0, vkeys=0; #pragma unroll for (int s=0;s<3;++s) { if (fill.next()) { ++pending; + vkeys|=fill.keys()<<(8*s); if (local_tid==0) load_operand(base+VOffset+s*OperandBytes,&vmap,fill,head_row,operand_full(1,s)); } } @@ -555,11 +578,14 @@ __global__ __launch_bounds__(Threads,1) void forward(const Vsa128Args args, uint32_t vaddr[4]; #pragma unroll for (int j=0;j<4;++j) vaddr[j]=base+VOffset+stage*OperandBytes+vlane[j]; - pv(o,pf,vaddr); + const int keys=(vkeys>>(8*stage))&255; + if (keys(o,pf,vaddr,keys,lane); + else pv(o,pf,vaddr,keys,lane); } // All three V slots were prefetched. Refill exactly the slot just // read with i+3, not A3's i+2 slot (stage+2)%3. if (fill.next()) { + vkeys=(vkeys&~(255<<(8*stage)))|(fill.keys()<<(8*stage)); const uint32_t counter=operand_counter(1,stage); if (last_reader(counter,lane)) { if (lane==0) { @@ -596,20 +622,26 @@ inline cudaError_t launch(const Vsa128Args& args,cudaStream_t stream) { CUtensorMap kmap,vmap; if (split128::encode_map(&kmap,args.k,rows)!=CUDA_SUCCESS || split128::encode_map(&vmap,args.v,rows)!=CUDA_SUCCESS) return cudaErrorInvalidValue; - cudaError_t err=cudaFuncSetAttribute(split128::forward,cudaFuncAttributeMaxDynamicSharedMemorySize,split128::SharedBytes); - if (err!=cudaSuccess) return err; // One CTA per SM fits (99 KB of shared memory each): the SMs are the - // CTAs that run at once. A device property, read once per device. + // CTAs that run at once. Read once per device, after the per-device + // shared-memory attribute is set; 0 means not yet. Threads that race on a + // device's first launch each set the same attribute and store the same + // count, and a thread never sees another device's count. + constexpr int MaxDevices=64; + static std::atomic sm_count[MaxDevices]; int device=0; - err=cudaGetDevice(&device); + cudaError_t err=cudaGetDevice(&device); if (err!=cudaSuccess) return err; - static int slots_device=-1, slots=1; - if (device!=slots_device) { + if (device<0 || device>=MaxDevices) return cudaErrorInvalidDevice; + int slots=sm_count[device].load(std::memory_order_acquire); + if (slots==0) { + err=cudaFuncSetAttribute(split128::forward,cudaFuncAttributeMaxDynamicSharedMemorySize,split128::SharedBytes); + if (err!=cudaSuccess) return err; int sms=0; err=cudaDeviceGetAttribute(&sms,cudaDevAttrMultiProcessorCount,device); if (err!=cudaSuccess) return err; slots=sms>0 ? sms : 1; - slots_device=device; + sm_count[device].store(slots,std::memory_order_release); } split128::forward<<>>(args,kmap,vmap,slots); return cudaGetLastError(); 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 8940774d7..04cdf0cc0 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm121.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm121.py @@ -28,6 +28,22 @@ _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. +_MAX_BATCH_HEADS = 65535 +_MAX_ROWS = 2**31 - 1 +_MAX_BLOCKS = 2**22 - 1 + + +def within_launch_limits(q: torch.Tensor, num_blocks: int) -> bool: + """True iff one launch can cover q's [B, H, S, 128] shape with num_blocks blocks. + + The part of ``is_supported`` that depends on the size of the call, not on the device, dtype or + layout, so callers that cache ``is_supported`` per device must still check it per call. + ``map_to_index`` makes q2k_idx's last dimension num_blocks, which this bounds too. + """ + batch_heads = q.shape[0] * q.shape[1] + return batch_heads <= _MAX_BATCH_HEADS and batch_heads * q.shape[2] <= _MAX_ROWS and num_blocks <= _MAX_BLOCKS def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: @@ -37,12 +53,14 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: * the extension carries the sm_121a kernel and q is on a GB10 (compute capability 12.1); * q is a contiguous 4-D bf16 BHSD tensor with head_dim 128 and seqlen == num_blocks * 128, num_blocks = variable_block_sizes.numel(); any num_blocks, odd included; - * variable_block_sizes is an integer CUDA tensor. + * variable_block_sizes is an integer CUDA tensor; + * the call is within the launch limits (``within_launch_limits``): batch * heads <= 65535, + batch * heads * seqlen < 2^31 and num_blocks < 2^22. What the kernel accepts beyond that, and is tested to handle: per-row counts q2k_num in [0, max_kv], non-uniform across rows (a row with count 0 gets exactly-zero output); q2k_idx entries are only read below each row's count (map_to_index's -1 padding is never read); any variable_block_sizes in [0, 128] (keys at or past a block's count are masked, and query rows - past it are written zero). + past it are written zero); padding rows of q, k and v that hold anything, NaN and Inf included. """ if not _HAS_VSA_SM121 or not q.is_cuda: return False @@ -55,7 +73,7 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: return False if not variable_block_sizes.is_cuda or variable_block_sizes.dtype not in (torch.int32, torch.int64): return False - return True + return within_launch_limits(q, num_blocks) @torch.library.custom_op( diff --git a/fastvideo-kernel/tests/test_block_sparse_sm121.py b/fastvideo-kernel/tests/test_block_sparse_sm121.py index f7c8cc754..0d9d28c36 100644 --- a/fastvideo-kernel/tests/test_block_sparse_sm121.py +++ b/fastvideo-kernel/tests/test_block_sparse_sm121.py @@ -135,6 +135,37 @@ def test_deterministic(): assert torch.equal(first, second) +@pytest.mark.parametrize("fill", [float("nan"), float("inf")]) +def test_padding_keys_holding_nan_or_inf_do_not_leak(fill): + """Padding rows of K and V may hold anything: P is zero there, and 0 * NaN must not reach the output.""" + q, k, v, idx, num, vbs = make_case(num_blocks=6, topk=3, ragged=True, seed=23) + assert (vbs % 64 != 0).any(), "the case must have partial 64-key chunks" + expected, has = reference(q, k, v, idx, num, vbs) + valid = (torch.arange(BLOCK, device="cuda")[None] < vbs[:, None]).reshape(-1) + for t in (q, k, v): + t[:, :, ~valid] = fill + out = run_with_nan_pool(q, k, v, idx, num, vbs) + assert torch.isfinite(out).all() + assert (out[~has] == 0).all() + torch.testing.assert_close(out.float()[has], expected[has], 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): + return torch.empty(shape, device="meta", dtype=torch.bfloat16) + + assert vsa.within_launch_limits(meta(1, 65535, 2 * BLOCK, HEAD_DIM), 2) + assert not vsa.within_launch_limits(meta(2, 32768, 2 * BLOCK, HEAD_DIM), 2) + assert not vsa.within_launch_limits(meta(1, 2, 2**30, HEAD_DIM), 2**23) + assert not vsa.within_launch_limits(meta(1, 1, 2**22 * BLOCK, HEAD_DIM), 2**22) + assert vsa.within_launch_limits(meta(1, 1, (2**22 - 1) * BLOCK, HEAD_DIM), 2**22 - 1) + q, _, _, _, _, vbs = make_case(num_blocks=2, topk=1, heads=4) + assert vsa.is_supported(q, vbs) + monkeypatch.setattr(vsa, "_MAX_BATCH_HEADS", 3) + assert not vsa.is_supported(q, vbs) + + def test_is_supported_rejects_what_it_cannot_run(): q, k, v, idx, num, vbs = make_case(num_blocks=2, topk=1) assert vsa.is_supported(q, vbs) diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index 0a13490ea..5c332646d 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -56,10 +56,10 @@ Tile 128 prefers the same ptx_blackwell forward, which carries a 128-token block instantiation. It needs no opt-in, pads an odd tile count with the same zero-valid partner tile, and runs no-grad forwards only. On a GB10 (DGX Spark, sm_121), where that kernel cannot run, the 128-token sm_121a forward -(``fastvideo_kernel.block_sparse_attn_sm121``) takes the same padded call. When -neither the extension nor the device can run one of them, each 128-token tile -splits into two contiguous 64-token halves on the Triton kernels; if those are -not installed either, the call fails closed. +(``fastvideo_kernel.block_sparse_attn_sm121``) takes the call, which needs no +partner tile. When neither the extension nor the device can run one of them, +each 128-token tile splits into two contiguous 64-token halves on the Triton +kernels; if those are not installed either, the call fails closed. """ import functools @@ -569,9 +569,21 @@ def _is_datacenter_blackwell(device: torch.device) -> bool: return device.type == "cuda" and torch.cuda.get_device_capability(device)[0] == 10 -def _sm121_can_run(query_bhsd: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: - """Whether the GB10 (sm_121a) 128-token forward can take a tile-128 call (needs map_to_index too).""" - return _sm121 is not None and map_to_index is not None and _sm121.is_supported(query_bhsd, variable_block_sizes) +def _is_gb10(device: torch.device) -> bool: + """sm_121 (DGX Spark), the device the sm_121a tile-128 kernel targets.""" + return device.type == "cuda" and torch.cuda.get_device_capability(device) == (12, 1) + + +def _sm121_unavailable_reason(query_bhsd: torch.Tensor, variable_block_sizes: torch.Tensor) -> str | None: + """Why the GB10 (sm_121a) 128-token forward cannot run here, or None if it can.""" + if _sm121 is None: + return "fastvideo_kernel.block_sparse_attn_sm121 is not installed" + if map_to_index is None: + return "fastvideo_kernel.triton_kernels.index (map_to_index) is not importable" + if not _sm121.is_supported(query_bhsd, variable_block_sizes): + return ("block_sparse_attn_sm121.is_supported returned False (needs a GB10 and a fastvideo-kernel built " + "with TORCH_CUDA_ARCH_LIST containing 12.1a and CUDA >= 13, bf16, head_dim 128)") + return None @functools.cache @@ -623,15 +635,15 @@ class MiniMaxH3VSAImpl(AttentionImpl): # request-time env/probe/fallback behavior; only Dynamo capture reads # the prepared, static route. self._regional_compile_ptx_blackwell_enabled: bool | None = None - # Tile-128 route per (device, dtype, head size): None when the CUDA - # kernel can run it, else why not. The kernel's predicate otherwise - # depends only on the tile-128 buffer contract (contiguous BHSD, - # 128-token blocks, an even tile count with the partner tile, integer - # tile sizes), which every call meets, so it is evaluated once. - self._tile128_route: dict[tuple[torch.device, torch.dtype, int], str | None] = {} - # Per the same key, when ptx_blackwell cannot run it: whether the GB10 - # (sm_121a) 128-token forward can, under the same contract. - self._tile128_sm121: dict[tuple[torch.device, torch.dtype, int], bool] = {} + # Tile-128 route per (device, dtype, head size): the kernel that runs + # it ("ptx_blackwell", "sm121" or "split", see _tile128_route) and why + # the faster ones cannot (None for ptx_blackwell). The kernels' + # predicates otherwise depend only on the tile-128 buffer contract + # (contiguous BHSD, 128-token blocks, the partner tile where + # ptx_blackwell needs one, integer tile sizes), which every call meets, + # so they are evaluated once; only the sm121 kernel's launch-size + # limits are checked per call. + self._tile128_routes: dict[tuple[torch.device, torch.dtype, int], tuple[str, str | None]] = {} # Optional replacement of the fine block-sparse kernel only (selection, # mask and the gated compression branch are unchanged). Called as # ``fn(query, key, value, mask, metadata) -> [B, S, H, D]`` on the @@ -695,6 +707,45 @@ class MiniMaxH3VSAImpl(AttentionImpl): logger.warning_once(f"VSA-H3 regional compile is unavailable and will stay eager: {reason}") return reason + def _tile128_route(self, device: torch.device, dtype: torch.dtype, head_size: int, + compiling: bool) -> tuple[str, str | None]: + """The no-grad tile-128 route on this device: ``(kernel, reason)``. + + ``kernel`` is "ptx_blackwell" (sm_100/sm_103, an even tile count), + "sm121" (GB10, any tile count) or "split" (each tile as two halves on + the Triton tile-64 kernels); ``reason`` says why the faster kernels + cannot run, None for ptx_blackwell. Resolved once per key on a + two-tile probe, before ``tile()`` decides whether to add the partner + tile only ptx_blackwell needs. + """ + route_key = (device, dtype, head_size) + route = self._tile128_routes.get(route_key) + if route is not None: + return route + probe_query = torch.empty((1, 1, 2 * 128, head_size), device=device, dtype=dtype) + probe_block_sizes = torch.full((2, ), 128, device=device, dtype=torch.int32) + reason = _ptx_blackwell_unavailable_reason(_ptx_blackwell, probe_query, probe_block_sizes, grad_mode=False) + if reason is None and map_to_index is None: + reason = "fastvideo_kernel.triton_kernels.index (map_to_index) is not importable" + if reason is None: + route = ("ptx_blackwell", None) + else: + sm121_reason = _sm121_unavailable_reason(probe_query, probe_block_sizes) + if sm121_reason is None: + route = ("sm121", reason) + else: + route = ("split", f"ptx_blackwell: {reason}; sm_121a: {sm121_reason}") + # These GPUs should have a native kernel; a broken build would + # otherwise only show up as a slow run. + if not compiling and _is_datacenter_blackwell(device): + logger.warning_once(f"VSA-H3 tile 128 on a data-center Blackwell GPU is falling back to " + f"the slower Triton tile-64 kernels: {reason}.") + elif not compiling and _is_gb10(device): + logger.warning_once(f"VSA-H3 tile 128 on a GB10 is falling back to the slower Triton tile-64 " + f"kernels: {sm121_reason}.") + self._tile128_routes[route_key] = route + return route + def tile(self, x: torch.Tensor, attn_metadata: MiniMaxH3VSAMetadata) -> torch.Tensor: """Scatter rows into the padded tile buffer (pad positions stay zero). @@ -703,9 +754,9 @@ class MiniMaxH3VSAImpl(AttentionImpl): sites in ``forward()`` read it immediately). A grad-tracking forward instead receives a fresh buffer and leaves the holder untouched, so the builder never retains autograd state across steps. Odd no-grad - ptx_blackwell requests (tile 128 always, tile 64 when opted in) carry one - additional all-zero tile internally; metadata and all observable - outputs retain the logical geometry. + ptx_blackwell requests (tile 128 when it routes there, tile 64 when + opted in) carry one additional all-zero tile internally; metadata and + all observable outputs retain the logical geometry. """ compiling = torch.compiler.is_compiling() regional_compiling = compiling and self._regional_compile_ptx_blackwell_enabled is True @@ -722,9 +773,12 @@ class MiniMaxH3VSAImpl(AttentionImpl): ptx_blackwell_requested = False else: ptx_blackwell_requested = envs.FASTVIDEO_VSA_PTX_BLACKWELL.get() - # Tile 128 has no route other than the ptx_blackwell CUDA kernel. - ptx_blackwell_route = attn_metadata.tile_elems == 128 or (attn_metadata.tile_elems == 64 - and ptx_blackwell_requested) + # Tile 128 picks its kernel by device; only ptx_blackwell needs the partner tile. + if 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: + ptx_blackwell_route = attn_metadata.tile_elems == 64 and ptx_blackwell_requested if regional_compiling: # Pad arithmetically so odd and even tile counts share one graph. needs_ptx_blackwell_pair, kernel_tiles = False, n_tiles + n_tiles % 2 @@ -937,32 +991,25 @@ class MiniMaxH3VSAImpl(AttentionImpl): split_tile128 = False if tile_elems == 128: # The CUDA kernel's 128-token instantiation needs no opt-in; on a GB10 - # its sm_121a counterpart runs the same call; without either, the + # its sm_121a counterpart runs the call; without either, the # Triton tile-64 kernels run each tile as two halves. if grad_mode: raise NotImplementedError("VSA-H3 tile 128 supports no-grad inference forwards only.") - route_key = (q_bhsd.device, q_bhsd.dtype, int(q_bhsd.shape[-1])) - if route_key not in self._tile128_route: - reason = _ptx_blackwell_unavailable_reason(_ptx_blackwell, q_bhsd, - ptx_blackwell_variable_block_sizes, grad_mode) - if reason is None and map_to_index is None: - reason = "fastvideo_kernel.triton_kernels.index (map_to_index) is not importable" - if reason is not None: - self._tile128_sm121[route_key] = _sm121_can_run(q_bhsd, ptx_blackwell_variable_block_sizes) - if reason is not None and not compiling and _is_datacenter_blackwell(q_bhsd.device): - # These GPUs should have the native kernel; a broken build - # would otherwise only show up as a slow run. - logger.warning_once(f"VSA-H3 tile 128 on a data-center Blackwell GPU is falling back to " - f"the slower Triton tile-64 kernels: {reason}.") - self._tile128_route[route_key] = reason - reason = self._tile128_route[route_key] - if reason is None: + kernel, reason = self._tile128_route(q_bhsd.device, q_bhsd.dtype, int(q_bhsd.shape[-1]), compiling) + if kernel == "sm121" and not _sm121.within_launch_limits(q_bhsd, n_tiles + pair_tiles): + # Too large for one launch (batch * heads, rows or tiles); this call only. + if not compiling: + 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": use_ptx_blackwell = True - elif self._tile128_sm121.get(route_key, False): + elif kernel == "sm121": use_sm121 = True elif block_sparse_attn_64_bhsd is None: - raise RuntimeError(f"VSA-H3 tile 128 requires the ptx_blackwell block-sparse kernel or " - f"the Triton tile-64 kernels: {reason}.") + raise RuntimeError(f"VSA-H3 tile 128 requires the ptx_blackwell block-sparse kernel " + f"(sm_100/sm_103), the GB10 block-sparse kernel (fastvideo-kernel built for " + f"sm_121a) or the Triton tile-64 kernels: {reason}.") else: split_tile128 = True elif regional_compiling: @@ -1026,8 +1073,9 @@ class MiniMaxH3VSAImpl(AttentionImpl): need_lse=False, ) elif use_sm121: - # The GB10 counterpart of the call above: the same padded tiles, - # mask and per-row counts, output only. + # The GB10 counterpart of the call above: the same mask and + # per-row counts, output only. It takes any tile count, so + # tile() adds no partner tile on this route. if not compiling: logger.info_once("MiniMax-H3 VSA tile-128 forward: using the GB10 (sm_121a) block-sparse kernel") q2k_idx, q2k_num = map_to_index(ptx_blackwell_mask) 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 8f99a945a..51dce050d 100644 --- a/fastvideo/tests/attention/test_vsa_h3_heads_first_tile.py +++ b/fastvideo/tests/attention/test_vsa_h3_heads_first_tile.py @@ -38,6 +38,8 @@ def _metadata(seq: int, n_tiles: int, tile_elems: int) -> SimpleNamespace: def _impl() -> vsa.MiniMaxH3VSAImpl: impl = object.__new__(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 = {} 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 a9e854ddc..dfcfae709 100644 --- a/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py +++ b/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py @@ -427,13 +427,17 @@ class _FakePtxBlackwell: class _FakeSm121: """Stands in for fastvideo_kernel.block_sparse_attn_sm121 (the GB10 128-token forward).""" - def __init__(self, supported=True): + def __init__(self, supported=True, within_limits=True): self.supported = supported + self.within_limits = within_limits self.calls = [] def is_supported(self, q, variable_block_sizes): return self.supported + def within_launch_limits(self, q, num_blocks): + return self.within_limits + def block_sparse_attn_sm121(self, q, k, v, q2k_idx, q2k_num, variable_block_sizes): self.calls.append(dict(q=q, q2k_num=q2k_num, vbs=variable_block_sizes)) return q.clone(), None @@ -499,6 +503,7 @@ def test_forward_tile128_even_tile_count_no_partner_tile(tile128): def test_forward_tile128_falls_back_to_triton_when_cuda_kernel_unavailable(tile128, monkeypatch): meta = _build_r2v(sparsity=0.5, tile_size=128) n_tiles = meta.variable_block_sizes.numel() + assert n_tiles % 2 == 1 tile128.supported = False calls = [] @@ -507,7 +512,9 @@ def test_forward_tile128_falls_back_to_triton_when_cuda_kernel_unavailable(tile1 return torch.zeros_like(q), None monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", triton) - _, out = _run_tile128(meta) + tiled, out = _run_tile128(meta) + # The split route needs no partner tile either. + assert tiled.shape[1] == n_tiles * 128 assert len(calls) == 1 q, mask, vbs = calls[0] assert q.shape[2] == n_tiles * 128 @@ -519,19 +526,35 @@ def test_forward_tile128_falls_back_to_triton_when_cuda_kernel_unavailable(tile1 def test_forward_tile128_uses_sm121_kernel_when_ptx_blackwell_unavailable(tile128, monkeypatch): - """On a GB10 the sm_121a forward takes the ptx_blackwell call: the padded partner tile and per-row counts.""" + """On a GB10 the sm_121a forward takes the call with the logical tile count: no partner tile, even when odd.""" 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 tile128.supported = False sm121 = _FakeSm121() monkeypatch.setattr(vsa_h3, "_sm121", sm121) - _, out = _run_tile128(meta) + tiled, out = _run_tile128(meta) + assert tiled.shape[1] == n_tiles * 128 assert tile128.calls == [] and len(sm121.calls) == 1 call = sm121.calls[0] - assert call["q"].shape[2] == (n_tiles + 1) * 128 - assert call["vbs"].dtype == torch.int32 and call["vbs"].numel() == n_tiles + 1 and call["vbs"][-1] == 0 - assert (call["q2k_num"][..., n_tiles] == 0).all() + assert call["q"].shape[2] == n_tiles * 128 + assert call["vbs"].dtype == torch.int32 and call["vbs"].numel() == n_tiles + assert call["q2k_num"].shape[-1] == n_tiles + assert out.shape == (1, n_tiles * 128, _HEADS, _DIM) + + +def test_forward_tile128_sm121_call_past_launch_limits_runs_split(tile128, monkeypatch): + """A call too large for one GB10 launch takes the Triton split route instead of failing in the binding.""" + meta = _build_r2v(sparsity=0.5, tile_size=128, ref_keep_rate=0.25) + n_tiles = meta.variable_block_sizes.numel() + tile128.supported = False + sm121 = _FakeSm121(within_limits=False) + monkeypatch.setattr(vsa_h3, "_sm121", sm121) + 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]) + _, out = _run_tile128(meta) + assert sm121.calls == [] and len(calls) == 1 and calls[0][2] == n_tiles * 128 assert out.shape == (1, n_tiles * 128, _HEADS, _DIM) @@ -555,26 +578,32 @@ def test_forward_tile128_falls_back_to_triton_when_sm121_unsupported(tile128, mo assert len(calls) == 1 and sm121.calls == [] -@pytest.mark.parametrize("datacenter", [False, True]) -def test_forward_tile128_fallback_warns_only_on_datacenter_blackwell(tile128, monkeypatch, datacenter): +@pytest.mark.parametrize("device_kind", ["other", "datacenter", "gb10"]) +def test_forward_tile128_fallback_warns_only_where_a_native_kernel_belongs(tile128, monkeypatch, device_kind): meta = _build_r2v(sparsity=0.5, tile_size=128) tile128.supported = False warnings = [] monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", lambda q, k, v, mask, vbs: (torch.zeros_like(q), None)) - monkeypatch.setattr(vsa_h3, "_is_datacenter_blackwell", lambda device: datacenter) + monkeypatch.setattr(vsa_h3, "_is_datacenter_blackwell", lambda device: device_kind == "datacenter") + monkeypatch.setattr(vsa_h3, "_is_gb10", lambda device: device_kind == "gb10") monkeypatch.setattr(vsa_h3.logger, "warning_once", warnings.append) _run_tile128(meta) - assert len(warnings) == int(datacenter) - if datacenter: - assert "falling back to the slower Triton tile-64 kernels" in warnings[0] + assert len(warnings) == int(device_kind != "other") + if device_kind == "datacenter": + assert "data-center Blackwell GPU is falling back to the slower Triton tile-64 kernels" in warnings[0] + elif device_kind == "gb10": + assert "on a GB10 is falling back to the slower Triton tile-64 kernels" in warnings[0] + assert "block_sparse_attn_sm121 is not installed" in warnings[0] def test_forward_tile128_unavailable_cuda_kernel_and_triton(tile128, monkeypatch): meta = _build_r2v(sparsity=0.5, tile_size=128) tile128.supported = False monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", None) - with pytest.raises(RuntimeError, match="tile 128 requires the ptx_blackwell block-sparse kernel"): + with pytest.raises(RuntimeError, match="tile 128 requires the ptx_blackwell block-sparse kernel") as error: _run_tile128(meta) + assert "GB10 block-sparse kernel (fastvideo-kernel built for sm_121a)" in str(error.value) + assert "sm_121a: fastvideo_kernel.block_sparse_attn_sm121 is not installed" in str(error.value) monkeypatch.setattr(vsa_h3, "_ptx_blackwell", None) with pytest.raises(RuntimeError, match="not installed"): _run_tile128(meta) @@ -632,28 +661,44 @@ def test_build_tile128_uses_4x4x8_tiles(): @pytest.mark.parametrize("sparsity,ref_keep_rate", [(0.5, 0.25), (0.9, 0.1)]) -def test_forward_tile128_real_ptx_blackwell_kernel_matches_oracle(env_overrides, sparsity, ref_keep_rate): +def test_forward_tile128_real_ptx_blackwell_kernel_matches_oracle(env_overrides, monkeypatch, sparsity, ref_keep_rate): """The real 128-token CUDA forward on a multi-region layout against an FP32 masked-SDPA oracle.""" if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}: pytest.skip("requires an sm_100a/sm_103a GPU (B200, B300, GB200, GB300)") forwards = getattr(vsa_h3._ptx_blackwell, "_FWD_BY_BLOCK", {}) if vsa_h3._ptx_blackwell is not None else {} if forwards.get(128) is None or vsa_h3.map_to_index is None: pytest.skip("requires a fastvideo_kernel build with the 128-token ptx_blackwell forward") - _check_real_tile128_forward_against_oracle(env_overrides, sparsity, ref_keep_rate) + _check_real_tile128_forward_against_oracle(env_overrides, monkeypatch, sparsity, ref_keep_rate, + kernel_module=vsa_h3._ptx_blackwell, + entry="block_sparse_attn_ptx_blackwell", + route="ptx_blackwell") @pytest.mark.parametrize("sparsity,ref_keep_rate", [(0.5, 0.25), (0.9, 0.1)]) -def test_forward_tile128_real_sm121_kernel_matches_oracle(env_overrides, sparsity, ref_keep_rate): +def test_forward_tile128_real_sm121_kernel_matches_oracle(env_overrides, monkeypatch, sparsity, ref_keep_rate): """The real GB10 (sm_121a) 128-token forward on a multi-region layout against an FP32 masked-SDPA oracle.""" if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (12, 1): pytest.skip("requires a GB10 GPU (DGX Spark, sm_121)") if vsa_h3._sm121 is None or not vsa_h3._sm121._HAS_VSA_SM121 or vsa_h3.map_to_index is None: pytest.skip("requires a fastvideo_kernel build with the sm_121a block-sparse forward") - _check_real_tile128_forward_against_oracle(env_overrides, sparsity, ref_keep_rate) + _check_real_tile128_forward_against_oracle(env_overrides, monkeypatch, sparsity, ref_keep_rate, + kernel_module=vsa_h3._sm121, + entry="block_sparse_attn_sm121", + route="sm121") -def _check_real_tile128_forward_against_oracle(env_overrides, sparsity, ref_keep_rate): +def _check_real_tile128_forward_against_oracle(env_overrides, monkeypatch, sparsity, ref_keep_rate, *, kernel_module, + entry, route): + """Also proves the route: the Triton split route matches the oracle too, so a silent fallback must fail.""" env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(None)) + launches = [] + real_entry = getattr(kernel_module, entry) + + def spy(*args, **kwargs): + launches.append(args[0].shape) + return real_entry(*args, **kwargs) + + monkeypatch.setattr(kernel_module, entry, spy) device = torch.device("cuda") meta = _build_r2v(sparsity=sparsity, tile_size=128, ref_keep_rate=ref_keep_rate, device=device) n_tiles = meta.variable_block_sizes.numel() @@ -663,9 +708,12 @@ def _check_real_tile128_forward_against_oracle(env_overrides, sparsity, ref_keep raw = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, device=device, dtype=torch.bfloat16) with torch.inference_mode(): tiled = impl.preprocess_qkv(raw, meta) - assert tiled.shape[1] == (n_tiles + n_tiles % 2) * 128 + assert impl._tile128_routes[(raw.device, torch.bfloat16, _DIM)][0] == route + # Only ptx_blackwell needs an even tile count. + assert tiled.shape[1] == (n_tiles + (n_tiles % 2 if route == "ptx_blackwell" else 0)) * 128 query, key, value = tiled.chunk(3, dim=0) actual = impl.postprocess_output(impl.forward(query, key, value, None, meta), meta) + assert len(launches) == 1, f"the {route} kernel must run exactly once, ran {len(launches)} times" # The mask the backend selected: pooled scores of the logical tiles. logical_query, logical_key, logical_value = (t[:, :logical] for t in (query, key, value)) scores = torch.matmul(_pool_tiles(logical_query, meta.variable_block_sizes, 128),