From 3196835913ce9e97b23888b32643a1207cc5ee70 Mon Sep 17 00:00:00 2001 From: lishunyang12 Date: Thu, 3 Sep 2026 08:29:48 +0000 Subject: [PATCH] feat(vsa): add native sm103a support Signed-off-by: lishunyang12 --- .github/workflows/publish-kernel.yml | 4 ++-- fastvideo-kernel/CMakeLists.txt | 24 +++++++++---------- fastvideo-kernel/README.md | 4 ++-- fastvideo-kernel/build.sh | 4 ++++ .../attention/block_sparse_kernel_sm100a.cuh | 22 ++++++++--------- .../attention/block_sparse_launch_sm100a.cuh | 2 +- .../csrc/attention/block_sparse_sm100a.cu | 2 +- .../csrc/attention/primitives.cuh | 3 +-- fastvideo-kernel/csrc/common_extension.cpp | 4 ++-- .../fastvideo_kernel/block_sparse_attn.py | 19 ++++++++------- .../block_sparse_attn_sm100a.py | 15 +++++++----- .../python/fastvideo_kernel/vsa_utils.py | 3 ++- .../backends/video_sparse_attn_h3.py | 20 ++++++++-------- .../attention/test_vsa_h3_sm100a_route.py | 21 ++++++++-------- tests/bench_block_sparse_sm100a.py | 11 +++++---- tests/test_block_sparse_sm100a.py | 6 ++--- tests/test_block_sparse_sm100a_dispatch.py | 6 ++--- tests/test_fasth3_packaging.py | 5 +++- 18 files changed, 94 insertions(+), 81 deletions(-) diff --git a/.github/workflows/publish-kernel.yml b/.github/workflows/publish-kernel.yml index e73956c700..5e89ecf2cb 100644 --- a/.github/workflows/publish-kernel.yml +++ b/.github/workflows/publish-kernel.yml @@ -63,7 +63,7 @@ jobs: torch-cuda-short: 'cu130' platform: # x86_64 builds the full cu126 + cu130 set. cu130 ships the - # data-center Blackwell sm_100a VSA and consumer sm_120a FP4 + # data-center Blackwell sm_100a/sm_103a VSA and consumer sm_120a FP4 # kernels. - os: ubuntu-22.04 arch: x86_64 @@ -169,7 +169,7 @@ jobs: # covers sm_120a; turbodiffusion covers sm_100a+sm_120a. The sm_100 FP4 # forward is the FA4 CuTe DSL path in the fastvideo package (PR #1221), # JIT-compiled at runtime — not built into this wheel. - # * x86_64 cu130 = Hopper TK + data-center Blackwell sm_100a VSA + # * x86_64 cu130 = Hopper TK + data-center Blackwell sm_100a/sm_103a VSA # + consumer Blackwell sm_120a FP4. # * x86_64 cu126 = Hopper TK only (older drivers; CUDA < 12.8 has no FP4). # The per-arch split in CMakeLists pins the FP4 targets to sm_120a and builds diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index fdb0284260..ac4b30aea7 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -72,7 +72,7 @@ if(NOT GPU_BACKEND STREQUAL "ROCM") else() # Best-effort probe of the visible GPU (mirrors build.sh detect_with_torch). execute_process( - COMMAND "${Python_EXECUTABLE}" -c "import torch; assert torch.cuda.is_available(); mj, mn = torch.cuda.get_device_capability(0); print(f'{mj}.{mn}a' if (mj, mn) in ((9, 0), (12, 0)) else f'{mj}.{mn}')" + COMMAND "${Python_EXECUTABLE}" -c "import torch; assert torch.cuda.is_available(); mj, mn = torch.cuda.get_device_capability(0); print(f'{mj}.{mn}a' if (mj, mn) in ((9, 0), (10, 0), (10, 3), (12, 0)) else f'{mj}.{mn}')" OUTPUT_VARIABLE _FV_ARCH_LIST OUTPUT_STRIP_TRAILING_WHITESPACE RESULT_VARIABLE _fv_detect_rc @@ -382,11 +382,11 @@ if(BUILD_CXX_KERNELS) # Combined FastVideo Extension # Using name 'fastvideo_kernel_ops' to distinguish from the python package namespace # --------------------------------------------------------------------------- - # VSA block-sparse attention forward, Blackwell (sm_100a) only. + # VSA block-sparse attention forward, data-center Blackwell (sm_100a/sm_103a). # - # NOTE the "a" suffix: -arch=sm_100a is NOT enough -- it emits a plain sm_100 target and - # ptxas rejects every tcgen05 / setmaxnreg instruction. The explicit gencode spelling - # below is required, and matches the 10.0a entry in TORCH_CUDA_ARCH_LIST. + # NOTE the "a" suffix: these kernels use architecture-conditional tcgen05 / setmaxnreg + # instructions. Explicit gencode entries are required for both the GB200 (sm_100a) and + # B300/GB300 (sm_103a) images. The source body is shared; ptxas specializes each image. # # Built for 64-token sparse blocks -- FastVideo's default (4,4,4) tiling, so no # tile-size change and no top-k granularity change is needed. VSA_BHSD @@ -396,21 +396,21 @@ if(BUILD_CXX_KERNELS) # exported (build.sh, and `pip install` with it set) the branch above only prints it -- # the cmake variable stays empty, so testing that alone silently skips the kernel and # leaves a build that succeeds with the op missing. - set(ENABLE_VSA_SM100A OFF) + set(ENABLE_VSA_SM100_FAMILY OFF) set(_VSA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}") if(NOT _VSA_ARCH_LIST AND DEFINED ENV{TORCH_CUDA_ARCH_LIST}) set(_VSA_ARCH_LIST "$ENV{TORCH_CUDA_ARCH_LIST}") endif() - if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(10\\.0a|100a|sm_100a)([; ,]|$)") - set(ENABLE_VSA_SM100A ON) + if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(10\\.(0|3)a|10(0|3)a|sm_10(0|3)a)([; ,]|$)") + set(ENABLE_VSA_SM100_FAMILY ON) endif() - if(ENABLE_VSA_SM100A) - message(STATUS "fastvideo-kernel: building block_sparse_sm100a (Blackwell, 64- and 128-token blocks)") + if(ENABLE_VSA_SM100_FAMILY) + message(STATUS "fastvideo-kernel: building block_sparse_sm100a (sm_100a/sm_103a, 64- and 128-token blocks)") list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_sm100a.cu csrc/attention/block_sparse_blk128_sm100a.cu) set_source_files_properties(csrc/attention/block_sparse_sm100a.cu csrc/attention/block_sparse_blk128_sm100a.cu PROPERTIES - COMPILE_OPTIONS "-gencode;arch=compute_100a,code=sm_100a;-DVSA_BHSD=true") + COMPILE_OPTIONS "-gencode;arch=compute_100a,code=sm_100a;-gencode;arch=compute_103a,code=sm_103a;-DVSA_BHSD=true") endif() Python_add_library(fastvideo_kernel_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI @@ -428,7 +428,7 @@ if(BUILD_CXX_KERNELS) # Build compile definitions list set(COMPILE_DEFS TORCH_EXTENSION_NAME=fastvideo_kernel_ops) - if(ENABLE_VSA_SM100A) + if(ENABLE_VSA_SM100_FAMILY) list(APPEND COMPILE_DEFS TK_COMPILE_BLOCK_SPARSE_VSA_SM100A) endif() if(ENABLE_TK_KERNELS) diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index 2e14b77495..c321df5c47 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -10,7 +10,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*_sm100a.cu` | Blackwell `sm_100a` only | ON iff `10.0a` is in `TORCH_CUDA_ARCH_LIST` | +| same extension, optional part | MiniMax-H3 block-sparse VSA forward (64- and 128-token blocks) | `csrc/attention/block_sparse*_sm100a.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 | 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 Blackwell `sm_120a` only, CUDA ≥ 12.8 | `FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER` (AUTO = ON iff `12.0a` is in the arch list) | @@ -37,7 +37,7 @@ Runtime-JIT kernels (no build step, ship in every wheel/image): Notes: - No Docker image ships the FP4 kernels; only the x86_64/aarch64 cu130 wheels do. -- Both cu130 PyPI wheels ship the MiniMax-H3 sm_100a VSA forward. +- Both cu130 PyPI wheels ship native sm_100a and sm_103a images for the MiniMax-H3 VSA forward. - On arm64 images (GH200 included) STA/VSA run on the Triton fallbacks, since TK never builds on aarch64. - Ulysses AUTO builds only when CMake finds a NCCL library and device-API headers with the 2.29 initializers. Use `-DFASTVIDEO_KERNEL_BUILD_ULYSSES_A2A=ON` to require it or `OFF` to test the portable build. diff --git a/fastvideo-kernel/build.sh b/fastvideo-kernel/build.sh index c7ed847c36..972e0bf3f1 100755 --- a/fastvideo-kernel/build.sh +++ b/fastvideo-kernel/build.sh @@ -139,6 +139,10 @@ if [ "${GPU_BACKEND}" = "CUDA" ]; then if [ -z "${TORCH_CUDA_ARCH_LIST:-}" ]; then if [ "${cc_major}" = "9" ] && [ "${cc_minor}" = "0" ]; then export TORCH_CUDA_ARCH_LIST="9.0a" + elif [ "${cc_major}" = "10" ] && { [ "${cc_minor}" = "0" ] || [ "${cc_minor}" = "3" ]; }; then + # Data-center Blackwell VSA uses architecture-conditional tcgen05 + # instructions and therefore requires the 'a' target. + export TORCH_CUDA_ARCH_LIST="${cc_major}.${cc_minor}a" elif [ "${cc_major}" = "12" ] && [ "${cc_minor}" = "0" ]; then # Blackwell sm_120 needs the arch-conditional 'a' suffix so CMake's # AUTO gate (matches 12.0a/120a/sm_120a) builds the attn_qat_infer diff --git a/fastvideo-kernel/csrc/attention/block_sparse_kernel_sm100a.cuh b/fastvideo-kernel/csrc/attention/block_sparse_kernel_sm100a.cuh index 3113b03b00..fa17cf251b 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_kernel_sm100a.cuh +++ b/fastvideo-kernel/csrc/attention/block_sparse_kernel_sm100a.cuh @@ -1,4 +1,5 @@ -// block_sparse_kernel_sm100a.cuh -- VSA block-sparse FMHA forward (per-q-block top-k), sm_100a. +// block_sparse_kernel_sm100a.cuh -- VSA block-sparse FMHA forward (per-q-block top-k), +// data-center Blackwell sm_100a/sm_103a. The filename is retained for API compatibility. // Warp-specialized: load / MMA (tcgen05) / softmax / correction / epilogue / scheduler. // Writes O and, when asked, the log-sum-exp the backward consumes. // @@ -202,15 +203,14 @@ fmha_context_bf16_gen_kernel(const __grid_constant__ CUtensorMap tmap_q, const int* __restrict__ variable_block_sizes, float* __restrict__ lse_out) { // Multi-arch builds: torch's cmake appends -gencode for EVERY entry of -// TORCH_CUDA_ARCH_LIST to this TU on top of the pinned compute_100a pass, and -// tcgen05/setmaxnreg do not exist outside sm_100a -- ptxas rejects the sm_120a -// (or plain sm_100) pass outright. Keep the body only where it can compile: -// the host pass (no __CUDA_ARCH__, needed for launch plumbing) and the -// sm_100a device pass (arch 1000 WITH the family-specific feature set that the -// "a" suffix defines). Every other device pass gets an empty stub; the Python -// is_supported() / host launcher never dispatch here off sm_100, so the stub -// is unreachable at runtime. -#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) +// TORCH_CUDA_ARCH_LIST to this TU on top of the pinned data-center Blackwell +// passes. tcgen05/setmaxnreg require an architecture-conditional target, so +// plain sm_100/sm_103 and consumer sm_120 passes cannot compile this body. +// Keep it for the host pass (needed for launch plumbing) plus sm_100a/sm_103a; +// every other device pass gets an unreachable empty stub. +#if !defined(__CUDA_ARCH__) || \ + ((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \ + (__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL))) const int total_workitems = num_samples * num_heads * packed_mtiles_per_seq; @@ -1070,7 +1070,7 @@ fmha_context_bf16_gen_kernel(const __grid_constant__ CUtensorMap tmap_q, } __syncthreads(); if (warp_id == 0) tcgen05_dealloc<1>(tmem_base, TMEM_TOTAL); -#endif // host pass or sm_100a device pass (multi-arch guard; see note at the top of the body) +#endif // host pass or sm_100a/sm_103a device pass (see multi-arch note above) } } // namespace VSA_NAMESPACE diff --git a/fastvideo-kernel/csrc/attention/block_sparse_launch_sm100a.cuh b/fastvideo-kernel/csrc/attention/block_sparse_launch_sm100a.cuh index 7b30dfb411..bdb9245888 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_launch_sm100a.cuh +++ b/fastvideo-kernel/csrc/attention/block_sparse_launch_sm100a.cuh @@ -1,7 +1,7 @@ #ifndef BLOCK_SPARSE_VSA_LAUNCH_SM100A_CUH #define BLOCK_SPARSE_VSA_LAUNCH_SM100A_CUH -// Launch surface for the sm_100a VSA block-sparse FMHA forward. +// Launch surface for the sm_100a/sm_103a VSA block-sparse FMHA forward. // // Everything a caller needs: a POD argument struct, a predicate saying whether this build can // run those arguments, and one launch entry point. The benchmark in diff --git a/fastvideo-kernel/csrc/attention/block_sparse_sm100a.cu b/fastvideo-kernel/csrc/attention/block_sparse_sm100a.cu index c7bee3420a..e9ade77312 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_sm100a.cu +++ b/fastvideo-kernel/csrc/attention/block_sparse_sm100a.cu @@ -1,4 +1,4 @@ -// block_sparse_sm100a.cu -- torch binding for the sm_100a VSA block-sparse FMHA forward. +// block_sparse_sm100a.cu -- torch binding for the sm_100a/sm_103a VSA block-sparse FMHA forward. // // Forward only: returns (out, lse) so FastVideo's existing Triton backward keeps working // unchanged. lse is exactly the M tensor triton_block_sparse_attn_forward writes -- diff --git a/fastvideo-kernel/csrc/attention/primitives.cuh b/fastvideo-kernel/csrc/attention/primitives.cuh index ac0cef77fd..8075bb908d 100644 --- a/fastvideo-kernel/csrc/attention/primitives.cuh +++ b/fastvideo-kernel/csrc/attention/primitives.cuh @@ -1,4 +1,4 @@ -// primitives.cuh -- device primitives for the sm_100a VSA block-sparse attention +// primitives.cuh -- device primitives for the sm_100a/sm_103a VSA block-sparse attention // forward: tcgen05 (alloc / mma / ld / st / commit / wait / fence), TMA load / store / // tensormap, mbarrier, cluster launch control, setmaxnreg, fast math, and the FMHA helpers. // @@ -874,4 +874,3 @@ __device__ __forceinline__ void sts_f32(uint32_t smem_addr, float val) { asm volatile("st.shared.f32 [%0], %1;" :: "r"(smem_addr), "f"(val) : "memory"); } - diff --git a/fastvideo-kernel/csrc/common_extension.cpp b/fastvideo-kernel/csrc/common_extension.cpp index 33cf29654c..0243b70480 100644 --- a/fastvideo-kernel/csrc/common_extension.cpp +++ b/fastvideo-kernel/csrc/common_extension.cpp @@ -50,10 +50,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { #ifdef TK_COMPILE_BLOCK_SPARSE_VSA_SM100A m.def("block_sparse_sm100a_fwd", torch::wrap_pybind_function(block_sparse_sm100a_fwd), - "VSA block-sparse attention forward, 64-token blocks (Blackwell sm100a)"); + "VSA block-sparse attention forward, 64-token blocks (Blackwell sm100a/sm103a)"); m.def("block_sparse_sm100a_blk128_fwd", torch::wrap_pybind_function(block_sparse_sm100a_blk128_fwd), - "VSA block-sparse attention forward, 128-token blocks (Blackwell sm100a)"); + "VSA block-sparse attention forward, 128-token blocks (Blackwell sm100a/sm103a)"); #endif #ifdef TK_COMPILE_ST_ATTN diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py index 5aab846d86..34b975e4a0 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py @@ -50,9 +50,9 @@ def _force_tk() -> bool: def _force_sm100a() -> bool: - """True iff the sm_100a (Blackwell) forward is explicitly opted into. + """True iff the data-center Blackwell forward is explicitly opted into. - Opt-in only (same env the H3 backend honors): the sm_100a extension is + Opt-in only (same legacy-named env the H3 backend honors): the extension is forward-only, so this routing pairs it with the Triton backward -- its lse is already in Triton's M format. Honored only when ``block_sparse_attn_sm100a.is_supported`` passes. Unsupported 64-token @@ -370,10 +370,10 @@ def _backward_sm90(ctx, grad_o, grad_lse): block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_context_sm90) # --------------------------------------------------------------------------- -# SM100A backend custom op (index-native) +# Data-center Blackwell backend custom op (index-native; legacy sm100a API name) # -# Forward runs the sm_100a CUDA extension; backward reuses the Triton kernels. -# The sm_100a forward emits lse in exactly Triton's M format (max*log2e + +# Forward runs the sm_100a/sm_103a CUDA extension; backward reuses the Triton kernels. +# The native forward emits lse in exactly Triton's M format (max*log2e + # log2(l)), so the pairing needs no conversion. The Triton backward is # hardcoded to 64-token blocks, hence the block-size assert below. # --------------------------------------------------------------------------- @@ -428,7 +428,7 @@ def _backward_sm100a(ctx, grad_o, grad_M): block = q.shape[2] // variable_block_sizes.numel() if block != 64: raise RuntimeError( - "block_sparse_attn_sm100a backward pairs the sm_100a forward with the " + "block_sparse_attn_sm100a backward pairs the sm_100a/sm_103a forward with the " f"Triton backward, which is hardcoded to 64-token blocks; got {block}. " "Run 128-token-block metadata without grad, or use the Triton forward.") dq, dk, dv = block_sparse_attn_backward_triton(grad_o, q, k, v, o, M, q2k_idx, @@ -464,7 +464,8 @@ def block_sparse_attn_from_indices( # Backend resolution: # - FASTVIDEO_VSA_TRITON forces Triton everywhere. - # - FASTVIDEO_VSA_SM100A opts into the sm_100a forward (Triton backward). + # - FASTVIDEO_VSA_SM100A opts into the data-center Blackwell forward + # (Triton backward). The environment name is retained for compatibility. # Unsupported 64-token metadata falls through; unsupported 128-token # metadata raises because Triton cannot consume it. # - FASTVIDEO_VSA_TK requests sm_90 TK; honored only when it's actually @@ -478,9 +479,9 @@ def block_sparse_attn_from_indices( variable_block_sizes) if _infer_block_size(q, variable_block_sizes) == 128: raise NotImplementedError( - "128-token block-sparse attention requires the sm_100a forward; " + "128-token block-sparse attention requires the sm_100a/sm_103a forward; " "the Triton fallback only supports 64-token blocks, and the " - "sm_100a route is unavailable for this input.") + "native data-center Blackwell route is unavailable for this input.") use_sm90 = sm90_available elif _force_tk(): use_sm90 = sm90_available diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py index a254b6586c..e8347450a5 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py @@ -1,5 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 -"""sm_100a (Blackwell) CUDA block-sparse VSA forward. +"""Data-center Blackwell CUDA block-sparse VSA forward. + +The historical ``sm100a`` module and symbol names are retained for compatibility, but the +extension carries native sm_100a and sm_103a images and supports both device generations. A third backend behind the same VSA op as the Triton and CuTe-DSL paths. Forward only: it returns ``(out, lse)`` with ``lse`` in exactly the form ``triton_block_sparse_attn_forward`` @@ -31,7 +34,7 @@ _FWD_BY_BLOCK = {} _HAS_VSA_SM100A = False -_SM100 = (10, 0) +_SUPPORTED_COMPUTE_CAPABILITIES = {(10, 0), (10, 3)} HEAD_DIM = 128 # Must match the -DVSA_BHSD the extension was compiled with (see CMakeLists). BHSD = True @@ -52,7 +55,7 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: enough to sit on a per-layer dispatch path. What the kernel accepts (and is tested to handle): - * q/k/v: contiguous 4-D bf16 CUDA tensors on an sm_100 device, head_dim 128, laid out + * q/k/v: contiguous 4-D bf16 CUDA tensors on an sm_100/sm_103 device, head_dim 128, laid out as compiled (BHSD here); seqlen == num_blocks * block with an EVEN num_blocks (a CTA owns an adjacent pair of query blocks) and a 64- or 128-token build present. * q2k_num: any per-row counts in [0, max_kv], NON-uniform across rows included. Rows @@ -67,7 +70,7 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: """ if not _HAS_VSA_SM100A or not q.is_cuda: return False - if torch.cuda.get_device_capability(q.device) != _SM100: + if torch.cuda.get_device_capability(q.device) not in _SUPPORTED_COMPUTE_CAPABILITIES: return False if q.dtype != torch.bfloat16 or q.dim() != 4 or q.shape[-1] != HEAD_DIM: return False @@ -97,7 +100,7 @@ def _block_sparse_attn_sm100a_inference( q2k_num: torch.Tensor, variable_block_sizes: torch.Tensor, ) -> torch.Tensor: - """Opaque no-LSE launch used by the inference-only sm_100a route. + """Opaque no-LSE launch used by the inference-only sm_100a/sm_103a route. The extension is exposed as a raw pybind function rather than a dispatcher op. Calling it directly makes Dynamo descend through a Python/C++ boundary @@ -136,7 +139,7 @@ def _block_sparse_attn_sm100a_from_mask_inference( block_map: torch.Tensor, variable_block_sizes: torch.Tensor, ) -> torch.Tensor: - """Opaque mask compaction plus no-LSE sm_100a launch for inference. + """Opaque mask compaction plus no-LSE sm_100a/sm_103a launch for inference. H3 naturally produces a bool block map. Its Triton ``map_to_index`` call must live behind the same opaque boundary as the raw pybind launch; diff --git a/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py b/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py index 1c6f3e8a61..92e3118121 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py +++ b/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py @@ -14,7 +14,8 @@ import torch VSA_TILE_SIZE = (4, 4, 4) -# 128 is served by the sm_100a CUDA backend (block_sparse_attn_sm100a); 64 and 256 by +# 128 is served by the sm_100a/sm_103a CUDA backend (legacy API name +# block_sparse_attn_sm100a); 64 and 256 by # Triton and the CuTe-DSL path. A volume here only needs a backend that accepts it. _SUPPORTED_VSA_BLOCK_VOLUMES = (64, 128, 256) diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index dcfd54f3a9..e5e6f227b8 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -31,11 +31,11 @@ backward run the Triton block-sparse kernels directly (no expansion, ``FASTVIDEO_VSA_CUTEDSL`` does not apply). A third, opt-in route exists for the tile-64 FORWARD only: ``FASTVIDEO_VSA_SM100A=1`` sends no-grad -forwards through the sm_100a CUDA block-sparse kernel +forwards through the data-center Blackwell CUDA block-sparse kernel (``fastvideo_kernel.block_sparse_attn_sm100a``, upstream PR #1719 plus our per-q-tile ``q2k_num`` fix) when the extension is built, the device -is sm_100, and the geometry qualifies. The CUDA kernel assigns adjacent -pairs of query tiles to CTAs, so an odd logical tile count receives one +is sm_100 or sm_103, and the geometry qualifies. The CUDA kernel assigns +adjacent pairs of query tiles to CTAs, so an odd logical tile count receives one internal, zero-valid partner tile for the no-grad call only. Score search, the trained mask, gate-compress, and the returned packed sequence remain on the original logical tiles. Grad-tracking forwards and every backward stay @@ -61,8 +61,8 @@ map_to_index = None try: - # Optional: only present in fastvideo_kernel builds that carry the sm_100a - # CUDA block-sparse forward (upstream PR #1719). The module itself imports + # Optional: only present in fastvideo_kernel builds that carry the + # sm_100a/sm_103a CUDA block-sparse forward (upstream PR #1719). The module itself imports # fine without the compiled symbols (`_HAS_VSA_SM100A` is then False and # `is_supported` says no), so this only guards *module* availability. from fastvideo_kernel import block_sparse_attn_sm100a as _sm100a @@ -79,7 +79,7 @@ logger = init_logger(__name__) -# Opt-in switch for the sm_100a CUDA forward on the tile-64 no-grad path. +# Opt-in switch for the data-center Blackwell CUDA forward on the tile-64 no-grad path. VSA_SM100A_ENV = "FASTVIDEO_VSA_SM100A" VSA_H3_TILE_SIZE = (4, 8, 8) # 256 elements -> FA4 CuTe fastpath on sm10.x (default) @@ -411,7 +411,7 @@ def _build_block_mask( def _sm100a_unavailable_reason(sm100a_mod: Any, query_bhsd: torch.Tensor, variable_block_sizes: torch.Tensor, grad_mode: bool) -> str | None: - """Why the opt-in sm_100a forward route cannot run here, or None if it can. + """Why the opt-in data-center Blackwell route cannot run here, or None if it can. Pure decision logic, split out so the routing is unit-testable without a GPU or the compiled extension (tests substitute ``sm100a_mod``). Order @@ -420,9 +420,9 @@ def _sm100a_unavailable_reason(sm100a_mod: Any, query_bhsd: torch.Tensor, variab if sm100a_mod is None: return "fastvideo_kernel.block_sparse_attn_sm100a is not installed" if grad_mode: - return "inputs require grad and the sm_100a kernel is forward-only; grad paths keep Triton" + return "inputs require grad and the sm_100a/sm_103a kernel is forward-only; grad paths keep Triton" if not sm100a_mod.is_supported(query_bhsd, variable_block_sizes): - return ("block_sparse_attn_sm100a.is_supported returned False (needs an sm_100 device, a built " + return ("block_sparse_attn_sm100a.is_supported returned False (needs an sm_100 or sm_103 device, a built " "extension, bf16, head_dim 128, an even tile count, and integer tile sizes)") return None @@ -703,7 +703,7 @@ def forward( # type: ignore[override] # capture. Logging from this branch would itself break a # ``fullgraph=True`` forward. if not compiling: - logger.info_once("MiniMax-H3 VSA tile-64 forward: using the sm100a CUDA block-sparse kernel") + logger.info_once("MiniMax-H3 VSA tile-64 forward: using the sm100a/sm103a CUDA block-sparse kernel") if regional_compiling: # The compile-safe wrapper keeps both Triton mask # compaction and the raw pybind launch behind one diff --git a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py index 41cf6ebf62..2deaed9f47 100644 --- a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py +++ b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Checks for the VSA-H3 tile-64 sm_100a route and odd-tile transport. +"""Checks for the VSA-H3 tile-64 sm_100a/sm_103a route and odd-tile transport. The opt-in third kernel route (``FASTVIDEO_VSA_SM100A=1``) must (a) stay off by default, (b) engage only when the extension is present, the device qualifies, @@ -397,7 +397,7 @@ def test_env_on_routes_sm100a_with_index_metadata(routed, monkeypatch): assert call["q2k_idx"].shape[-1] == n_tiles and call["q2k_idx"].dtype == torch.int32 assert call["vbs"].dtype == torch.int32 assert call["need_lse"] is False - assert messages == ["MiniMax-H3 VSA tile-64 forward: using the sm100a CUDA block-sparse kernel"] + assert messages == ["MiniMax-H3 VSA tile-64 forward: using the sm100a/sm103a CUDA block-sparse kernel"] # BHSD kernel result comes back in the backend's BSHD layout assert out.shape == (1, n_tiles * 64, _HEADS, _DIM) @@ -420,7 +420,8 @@ def capture_log(level, msg, *args, **kwargs): assert fake_triton.calls == 0 assert len(fake_sm.calls) == 2 - assert records == [(logging.INFO, "MiniMax-H3 VSA tile-64 forward: using the sm100a CUDA block-sparse kernel", (), + assert records == [(logging.INFO, + "MiniMax-H3 VSA tile-64 forward: using the sm100a/sm103a CUDA block-sparse kernel", (), {"stacklevel": 2})] @@ -605,11 +606,11 @@ def test_forward_rejects_non_contract_transport_shapes(routed): def test_real_sm100a_odd_route_matches_triton_oracle(monkeypatch): - """Exercise the padded route through the actual GB200 extension.""" - if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): - pytest.skip("requires a GB200 (sm_100a) compute node") + """Exercise the padded route through the actual data-center Blackwell extension.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}: + pytest.skip("requires a GB200/B300 (sm_100a/sm_103a) compute node") if vsa_h3._sm100a is None or not vsa_h3._sm100a._HAS_VSA_SM100A: - pytest.skip("requires a fastvideo_kernel build containing sm_100a VSA") + pytest.skip("requires a fastvideo_kernel build containing data-center Blackwell VSA") device = torch.device("cuda") meta = _build_meta(device=device, sparsity=0.5) @@ -643,10 +644,10 @@ def reject_triton(*args, **kwargs): def test_real_sm100a_odd_preprocess_forward_postprocess_fullgraph(monkeypatch): """Capture #1745's odd transport buffer mutation with real Inductor.""" - if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): - pytest.skip("requires a GB200 (sm_100a) compute node") + if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}: + pytest.skip("requires a GB200/B300 (sm_100a/sm_103a) compute node") if vsa_h3._sm100a is None or not vsa_h3._sm100a._HAS_VSA_SM100A: - pytest.skip("requires a fastvideo_kernel build containing sm_100a VSA") + pytest.skip("requires a fastvideo_kernel build containing data-center Blackwell VSA") device = torch.device("cuda") monkeypatch.setenv(VSA_SM100A_ENV, "1") diff --git a/tests/bench_block_sparse_sm100a.py b/tests/bench_block_sparse_sm100a.py index 6ced276dc2..8e75f1b15a 100644 --- a/tests/bench_block_sparse_sm100a.py +++ b/tests/bench_block_sparse_sm100a.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""In-framework benchmark of the sm_100a VSA backend at real Wan shapes. +"""In-framework benchmark of the sm_100a/sm_103a VSA backend at real Wan shapes. Not the powers-of-two bench grid: this builds the metadata FastVideo actually builds, at the deployed latent shapes, and drives the same block_sparse_attn_from_indices entry point the @@ -39,10 +39,11 @@ def timed(fn, iters=20, warmup=5): def main(): - if torch.cuda.get_device_capability() != (10, 0): - print("not Blackwell; skipping") + if torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}: + print("not data-center Blackwell; skipping") return - print(f"{'case':<8} {'S_pad':>7} {'blocks':>7} {'topk':>5} {'sm100a ms':>10} " + arch = "sm103a" if torch.cuda.get_device_capability() == (10, 3) else "sm100a" + print(f"{'case':<8} {'S_pad':>7} {'blocks':>7} {'topk':>5} {f'{arch} ms':>10} " f"{'triton ms':>10} {'speedup':>8} selected") for label, latent, tile, heads in CASES: meta = build_vsa_metadata(latent, tile_size=tile, device="cuda") @@ -76,7 +77,7 @@ def main(): tri = timed(lambda: block_sparse_attn_triton(q, k, v, i64, n64, vbs64)) print(f"{label:<8} {S:>7} {nb:>7} {topk:>5} {ours:>10.3f} {tri:>10.3f} " - f"{tri / ours:>7.2f}x {'sm100a' if ok else 'FALLBACK'}") + f"{tri / ours:>7.2f}x {arch if ok else 'FALLBACK'}") if __name__ == "__main__": diff --git a/tests/test_block_sparse_sm100a.py b/tests/test_block_sparse_sm100a.py index 479a59affe..b449a2da8c 100644 --- a/tests/test_block_sparse_sm100a.py +++ b/tests/test_block_sparse_sm100a.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Correctness tests for the sm_100a CUDA block-sparse VSA forward. +"""Correctness tests for the sm_100a/sm_103a CUDA block-sparse VSA forward. Compared against an explicit PyTorch reference rather than the Triton kernel: Triton's block-sparse forward is hardcoded to 64-token blocks (BLOCK_M = BLOCK_N = 64) while this @@ -24,9 +24,9 @@ BLOCK_SIZES = [64, 128] pytestmark = pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0) + not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)} or not vsa._HAS_VSA_SM100A, - reason="requires Blackwell (sm_100a) and a built fastvideo_kernel extension", + reason="requires data-center Blackwell (sm_100a/sm_103a) and a built fastvideo_kernel extension", ) diff --git a/tests/test_block_sparse_sm100a_dispatch.py b/tests/test_block_sparse_sm100a_dispatch.py index 75b7723ee7..5697b488d2 100644 --- a/tests/test_block_sparse_sm100a_dispatch.py +++ b/tests/test_block_sparse_sm100a_dispatch.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Routing tests for the opt-in sm_100a dispatch in ``block_sparse_attn_from_indices``. +"""Routing tests for the opt-in sm_100a/sm_103a dispatch in ``block_sparse_attn_from_indices``. ``FASTVIDEO_VSA_SM100A=1`` routes the forward to the sm_100a extension when ``block_sparse_attn_sm100a.is_supported`` passes, pairing it with the Triton @@ -20,9 +20,9 @@ ENV = "FASTVIDEO_VSA_SM100A" pytestmark = pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0) + not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)} or not vsa._HAS_VSA_SM100A, - reason="requires Blackwell (sm_100a) and a built fastvideo_kernel extension", + reason="requires data-center Blackwell (sm_100a/sm_103a) and a built fastvideo_kernel extension", ) diff --git a/tests/test_fasth3_packaging.py b/tests/test_fasth3_packaging.py index 6162ffbc87..c3d2ccdc4c 100644 --- a/tests/test_fasth3_packaging.py +++ b/tests/test_fasth3_packaging.py @@ -30,9 +30,12 @@ def test_fasth3_extra_and_root_kernel_pin_match_source_release(): } in kernel_sources -def test_kernel_release_matrix_can_publish_sm100a_wheels(): +def test_kernel_release_matrix_can_publish_data_center_blackwell_wheels(): workflow = (REPO_ROOT / ".github" / "workflows" / "publish-kernel.yml").read_text(encoding="utf-8") + cmake = (REPO_ROOT / "fastvideo-kernel" / "CMakeLists.txt").read_text(encoding="utf-8") assert 'TORCH_CUDA_ARCH_LIST="9.0a;10.0a;12.0a"' in workflow assert 'TORCH_CUDA_ARCH_LIST="10.0a;12.0a"' in workflow + assert "arch=compute_100a,code=sm_100a" in cmake + assert "arch=compute_103a,code=sm_103a" in cmake assert "patchelf==0.17.2.4" in workflow