[kernel] VSA kernel: one sm_100a / sm_103a image per listed arch; un-gate backward on sm_103a

.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

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)

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<DQ_DTYPE>;
   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 <bool BHSD = false, typename DQ_DTYPE = float>
 __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;

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

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 @@
     _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

tests/test_block_sparse_bwd_sm100a.py

@@ -31,9 +31,9 @@
 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",
 )
 

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