From aef4e9b3b1423be9d5e39ca463fb085b85c221f6 Mon Sep 17 00:00:00 2001 From: lpc0220 Date: Sat, 12 Sep 2026 16:49:36 -0700 Subject: [PATCH] [kernel] VSA kernel: one sm_100a / sm_103a image per listed arch; un-gate backward on sm_103a (#1833) Co-authored-by: Claude Fable 5.1 --- .github/workflows/publish-kernel.yml | 4 +-- fastvideo-kernel/CMakeLists.txt | 36 +++++++++++++------ .../block_sparse_bwd_kernel_sm100a.cuh | 16 ++++++--- fastvideo-kernel/csrc/common_extension.cpp | 2 +- .../block_sparse_attn_bwd_sm100a.py | 6 ++-- tests/test_block_sparse_bwd_sm100a.py | 4 +-- tests/test_block_sparse_sm100a_dispatch.py | 5 +-- 7 files changed, 48 insertions(+), 25 deletions(-) diff --git a/.github/workflows/publish-kernel.yml b/.github/workflows/publish-kernel.yml index 5e89ecf2c..6c4fea1ca 100644 --- a/.github/workflows/publish-kernel.yml +++ b/.github/workflows/publish-kernel.yml @@ -176,11 +176,11 @@ jobs: # the main extension for the full arch list. CMAKE_BUILD_PARALLEL_LEVEL caps # Ninja so heavy CUTLASS/TK template TUs don't OOM the 16 GB runner (exit 143). if [ "${{ matrix.platform.arch }}" = "aarch64" ]; then - export TORCH_CUDA_ARCH_LIST="10.0a;12.0a" + export TORCH_CUDA_ARCH_LIST="10.0a;10.3a;12.0a" export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=OFF -DFASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON" export CMAKE_BUILD_PARALLEL_LEVEL=1 elif [ "${{ matrix.torch-cuda.torch-cuda-short }}" = "cu130" ]; then - export TORCH_CUDA_ARCH_LIST="9.0a;10.0a;12.0a" + export TORCH_CUDA_ARCH_LIST="9.0a;10.0a;10.3a;12.0a" export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DFASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON -DCMAKE_CUDA_ARCHITECTURES=90a" # A single FP4 TU (attn_qat_infer) can use ~8-12 GB on its own, so serialize. export CMAKE_BUILD_PARALLEL_LEVEL=1 diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index a7f968242..405d7b751 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -396,27 +396,41 @@ 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_SM100_FAMILY OFF) + set(ENABLE_VSA_SM100A OFF) + set(ENABLE_VSA_SM103A 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\\.(0|3)a|10(0|3)a|sm_10(0|3)a)([; ,]|$)") - set(ENABLE_VSA_SM100_FAMILY ON) + if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(10\\.0a|100a|sm_100a)([; ,]|$)") + set(ENABLE_VSA_SM100A ON) endif() - if(ENABLE_VSA_SM100_FAMILY) - message(STATUS "fastvideo-kernel: building block_sparse_sm100a (sm_100a/sm_103a, 64- and 128-token blocks)") + if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(10\\.3a|103a|sm_103a)([; ,]|$)") + set(ENABLE_VSA_SM103A ON) + endif() + set(_VSA_FWD_GENCODE "") + set(_VSA_FWD_ARCHS "") + if(ENABLE_VSA_SM100A) + list(APPEND _VSA_FWD_GENCODE "-gencode;arch=compute_100a,code=sm_100a") + list(APPEND _VSA_FWD_ARCHS "sm_100a") + endif() + if(ENABLE_VSA_SM103A) + list(APPEND _VSA_FWD_GENCODE "-gencode;arch=compute_103a,code=sm_103a") + list(APPEND _VSA_FWD_ARCHS "sm_103a") + endif() + if(_VSA_FWD_GENCODE) + message(STATUS "fastvideo-kernel: building block_sparse_sm100a (${_VSA_FWD_ARCHS}; 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;-gencode;arch=compute_103a,code=sm_103a;-DVSA_BHSD=true") - # VSA block-sparse attention BACKWARD, 64-token blocks. sm_100a only for now: validated - # on GB200, not yet on B300/GB300, so no sm_103a image is built and the Python side keeps - # the Triton backward for sm_103a devices. + COMPILE_OPTIONS "${_VSA_FWD_GENCODE};-DVSA_BHSD=true") + endif() + # VSA block-sparse attention BACKWARD, 64-token blocks: the same images as the forward. + if(_VSA_FWD_GENCODE) list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_bwd_sm100a.cu) set_source_files_properties(csrc/attention/block_sparse_bwd_sm100a.cu PROPERTIES - COMPILE_OPTIONS "-gencode;arch=compute_100a,code=sm_100a;-DVSA_BHSD=true") + COMPILE_OPTIONS "${_VSA_FWD_GENCODE};-DVSA_BHSD=true") endif() Python_add_library(fastvideo_kernel_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI @@ -434,7 +448,7 @@ if(BUILD_CXX_KERNELS) # Build compile definitions list set(COMPILE_DEFS TORCH_EXTENSION_NAME=fastvideo_kernel_ops) - if(ENABLE_VSA_SM100_FAMILY) + if(ENABLE_VSA_SM100A OR ENABLE_VSA_SM103A) list(APPEND COMPILE_DEFS TK_COMPILE_BLOCK_SPARSE_VSA_SM100A) endif() if(ENABLE_TK_KERNELS) diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh b/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh index 86d738cac..3d3014ff0 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh @@ -123,7 +123,9 @@ __global__ void __cluster_dims__(1, 1, 1) __launch_bounds__(N_WARPS * 32, 1) vsa const int* __restrict__ k2q_num, const int* __restrict__ workitem_remap, const int* __restrict__ variable_block_sizes, int max_q_blocks, int num_samples, int num_heads, int seqlen, float scale_log2, float sm_scale) { -#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) +#if !defined(__CUDA_ARCH__) || \ + ((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \ + (__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL))) using DQ = DQConfig; const int num_kv_blocks_per_seq = seqlen / BLOCK; @@ -863,7 +865,9 @@ __global__ void __launch_bounds__(ORDER_THREADS, 1) vsa_bwd_order_kernel(const int* __restrict__ k2q_idx, const int* __restrict__ k2q_num, int max_q_blocks, int num_kv_blocks_per_seq, int order_bin, bool snake, int* __restrict__ order_out) { -#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) +#if !defined(__CUDA_ARCH__) || \ + ((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \ + (__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL))) extern __shared__ int order_smem[]; int* sbin = order_smem; int* smid = order_smem + num_kv_blocks_per_seq; @@ -908,7 +912,9 @@ __global__ void __launch_bounds__(256, 1) __nv_bfloat16* __restrict__ dk, __nv_bfloat16* __restrict__ dv, const int* __restrict__ k2q_num, int num_samples, int num_heads, int seqlen) { -#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) +#if !defined(__CUDA_ARCH__) || \ + ((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \ + (__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL))) __shared__ __align__(128) __nv_bfloat16 tile[PRE_TOKENS][SUB_COLS_BF16 + 4]; const int num_kv_blocks_per_seq = seqlen / BLOCK; const int token_block_id = (int)blockIdx.x; @@ -1038,7 +1044,9 @@ template __global__ void __launch_bounds__(128, 1) vsa_bwd_postprocess_kernel(const DQ_DTYPE* __restrict__ dqaccum, __nv_bfloat16* __restrict__ dq, int num_heads, int seqlen, float sm_scale) { -#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) +#if !defined(__CUDA_ARCH__) || \ + ((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \ + (__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL))) const int q_block_id = (int)blockIdx.x; const int batch_head = (int)blockIdx.y; const int batch = batch_head / num_heads, head = batch_head % num_heads; diff --git a/fastvideo-kernel/csrc/common_extension.cpp b/fastvideo-kernel/csrc/common_extension.cpp index f58a39711..45d5744bb 100644 --- a/fastvideo-kernel/csrc/common_extension.cpp +++ b/fastvideo-kernel/csrc/common_extension.cpp @@ -60,7 +60,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "VSA block-sparse attention forward, 128-token blocks (Blackwell sm100a/sm103a)"); m.def("block_sparse_sm100a_bwd", torch::wrap_pybind_function(block_sparse_sm100a_bwd), - "VSA block-sparse attention backward, 64-token blocks (Blackwell sm100a)"); + "VSA block-sparse attention backward, 64-token blocks (Blackwell sm100a/sm103a)"); #endif #ifdef TK_COMPILE_ST_ATTN diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py index 76d338cd6..ea073d868 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""sm_100a (Blackwell) CUDA block-sparse VSA backward. +"""sm_100a/sm_103a (data-center Blackwell) CUDA block-sparse VSA backward. Companion of ``block_sparse_attn_sm100a`` (the forward): consumes the forward's ``lse`` in the Triton "M format" (``max(qk * sm_scale * log2e) + log2(l)``, ``[B, H, S]`` fp32) unchanged and @@ -23,7 +23,7 @@ except ImportError: # pragma: no cover - extension not built _BWD = None _HAS_VSA_BWD_SM100A = False -_SM100 = (10, 0) +_SUPPORTED_COMPUTE_CAPABILITIES = {(10, 0), (10, 3)} HEAD_DIM = 128 BLOCK = 64 # Must match the -DVSA_BHSD the extension was compiled with (FastVideo builds with true). @@ -58,7 +58,7 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: """ if not _HAS_VSA_BWD_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 diff --git a/tests/test_block_sparse_bwd_sm100a.py b/tests/test_block_sparse_bwd_sm100a.py index a2a7c1351..b9daea240 100644 --- a/tests/test_block_sparse_bwd_sm100a.py +++ b/tests/test_block_sparse_bwd_sm100a.py @@ -31,9 +31,9 @@ REL_MAX_TOL = 1e-2 # max|got - ref| / max|ref| MEAN_ABS_TOL = 1e-3 # mean|got - ref| 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 bwd._HAS_VSA_BWD_SM100A, - reason="requires a compute-capability (10, 0) GPU (sm_100a) and a fastvideo_kernel " + reason="requires a data-center Blackwell GPU (sm_100a/sm_103a) and a fastvideo_kernel " "extension built with block_sparse_sm100a_bwd", ) diff --git a/tests/test_block_sparse_sm100a_dispatch.py b/tests/test_block_sparse_sm100a_dispatch.py index 9bfdc719c..122aec2ea 100644 --- a/tests/test_block_sparse_sm100a_dispatch.py +++ b/tests/test_block_sparse_sm100a_dispatch.py @@ -175,8 +175,9 @@ def test_backward_uses_sm100a_kernel_when_built(monkeypatch): """Guards against a silent Triton fallback: with the op built on an sm_100a device the CUDA backward must be the one that runs.""" from fastvideo_kernel import block_sparse_attn_bwd_sm100a as vsa_bwd - if not vsa_bwd._HAS_VSA_BWD_SM100A or torch.cuda.get_device_capability() != (10, 0): - pytest.skip("sm_100a backward not built for this device") + if (not vsa_bwd._HAS_VSA_BWD_SM100A + or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}): + pytest.skip("sm_100a/sm_103a backward not built for this device") _, _, sm100a_backward = _grads_sm100a_route_vs_triton(monkeypatch) assert sm100a_backward