[feat] Add native SM103a VSA support (#1812)

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

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

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 @@ except ImportError:  # pragma: no cover - extension not built
     _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;

fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py

@@ -14,7 +14,8 @@ import math
 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)