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

Co-authored-by: root <root@ptyche0124.ptyche.clusters.nvidia.com>
2026-10-17 21:10:19 -07:00
parent 377c79fa7e
commit d4ed61cc7c
75 changed files with 613 additions and 586 deletions
@@ -95,7 +95,7 @@ python examples/inference/basic/basic_fasth3.py --prompt "your prompt"
 ```
 The default checkpoint, [FastH3 Preview v0.2](https://huggingface.co/FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2), is public on the Hub under the MiniMax H3 Community License. Review its model card and license before use or redistribution.

-The default `all` profile is the fastest measured four-GPU Preview recipe on GB200. It selects VSA sparsity 0.9 with 64-token tiles and the sm_100a sparse kernel, enables FA4 for eligible non-VSA paths, regionally compiles and replicates the sparse DiT, compiles and temporally parallelizes the video VAE with the `gather` strategy, and pins CPU-offloaded component memory. It also pins the benchmark protocol: five sigma-grid points (exactly four DiT forwards), one excluded seed-999 warmup, then three timed seed-1000 requests with distinct output paths.
+The default `all` profile is the fastest measured four-GPU Preview recipe on GB200. It selects VSA sparsity 0.9 with 64-token tiles and the ptx_blackwell sparse kernel, enables FA4 for eligible non-VSA paths, regionally compiles and replicates the sparse DiT, compiles and temporally parallelizes the video VAE with the `gather` strategy, and pins CPU-offloaded component memory. It also pins the benchmark protocol: five sigma-grid points (exactly four DiT forwards), one excluded seed-999 warmup, then three timed seed-1000 requests with distinct output paths.

 The equivalent explicit command is:

@@ -107,7 +107,7 @@ python examples/inference/basic/basic_fasth3.py \
  --steps 5 \
  --vsa-sparsity 0.9 \
  --vsa-tile-size 64 \
-  --vsa-kernel sm100a \
+  --vsa-kernel ptx_blackwell \
  --compile-vae \
  --parallel-vae \
  --replicated-dit \
@@ -143,13 +143,13 @@ Pass `--no-inference-torch-compile` to recover the eager sparse-DiT route.

 #### Hopper (H100 80 GB)

-Hopper has no FA4 and no sm_100a VSA kernel, so the GB200 profile above does
+Hopper has no FA4 and no ptx_blackwell VSA kernel, so the GB200 profile above does
 not apply there. The measured Hopper route keeps the same checkpoint, four
 forwards and 90% VSA sparsity and changes three things: the 64-token VSA path
 runs on the ThunderKittens sm_90a kernel, which fastvideo-kernel picks on sm_90
 whenever it is built with
 `cd fastvideo-kernel && CMAKE_ARGS='-DFASTVIDEO_KERNEL_BUILD_TK=ON' ./build.sh`
-(`--vsa-kernel tk` keeps the sm100a route off and fails early when that
+(`--vsa-kernel tk` keeps the ptx_blackwell route off and fails early when that
 extension is missing),
 the DiT is FSDP-sharded across the GPUs instead of replicated
 (`--no-replicated-dit`; a replicated 66 GB BF16 DiT does not fit 80 GB), and
@@ -216,7 +216,7 @@ runs a Parallel Decoding Distillation Ref2VA student in eight transformer
 forwards. The export carries only its `transformer_ref`, scheduler configs,
 and `fastvideo_inference.json`; the script links them with the base
 MiniMax-H3 components into one local model directory. Its 128-token VSA tiles
-need the sm_100a/sm_103a kernel:
+need the ptx_blackwell kernel:

 ```bash
 python examples/inference/basic/basic_fasth3_omniref_pdd.py \
@@ -1,4 +1,4 @@
-#   FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 \
+#   FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \
 #     FLASHINFER_CUDA_ARCH_LIST=12.0a FASTVIDEO_STAGE_LOGGING=1 \
 #     fastvideo generate --config examples/inference/basic/basic_compacth3_rtx5090.yaml
 generator:
@@ -1,4 +1,4 @@
-#   FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 \
+#   FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \
 #     FLASHINFER_CUDA_ARCH_LIST=12.0a FASTVIDEO_STAGE_LOGGING=1 \
 #     fastvideo generate --config examples/inference/basic/basic_compacth3_rtx_pro6000.yaml
 generator:
@@ -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():
@@ -19,7 +19,7 @@
 # request.sampling below is the measured 4xH100 shape (960x544, 345 frames =
 # 14.375 s at 24 FPS). Legal H3 frame counts are 17n+5, max 345.
 #
-#   FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 FASTVIDEO_VSA_TK=1 \
+#   FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 FASTVIDEO_VSA_TK=1 \
 #     FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING=1 \
 #     fastvideo generate --config examples/inference/basic/basic_fasth3_h100.yaml
 generator:
@@ -4,7 +4,7 @@
 This is the LoRA counterpart of ``basic_fasth3.py``. Both routes share the
 same performance profile: four DiT forwards, regional fullgraph DiT compile,
 H3 fusions, compiled and sequence-parallel video VAE decode, replicated DiT,
-pinned CPU offload, FA4, and the sm100a tile-64 kernel for VSA adapters.
+pinned CPU offload, FA4, and the ptx_blackwell tile-64 kernel for VSA adapters.

 The FastH3 adapters include low-rank factors plus exact dense deltas. Some also
 provide the VSA compression gate that is absent from the base checkpoint. Pass
@@ -33,7 +33,7 @@ from fastvideo.api import (  # noqa: E402
 os.environ.update({
    "FASTVIDEO_FA4": "1",
    "FASTVIDEO_MINIMAX_H3_FUSIONS": "all",
-    "FASTVIDEO_VSA_SM100A": "0",
+    "FASTVIDEO_VSA_PTX_BLACKWELL": "0",
 })
 os.environ.pop("FASTVIDEO_INFERENCE_TORCH_COMPILE", None)

@@ -15,7 +15,7 @@ export FASTVIDEO_FA4=1
 export FASTVIDEO_INFERENCE_TORCH_COMPILE=0
 export FASTVIDEO_MINIMAX_H3_FUSIONS=all
 export FASTVIDEO_NVTX_PROFILE=1
-export FASTVIDEO_VSA_SM100A=0
+export FASTVIDEO_VSA_PTX_BLACKWELL=0
 export PYTHONUNBUFFERED=1
 export PYTHONPATH="${WORKTREE_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"

@@ -1,6 +1,6 @@
 # FastH3 on one DGX Spark (one GB10). Install: docs/getting_started/installation/spark.md
 #
-# GB10 has no FA4 / sm_100a VSA kernel. Keep the env prefix below, or pass
+# GB10 has no FA4 / ptx_blackwell VSA kernel. Keep the env prefix below, or pass
 # equivalent CLI flags on examples/inference/basic/basic_fasth3.py
 # (`--vsa-kernel triton --no-fa4 --num-gpus 1`).
 #
@@ -9,7 +9,7 @@
 # counts are 17n+5, max 345 (15 s). Native 16:9 sizes include 832x480 and
 # 1344x768. 1024x576 is a common recipe-matched 16:9.
 #
-#   FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 \
+#   FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 \
 #     FASTVIDEO_STAGE_LOGGING=1 \
 #     fastvideo generate --config examples/inference/basic/basic_fasth3_spark.yaml
 generator:
@@ -8,7 +8,7 @@
 #   source examples/inference/optimizations/spark_pair_env.sh
 #   export RAY_ADDRESS=<qsfp-head>:6379
 #   export FASTVIDEO_HOST_IP=<qsfp-head>
-#   FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 \
+#   FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 \
 #     FASTVIDEO_VAE_PARALLEL_DECODE=1 FASTVIDEO_STAGE_LOGGING=1 \
 #     fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pair.yaml
 generator:
@@ -8,7 +8,7 @@
 #
 # FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cutlass \
 # FASTVIDEO_H3_VAE_TILE_BATCH=1 \
-# FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 \
+# FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \
 # FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING=1 \
 # fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pair_pruned_nvfp4.yaml
 generator:
@@ -8,7 +8,7 @@
 #
 # FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cutlass \
 # FASTVIDEO_H3_VAE_TILE_BATCH=1 \
-# FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 \
+# FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \
 # FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING=1 \
 # fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pair_v2_nvfp4.yaml
 generator:
@@ -7,7 +7,7 @@
 #
 # FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cutlass \
 # FASTVIDEO_H3_VAE_TILE_BATCH=1 \
-# FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 \
+# FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \
 # FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING=1 \
 # fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pruned_nvfp4.yaml
 generator:
@@ -7,7 +7,7 @@
 #
 # FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cutlass \
 # FASTVIDEO_H3_VAE_TILE_BATCH=1 \
-# FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 \
+# FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \
 # FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING=1 \
 # fastvideo generate --config examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml
 generator: