feat(vsa): add native sm103a support

Signed-off-by: lishunyang12 <lishunyang12@163.com>

Author: lishunyang12

.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

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)

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.

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

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

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

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

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");
 }
-

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

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