[refactor] VSA: rename the sm100a block-sparse kernel to ptx_blackwell (#1956)

Co-authored-by: root <root@ptyche0124.ptyche.clusters.nvidia.com>

examples/inference/basic/basic_fasth3.py

@@ -20,6 +20,7 @@ import importlib.util
 import os
 import statistics
 import time
+import warnings
 from collections.abc import Sequence
 from pathlib import Path
 
@@ -94,11 +95,11 @@ def build_parser(description: str | None = None) -> argparse.ArgumentParser:
                         default=64,
                         help="VSA-H3 tile size; 64 is the checkpoint's trained and measured geometry")
     parser.add_argument("--vsa-kernel",
-                        choices=("triton", "sm100a", "tk"),
-                        default="sm100a",
-                        help="tile-64 sparse kernel; sm100a is the measured GB200 route and requires a compatible "
-                        "fastvideo-kernel build; tk keeps the sm100a route off and requires the ThunderKittens "
-                        "sm_90a extension, which fastvideo-kernel already picks on sm_90 whenever it is built")
+                        choices=("triton", "ptx_blackwell", "tk", "sm100a"),
+                        default="ptx_blackwell",
+                        help="tile-64 sparse kernel; ptx_blackwell requires a compatible fastvideo-kernel build; "
+                        "tk keeps the ptx_blackwell route off and requires the ThunderKittens sm_90a extension, "
+                        "which fastvideo-kernel already picks on sm_90 whenever it is built")
     parser.add_argument("--fa4",
                         action=argparse.BooleanOptionalAction,
                         default=True,
@@ -174,6 +175,9 @@ def validate_args(parser: argparse.ArgumentParser, args: argparse.Namespace) ->
         parser.error("--vsa-sparsity must be in [0, 1)")
     if args.compile_mode is not None and args.inference_torch_compile:
         parser.error("--compile-mode cannot be combined with regional compile; pass --no-inference-torch-compile")
+    if args.vsa_kernel == "sm100a":
+        warnings.warn("--vsa-kernel sm100a is deprecated; use --vsa-kernel ptx_blackwell", FutureWarning, stacklevel=2)
+        args.vsa_kernel = "ptx_blackwell"
     return args
 
 
@@ -203,7 +207,7 @@ def profile_environment(args: argparse.Namespace) -> dict[str, str | None]:
     use_vsa = _uses_vsa(args)
     return {
         "FASTVIDEO_ATTENTION_BACKEND": "VIDEO_SPARSE_ATTN_H3" if use_vsa else "FLASH_ATTN",
-        "FASTVIDEO_VSA_SM100A": "1" if use_vsa and args.vsa_kernel == "sm100a" else "0",
+        "FASTVIDEO_VSA_PTX_BLACKWELL": "1" if use_vsa and args.vsa_kernel == "ptx_blackwell" else "0",
         "FASTVIDEO_VSA_TK": "1" if use_vsa and args.vsa_kernel == "tk" else "0",
         "FASTVIDEO_VSA_CUTEDSL": "0",
         # A non-empty output path enables the diagnostic probe.
@@ -239,12 +243,12 @@ def _fa4_is_installed() -> bool:
         return False
 
 
-def _sm100a_kernel_is_installed() -> bool:
+def _ptx_blackwell_kernel_is_installed() -> bool:
     try:
-        from fastvideo_kernel import block_sparse_attn_sm100a
+        from fastvideo_kernel import block_sparse_attn_ptx_blackwell
     except ImportError:
         return False
-    return bool(getattr(block_sparse_attn_sm100a, "_HAS_VSA_SM100A", False))
+    return bool(getattr(block_sparse_attn_ptx_blackwell, "_HAS_VSA_PTX_BLACKWELL", False))
 
 
 def _sm90_kernel_is_installed() -> bool:
@@ -262,9 +266,9 @@ def validate_profile_dependencies(args: argparse.Namespace) -> None:
         raise RuntimeError(
             "FastH3's FA4 profile requires the pinned flash-attn-4 package. Install it with "
             "`UV_TORCH_BACKEND=cu130 uv pip install -e \".[fasth3]\"`, or pass --no-fa4.")
-    if _uses_vsa(args) and args.vsa_kernel == "sm100a" and not _sm100a_kernel_is_installed():
+    if _uses_vsa(args) and args.vsa_kernel == "ptx_blackwell" and not _ptx_blackwell_kernel_is_installed():
         raise RuntimeError(
-            "FastH3's sm100a profile requires fastvideo-kernel 0.3.4 built with the Blackwell VSA extension. "
+            "FastH3's ptx_blackwell profile requires fastvideo-kernel 0.3.4 built with the Blackwell VSA extension. "
             "Install this checkout with `UV_TORCH_BACKEND=cu130 uv pip install -e \".[fasth3]\"` (or run "
             "`cd fastvideo-kernel && ./build.sh`), or pass --vsa-kernel triton.")
     if _uses_vsa(args) and args.vsa_kernel == "tk" and not _sm90_kernel_is_installed():