From d4ed61cc7c81314398190bf382148b6c2c4bcb2d Mon Sep 17 00:00:00 2001 From: lpc0220 Date: Sat, 17 Oct 2026 21:10:19 -0700 Subject: [PATCH] [refactor] VSA: rename the sm100a block-sparse kernel to ptx_blackwell (#1956) Co-authored-by: root --- .../dreamverse/minimax_h3_generation.py | 2 +- .../tests/test_minimax_h3_generation.py | 4 +- apps/infinite_livestream/README.md | 2 +- .../infinite_livestream/backend.py | 20 +- .../configs/infinite_livestream.yaml | 4 +- docs/assets/cookbook-recipes.json | 10 +- docs/contributing/env_vars.md | 2 +- docs/cookbook/minimax-h3.md | 4 +- docs/cookbook/openai-api.md | 4 +- docs/cookbook/recipes/fasth3-trim.yaml | 2 +- .../installation/spark_pair.md | 12 +- .../installation/spark_performance.md | 2 +- docs/inference/fasth3-distilled.md | 6 +- docs/inference/optimizations.md | 8 +- examples/inference/basic/README.md | 10 +- .../basic/basic_compacth3_rtx5090.yaml | 2 +- .../basic/basic_compacth3_rtx_pro6000.yaml | 2 +- examples/inference/basic/basic_fasth3.py | 26 ++- .../inference/basic/basic_fasth3_h100.yaml | 2 +- .../basic/basic_fasth3_lora_preview.py | 2 +- .../basic/basic_fasth3_simplified.py | 2 +- .../basic/basic_fasth3_simplified_profile.sh | 2 +- .../inference/basic/basic_fasth3_spark.yaml | 4 +- .../basic/basic_fasth3_spark_pair.yaml | 2 +- .../basic_fasth3_spark_pair_pruned_nvfp4.yaml | 2 +- .../basic_fasth3_spark_pair_v2_nvfp4.yaml | 2 +- .../basic_fasth3_spark_pruned_nvfp4.yaml | 2 +- .../basic/basic_fasth3_spark_v2_nvfp4.yaml | 2 +- .../lora/minimax_h3_lora_inference.py | 8 +- .../serving/openai_compacth3_rtx5090.yaml | 2 +- .../serving/openai_compacth3_rtx_pro6000.yaml | 2 +- examples/serving/openai_fasth3_spark.yaml | 4 +- fastvideo-kernel/CMakeLists.txt | 20 +- fastvideo-kernel/README.md | 2 +- ...se_sm100a.cu => block_sparse_blackwell.cu} | 20 +- .../block_sparse_blk128_blackwell.cu | 7 + .../attention/block_sparse_blk128_sm100a.cu | 7 - ...m100a.cu => block_sparse_bwd_blackwell.cu} | 24 +-- .../block_sparse_bwd_blk128_blackwell.cu | 7 + .../block_sparse_bwd_blk128_sm100a.cu | 7 - ... => block_sparse_bwd_kernel_blackwell.cuh} | 8 +- ...ck_sparse_bwd_kernel_blk128_blackwell.cuh} | 8 +- ... => block_sparse_bwd_launch_blackwell.cuh} | 12 +- ...ck_sparse_bwd_launch_blk128_blackwell.cuh} | 12 +- ....cuh => block_sparse_kernel_blackwell.cuh} | 8 +- ....cuh => block_sparse_launch_blackwell.cuh} | 14 +- fastvideo-kernel/csrc/common_extension.cpp | 28 +-- .../fastvideo_kernel/block_sparse_attn.py | 94 ++++---- ...=> block_sparse_attn_bwd_ptx_blackwell.py} | 40 ++-- ....py => block_sparse_attn_ptx_blackwell.py} | 50 ++--- .../python/fastvideo_kernel/vsa_utils.py | 4 +- ....py => test_block_sparse_blackwell_clc.py} | 12 +- ...> test_block_sparse_blackwell_dispatch.py} | 90 ++++---- ....py => test_block_sparse_bwd_blackwell.py} | 22 +- .../backends/video_sparse_attn_h3.py | 203 +++++++++--------- fastvideo/envs.py | 7 +- fastvideo/fastvideo_args.py | 6 +- fastvideo/models/dits/minimax_h3.py | 2 +- fastvideo/models/loader/fsdp_load.py | 8 +- .../minimax_h3/stages/minimax_h3_denoising.py | 2 +- ...oute.py => test_vsa_h3_blackwell_route.py} | 189 ++++++++-------- .../attention/test_vsa_h3_heads_first_tile.py | 10 +- .../attention/test_vsa_h3_ref2va_regions.py | 32 +-- .../inference/test_basic_fasth3_profile.py | 12 +- .../test_inference_regional_compile.py | 18 +- ...vation_checkpoint_ops_retains_attention.py | 2 +- .../test_minimax_h3_tile_first.py | 2 +- .../train/utils/activation_checkpoint.py | 2 +- fastvideo/training/activation_checkpoint.py | 2 +- pyproject.toml | 2 +- .../benchmarks/minimax_h3_4090/bench_pod.py | 2 +- scripts/benchmarks/minimax_h3_pro6000/app.py | 2 +- ...00a.py => bench_block_sparse_blackwell.py} | 8 +- tests/local_tests/test_cookbook_serving.py | 2 +- ...100a.py => test_block_sparse_blackwell.py} | 30 +-- 75 files changed, 613 insertions(+), 586 deletions(-) rename fastvideo-kernel/csrc/attention/{block_sparse_sm100a.cu => block_sparse_blackwell.cu} (83%) create mode 100644 fastvideo-kernel/csrc/attention/block_sparse_blk128_blackwell.cu delete mode 100644 fastvideo-kernel/csrc/attention/block_sparse_blk128_sm100a.cu rename fastvideo-kernel/csrc/attention/{block_sparse_bwd_sm100a.cu => block_sparse_bwd_blackwell.cu} (89%) create mode 100644 fastvideo-kernel/csrc/attention/block_sparse_bwd_blk128_blackwell.cu delete mode 100644 fastvideo-kernel/csrc/attention/block_sparse_bwd_blk128_sm100a.cu rename fastvideo-kernel/csrc/attention/{block_sparse_bwd_kernel_sm100a.cuh => block_sparse_bwd_kernel_blackwell.cuh} (99%) rename fastvideo-kernel/csrc/attention/{block_sparse_bwd_kernel_sm100a_blk128.cuh => block_sparse_bwd_kernel_blk128_blackwell.cuh} (99%) rename fastvideo-kernel/csrc/attention/{block_sparse_bwd_launch_sm100a.cuh => block_sparse_bwd_launch_blackwell.cuh} (97%) rename fastvideo-kernel/csrc/attention/{block_sparse_bwd_launch_sm100a_blk128.cuh => block_sparse_bwd_launch_blk128_blackwell.cuh} (95%) rename fastvideo-kernel/csrc/attention/{block_sparse_kernel_sm100a.cuh => block_sparse_kernel_blackwell.cuh} (99%) rename fastvideo-kernel/csrc/attention/{block_sparse_launch_sm100a.cuh => block_sparse_launch_blackwell.cuh} (95%) rename fastvideo-kernel/python/fastvideo_kernel/{block_sparse_attn_bwd_sm100a.py => block_sparse_attn_bwd_ptx_blackwell.py} (80%) rename fastvideo-kernel/python/fastvideo_kernel/{block_sparse_attn_sm100a.py => block_sparse_attn_ptx_blackwell.py} (82%) rename fastvideo-kernel/tests/{test_block_sparse_sm100a_clc.py => test_block_sparse_blackwell_clc.py} (88%) rename fastvideo-kernel/tests/{test_block_sparse_sm100a_dispatch.py => test_block_sparse_blackwell_dispatch.py} (77%) rename fastvideo-kernel/tests/{test_block_sparse_bwd_sm100a.py => test_block_sparse_bwd_blackwell.py} (93%) rename fastvideo/tests/attention/{test_vsa_h3_sm100a_route.py => test_vsa_h3_blackwell_route.py} (76%) rename tests/{bench_block_sparse_sm100a.py => bench_block_sparse_blackwell.py} (92%) rename tests/{test_block_sparse_sm100a.py => test_block_sparse_blackwell.py} (92%) diff --git a/apps/dreamverse/dreamverse/minimax_h3_generation.py b/apps/dreamverse/dreamverse/minimax_h3_generation.py index 4861ef2bc..ea7fcef43 100644 --- a/apps/dreamverse/dreamverse/minimax_h3_generation.py +++ b/apps/dreamverse/dreamverse/minimax_h3_generation.py @@ -48,7 +48,7 @@ class MiniMaxH3GenerationBackend: "FASTVIDEO_ATTENTION_BACKEND": attention_backend, "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) diff --git a/apps/dreamverse/dreamverse/tests/test_minimax_h3_generation.py b/apps/dreamverse/dreamverse/tests/test_minimax_h3_generation.py index c8d77c86e..b7c11f6d4 100644 --- a/apps/dreamverse/dreamverse/tests/test_minimax_h3_generation.py +++ b/apps/dreamverse/dreamverse/tests/test_minimax_h3_generation.py @@ -73,7 +73,7 @@ def test_initialize_builds_vsa_datafree_fasth3_generator(monkeypatch): monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "test-attention") monkeypatch.setenv("FASTVIDEO_FA4", "0") monkeypatch.setenv("FASTVIDEO_MINIMAX_H3_FUSIONS", "0") - monkeypatch.setenv("FASTVIDEO_VSA_SM100A", "1") + monkeypatch.setenv("FASTVIDEO_VSA_PTX_BLACKWELL", "1") monkeypatch.setenv("FASTVIDEO_INFERENCE_TORCH_COMPILE", "1") backend = MiniMaxH3GenerationBackend(gpu_id=0) @@ -108,7 +108,7 @@ def test_initialize_builds_vsa_datafree_fasth3_generator(monkeypatch): assert os.environ["FASTVIDEO_ATTENTION_BACKEND"] == "VIDEO_SPARSE_ATTN_H3" assert os.environ["FASTVIDEO_FA4"] == "1" assert os.environ["FASTVIDEO_MINIMAX_H3_FUSIONS"] == "all" - assert os.environ["FASTVIDEO_VSA_SM100A"] == "0" + assert os.environ["FASTVIDEO_VSA_PTX_BLACKWELL"] == "0" assert "FASTVIDEO_INFERENCE_TORCH_COMPILE" not in os.environ diff --git a/apps/infinite_livestream/README.md b/apps/infinite_livestream/README.md index af1add185..608af9a07 100644 --- a/apps/infinite_livestream/README.md +++ b/apps/infinite_livestream/README.md @@ -276,7 +276,7 @@ any GPU work begins. The model directory needs `transformer`, `text_encoder`, and `modular_model_index.json`. See [Download weights](#download-weights) for the complete checkpoint. -**`FastH3's sm100a route needs fastvideo-kernel built with the Blackwell VSA +**`FastH3's ptx_blackwell route needs fastvideo-kernel built with the Blackwell VSA extension`.** Startup checks for the fast sparse-attention kernel before loading any weights. Follow [Install](#install) to build it, or set `inference.vsa_kernel: triton` in your config to use the slower fallback. diff --git a/apps/infinite_livestream/infinite_livestream/backend.py b/apps/infinite_livestream/infinite_livestream/backend.py index cd4f2b501..d9bd445fc 100644 --- a/apps/infinite_livestream/infinite_livestream/backend.py +++ b/apps/infinite_livestream/infinite_livestream/backend.py @@ -193,6 +193,14 @@ class FastH3Backend: # --------------------------------------------------------------- profile + def _vsa_kernel(self) -> str: + """inference.vsa_kernel; the pre-rename value sm100a still selects ptx_blackwell, with a warning.""" + kernel = str(self._config.inference.get("vsa_kernel", "ptx_blackwell")) + if kernel == "sm100a": + logger.warning("inference.vsa_kernel: sm100a is deprecated; use ptx_blackwell") + return "ptx_blackwell" + return kernel + def _apply_profile_environment(self) -> None: """Set the FastH3 profile environment, as the reference CLI does. @@ -201,11 +209,11 @@ class FastH3Backend: experiment settings cannot silently change what gets served. """ cfg = self._config.inference - vsa_kernel = str(cfg.get("vsa_kernel", "sm100a")) + vsa_kernel = self._vsa_kernel() fusions = "all" if bool(cfg.get("h3_fusions", True)) else "0" environment: dict[str, str | None] = { "FASTVIDEO_ATTENTION_BACKEND": "VIDEO_SPARSE_ATTN_H3", - "FASTVIDEO_VSA_SM100A": "1" if vsa_kernel == "sm100a" else "0", + "FASTVIDEO_VSA_PTX_BLACKWELL": "1" if vsa_kernel == "ptx_blackwell" else "0", "FASTVIDEO_VSA_CUTEDSL": "0", # A non-empty path enables the diagnostic probe; it must stay unset. "FASTVIDEO_H3_VSA_PROBE": None, @@ -241,15 +249,15 @@ class FastH3Backend: if not present: raise RuntimeError("FastH3's FA4 route needs the pinned flash-attn-4 package. Install it, " "or set inference.fa4: false in your config.") - if str(cfg.get("vsa_kernel", "sm100a")) == "sm100a": + if self._vsa_kernel() == "ptx_blackwell": try: - from fastvideo_kernel import block_sparse_attn_sm100a + from fastvideo_kernel import block_sparse_attn_ptx_blackwell except ImportError: present = False else: - present = bool(getattr(block_sparse_attn_sm100a, "_HAS_VSA_SM100A", False)) + present = bool(getattr(block_sparse_attn_ptx_blackwell, "_HAS_VSA_PTX_BLACKWELL", False)) if not present: - raise RuntimeError("FastH3's sm100a route needs fastvideo-kernel built with the Blackwell VSA " + raise RuntimeError("FastH3's ptx_blackwell route needs fastvideo-kernel built with the Blackwell VSA " "extension. Install a matching wheel, or set inference.vsa_kernel: triton.") def _generator_config(self) -> Any: diff --git a/apps/infinite_livestream/infinite_livestream/configs/infinite_livestream.yaml b/apps/infinite_livestream/infinite_livestream/configs/infinite_livestream.yaml index 7bbbf308a..1a0b156a9 100644 --- a/apps/infinite_livestream/infinite_livestream/configs/infinite_livestream.yaml +++ b/apps/infinite_livestream/infinite_livestream/configs/infinite_livestream.yaml @@ -26,8 +26,8 @@ inference: vsa_sparsity: 0.9 vsa_tile_size: 64 - # sm100a is the Blackwell VSA kernel; `triton` is the ~2.5x slower fallback. - vsa_kernel: sm100a + # ptx_blackwell is the Blackwell VSA kernel; `triton` is the ~2.5x slower fallback. + vsa_kernel: ptx_blackwell fa4: true h3_fusions: true inference_torch_compile: true diff --git a/docs/assets/cookbook-recipes.json b/docs/assets/cookbook-recipes.json index 0038beda5..1bf270590 100644 --- a/docs/assets/cookbook-recipes.json +++ b/docs/assets/cookbook-recipes.json @@ -552,10 +552,10 @@ "serving": { "source": "examples/serving/openai_fasth3_spark.yaml", "install": "UV_TORCH_BACKEND=cu130 uv pip install -e .", - "env": "FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3", + "env": "FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3", "audio": true }, - "command": "FASTVIDEO_VSA_SM100A=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", + "command": "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", "gpu_types": ["NVIDIA"], "hardware": { "platform": "cuda", @@ -567,7 +567,7 @@ "expected_artifact": "MP4 under outputs/fasth3_spark/", "modes": ["T2VA", "1-Spark"], "limitations": [ - "Install from the DGX Spark guide, not the generic CUDA extra. GB10 has no FA4 / sm_100a VSA kernel; keep FASTVIDEO_FA4=0 and FASTVIDEO_VSA_SM100A=0.", + "Install from the DGX Spark guide, not the generic CUDA extra. GB10 has no FA4 / sm_100a VSA kernel; keep FASTVIDEO_FA4=0 and FASTVIDEO_VSA_PTX_BLACKWELL=0.", "Legal num_frames values are 17n+5, capped at 345 (15 s). Native 16:9 sizes include 832x480 and 1344x768.", "Lazy module load reloads Qwen3-VL and the DiT between phases of each request. Do not pass --no-lazy-module-load on this box.", "A 345-frame request on one Spark can OOM. Prefer 124 or 243 frames, TAEH3 decode, or two Sparks over QSFP." @@ -585,7 +585,7 @@ "summary": "Run one FastH3 clip across two GB10s with Ray sequence parallel over QSFP RoCE. Sequential load and lazy module load stay on because SP replicates the DiT on each node.", "model": "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree", "source": "examples/inference/basic/basic_fasth3_spark_pair.yaml", - "command": "source examples/inference/optimizations/spark_pair_env.sh && FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_VAE_PARALLEL_DECODE=1 fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pair.yaml", + "command": "source examples/inference/optimizations/spark_pair_env.sh && FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_VAE_PARALLEL_DECODE=1 fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pair.yaml", "gpu_types": ["NVIDIA"], "hardware": { "platform": "cuda", @@ -842,7 +842,7 @@ "summary": "Eight-step video and synchronized audio with the released NVFP4 transformer, text encoder, and lightweight VAE.", "model": "FastVideo/FastVideo-FastH3-Trim-8-Step-NVFP4", "source": "examples/inference/basic/basic_fasth3_spark_pruned_nvfp4.yaml", - "command": "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_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING=1 fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pruned_nvfp4.yaml", + "command": "FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cutlass FASTVIDEO_H3_VAE_TILE_BATCH=1 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", "gpu_types": [ "NVIDIA" ], diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index 7f7562272..98ed6b0f1 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -194,7 +194,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry. | `FASTVIDEO_ATTENTION_BACKEND` | str | unset | attention | Attention backend, as an AttentionBackendEnum name such as TORCH_SDPA, FLASH_ATTN, VIDEO_SPARSE_ATTN, SAGE_ATTN, or SAGE_ATTN_THREE. FastVideoArgs uses it when FastVideoArgs.attention_backend is unset. | | `FASTVIDEO_FA4` | bool | `0` | attention | The FLASH_ATTN backend uses FlashAttention-4 (flash_attn.cute) instead of FA3 or FA2. | | `FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN` | bool | `0` | attention | MiniMax-H3 dense DiT self-attention uses the FlashAttention-4 packed-varlen entry point. This changes the floating-point reduction order, so it is an inference-only opt-in. | -| `FASTVIDEO_VSA_SM100A` | bool | `0` | attention | VIDEO_SPARSE_ATTN_H3 sends no-grad tile-64 forwards to the data-center Blackwell (sm_100a) kernel. fastvideo-kernel reads the same variable with the same rule. | +| `FASTVIDEO_VSA_PTX_BLACKWELL` | bool | `0` | attention | VIDEO_SPARSE_ATTN_H3 sends no-grad tile-64 forwards to the data-center Blackwell (sm_100a / sm_103a) kernel. fastvideo-kernel reads the same variable with the same rule. Deprecated names: `FASTVIDEO_VSA_SM100A`. | | `FASTVIDEO_VSA_TRITON` | bool | `0` | attention | Force the Triton MiniMax-H3 sparse attention kernel. fastvideo-kernel reads the same variable. | | `FASTVIDEO_NVFP4_FA4` | bool | `0` | attention | FlashAttention-4 quantizes Q and K to NVFP4. An explicit nvfp4_fa4 attention implementation argument takes precedence. | | `FASTVIDEO_FA4_PV_MODE` | one of bf16, fp8 | `bf16` | attention | V dtype on the NVFP4 FlashAttention-4 path (FLASH_ATTN with nvfp4_fa4, ATTN_QAT_INFER on sm_100a/sm_103a): bf16 keeps V in BF16; fp8 casts V to float8 e4m3 without scaling. An explicit fa4_pv_mode attention implementation argument takes precedence. | diff --git a/docs/cookbook/minimax-h3.md b/docs/cookbook/minimax-h3.md index 002f3ef97..640baf26c 100644 --- a/docs/cookbook/minimax-h3.md +++ b/docs/cookbook/minimax-h3.md @@ -101,8 +101,8 @@ cd FastVideo
  • The MLX source runtime supports T2VA, optional temporal --fast, optional spatial --fast-spatial, and opt-in VSA on --include-vsa checkpoints. FastH3 V2 MLX converts with --include-vsa and runs eight forwards. FL2VA, Ref2VA, and two-pass refinement are not wired.
  • GPU count and VAE decode backend are configurable in the builder above for FastH3 CUDA recipes. Only the value shown by default has a recorded run; other supported values are unmeasured here.
  • DGX Spark is a runtime on FastH3 V1, not a separate family card. Select NVIDIA DGX Spark, then 1 Spark or 2 Sparks. The CUDA GPU-count knob does not apply to Spark.
    • -
  • CompactH3 NVFP4 is one Blackwell GPU. RTX 5090 (32 GB) parks the encoder in pinned host RAM. RTX PRO 6000 Blackwell (96 GB) keeps encoder, DiT, and VAE resident. Keep FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER, FASTVIDEO_FA4=0, FASTVIDEO_VSA_SM100A=0, and FLASHINFER_CUDA_ARCH_LIST=12.0a. On PRO 6000 enable VAE compile and leave DiT inference_torch_compile off. CompactH3 is a dense prune; do not enable VIDEO_SPARSE_ATTN_H3 until a VSA-trained student exists.
    • -
  • GB10 has no FA4 / sm_100a VSA kernel. Keep FASTVIDEO_FA4=0 and FASTVIDEO_VSA_SM100A=0. Legal num_frames values are 17n+5, capped at 362 (15.08 s). A 345-frame request on one Spark can OOM. Native 16:9 sizes include 832×480 and 1344×768.
    • +
  • CompactH3 NVFP4 is one Blackwell GPU. RTX 5090 (32 GB) parks the encoder in pinned host RAM. RTX PRO 6000 Blackwell (96 GB) keeps encoder, DiT, and VAE resident. Keep FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER, FASTVIDEO_FA4=0, FASTVIDEO_VSA_PTX_BLACKWELL=0, and FLASHINFER_CUDA_ARCH_LIST=12.0a. On PRO 6000 enable VAE compile and leave DiT inference_torch_compile off. CompactH3 is a dense prune; do not enable VIDEO_SPARSE_ATTN_H3 until a VSA-trained student exists.
    • +
  • GB10 has no FA4 / ptx_blackwell VSA kernel. Keep FASTVIDEO_FA4=0 and FASTVIDEO_VSA_PTX_BLACKWELL=0. Legal num_frames values are 17n+5, capped at 362 (15.08 s). A 345-frame request on one Spark can OOM. Native 16:9 sizes include 832×480 and 1344×768.
  • Gated or missing checkpoints: run huggingface-cli login and confirm you accepted the model's license on Hugging Face.
    • diff --git a/docs/cookbook/openai-api.md b/docs/cookbook/openai-api.md index 530efa24e..ca0c524b9 100644 --- a/docs/cookbook/openai-api.md +++ b/docs/cookbook/openai-api.md @@ -56,11 +56,11 @@ After model loading completes, the response is `{"status":"ok"}`. ### NVIDIA DGX Spark Complete the [DGX Spark installation](../getting_started/installation/spark.md) -before running these commands. GB10 has no FA4 / sm_100a VSA kernel: +before running these commands. GB10 has no FA4 / ptx_blackwell VSA kernel: ```bash UV_TORCH_BACKEND=cu130 uv pip install -e . -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 serve --config examples/serving/openai_fasth3_spark.yaml --server.host 127.0.0.1 ``` diff --git a/docs/cookbook/recipes/fasth3-trim.yaml b/docs/cookbook/recipes/fasth3-trim.yaml index 62b68c69c..8ec8e242f 100644 --- a/docs/cookbook/recipes/fasth3-trim.yaml +++ b/docs/cookbook/recipes/fasth3-trim.yaml @@ -12,7 +12,7 @@ deployments: FASTVIDEO_NVFP4_MM_BACKEND: cutlass FASTVIDEO_H3_VAE_TILE_BATCH: "1" FASTVIDEO_VSA_TRITON: "1" - FASTVIDEO_VSA_SM100A: "0" + FASTVIDEO_VSA_PTX_BLACKWELL: "0" FASTVIDEO_FA4: "0" FASTVIDEO_ATTENTION_BACKEND: VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING: "1" diff --git a/docs/getting_started/installation/spark_pair.md b/docs/getting_started/installation/spark_pair.md index 928bfe91b..e6dab20ca 100644 --- a/docs/getting_started/installation/spark_pair.md +++ b/docs/getting_started/installation/spark_pair.md @@ -106,12 +106,12 @@ Check `ray status` on the head: `0.0/2.0 GPU` idle. Run the driver on the **head**, same venv, same QSFP IP. -`basic_fasth3.py` defaults target a four-GPU GB200 profile: 768×1344, `sm100a` +`basic_fasth3.py` defaults target a four-GPU GB200 profile: 768×1344, `ptx_blackwell` VSA, FA4, four GPUs. On Sparks you must override the kernel flags. Height, width, frames, steps, seed, and prompt are yours. Change them. Legal `num_frames` values are `17n+5`, capped at 362. -GB10 has no FA4 / sm_100a VSA kernel, so `--vsa-kernel triton --no-fa4` stays +GB10 has no FA4 / ptx_blackwell VSA kernel, so `--vsa-kernel triton --no-fa4` stays required on this box. `--execution-backend ray` is optional when `RAY_ADDRESS` is already set. @@ -141,7 +141,7 @@ Config-first equivalent. Edit the YAML the same way, `request.sampling` is not locked: ```bash -FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 \ +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 @@ -158,7 +158,7 @@ settings, then run one of these configs from the head: ```bash export FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cutlass export FASTVIDEO_H3_VAE_TILE_BATCH=1 FASTVIDEO_VSA_TRITON=1 -export FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 +export FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING=1 fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pair_v2_nvfp4.yaml # Or basic_fasth3_spark_pair_pruned_nvfp4.yaml for FastH3 Trim. @@ -207,8 +207,8 @@ frames. The first VAE decode still pays `torch.compile`. Later `generate()` calls in the same workers are cheaper. GB10 regional DiT compile stays off because the -sm_100a VSA kernel is not on this chip, so denoise is slower than a GB200 -`sm100a` run at the same geometry. +`ptx_blackwell` kernel is not on this chip, so denoise is slower than a GB200 +run at the same geometry. ## Troubleshooting diff --git a/docs/getting_started/installation/spark_performance.md b/docs/getting_started/installation/spark_performance.md index 9b4974a86..2a4740bf6 100644 --- a/docs/getting_started/installation/spark_performance.md +++ b/docs/getting_started/installation/spark_performance.md @@ -262,7 +262,7 @@ It does not use frame dropping or spatial upscaling. 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 \ nice -n 19 fastvideo generate \ diff --git a/docs/inference/fasth3-distilled.md b/docs/inference/fasth3-distilled.md index d77f68bd5..5c9b9d958 100644 --- a/docs/inference/fasth3-distilled.md +++ b/docs/inference/fasth3-distilled.md @@ -19,7 +19,7 @@ other flag works unchanged: ```bash python examples/inference/basic/basic_fasth3_8step.py \ --prompt 'A slow cinematic drone shot glides over a coastal town; gulls call over the harbor.' \ - --num-gpus 4 --vsa-kernel sm100a \ + --num-gpus 4 --vsa-kernel ptx_blackwell \ --profile strict --no-inference-torch-compile --no-compile-vae \ --height 768 --width 1344 --num-frames 124 \ --output outputs/fasth3-8step @@ -35,8 +35,8 @@ H3 clean-time values (`1 - sigma`). A uniform nine-point grid is not a substitut for those rungs. Compilation and H3 fusions are disabled above to establish an eager reference; -they can be evaluated separately. On hardware without the sm100a extension, -use `--vsa-kernel triton`; compare outputs and performance before adopting that +they can be evaluated separately. Where the `ptx_blackwell` kernel is unavailable +(non-Blackwell GPUs, or a build without it), use `--vsa-kernel triton`; compare outputs and performance before adopting that backend. This recipe is T2AV-only, not a distilled `transformer_ref` model. ## Export metadata and validation diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index 763db7335..d8a4d91d4 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -399,8 +399,8 @@ loader wraps each `_compile_conditions` block in loads. The ordinary compile path keeps its historical compiler-disabled attention boundary by default; regional compile opts in only the compatible attention instances owned by this transformer. MiniMax-H3 VSA is supported -only by the inference-only sm_100a tile-64 route -(`FASTVIDEO_VSA_SM100A=1` and `VSA_tile_size=64`); the loader probes that +only by the inference-only ptx_blackwell tile-64 route +(`FASTVIDEO_VSA_PTX_BLACKWELL=1` and `VSA_tile_size=64`); the loader probes that route before capture and keeps the transformer eager when the kernel or device is unsupported. Legacy VSA, MiniMax-H3 tile-256 VSA, and the explicit `FASTVIDEO_DISABLE_ATTENTION_COMPILE=1` escape hatch keep the transformer @@ -450,7 +450,7 @@ regional MiniMax-H3 compile as an opt-in performance experiment, not an eager-parity-safe mode. The same caveat applies to sparse MiniMax-H3 regional compile. Its mask -compaction, sm_100a launch, trained compression gates, and inference-only H3 +compaction, ptx_blackwell launch, trained compression gates, and inference-only H3 fusions are fullgraph-compatible, but compilation can still change model numerics. The FastH3 `all` profile enables regional compile by default because it is the fastest measured route at 124, 243, and 345 frames. Use @@ -479,7 +479,7 @@ MS-SSIM gate on *your* config, especially when combining PyTorch rejects `mode` together with `options`. The generic compile path can accept `mode`, but CUDA-graph compatibility remains backend- and shape-dependent. FA2/FA3 inference and FA4 expose traceable custom-op - boundaries. MiniMax-H3's sm_100a tile-64 inference route is the only VSA + boundaries. MiniMax-H3's ptx_blackwell tile-64 inference route is the only VSA path in the regional support envelope; other VSA paths and the FA3 grad-enabled path remain outside it. Use the default inductor mode shown above unless your exact configuration has its own gate. diff --git a/examples/inference/basic/README.md b/examples/inference/basic/README.md index 131210ea6..a3e063349 100644 --- a/examples/inference/basic/README.md +++ b/examples/inference/basic/README.md @@ -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 \ diff --git a/examples/inference/basic/basic_compacth3_rtx5090.yaml b/examples/inference/basic/basic_compacth3_rtx5090.yaml index 46ce251b8..e290a1347 100644 --- a/examples/inference/basic/basic_compacth3_rtx5090.yaml +++ b/examples/inference/basic/basic_compacth3_rtx5090.yaml @@ -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: diff --git a/examples/inference/basic/basic_compacth3_rtx_pro6000.yaml b/examples/inference/basic/basic_compacth3_rtx_pro6000.yaml index 6c92ae7e8..adafec2b1 100644 --- a/examples/inference/basic/basic_compacth3_rtx_pro6000.yaml +++ b/examples/inference/basic/basic_compacth3_rtx_pro6000.yaml @@ -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: diff --git a/examples/inference/basic/basic_fasth3.py b/examples/inference/basic/basic_fasth3.py index 54d7ad100..91314f8ca 100644 --- a/examples/inference/basic/basic_fasth3.py +++ b/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(): diff --git a/examples/inference/basic/basic_fasth3_h100.yaml b/examples/inference/basic/basic_fasth3_h100.yaml index a84636d4c..dbe0cd83e 100644 --- a/examples/inference/basic/basic_fasth3_h100.yaml +++ b/examples/inference/basic/basic_fasth3_h100.yaml @@ -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: diff --git a/examples/inference/basic/basic_fasth3_lora_preview.py b/examples/inference/basic/basic_fasth3_lora_preview.py index 7db20a2df..7d2e775ea 100644 --- a/examples/inference/basic/basic_fasth3_lora_preview.py +++ b/examples/inference/basic/basic_fasth3_lora_preview.py @@ -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 diff --git a/examples/inference/basic/basic_fasth3_simplified.py b/examples/inference/basic/basic_fasth3_simplified.py index b53bd655c..81ebb1d1f 100644 --- a/examples/inference/basic/basic_fasth3_simplified.py +++ b/examples/inference/basic/basic_fasth3_simplified.py @@ -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) diff --git a/examples/inference/basic/basic_fasth3_simplified_profile.sh b/examples/inference/basic/basic_fasth3_simplified_profile.sh index c20981cd3..2430615dd 100755 --- a/examples/inference/basic/basic_fasth3_simplified_profile.sh +++ b/examples/inference/basic/basic_fasth3_simplified_profile.sh @@ -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}}" diff --git a/examples/inference/basic/basic_fasth3_spark.yaml b/examples/inference/basic/basic_fasth3_spark.yaml index 3cbff23d0..46b8f9848 100644 --- a/examples/inference/basic/basic_fasth3_spark.yaml +++ b/examples/inference/basic/basic_fasth3_spark.yaml @@ -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: diff --git a/examples/inference/basic/basic_fasth3_spark_pair.yaml b/examples/inference/basic/basic_fasth3_spark_pair.yaml index 4f81a1007..eb1d165ef 100644 --- a/examples/inference/basic/basic_fasth3_spark_pair.yaml +++ b/examples/inference/basic/basic_fasth3_spark_pair.yaml @@ -8,7 +8,7 @@ # source examples/inference/optimizations/spark_pair_env.sh # export RAY_ADDRESS=:6379 # export FASTVIDEO_HOST_IP= -# 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: diff --git a/examples/inference/basic/basic_fasth3_spark_pair_pruned_nvfp4.yaml b/examples/inference/basic/basic_fasth3_spark_pair_pruned_nvfp4.yaml index 5d59fcfb1..9f75d084b 100644 --- a/examples/inference/basic/basic_fasth3_spark_pair_pruned_nvfp4.yaml +++ b/examples/inference/basic/basic_fasth3_spark_pair_pruned_nvfp4.yaml @@ -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: diff --git a/examples/inference/basic/basic_fasth3_spark_pair_v2_nvfp4.yaml b/examples/inference/basic/basic_fasth3_spark_pair_v2_nvfp4.yaml index 4d1a0d55e..3e76091b2 100644 --- a/examples/inference/basic/basic_fasth3_spark_pair_v2_nvfp4.yaml +++ b/examples/inference/basic/basic_fasth3_spark_pair_v2_nvfp4.yaml @@ -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: diff --git a/examples/inference/basic/basic_fasth3_spark_pruned_nvfp4.yaml b/examples/inference/basic/basic_fasth3_spark_pruned_nvfp4.yaml index 5ad449512..d9b86d72f 100644 --- a/examples/inference/basic/basic_fasth3_spark_pruned_nvfp4.yaml +++ b/examples/inference/basic/basic_fasth3_spark_pruned_nvfp4.yaml @@ -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: diff --git a/examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml b/examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml index 27531db0e..023010956 100644 --- a/examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml +++ b/examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml @@ -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: diff --git a/examples/inference/lora/minimax_h3_lora_inference.py b/examples/inference/lora/minimax_h3_lora_inference.py index 9a0c6c78a..a4cfeb853 100644 --- a/examples/inference/lora/minimax_h3_lora_inference.py +++ b/examples/inference/lora/minimax_h3_lora_inference.py @@ -33,6 +33,7 @@ import argparse import json import math import os +import warnings from collections.abc import Sequence from pathlib import Path @@ -82,9 +83,12 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: help="video sparse attention; inferred from the adapter when omitted") parser.add_argument("--vsa-sparsity", type=float, default=0.9) parser.add_argument("--vsa-tile-size", type=int, choices=(64, 256), default=64) - parser.add_argument("--vsa-kernel", choices=("triton", "sm100a"), default="sm100a") + parser.add_argument("--vsa-kernel", choices=("triton", "ptx_blackwell", "sm100a"), default="ptx_blackwell") parser.add_argument("--fa4", action=argparse.BooleanOptionalAction, default=True) args = parser.parse_args(argv) + 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" if not args.prompt and not args.prompts_file: parser.error("pass --prompt or --prompts-file") # Whether VSA is required is a property of the adapter, not of having one at all -- @@ -110,7 +114,7 @@ def configure_environment(args: argparse.Namespace) -> None: """ env: dict[str, str | None] = { "FASTVIDEO_ATTENTION_BACKEND": "VIDEO_SPARSE_ATTN_H3" if args.vsa else "FLASH_ATTN", - "FASTVIDEO_VSA_SM100A": "1" if (args.vsa and args.vsa_kernel == "sm100a") else "0", + "FASTVIDEO_VSA_PTX_BLACKWELL": "1" if (args.vsa and args.vsa_kernel == "ptx_blackwell") else "0", "FASTVIDEO_VSA_CUTEDSL": "0", "FASTVIDEO_H3_VSA_PROBE": None, "FASTVIDEO_DISABLE_ATTENTION_COMPILE": "0", diff --git a/examples/serving/openai_compacth3_rtx5090.yaml b/examples/serving/openai_compacth3_rtx5090.yaml index 136fc0af9..7d38bc849 100644 --- a/examples/serving/openai_compacth3_rtx5090.yaml +++ b/examples/serving/openai_compacth3_rtx5090.yaml @@ -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 serve --config examples/serving/openai_compacth3_rtx5090.yaml generator: diff --git a/examples/serving/openai_compacth3_rtx_pro6000.yaml b/examples/serving/openai_compacth3_rtx_pro6000.yaml index 0f947feec..97f1351e9 100644 --- a/examples/serving/openai_compacth3_rtx_pro6000.yaml +++ b/examples/serving/openai_compacth3_rtx_pro6000.yaml @@ -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 serve --config examples/serving/openai_compacth3_rtx_pro6000.yaml generator: diff --git a/examples/serving/openai_fasth3_spark.yaml b/examples/serving/openai_fasth3_spark.yaml index f240e4346..56f9a9804 100644 --- a/examples/serving/openai_fasth3_spark.yaml +++ b/examples/serving/openai_fasth3_spark.yaml @@ -1,8 +1,8 @@ # OpenAI-compatible FastH3 server for one DGX Spark (one GB10). # Install: docs/getting_started/installation/spark.md # -# GB10 has no FA4 / sm_100a VSA kernel. Prefix the serve command with: -# FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 +# GB10 has no FA4 / ptx_blackwell VSA kernel. Prefix the serve command with: +# FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 # # default_request.sampling is an example. Legal H3 frame counts are 17n+5, # max 345 (15 s). A 345-frame request on one Spark can OOM; prefer 124 or 243, diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index afd516493..b84b32327 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -479,19 +479,19 @@ if(BUILD_CXX_KERNELS) 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 + message(STATUS "fastvideo-kernel: building block_sparse_ptx_blackwell (${_VSA_FWD_ARCHS}; 64- and 128-token blocks)") + list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_blackwell.cu + csrc/attention/block_sparse_blk128_blackwell.cu) + set_source_files_properties(csrc/attention/block_sparse_blackwell.cu + csrc/attention/block_sparse_blk128_blackwell.cu PROPERTIES COMPILE_OPTIONS "${_VSA_FWD_GENCODE};-DVSA_BHSD=true") endif() # VSA block-sparse attention BACKWARD, 64- and 128-token blocks: the same images as the forward. if(_VSA_FWD_GENCODE) - list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_bwd_sm100a.cu - csrc/attention/block_sparse_bwd_blk128_sm100a.cu) - set_source_files_properties(csrc/attention/block_sparse_bwd_sm100a.cu - csrc/attention/block_sparse_bwd_blk128_sm100a.cu PROPERTIES + list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_bwd_blackwell.cu + csrc/attention/block_sparse_bwd_blk128_blackwell.cu) + set_source_files_properties(csrc/attention/block_sparse_bwd_blackwell.cu + csrc/attention/block_sparse_bwd_blk128_blackwell.cu PROPERTIES COMPILE_OPTIONS "${_VSA_FWD_GENCODE};-DVSA_BHSD=true") endif() @@ -511,7 +511,7 @@ if(BUILD_CXX_KERNELS) # Build compile definitions list set(COMPILE_DEFS TORCH_EXTENSION_NAME=fastvideo_kernel_ops) if(ENABLE_VSA_SM100A OR ENABLE_VSA_SM103A) - list(APPEND COMPILE_DEFS TK_COMPILE_BLOCK_SPARSE_VSA_SM100A) + list(APPEND COMPILE_DEFS TK_COMPILE_BLOCK_SPARSE_VSA_PTX_BLACKWELL) endif() if(ENABLE_TK_KERNELS) list(APPEND COMPILE_DEFS TK_COMPILE_ST_ATTN TK_COMPILE_BLOCK_SPARSE) diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index 23bc42fbf..b5f08734f 100644 --- a/fastvideo-kernel/README.md +++ b/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` | Data-center Blackwell `sm_100a`/`sm_103a` | ON iff `10.0a` or `10.3a` 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*.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/workstation Blackwell `sm_120a` (CUDA ≥ 12.8) + `sm_121a` (DGX Spark GB10; CUDA ≥ 13) | `FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER` (AUTO = ON iff `12.0a` or `12.1a` is in the arch list) | diff --git a/fastvideo-kernel/csrc/attention/block_sparse_sm100a.cu b/fastvideo-kernel/csrc/attention/block_sparse_blackwell.cu similarity index 83% rename from fastvideo-kernel/csrc/attention/block_sparse_sm100a.cu rename to fastvideo-kernel/csrc/attention/block_sparse_blackwell.cu index e9ade7731..59580a348 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_sm100a.cu +++ b/fastvideo-kernel/csrc/attention/block_sparse_blackwell.cu @@ -1,4 +1,4 @@ -// block_sparse_sm100a.cu -- torch binding for the sm_100a/sm_103a VSA block-sparse FMHA forward. +// block_sparse_blackwell.cu -- torch binding for the Blackwell 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 -- @@ -13,7 +13,7 @@ #include #include -#include "block_sparse_launch_sm100a.cuh" +#include "block_sparse_launch_blackwell.cuh" namespace { @@ -40,17 +40,17 @@ void check_index(const torch::Tensor& t, const char* name) { } // namespace -// The exported symbol carries the block size: block_sparse_sm100a_fwd is the 64-token build, -// block_sparse_sm100a_blk128_fwd the 128-token one (block_sparse_blk128_sm100a.cu re-includes +// The exported symbol carries the block size: block_sparse_ptx_blackwell_fwd is the 64-token build, +// block_sparse_ptx_blackwell_blk128_fwd the 128-token one (block_sparse_blk128_blackwell.cu re-includes // this file with VSA_BLK128 set). The python backend picks by the metadata's block size. #if VSA_BLK128 -#define BLOCK_SPARSE_SM100A_FWD block_sparse_sm100a_blk128_fwd +#define BLOCK_SPARSE_PTX_BLACKWELL_FWD block_sparse_ptx_blackwell_blk128_fwd #else -#define BLOCK_SPARSE_SM100A_FWD block_sparse_sm100a_fwd +#define BLOCK_SPARSE_PTX_BLACKWELL_FWD block_sparse_ptx_blackwell_fwd #endif // Returns {out} or {out, lse}. Layout of out matches the inputs. -std::vector BLOCK_SPARSE_SM100A_FWD(torch::Tensor q, torch::Tensor k, +std::vector BLOCK_SPARSE_PTX_BLACKWELL_FWD(torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional v_t, torch::Tensor q2k_idx, @@ -100,14 +100,14 @@ std::vector BLOCK_SPARSE_SM100A_FWD(torch::Tensor q, torch::Tenso // Report an unsupported regime loudly rather than returning plausible-looking wrong values. TORCH_CHECK(block_sparse_supported(a) == cudaSuccess, - "block_sparse_sm100a: unsupported configuration -- requires head_dim==", + "block_sparse_ptx_blackwell: unsupported configuration -- requires head_dim==", HEAD_DIM, ", an even num_blocks, seqlen == num_blocks*", BLOCK, ", and a variable_block_sizes tensor. Got head_dim=", D, " num_blocks=", num_blocks, " seqlen=", S); - const cudaError_t err = launch_block_sparse_sm100a(a, at::cuda::getCurrentCUDAStream()); + const cudaError_t err = launch_block_sparse_ptx_blackwell(a, at::cuda::getCurrentCUDAStream()); TORCH_CHECK(err == cudaSuccess, - "block_sparse_sm100a launch failed: ", cudaGetErrorString(err)); + "block_sparse_ptx_blackwell launch failed: ", cudaGetErrorString(err)); if (need_lse) return {out, lse}; return {out}; diff --git a/fastvideo-kernel/csrc/attention/block_sparse_blk128_blackwell.cu b/fastvideo-kernel/csrc/attention/block_sparse_blk128_blackwell.cu new file mode 100644 index 000000000..c1efa5a5d --- /dev/null +++ b/fastvideo-kernel/csrc/attention/block_sparse_blk128_blackwell.cu @@ -0,0 +1,7 @@ +// block_sparse_blk128_blackwell.cu -- the 128-token-block instantiation of the torch binding. +// +// Same source as block_sparse_blackwell.cu with VSA_BLK128 set: the kernel and launch land in +// namespace vsa_blk128 (distinct symbols, no ODR clash with the blk64 objects) and the +// exported entry point becomes block_sparse_ptx_blackwell_blk128_fwd. +#define VSA_BLK128 true +#include "block_sparse_blackwell.cu" diff --git a/fastvideo-kernel/csrc/attention/block_sparse_blk128_sm100a.cu b/fastvideo-kernel/csrc/attention/block_sparse_blk128_sm100a.cu deleted file mode 100644 index cb1c2eefe..000000000 --- a/fastvideo-kernel/csrc/attention/block_sparse_blk128_sm100a.cu +++ /dev/null @@ -1,7 +0,0 @@ -// block_sparse_blk128_sm100a.cu -- the 128-token-block instantiation of the torch binding. -// -// Same source as block_sparse_sm100a.cu with VSA_BLK128 set: the kernel and launch land in -// namespace vsa_blk128 (distinct symbols, no ODR clash with the blk64 objects) and the -// exported entry point becomes block_sparse_sm100a_blk128_fwd. -#define VSA_BLK128 true -#include "block_sparse_sm100a.cu" diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_sm100a.cu b/fastvideo-kernel/csrc/attention/block_sparse_bwd_blackwell.cu similarity index 89% rename from fastvideo-kernel/csrc/attention/block_sparse_bwd_sm100a.cu rename to fastvideo-kernel/csrc/attention/block_sparse_bwd_blackwell.cu index 8397ebf56..1337ddb09 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_bwd_sm100a.cu +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_blackwell.cu @@ -1,6 +1,6 @@ -// block_sparse_bwd_sm100a.cu -- torch binding for the sm_100a VSA block-sparse FMHA backward. +// block_sparse_bwd_blackwell.cu -- torch binding for the Blackwell VSA block-sparse FMHA backward. // -// Pairs with block_sparse_sm100a_fwd. Inputs are the forward's operands plus its output o and +// Pairs with block_sparse_ptx_blackwell_fwd. Inputs are the forward's operands plus its output o and // its lse; lse is the Triton "M format" tensor the forward returns -- [B, H, S] fp32, // M = max(qk * sm_scale * log2e) + log2(l) -- and is consumed as-is. Sparsity arrives as // FastVideo's k2q metadata (fastvideo_kernel.triton_kernels.index.invert_indices): for every @@ -10,8 +10,8 @@ // // The layout is fixed at compile time: VSA_BHSD true -> [B, H, S, 128] (FastVideo's build), // false -> [B, S, H, 128] (repo native; the kernel addresses it as [B*S tokens, H, 128]). -// VSA_BLK128 picks the kernel: false -> 64-token sparse blocks (block_sparse_sm100a_bwd), true -> -// 128-token (block_sparse_sm100a_blk128_bwd; block_sparse_bwd_blk128_sm100a.cu re-includes this +// VSA_BLK128 picks the kernel: false -> 64-token sparse blocks (block_sparse_ptx_blackwell_bwd), true -> +// 128-token (block_sparse_ptx_blackwell_blk128_bwd; block_sparse_bwd_blk128_blackwell.cu re-includes this // file with it set). The python backend picks by the metadata's block size. #include @@ -24,11 +24,11 @@ #define VSA_BLK128 false #endif #if VSA_BLK128 -#include "block_sparse_bwd_launch_sm100a_blk128.cuh" -#define BLOCK_SPARSE_SM100A_BWD block_sparse_sm100a_blk128_bwd +#include "block_sparse_bwd_launch_blk128_blackwell.cuh" +#define BLOCK_SPARSE_PTX_BLACKWELL_BWD block_sparse_ptx_blackwell_blk128_bwd #else -#include "block_sparse_bwd_launch_sm100a.cuh" -#define BLOCK_SPARSE_SM100A_BWD block_sparse_sm100a_bwd +#include "block_sparse_bwd_launch_blackwell.cuh" +#define BLOCK_SPARSE_PTX_BLACKWELL_BWD block_sparse_ptx_blackwell_bwd #endif namespace { @@ -61,7 +61,7 @@ __nv_bfloat16* bf16_ptr(const torch::Tensor& t) { } // namespace // Returns {dq, dk, dv}: bf16, each with the shape and layout of q, k, v respectively. -std::vector BLOCK_SPARSE_SM100A_BWD(torch::Tensor grad_o, torch::Tensor q, +std::vector BLOCK_SPARSE_PTX_BLACKWELL_BWD(torch::Tensor grad_o, torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, torch::Tensor lse, torch::Tensor k2q_idx, torch::Tensor k2q_num, @@ -166,7 +166,7 @@ std::vector BLOCK_SPARSE_SM100A_BWD(torch::Tensor grad_o, torch:: // Report an unsupported regime loudly rather than returning plausible-looking wrong values. TORCH_CHECK(block_sparse_bwd_supported(a) == cudaSuccess, - "block_sparse_sm100a_bwd (", BLOCK, + "block_sparse_ptx_blackwell_bwd (", BLOCK, "-token blocks): unsupported configuration -- requires head_dim==", HEAD_DIM, ", seqlen == num_kv_blocks_per_seq*", BLOCK, " with seqlen % 128 == 0, batch*heads <= 65535, " @@ -174,8 +174,8 @@ std::vector BLOCK_SPARSE_SM100A_BWD(torch::Tensor grad_o, torch:: D, " num_kv_blocks_per_seq=", num_kv_blocks_per_seq, " seqlen=", S, " max_q_blocks=", max_q_blocks, " sm_scale=", sm_scale); - const cudaError_t err = launch_block_sparse_bwd_sm100a(a, at::cuda::getCurrentCUDAStream()); - TORCH_CHECK(err == cudaSuccess, "block_sparse_sm100a_bwd (", BLOCK, + const cudaError_t err = launch_block_sparse_bwd_ptx_blackwell(a, at::cuda::getCurrentCUDAStream()); + TORCH_CHECK(err == cudaSuccess, "block_sparse_ptx_blackwell_bwd (", BLOCK, "-token blocks) launch failed: ", cudaGetErrorString(err)); return {dq, dk, dv}; diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_blk128_blackwell.cu b/fastvideo-kernel/csrc/attention/block_sparse_bwd_blk128_blackwell.cu new file mode 100644 index 000000000..d73fc97ab --- /dev/null +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_blk128_blackwell.cu @@ -0,0 +1,7 @@ +// block_sparse_bwd_blk128_blackwell.cu -- the 128-token-block instantiation of the torch binding. +// +// Same source as block_sparse_bwd_blackwell.cu with VSA_BLK128 set: the kernel and launch land in +// namespace vsa_bwd_blk128 (distinct symbols, no ODR clash with the blk64 objects) and the +// exported entry point becomes block_sparse_ptx_blackwell_blk128_bwd. +#define VSA_BLK128 true +#include "block_sparse_bwd_blackwell.cu" diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_blk128_sm100a.cu b/fastvideo-kernel/csrc/attention/block_sparse_bwd_blk128_sm100a.cu deleted file mode 100644 index 2d04613a1..000000000 --- a/fastvideo-kernel/csrc/attention/block_sparse_bwd_blk128_sm100a.cu +++ /dev/null @@ -1,7 +0,0 @@ -// block_sparse_bwd_blk128_sm100a.cu -- the 128-token-block instantiation of the torch binding. -// -// Same source as block_sparse_bwd_sm100a.cu with VSA_BLK128 set: the kernel and launch land in -// namespace vsa_bwd_blk128 (distinct symbols, no ODR clash with the blk64 objects) and the -// exported entry point becomes block_sparse_sm100a_blk128_bwd. -#define VSA_BLK128 true -#include "block_sparse_bwd_sm100a.cu" diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh b/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_blackwell.cuh similarity index 99% rename from fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh rename to fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_blackwell.cuh index 3d3014ff0..e4a6fe2dc 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_blackwell.cuh @@ -1,9 +1,9 @@ -// block_sparse_bwd_kernel_sm100a.cuh -- VSA block-sparse attention BACKWARD (blk64 -// one-pass), sm_100a. Warp-specialized: load / MMA (tcgen05) / softmax (P^T, dS^T) / +// block_sparse_bwd_kernel_blackwell.cuh -- VSA block-sparse attention BACKWARD (blk64 +// one-pass), Blackwell. Warp-specialized: load / MMA (tcgen05) / softmax (P^T, dS^T) / // epilogue (dQ drain) / scheduler. Three kernels: preprocess (Delta, Q^T, dO^T, dqaccum // zero), main (dK, dV, dQ partials), postprocess (dQ unscramble + scale). -#ifndef BLOCK_SPARSE_VSA_BWD_KERNEL_SM100A_CUH -#define BLOCK_SPARSE_VSA_BWD_KERNEL_SM100A_CUH +#ifndef BLOCK_SPARSE_VSA_BWD_KERNEL_PTX_BLACKWELL_CUH +#define BLOCK_SPARSE_VSA_BWD_KERNEL_PTX_BLACKWELL_CUH #include #include diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a_blk128.cuh b/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_blk128_blackwell.cuh similarity index 99% rename from fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a_blk128.cuh rename to fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_blk128_blackwell.cuh index 1ca9bcf4e..3e6780d07 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a_blk128.cuh +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_blk128_blackwell.cuh @@ -1,9 +1,9 @@ -// block_sparse_bwd_kernel_sm100a_blk128.cuh -- VSA block-sparse attention BACKWARD, -// 128-token blocks, sm_100a (one CTA per kv block). Warp-specialized: load / +// block_sparse_bwd_kernel_blk128_blackwell.cuh -- VSA block-sparse attention BACKWARD, +// 128-token blocks, Blackwell (one CTA per kv block). Warp-specialized: load / // MMA (tcgen05) / softmax (P^T, dS^T) / epilogue (dQ drain). Three kernels: preprocess // (Delta, dqaccum zero), main (dK, dV, dQ partials), postprocess (dQ unscramble + scale). -#ifndef BLOCK_SPARSE_VSA_BWD_KERNEL_SM100A_BLK128_CUH -#define BLOCK_SPARSE_VSA_BWD_KERNEL_SM100A_BLK128_CUH +#ifndef BLOCK_SPARSE_VSA_BWD_KERNEL_PTX_BLACKWELL_BLK128_CUH +#define BLOCK_SPARSE_VSA_BWD_KERNEL_PTX_BLACKWELL_BLK128_CUH #include #include diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a.cuh b/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_blackwell.cuh similarity index 97% rename from fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a.cuh rename to fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_blackwell.cuh index 559152eac..d936b28ae 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a.cuh +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_blackwell.cuh @@ -1,13 +1,13 @@ -// block_sparse_bwd_launch_sm100a.cuh -- host surface of the VSA block-sparse backward drop: +// block_sparse_bwd_launch_blackwell.cuh -- host surface of the VSA block-sparse backward drop: // argument struct, workspace sizes, the support predicate and the stream-chained launch // (preprocess -> order -> main -> postprocess). Tensor maps are encoded per call (no static // cache: a torch caller hands us fresh pointers every time). -#ifndef BLOCK_SPARSE_VSA_BWD_LAUNCH_SM100A_CUH -#define BLOCK_SPARSE_VSA_BWD_LAUNCH_SM100A_CUH +#ifndef BLOCK_SPARSE_VSA_BWD_LAUNCH_PTX_BLACKWELL_CUH +#define BLOCK_SPARSE_VSA_BWD_LAUNCH_PTX_BLACKWELL_CUH #include #include -#include "block_sparse_bwd_kernel_sm100a.cuh" +#include "block_sparse_bwd_kernel_blackwell.cuh" #ifndef VSA_BHSD #define VSA_BHSD false @@ -230,7 +230,7 @@ __host__ inline cudaError_t launch_main(const BlockSparseVsaBwdArgs& args, const } } -__host__ inline cudaError_t launch_block_sparse_bwd_sm100a(const BlockSparseVsaBwdArgs& args, +__host__ inline cudaError_t launch_block_sparse_bwd_ptx_blackwell(const BlockSparseVsaBwdArgs& args, cudaStream_t stream) { const cudaError_t supported = block_sparse_bwd_supported(args); if (supported != cudaSuccess) { @@ -327,4 +327,4 @@ __host__ inline cudaError_t launch_block_sparse_bwd_sm100a(const BlockSparseVsaB using namespace vsa_bwd_blk64; -#endif // BLOCK_SPARSE_VSA_BWD_LAUNCH_SM100A_CUH +#endif // BLOCK_SPARSE_VSA_BWD_LAUNCH_PTX_BLACKWELL_CUH diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a_blk128.cuh b/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_blk128_blackwell.cuh similarity index 95% rename from fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a_blk128.cuh rename to fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_blk128_blackwell.cuh index 9cacbda24..82fd6d28f 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a_blk128.cuh +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_blk128_blackwell.cuh @@ -1,12 +1,12 @@ -// block_sparse_bwd_launch_sm100a_blk128.cuh -- host surface of the 128-token VSA block-sparse +// block_sparse_bwd_launch_blk128_blackwell.cuh -- host surface of the 128-token VSA block-sparse // backward drop: argument struct, workspace sizes, the support predicate and the stream-chained // launch (preprocess -> main -> postprocess, each a programmatic dependent launch of the one // before). Tensor maps are encoded per call (a torch caller hands us fresh pointers every time). -#ifndef BLOCK_SPARSE_VSA_BWD_LAUNCH_SM100A_BLK128_CUH -#define BLOCK_SPARSE_VSA_BWD_LAUNCH_SM100A_BLK128_CUH +#ifndef BLOCK_SPARSE_VSA_BWD_LAUNCH_PTX_BLACKWELL_BLK128_CUH +#define BLOCK_SPARSE_VSA_BWD_LAUNCH_PTX_BLACKWELL_BLK128_CUH #include -#include "block_sparse_bwd_kernel_sm100a_blk128.cuh" +#include "block_sparse_bwd_kernel_blk128_blackwell.cuh" #ifndef VSA_BHSD #define VSA_BHSD false @@ -149,7 +149,7 @@ __host__ inline cudaLaunchConfig_t pdl_launch_config(dim3 grid, dim3 block, size return cfg; } -__host__ inline cudaError_t launch_block_sparse_bwd_sm100a(const BlockSparseVsaBwdArgs& args, +__host__ inline cudaError_t launch_block_sparse_bwd_ptx_blackwell(const BlockSparseVsaBwdArgs& args, cudaStream_t stream) { const cudaError_t supported = block_sparse_bwd_supported(args); if (supported != cudaSuccess) { @@ -228,4 +228,4 @@ __host__ inline cudaError_t launch_block_sparse_bwd_sm100a(const BlockSparseVsaB using namespace vsa_bwd_blk128; -#endif // BLOCK_SPARSE_VSA_BWD_LAUNCH_SM100A_BLK128_CUH +#endif // BLOCK_SPARSE_VSA_BWD_LAUNCH_PTX_BLACKWELL_BLK128_CUH diff --git a/fastvideo-kernel/csrc/attention/block_sparse_kernel_sm100a.cuh b/fastvideo-kernel/csrc/attention/block_sparse_kernel_blackwell.cuh similarity index 99% rename from fastvideo-kernel/csrc/attention/block_sparse_kernel_sm100a.cuh rename to fastvideo-kernel/csrc/attention/block_sparse_kernel_blackwell.cuh index d2af53a1a..8876e99bc 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_kernel_sm100a.cuh +++ b/fastvideo-kernel/csrc/attention/block_sparse_kernel_blackwell.cuh @@ -1,9 +1,9 @@ -// 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. +// block_sparse_kernel_blackwell.cuh -- VSA block-sparse FMHA forward (per-q-block top-k), +// data-center Blackwell. // Warp-specialized: load / MMA (tcgen05) / softmax / correction / epilogue / scheduler. // Writes O and, when asked, the log-sum-exp the backward consumes. -#ifndef BLOCK_SPARSE_VSA_KERNEL_SM100A_CUH -#define BLOCK_SPARSE_VSA_KERNEL_SM100A_CUH +#ifndef BLOCK_SPARSE_VSA_KERNEL_PTX_BLACKWELL_CUH +#define BLOCK_SPARSE_VSA_KERNEL_PTX_BLACKWELL_CUH #include #include diff --git a/fastvideo-kernel/csrc/attention/block_sparse_launch_sm100a.cuh b/fastvideo-kernel/csrc/attention/block_sparse_launch_blackwell.cuh similarity index 95% rename from fastvideo-kernel/csrc/attention/block_sparse_launch_sm100a.cuh rename to fastvideo-kernel/csrc/attention/block_sparse_launch_blackwell.cuh index bdb924588..53ba27c35 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_launch_sm100a.cuh +++ b/fastvideo-kernel/csrc/attention/block_sparse_launch_blackwell.cuh @@ -1,18 +1,18 @@ -#ifndef BLOCK_SPARSE_VSA_LAUNCH_SM100A_CUH -#define BLOCK_SPARSE_VSA_LAUNCH_SM100A_CUH +#ifndef BLOCK_SPARSE_VSA_LAUNCH_PTX_BLACKWELL_CUH +#define BLOCK_SPARSE_VSA_LAUNCH_PTX_BLACKWELL_CUH -// Launch surface for the sm_100a/sm_103a VSA block-sparse FMHA forward. +// Launch surface for the Blackwell 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 -// block_sparse_bench_sm100a.cu and the torch binding both go through here, so there is +// block_sparse_bench.cu and the torch binding both go through here, so there is // one tensormap construction and one launch configuration rather than two that can drift. // // Two compile-time knobs select the four builds: // VSA_BLK128 false -> 64-token sparse blocks, true -> 128-token // VSA_BHSD false -> [token][head][dim] (BSHD), true -> [batch][head][token][dim] (BHSD) -#include "block_sparse_kernel_sm100a.cuh" +#include "block_sparse_kernel_blackwell.cuh" namespace VSA_NAMESPACE { @@ -53,7 +53,7 @@ __host__ inline cudaError_t block_sparse_supported(const BlockSparseVsaArgs& a) return cudaSuccess; } -__host__ inline cudaError_t launch_block_sparse_sm100a(const BlockSparseVsaArgs& a, +__host__ inline cudaError_t launch_block_sparse_ptx_blackwell(const BlockSparseVsaArgs& a, cudaStream_t stream) { const cudaError_t sup = block_sparse_supported(a); if (sup != cudaSuccess) return sup; @@ -198,4 +198,4 @@ __host__ inline cudaError_t launch_block_sparse_sm100a(const BlockSparseVsaArgs& // only ever sees the one configuration its VSA_BLK128 selected. using namespace VSA_NAMESPACE; -#endif // BLOCK_SPARSE_VSA_LAUNCH_SM100A_CUH +#endif // BLOCK_SPARSE_VSA_LAUNCH_PTX_BLACKWELL_CUH diff --git a/fastvideo-kernel/csrc/common_extension.cpp b/fastvideo-kernel/csrc/common_extension.cpp index 9f839dd62..1d0462808 100644 --- a/fastvideo-kernel/csrc/common_extension.cpp +++ b/fastvideo-kernel/csrc/common_extension.cpp @@ -33,20 +33,20 @@ void register_rms_norm(pybind11::module_ &); void register_layer_norm(pybind11::module_ &); void register_gemm(pybind11::module_ &); -#ifdef TK_COMPILE_BLOCK_SPARSE_VSA_SM100A -extern std::vector block_sparse_sm100a_fwd( +#ifdef TK_COMPILE_BLOCK_SPARSE_VSA_PTX_BLACKWELL +extern std::vector block_sparse_ptx_blackwell_fwd( torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional v_t, torch::Tensor q2k_idx, torch::Tensor q2k_num, torch::Tensor variable_block_sizes, double sm_scale, bool need_lse); -extern std::vector block_sparse_sm100a_blk128_fwd( +extern std::vector block_sparse_ptx_blackwell_blk128_fwd( torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional v_t, torch::Tensor q2k_idx, torch::Tensor q2k_num, torch::Tensor variable_block_sizes, double sm_scale, bool need_lse); -extern std::vector block_sparse_sm100a_bwd( +extern std::vector block_sparse_ptx_blackwell_bwd( torch::Tensor grad_o, torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, torch::Tensor lse, torch::Tensor k2q_idx, torch::Tensor k2q_num, torch::Tensor variable_block_sizes, double sm_scale); -extern std::vector block_sparse_sm100a_blk128_bwd( +extern std::vector block_sparse_ptx_blackwell_blk128_bwd( torch::Tensor grad_o, torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, torch::Tensor lse, torch::Tensor k2q_idx, torch::Tensor k2q_num, torch::Tensor variable_block_sizes, double sm_scale); @@ -55,18 +55,18 @@ extern std::vector block_sparse_sm100a_blk128_bwd( PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "FastVideo CUDA Kernels"; -#ifdef TK_COMPILE_BLOCK_SPARSE_VSA_SM100A - m.def("block_sparse_sm100a_fwd", - torch::wrap_pybind_function(block_sparse_sm100a_fwd), +#ifdef TK_COMPILE_BLOCK_SPARSE_VSA_PTX_BLACKWELL + m.def("block_sparse_ptx_blackwell_fwd", + torch::wrap_pybind_function(block_sparse_ptx_blackwell_fwd), "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), + m.def("block_sparse_ptx_blackwell_blk128_fwd", + torch::wrap_pybind_function(block_sparse_ptx_blackwell_blk128_fwd), "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), + m.def("block_sparse_ptx_blackwell_bwd", + torch::wrap_pybind_function(block_sparse_ptx_blackwell_bwd), "VSA block-sparse attention backward, 64-token blocks (Blackwell sm100a/sm103a)"); - m.def("block_sparse_sm100a_blk128_bwd", - torch::wrap_pybind_function(block_sparse_sm100a_blk128_bwd), + m.def("block_sparse_ptx_blackwell_blk128_bwd", + torch::wrap_pybind_function(block_sparse_ptx_blackwell_blk128_bwd), "VSA block-sparse attention backward, 128-token blocks (Blackwell sm100a/sm103a)"); #endif diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py index b463932ca..a11ce2f2a 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py @@ -49,18 +49,21 @@ def _force_tk() -> bool: return os.environ.get("FASTVIDEO_VSA_TK", "0") == "1" -def _force_sm100a() -> bool: +def _force_ptx_blackwell() -> bool: """True iff the data-center Blackwell forward is explicitly opted into. - Opt-in only (same legacy-named env the H3 backend honors). Honored only when - ``block_sparse_attn_sm100a.is_supported`` passes; the backward then runs the - sm_100a/sm_103a CUDA backward when ``block_sparse_attn_bwd_sm100a.is_supported`` + Opt-in only (same env the H3 backend honors). Honored only when + ``block_sparse_attn_ptx_blackwell.is_supported`` passes; the backward then runs the + ptx_blackwell backward when ``block_sparse_attn_bwd_ptx_blackwell.is_supported`` passes and Triton otherwise (the forward's lse is already in Triton's M format). Unsupported 64-token metadata falls through to the default selection; unsupported 128-token metadata raises because Triton has no - compatible fallback. ``FASTVIDEO_VSA_TRITON`` still wins. + compatible fallback. ``FASTVIDEO_VSA_TRITON`` still wins. The deprecated name + ``FASTVIDEO_VSA_SM100A`` is read when ``FASTVIDEO_VSA_PTX_BLACKWELL`` is unset. """ - return _env_bool("FASTVIDEO_VSA_SM100A") + if "FASTVIDEO_VSA_PTX_BLACKWELL" not in os.environ and "FASTVIDEO_VSA_SM100A" in os.environ: + return _env_bool("FASTVIDEO_VSA_SM100A") + return _env_bool("FASTVIDEO_VSA_PTX_BLACKWELL") def _env_bool(name: str) -> bool: @@ -79,12 +82,12 @@ def _env_bool(name: str) -> bool: "or an empty string") -def _sm100a_is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: +def _ptx_blackwell_is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: try: - from fastvideo_kernel import block_sparse_attn_sm100a as vsa_sm100a + from fastvideo_kernel import block_sparse_attn_ptx_blackwell as vsa_ptx_blackwell except Exception: return False - return vsa_sm100a.is_supported(q, variable_block_sizes) + return vsa_ptx_blackwell.is_supported(q, variable_block_sizes) def _infer_block_size(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> int: @@ -386,10 +389,10 @@ def _backward_sm90(ctx, grad_o, grad_lse): block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_context_sm90) # --------------------------------------------------------------------------- -# Data-center Blackwell backend custom op (index-native; legacy sm100a API name) +# Data-center Blackwell backend custom op (index-native; the ptx_blackwell kernels) # -# Forward runs the sm_100a/sm_103a CUDA extension. Backward runs the sm_100a CUDA -# backward when block_sparse_attn_bwd_sm100a.is_supported passes (64- or 128-token +# Forward runs the ptx_blackwell extension. Backward runs the ptx_blackwell +# backward when block_sparse_attn_bwd_ptx_blackwell.is_supported passes (64- or 128-token # blocks, data-center Blackwell device, extension built with that block's op) and # the Triton kernels otherwise. The native forward emits lse in exactly Triton's M # format (max*log2e + log2(l)), so either pairing needs no conversion. The Triton @@ -399,11 +402,11 @@ block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_co @torch.library.custom_op( - "fastvideo_kernel::block_sparse_attn_backward_sm100a", + "fastvideo_kernel::block_sparse_attn_backward_ptx_blackwell", mutates_args=(), device_types="cuda", ) -def block_sparse_attn_backward_sm100a( +def block_sparse_attn_backward_ptx_blackwell( grad_o: torch.Tensor, q: torch.Tensor, k: torch.Tensor, @@ -414,19 +417,19 @@ def block_sparse_attn_backward_sm100a( q2k_num: torch.Tensor, variable_block_sizes: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - from fastvideo_kernel.block_sparse_attn_bwd_sm100a import ( - block_sparse_attn_backward_sm100a_from_k2q, ) + from fastvideo_kernel.block_sparse_attn_bwd_ptx_blackwell import ( + block_sparse_attn_backward_ptx_blackwell_from_k2q, ) num_kv_blocks = variable_block_sizes.numel() k2q_idx, k2q_num = _invert_indices_for_backward(q2k_idx, q2k_num, num_kv_blocks) - dq, dk, dv = block_sparse_attn_backward_sm100a_from_k2q( + dq, dk, dv = block_sparse_attn_backward_ptx_blackwell_from_k2q( grad_o.contiguous(), q.contiguous(), k.contiguous(), v.contiguous(), o.contiguous(), lse.contiguous(), k2q_idx, k2q_num, variable_block_sizes) return dq, dk, dv -@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_backward_sm100a") -def _block_sparse_attn_backward_sm100a_fake( +@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_backward_ptx_blackwell") +def _block_sparse_attn_backward_ptx_blackwell_fake( grad_o: torch.Tensor, q: torch.Tensor, k: torch.Tensor, @@ -440,20 +443,20 @@ def _block_sparse_attn_backward_sm100a_fake( return torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) -def _sm100a_backward_is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: +def _ptx_blackwell_backward_is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: try: - from fastvideo_kernel import block_sparse_attn_bwd_sm100a as vsa_bwd_sm100a + from fastvideo_kernel import block_sparse_attn_bwd_ptx_blackwell as vsa_bwd_ptx_blackwell except ImportError: # pragma: no cover - extension not built return False - return vsa_bwd_sm100a.is_supported(q, variable_block_sizes) + return vsa_bwd_ptx_blackwell.is_supported(q, variable_block_sizes) @torch.library.custom_op( - "fastvideo_kernel::block_sparse_attn_sm100a", + "fastvideo_kernel::block_sparse_attn_ptx_blackwell", mutates_args=(), device_types="cuda", ) -def block_sparse_attn_sm100a_op( +def block_sparse_attn_ptx_blackwell_op( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -461,15 +464,15 @@ def block_sparse_attn_sm100a_op( q2k_num: torch.Tensor, variable_block_sizes: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: - from fastvideo_kernel.block_sparse_attn_sm100a import block_sparse_attn_sm100a + from fastvideo_kernel.block_sparse_attn_ptx_blackwell import block_sparse_attn_ptx_blackwell - o, M = block_sparse_attn_sm100a(q, k, v, q2k_idx, q2k_num, variable_block_sizes, - need_lse=True) + o, M = block_sparse_attn_ptx_blackwell(q, k, v, q2k_idx, q2k_num, variable_block_sizes, + need_lse=True) return o, M -@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm100a") -def _block_sparse_attn_sm100a_fake( +@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_ptx_blackwell") +def _block_sparse_attn_ptx_blackwell_fake( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -486,32 +489,32 @@ def _block_sparse_attn_sm100a_fake( return o, M -def _setup_context_sm100a(ctx, inputs, output): +def _setup_context_ptx_blackwell(ctx, inputs, output): q, k, v, q2k_idx, q2k_num, variable_block_sizes = inputs o, M = output ctx.save_for_backward(q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes) -def _backward_sm100a(ctx, grad_o, grad_M): +def _backward_ptx_blackwell(ctx, grad_o, grad_M): q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes = ctx.saved_tensors - if _sm100a_backward_is_supported(q, variable_block_sizes): - dq, dk, dv = block_sparse_attn_backward_sm100a(grad_o, q, k, v, o, M, q2k_idx, - q2k_num, variable_block_sizes) + if _ptx_blackwell_backward_is_supported(q, variable_block_sizes): + dq, dk, dv = block_sparse_attn_backward_ptx_blackwell(grad_o, q, k, v, o, M, q2k_idx, + q2k_num, variable_block_sizes) else: block = q.shape[2] // variable_block_sizes.numel() if block != 64: raise RuntimeError( - "block_sparse_attn_sm100a backward: no sm_100a/sm_103a backward for " + "block_sparse_attn_ptx_blackwell backward: no ptx_blackwell backward for " f"{block}-token blocks on this build/device, and the Triton backward is " "hardcoded to 64-token blocks. Run this metadata without grad, or build " - "the extension with block_sparse_sm100a_blk128_bwd.") + "the extension with block_sparse_ptx_blackwell_blk128_bwd.") dq, dk, dv = block_sparse_attn_backward_triton(grad_o, q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes) return dq, dk, dv, None, None, None -block_sparse_attn_sm100a_op.register_autograd(_backward_sm100a, - setup_context=_setup_context_sm100a) +block_sparse_attn_ptx_blackwell_op.register_autograd(_backward_ptx_blackwell, + setup_context=_setup_context_ptx_blackwell) # --------------------------------------------------------------------------- # Public API @@ -538,9 +541,8 @@ def block_sparse_attn_from_indices( # Backend resolution: # - FASTVIDEO_VSA_TRITON forces Triton everywhere. - # - FASTVIDEO_VSA_SM100A opts into the data-center Blackwell forward (CUDA - # backward when supported, else Triton). The environment name is retained - # for compatibility. + # - FASTVIDEO_VSA_PTX_BLACKWELL opts into the data-center Blackwell forward (CUDA + # backward when supported, else Triton). # 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 @@ -548,13 +550,13 @@ def block_sparse_attn_from_indices( # - Otherwise: TK on sm_90 if available, else Triton. if _force_triton(): use_sm90 = False - elif _force_sm100a(): - if _sm100a_is_supported(q, variable_block_sizes): - return block_sparse_attn_sm100a_op(q, k, v, q2k_idx, q2k_num, - variable_block_sizes) + elif _force_ptx_blackwell(): + if _ptx_blackwell_is_supported(q, variable_block_sizes): + return block_sparse_attn_ptx_blackwell_op(q, k, v, q2k_idx, q2k_num, + variable_block_sizes) if _infer_block_size(q, variable_block_sizes) == 128: raise NotImplementedError( - "128-token block-sparse attention requires the sm_100a/sm_103a forward; " + "128-token block-sparse attention requires the ptx_blackwell forward; " "the Triton fallback only supports 64-token blocks, and the " "native data-center Blackwell route is unavailable for this input.") use_sm90 = sm90_available diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_ptx_blackwell.py similarity index 80% rename from fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py rename to fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_ptx_blackwell.py index cb2b1c7ca..895d8895d 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_ptx_blackwell.py @@ -1,12 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -"""sm_100a/sm_103a (data-center Blackwell) CUDA block-sparse VSA backward. +"""ptx_blackwell block-sparse VSA backward for data-center Blackwell (sm_100a / sm_103a). -Companion of ``block_sparse_attn_sm100a`` (the forward): consumes the forward's ``lse`` in the +Companion of ``block_sparse_attn_ptx_blackwell`` (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 FastVideo's k2q index metadata, returns ``(dq, dk, dv)`` in bf16 with the inputs' layout and the Triton backward's scaling (dq and dk carry sm_scale, dv does not). Two kernels, picked by the -metadata's block size: 64-token blocks (``block_sparse_sm100a_bwd``) and 128-token blocks -(``block_sparse_sm100a_blk128_bwd``); any other configuration falls back to Triton via +metadata's block size: 64-token blocks (``block_sparse_ptx_blackwell_bwd``) and 128-token blocks +(``block_sparse_ptx_blackwell_blk128_bwd``); any other configuration falls back to Triton via ``is_supported``. """ @@ -19,14 +19,14 @@ try: # (its __init__ is empty, so hasattr on the package fails with the kernel built and present). from fastvideo_kernel._C import fastvideo_kernel_ops as _C _BWD_BY_BLOCK = { - 64: getattr(_C, "block_sparse_sm100a_bwd", None), - 128: getattr(_C, "block_sparse_sm100a_blk128_bwd", None), + 64: getattr(_C, "block_sparse_ptx_blackwell_bwd", None), + 128: getattr(_C, "block_sparse_ptx_blackwell_blk128_bwd", None), } - _HAS_VSA_BWD_SM100A = any(_BWD_BY_BLOCK.values()) + _HAS_VSA_BWD_PTX_BLACKWELL = any(_BWD_BY_BLOCK.values()) except ImportError: # pragma: no cover - extension not built _C = None _BWD_BY_BLOCK = {} - _HAS_VSA_BWD_SM100A = False + _HAS_VSA_BWD_PTX_BLACKWELL = False _SUPPORTED_COMPUTE_CAPABILITIES = {(10, 0), (10, 3)} HEAD_DIM = 128 @@ -36,20 +36,20 @@ BHSD = True def set_extension(module) -> None: - """Use an already-loaded extension module exposing ``block_sparse_sm100a_bwd`` and / or - ``block_sparse_sm100a_blk128_bwd``. + """Use an already-loaded extension module exposing ``block_sparse_ptx_blackwell_bwd`` and / or + ``block_sparse_ptx_blackwell_blk128_bwd``. A standalone build of the binding .cu files (for example through ``torch.utils.cpp_extension.load`` with a ten-line pybind wrapper) can be injected here, so the backend can be exercised without rebuilding the fastvideo_kernel wheel. """ - global _C, _BWD_BY_BLOCK, _HAS_VSA_BWD_SM100A + global _C, _BWD_BY_BLOCK, _HAS_VSA_BWD_PTX_BLACKWELL _C = module _BWD_BY_BLOCK = { - 64: getattr(module, "block_sparse_sm100a_bwd", None), - 128: getattr(module, "block_sparse_sm100a_blk128_bwd", None), + 64: getattr(module, "block_sparse_ptx_blackwell_bwd", None), + 128: getattr(module, "block_sparse_ptx_blackwell_blk128_bwd", None), } - _HAS_VSA_BWD_SM100A = any(_BWD_BY_BLOCK.values()) + _HAS_VSA_BWD_PTX_BLACKWELL = any(_BWD_BY_BLOCK.values()) def _seqlen(q: torch.Tensor) -> int: @@ -77,7 +77,7 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: Per-row k2q counts may be anything in [0, num_q_blocks], including 0: unselected kv blocks get exactly-zero dk/dv rows. """ - if not _HAS_VSA_BWD_SM100A or not q.is_cuda: + if not _HAS_VSA_BWD_PTX_BLACKWELL or not q.is_cuda: return False if torch.cuda.get_device_capability(q.device) not in _SUPPORTED_COMPUTE_CAPABILITIES: return False @@ -99,7 +99,7 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: return True -def block_sparse_attn_backward_sm100a_from_k2q( +def block_sparse_attn_backward_ptx_blackwell_from_k2q( grad_o: torch.Tensor, q: torch.Tensor, k: torch.Tensor, @@ -119,7 +119,7 @@ def block_sparse_attn_backward_sm100a_from_k2q( block = _block_size(q, variable_block_sizes) backward = _BWD_BY_BLOCK.get(block) if backward is None: - raise RuntimeError(f"block_sparse_attn_backward_sm100a: no kernel for {block}-token " + raise RuntimeError(f"block_sparse_attn_backward_ptx_blackwell: no kernel for {block}-token " f"blocks (built: {sorted(b for b, f in _BWD_BY_BLOCK.items() if f)})") sm_scale = 1.0 / (q.shape[-1]**0.5) idx = k2q_idx.to(torch.int32).contiguous() @@ -130,7 +130,7 @@ def block_sparse_attn_backward_sm100a_from_k2q( return res[0], res[1], res[2] -def block_sparse_attn_backward_sm100a( +def block_sparse_attn_backward_ptx_blackwell( grad_o: torch.Tensor, q: torch.Tensor, k: torch.Tensor, @@ -158,5 +158,5 @@ def block_sparse_attn_backward_sm100a( if num.dim() != 3: num = num.view(batch, heads, -1) k2q_idx, k2q_num = invert_indices(idx, num, num_kv_blocks) - return block_sparse_attn_backward_sm100a_from_k2q(grad_o, q, k, v, o, lse, k2q_idx, k2q_num, - variable_block_sizes) + return block_sparse_attn_backward_ptx_blackwell_from_k2q(grad_o, q, k, v, o, lse, k2q_idx, k2q_num, + variable_block_sizes) diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_ptx_blackwell.py similarity index 82% rename from fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py rename to fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_ptx_blackwell.py index b2d2df6fe..a9cc6eda2 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_ptx_blackwell.py @@ -1,14 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 """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. +``ptx_blackwell`` is the kernel's name, not an architecture: 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. This module is the forward: it returns ``(out, lse)`` with ``lse`` in exactly the form ``triton_block_sparse_attn_forward`` writes -- ``max(qk * qk_scale) + log2(l)``, ``[B, H, S]`` -fp32 -- so both ``block_sparse_attn_backward_triton`` and the sm_100a CUDA backward -(``block_sparse_attn_bwd_sm100a``, 64- and 128-token blocks, sm_100a / sm_103a) run against it +fp32 -- so both ``block_sparse_attn_backward_triton`` and the ptx_blackwell backward +(``block_sparse_attn_bwd_ptx_blackwell``, 64- and 128-token blocks, sm_100a / sm_103a) run against it unchanged. The extension carries TWO instantiations of the kernel, for 64- and 128-token sparse blocks @@ -27,14 +27,14 @@ try: # kernel built and present. from fastvideo_kernel._C import fastvideo_kernel_ops as _C _FWD_BY_BLOCK = { - 64: getattr(_C, "block_sparse_sm100a_fwd", None), - 128: getattr(_C, "block_sparse_sm100a_blk128_fwd", None), + 64: getattr(_C, "block_sparse_ptx_blackwell_fwd", None), + 128: getattr(_C, "block_sparse_ptx_blackwell_blk128_fwd", None), } - _HAS_VSA_SM100A = any(_FWD_BY_BLOCK.values()) + _HAS_VSA_PTX_BLACKWELL = any(_FWD_BY_BLOCK.values()) except ImportError: # pragma: no cover - extension not built _C = None _FWD_BY_BLOCK = {} - _HAS_VSA_SM100A = False + _HAS_VSA_PTX_BLACKWELL = False _SUPPORTED_COMPUTE_CAPABILITIES = {(10, 0), (10, 3)} HEAD_DIM = 128 @@ -69,9 +69,9 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: (= q2k_idx.shape[-1]) must be >= 1, which the host launcher re-checks. * variable_block_sizes: per-KV-block valid-token counts in [0, block]; keys at or past a block's count are masked. Integer metadata is converted to int32/contiguous by - ``block_sparse_attn_sm100a`` itself, so int64 inputs merely cost a cast. + ``block_sparse_attn_ptx_blackwell`` itself, so int64 inputs merely cost a cast. """ - if not _HAS_VSA_SM100A or not q.is_cuda: + if not _HAS_VSA_PTX_BLACKWELL or not q.is_cuda: return False if torch.cuda.get_device_capability(q.device) not in _SUPPORTED_COMPUTE_CAPABILITIES: return False @@ -91,11 +91,11 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: @torch.library.custom_op( - "fastvideo_kernel::block_sparse_attn_sm100a_inference", + "fastvideo_kernel::block_sparse_attn_ptx_blackwell_inference", mutates_args=(), device_types="cuda", ) -def _block_sparse_attn_sm100a_inference( +def _block_sparse_attn_ptx_blackwell_inference( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -103,7 +103,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/sm_103a route. + """Opaque no-LSE launch used by the inference-only ptx_blackwell 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 @@ -117,8 +117,8 @@ def _block_sparse_attn_sm100a_inference( return res[0] -@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm100a_inference") -def _block_sparse_attn_sm100a_inference_fake( +@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_ptx_blackwell_inference") +def _block_sparse_attn_ptx_blackwell_inference_fake( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -131,18 +131,18 @@ def _block_sparse_attn_sm100a_inference_fake( @torch.library.custom_op( - "fastvideo_kernel::block_sparse_attn_sm100a_from_mask_inference", + "fastvideo_kernel::block_sparse_attn_ptx_blackwell_from_mask_inference", mutates_args=(), device_types="cuda", ) -def _block_sparse_attn_sm100a_from_mask_inference( +def _block_sparse_attn_ptx_blackwell_from_mask_inference( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, block_map: torch.Tensor, variable_block_sizes: torch.Tensor, ) -> torch.Tensor: - """Opaque mask compaction plus no-LSE sm_100a/sm_103a launch for inference. + """Opaque mask compaction plus no-LSE ptx_blackwell 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; @@ -152,7 +152,7 @@ def _block_sparse_attn_sm100a_from_mask_inference( from fastvideo_kernel.triton_kernels.index import map_to_index q2k_idx, q2k_num = map_to_index(block_map) - return _block_sparse_attn_sm100a_inference( + return _block_sparse_attn_ptx_blackwell_inference( q, k, v, @@ -162,8 +162,8 @@ def _block_sparse_attn_sm100a_from_mask_inference( ) -@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm100a_from_mask_inference") -def _block_sparse_attn_sm100a_from_mask_inference_fake( +@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_ptx_blackwell_from_mask_inference") +def _block_sparse_attn_ptx_blackwell_from_mask_inference_fake( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -173,7 +173,7 @@ def _block_sparse_attn_sm100a_from_mask_inference_fake( return torch.empty_like(q) -def block_sparse_attn_sm100a_from_mask( +def block_sparse_attn_ptx_blackwell_from_mask( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -181,7 +181,7 @@ def block_sparse_attn_sm100a_from_mask( variable_block_sizes: torch.Tensor, ) -> Tuple[torch.Tensor, None]: """Inference forward from a bool block map, with compaction kept opaque.""" - out = _block_sparse_attn_sm100a_from_mask_inference( + out = _block_sparse_attn_ptx_blackwell_from_mask_inference( q.contiguous(), k.contiguous(), v.contiguous(), @@ -191,7 +191,7 @@ def block_sparse_attn_sm100a_from_mask( return out, None -def block_sparse_attn_sm100a( +def block_sparse_attn_ptx_blackwell( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -212,7 +212,7 @@ def block_sparse_attn_sm100a( # This is the production inference path. The custom op keeps the raw # pybind launch opaque to Dynamo while its fake kernel carries output # metadata through full-graph capture. - return _block_sparse_attn_sm100a_inference(q, k, v, idx, num, vbs), None + return _block_sparse_attn_ptx_blackwell_inference(q, k, v, idx, num, vbs), None # Preserve the established LSE-producing path for correctness tests and # any future forward/backward pairing; only inference needs the opaque op. diff --git a/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py b/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py index 92e311812..a0b8dbde2 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py +++ b/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py @@ -14,8 +14,8 @@ import math import torch VSA_TILE_SIZE = (4, 4, 4) -# 128 is served by the sm_100a/sm_103a CUDA backend (legacy API name -# block_sparse_attn_sm100a); 64 and 256 by +# 128 is served by the ptx_blackwell backend +# (block_sparse_attn_ptx_blackwell); 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) diff --git a/fastvideo-kernel/tests/test_block_sparse_sm100a_clc.py b/fastvideo-kernel/tests/test_block_sparse_blackwell_clc.py similarity index 88% rename from fastvideo-kernel/tests/test_block_sparse_sm100a_clc.py rename to fastvideo-kernel/tests/test_block_sparse_blackwell_clc.py index 2a22aee23..e8047eaf2 100644 --- a/fastvideo-kernel/tests/test_block_sparse_sm100a_clc.py +++ b/fastvideo-kernel/tests/test_block_sparse_blackwell_clc.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""sm_100a VSA forward: CLC work-item scheduler under slot reuse. +"""ptx_blackwell VSA forward: CLC work-item scheduler under slot reuse. The persistent forward steals not-yet-launched CTAs with clusterlaunchcontrol.try_cancel and hands the 16-byte responses through a two-slot ring. The ring is only exercised when the grid @@ -13,13 +13,13 @@ and a tile written from the wrong work item as a value mismatch. import pytest import torch -from fastvideo_kernel import block_sparse_attn_sm100a as vsa +from fastvideo_kernel import block_sparse_attn_ptx_blackwell as vsa HEAD_DIM = 128 pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)} - or not vsa._HAS_VSA_SM100A, + or not vsa._HAS_VSA_PTX_BLACKWELL, reason="requires data-center Blackwell (sm_100a/sm_103a) and a built fastvideo_kernel extension", ) @@ -47,10 +47,10 @@ def test_clc_slot_reuse_never_drops_or_misplaces_a_tile(block, num_blocks): ctas = heads * num_blocks // 2 assert ctas >= 3 * torch.cuda.get_device_properties(0).multi_processor_count, ctas - out_g, lse_g = (t.clone() for t in vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs, need_lse=True)) + out_g, lse_g = (t.clone() for t in vsa.block_sparse_attn_ptx_blackwell(q, k, v, idx, num, vbs, need_lse=True)) assert torch.isfinite(out_g).all() and torch.isfinite(lse_g).all() for _ in range(3): - out, lse = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs, need_lse=True) + out, lse = vsa.block_sparse_attn_ptx_blackwell(q, k, v, idx, num, vbs, need_lse=True) assert torch.equal(out, out_g) and torch.equal(lse, lse_g), "forward is not bitwise deterministic" del out, lse @@ -64,7 +64,7 @@ def test_clc_slot_reuse_never_drops_or_misplaces_a_tile(block, num_blocks): poison_out = torch.full_like(q, nan) out_ptr, lse_ptr = poison_out.data_ptr(), poison_lse.data_ptr() del poison_out, poison_lse - out, lse = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs, need_lse=True) + out, lse = vsa.block_sparse_attn_ptx_blackwell(q, k, v, idx, num, vbs, need_lse=True) assert out.data_ptr() == out_ptr and lse.data_ptr() == lse_ptr, "allocator did not reuse the poisoned blocks" wrong = (out != out_g).view(tile_shape).any(-1) | (lse != lse_g).view(1, heads, num_blocks, block).any(-1) if bool(wrong.any()): diff --git a/fastvideo-kernel/tests/test_block_sparse_sm100a_dispatch.py b/fastvideo-kernel/tests/test_block_sparse_blackwell_dispatch.py similarity index 77% rename from fastvideo-kernel/tests/test_block_sparse_sm100a_dispatch.py rename to fastvideo-kernel/tests/test_block_sparse_blackwell_dispatch.py index 73cca33fa..fe0223b5b 100644 --- a/fastvideo-kernel/tests/test_block_sparse_sm100a_dispatch.py +++ b/fastvideo-kernel/tests/test_block_sparse_blackwell_dispatch.py @@ -1,16 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -"""Routing tests for the opt-in sm_100a/sm_103a dispatch in ``block_sparse_attn_from_indices``. +"""Routing tests for the opt-in ptx_blackwell dispatch in ``block_sparse_attn_from_indices``. -``FASTVIDEO_VSA_SM100A=1`` routes the forward to the sm_100a extension when -``block_sparse_attn_sm100a.is_supported`` passes. Its backward is the sm_100a -CUDA backward where that op is built and ``block_sparse_attn_bwd_sm100a.is_supported`` +``FASTVIDEO_VSA_PTX_BLACKWELL=1`` routes the forward to the ptx_blackwell kernel when +``block_sparse_attn_ptx_blackwell.is_supported`` passes. Its backward is the ptx_blackwell +backward where that op is built and ``block_sparse_attn_bwd_ptx_blackwell.is_supported`` passes (64- or 128-token blocks on a data-center Blackwell device), the Triton backward -otherwise (64-token blocks only); the sm_100a lse is already in Triton's M format, so +otherwise (64-token blocks only); the ptx_blackwell lse is already in Triton's M format, so either pairing needs no conversion. Everything else -- env unset, unsupported input, ``FASTVIDEO_VSA_TRITON`` override -- must keep the pre-existing selection, bit-for-bit. -Run with: python -m pytest tests/test_block_sparse_sm100a_dispatch.py -v +Run with: python -m pytest tests/test_block_sparse_blackwell_dispatch.py -v """ import importlib @@ -18,15 +18,15 @@ import importlib import pytest import torch -from fastvideo_kernel import block_sparse_attn_sm100a as vsa +from fastvideo_kernel import block_sparse_attn_ptx_blackwell as vsa from fastvideo_kernel.block_sparse_attn import block_sparse_attn_from_indices HEAD_DIM = 128 -ENV = "FASTVIDEO_VSA_SM100A" +ENV = "FASTVIDEO_VSA_PTX_BLACKWELL" pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)} - or not vsa._HAS_VSA_SM100A, + or not vsa._HAS_VSA_PTX_BLACKWELL, reason="requires data-center Blackwell (sm_100a/sm_103a) and a built fastvideo_kernel extension", ) @@ -39,7 +39,7 @@ def make_case(block, num_blocks=8, topk=4, heads=4, batch=1, seed=0, requires_gr requires_grad=requires_grad) for _ in range(3)) # from_indices takes Triton-shaped metadata: [B, H, Nq, Mk] / [B, H, Nq]. - # The sm_100a extension reads both flat, so the same tensors serve both paths. + # The ptx_blackwell extension reads both flat, so the same tensors serve both paths. idx = torch.empty((batch * heads * num_blocks, topk), dtype=torch.int32, device="cuda") for r in range(idx.shape[0]): idx[r] = torch.randperm(num_blocks, device="cuda")[:topk].to(torch.int32).sort().values @@ -57,11 +57,11 @@ def triton_reference(q, k, v, idx, num, vbs, monkeypatch): @pytest.mark.parametrize("block", [64, 128]) -def test_env_routes_to_sm100a(block, monkeypatch): - """With the env set on supported input, from_indices runs the sm_100a op: +def test_env_routes_to_ptx_blackwell(block, monkeypatch): + """With the env set on supported input, from_indices runs the ptx_blackwell op: bitwise-equal to the wrapper called directly, allclose to Triton.""" q, k, v, idx, num, vbs = make_case(block) - want_o, want_m = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs) + want_o, want_m = vsa.block_sparse_attn_ptx_blackwell(q, k, v, idx, num, vbs) monkeypatch.setenv(ENV, "1") got_o, got_m = block_sparse_attn_from_indices(q, k, v, idx, num, vbs) @@ -100,28 +100,28 @@ def test_unsupported_blk128_raises_not_implemented(monkeypatch): block_sparse_attn_from_indices(q, k, v, idx, num, vbs) -def test_force_triton_overrides_sm100a(monkeypatch): +def test_force_triton_overrides_ptx_blackwell(monkeypatch): monkeypatch.setenv(ENV, "1") monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1") q, k, v, idx, num, vbs = make_case(64) got = block_sparse_attn_from_indices(q, k, v, idx, num, vbs) monkeypatch.delenv("FASTVIDEO_VSA_TRITON") - sm100a_o, _ = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs) + ptx_blackwell_o, _ = vsa.block_sparse_attn_ptx_blackwell(q, k, v, idx, num, vbs) ref = triton_reference(q, k, v, idx, num, vbs, monkeypatch) assert torch.equal(got[0], ref[0]) - assert not torch.equal(got[0], sm100a_o) + assert not torch.equal(got[0], ptx_blackwell_o) -def _grads_sm100a_route_vs_triton(monkeypatch, vbs=None): - """Grads of (out**2).sum() through the sm_100a route vs the all-Triton route. +def _grads_ptx_blackwell_route_vs_triton(monkeypatch, vbs=None): + """Grads of (out**2).sum() through the ptx_blackwell route vs the all-Triton route. - On a device where the sm_100a backward is built and supported, the route must not enter + On a device where the ptx_blackwell backward is built and supported, the route must not enter the Triton backward at all (it is monkeypatched to fail); elsewhere the route pairs the - sm_100a forward with the Triton backward, as before. + ptx_blackwell forward with the Triton backward, as before. """ # The package exports a FUNCTION named block_sparse_attn; the module needs importlib. dispatch = importlib.import_module("fastvideo_kernel.block_sparse_attn") - from fastvideo_kernel import block_sparse_attn_bwd_sm100a as vsa_bwd + from fastvideo_kernel import block_sparse_attn_bwd_ptx_blackwell as vsa_bwd q, k, v, idx, num, default_vbs = make_case(64, requires_grad=True) vbs = default_vbs if vbs is None else vbs @@ -134,22 +134,22 @@ def _grads_sm100a_route_vs_triton(monkeypatch, vbs=None): monkeypatch.delenv("FASTVIDEO_VSA_TRITON") monkeypatch.setenv(ENV, "1") - sm100a_backward = vsa_bwd.is_supported(q, vbs) - if sm100a_backward: + ptx_blackwell_backward = vsa_bwd.is_supported(q, vbs) + if ptx_blackwell_backward: def no_triton_backward(*args, **kwargs): - raise AssertionError("Triton backward entered on a supported sm_100a input") + raise AssertionError("Triton backward entered on a supported ptx_blackwell input") monkeypatch.setattr(dispatch, "block_sparse_attn_backward_triton", no_triton_backward) out, _ = block_sparse_attn_from_indices(q, k, v, idx, num, vbs) out.float().square().sum().backward() got = [t.grad.float() for t in (q, k, v)] - return got, ref, sm100a_backward + return got, ref, ptx_blackwell_backward def _assert_grads_close(got, ref): # Two bf16 kernels against each other (not an fp32 reference): twice the tolerances the - # sm_100a backward holds against fp32 in tests/test_block_sparse_bwd_sm100a.py. + # ptx_blackwell backward holds against fp32 in tests/test_block_sparse_bwd_blackwell.py. for g, r, name in zip(got, ref, "qkv"): diff = (g - r).abs() rel_max = diff.max().item() / max(r.abs().max().item(), 1e-6) @@ -159,8 +159,8 @@ def _assert_grads_close(got, ref): def test_backward_matches_triton(monkeypatch): - """sm_100a route (CUDA backward where built, Triton backward otherwise) vs all-Triton.""" - got, ref, _ = _grads_sm100a_route_vs_triton(monkeypatch) + """ptx_blackwell route (its backward where built, Triton backward otherwise) vs all-Triton.""" + got, ref, _ = _grads_ptx_blackwell_route_vs_triton(monkeypatch) _assert_grads_close(got, ref) @@ -168,30 +168,30 @@ def test_backward_ragged_block_sizes_matches_triton(monkeypatch): """variable_block_sizes below 64 (padded kv rows) through the same two routes.""" torch.manual_seed(1) vbs = torch.randint(32, 65, (8, ), dtype=torch.int32, device="cuda") - got, ref, _ = _grads_sm100a_route_vs_triton(monkeypatch, vbs=vbs) + got, ref, _ = _grads_ptx_blackwell_route_vs_triton(monkeypatch, vbs=vbs) _assert_grads_close(got, ref) -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 +def test_backward_uses_ptx_blackwell_kernel_when_built(monkeypatch): + """Guards against a silent Triton fallback: with the op built on a data-center Blackwell device + the ptx_blackwell backward must be the one that runs.""" + from fastvideo_kernel import block_sparse_attn_bwd_ptx_blackwell as vsa_bwd + if (not vsa_bwd._HAS_VSA_BWD_PTX_BLACKWELL 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 + pytest.skip("ptx_blackwell backward not built for this device") + _, _, ptx_blackwell_backward = _grads_ptx_blackwell_route_vs_triton(monkeypatch) + assert ptx_blackwell_backward def test_backward_large_seq_matches_triton_and_is_deterministic(monkeypatch): """65536 tokens (1024 kv blocks, 12.5% density): the device-computed work order (nb >= 1024) and the sequence regime where a TMEM ordering bug corrupted dk/dv while every <= 16-block - test stayed green. sm_100a route vs all-Triton, and two sm_100a runs must agree to within + test stayed green. ptx_blackwell route vs all-Triton, and two ptx_blackwell runs must agree to within summation-order noise: invert_indices compacts each kv row's q list with atomics, so the kernel sums quads in a different order per call (measured run-to-run delta on GB200: rel_max 5e-3, mean|diff| 7e-6). The TMEM race this guards against gave rel_max 0.6-1.0 and mean|diff| 4e-2, an order of magnitude past the bounds below on both metrics.""" - from fastvideo_kernel import block_sparse_attn_bwd_sm100a as vsa_bwd + from fastvideo_kernel import block_sparse_attn_bwd_ptx_blackwell as vsa_bwd torch.manual_seed(0) block, num_blocks, heads, topk = 64, 1024, 4, 128 @@ -204,7 +204,7 @@ def test_backward_large_seq_matches_triton_and_is_deterministic(monkeypatch): num = torch.full((1, heads, num_blocks), topk, dtype=torch.int32, device="cuda") vbs = torch.full((num_blocks, ), block, dtype=torch.int32, device="cuda") if not vsa_bwd.is_supported(q, vbs): - pytest.skip("sm_100a backward not built for this device") + pytest.skip("ptx_blackwell backward not built for this device") def grads(route): qq, kk, vv = (t.clone().requires_grad_(True) for t in (q, k, v)) @@ -218,15 +218,15 @@ def test_backward_large_seq_matches_triton_and_is_deterministic(monkeypatch): return [t.grad.float() for t in (qq, kk, vv)] ref = grads("triton") - got = grads("sm100a") - again = grads("sm100a") + got = grads("ptx_blackwell") + again = grads("ptx_blackwell") _assert_grads_close(got, ref) for g1, g2, name in zip(got, again, "qkv"): diff = (g1 - g2).abs() rel_max = diff.max().item() / max(g1.abs().max().item(), 1e-6) mean_abs = diff.mean().item() assert rel_max <= 2e-2 and mean_abs <= 1e-4, \ - f"d{name}: two sm_100a runs differ beyond summation-order noise: " \ + f"d{name}: two ptx_blackwell runs differ beyond summation-order noise: " \ f"max|diff|/max|ref|={rel_max:.3e} mean|diff|={mean_abs:.3e}" @@ -252,10 +252,10 @@ def _blk128_triton_reference_grads(q, k, v, idx, num, vbs, monkeypatch): def test_blk128_backward(monkeypatch): - """128-token blocks: with the 128-token backward built the sm_100a route runs it (the Triton + """128-token blocks: with the 128-token backward built the ptx_blackwell route runs it (the Triton backward must not be entered) and its grads match the all-Triton 128 wrapper; without it the backward raises instead of silently falling back to a 64-token kernel.""" - from fastvideo_kernel import block_sparse_attn_bwd_sm100a as vsa_bwd + from fastvideo_kernel import block_sparse_attn_bwd_ptx_blackwell as vsa_bwd dispatch = importlib.import_module("fastvideo_kernel.block_sparse_attn") monkeypatch.setenv(ENV, "1") diff --git a/fastvideo-kernel/tests/test_block_sparse_bwd_sm100a.py b/fastvideo-kernel/tests/test_block_sparse_bwd_blackwell.py similarity index 93% rename from fastvideo-kernel/tests/test_block_sparse_bwd_sm100a.py rename to fastvideo-kernel/tests/test_block_sparse_bwd_blackwell.py index 832bef8df..b6ecd2fb6 100644 --- a/fastvideo-kernel/tests/test_block_sparse_bwd_sm100a.py +++ b/fastvideo-kernel/tests/test_block_sparse_bwd_blackwell.py @@ -1,14 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 -"""Correctness tests for the sm_100a CUDA block-sparse VSA backward, 64- and 128-token blocks. +"""Correctness tests for the ptx_blackwell block-sparse VSA backward, 64- and 128-token blocks. Reference: fp32 dense attention restricted to the selected blocks, keys past variable_block_sizes masked to -inf, differentiated with torch autograd on fp32 copies of the bf16 inputs with loss = (out * grad_o).sum(). The kernel is fed the reference's lse in Triton's M format (logsumexp * log2e) and the reference output rounded to bf16 as ``o`` -- exactly what -the sm_100a forward hands it in FastVideo. The k2q inversion is done here in torch (a stable +the ptx_blackwell forward hands it in FastVideo. The k2q inversion is done here in torch (a stable sort) so the tests do not need Triton. Every case runs once per built block size. -Run with: python -m pytest tests/test_block_sparse_bwd_sm100a.py -v +Run with: python -m pytest tests/test_block_sparse_bwd_blackwell.py -v """ import itertools @@ -17,7 +17,7 @@ import os import pytest import torch -from fastvideo_kernel import block_sparse_attn_bwd_sm100a as bwd +from fastvideo_kernel import block_sparse_attn_bwd_ptx_blackwell as bwd HEAD_DIM = 128 LOG2E = 1.4426950408889634 @@ -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() not in {(10, 0), (10, 3)} - or not bwd._HAS_VSA_BWD_SM100A, + or not bwd._HAS_VSA_BWD_PTX_BLACKWELL, reason="requires a data-center Blackwell GPU (sm_100a/sm_103a) and a fastvideo_kernel " - "extension built with block_sparse_sm100a_bwd / block_sparse_sm100a_blk128_bwd", + "extension built with block_sparse_ptx_blackwell_bwd / block_sparse_ptx_blackwell_blk128_bwd", ) @@ -159,8 +159,8 @@ def run_and_compare(block, num_blocks=8, topk=4, heads=4, batch=1, ragged=False, assert bwd.is_supported(q, vbs) o, lse, ref_dq, ref_dk, ref_dv = reference(block, q, k, v, grad_o, idx, num, vbs) k2q_idx, k2q_num = invert_indices_torch(idx, num, num_blocks) - dq, dk, dv = bwd.block_sparse_attn_backward_sm100a_from_k2q(grad_o, q, k, v, o, lse, - k2q_idx, k2q_num, vbs) + dq, dk, dv = bwd.block_sparse_attn_backward_ptx_blackwell_from_k2q(grad_o, q, k, v, o, lse, + k2q_idx, k2q_num, vbs) torch.cuda.synchronize() for name, got, ref in (("dq", dq, ref_dq), ("dk", dk, ref_dk), ("dv", dv, ref_dv)): assert got.dtype == torch.bfloat16, f"{name}: dtype {got.dtype}" @@ -273,6 +273,6 @@ def test_unsupported_is_rejected(block): lse = torch.zeros((q.shape[0], heads, vbs.numel() * block), dtype=torch.float32, device="cuda") with pytest.raises(RuntimeError): - bwd.block_sparse_attn_backward_sm100a_from_k2q(grad_o.float(), q.float(), k.float(), - v.float(), o.float(), lse, k2q_idx, - k2q_num, vbs) + bwd.block_sparse_attn_backward_ptx_blackwell_from_k2q(grad_o.float(), q.float(), k.float(), + v.float(), o.float(), lse, k2q_idx, + k2q_num, vbs) diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index 10741e394..5852f66ae 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -40,9 +40,9 @@ fallback and keeps identical mask semantics. At tile 64 the block map is already at the kernels' native 64-token granularity, so both forward and backward run the Triton block-sparse kernels directly (no expansion, ``FASTVIDEO_VSA_CUTEDSL`` does not apply). A third, opt-in route exists -for the tile-64 FORWARD only: ``FASTVIDEO_VSA_SM100A=1`` sends no-grad +for the tile-64 FORWARD only: ``FASTVIDEO_VSA_PTX_BLACKWELL=1`` sends no-grad forwards through the data-center Blackwell CUDA block-sparse kernel -(``fastvideo_kernel.block_sparse_attn_sm100a``, which reads a separate +(``fastvideo_kernel.block_sparse_attn_ptx_blackwell``, which reads a separate ``q2k_num`` key-tile count per query tile) when the extension is built, the device is sm_100 or sm_103, and the geometry qualifies. The CUDA kernel assigns adjacent pairs of query tiles to CTAs, so an odd logical tile count receives one @@ -52,7 +52,7 @@ the original logical tiles. Grad-tracking forwards and every backward stay on the Triton kernels. If the env is set but a precondition fails, the route logs one warning and falls back. -Tile 128 has exactly one implementation: the same sm_100a/sm_103a CUDA +Tile 128 has exactly one implementation: the same ptx_blackwell forward, which carries a 128-token block instantiation. It needs no opt-in, pads an odd tile count with the same zero-valid partner tile, runs no-grad forwards only, and fails closed (no Triton fallback) when the extension or @@ -77,12 +77,12 @@ except ImportError: try: # Optional: only present in fastvideo_kernel builds that carry the - # sm_100a/sm_103a CUDA block-sparse forward. The module itself imports - # fine without the compiled symbols (`_HAS_VSA_SM100A` is then False and + # ptx_blackwell block-sparse forward. The module itself imports + # fine without the compiled symbols (`_HAS_VSA_PTX_BLACKWELL` is then False and # `is_supported` says no), so this only guards *module* availability. - from fastvideo_kernel import block_sparse_attn_sm100a as _sm100a + from fastvideo_kernel import block_sparse_attn_ptx_blackwell as _ptx_blackwell except ImportError: - _sm100a = None + _ptx_blackwell = None import fastvideo.envs as envs from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata, @@ -98,7 +98,7 @@ from fastvideo.logger import init_logger logger = init_logger(__name__) # Opt-in switch for the data-center Blackwell CUDA forward on the tile-64 no-grad path. -VSA_SM100A_ENV = "FASTVIDEO_VSA_SM100A" +VSA_PTX_BLACKWELL_ENV = "FASTVIDEO_VSA_PTX_BLACKWELL" VSA_H3_TILE_SIZE = (4, 8, 8) # 256 elements -> FA4 CuTe fastpath on sm10.x (default) _TILE_ELEMS = math.prod(VSA_H3_TILE_SIZE) @@ -109,20 +109,20 @@ _TILE_ELEMS = math.prod(VSA_H3_TILE_SIZE) VSA_H3_TILE_SHAPES: dict[int, tuple[int, int, int]] = { _TILE_ELEMS: VSA_H3_TILE_SIZE, 64: (4, 4, 4), - # sm_100a/sm_103a CUDA forward only (128-token block instantiation). + # ptx_blackwell forward only (128-token block instantiation). 128: (4, 4, 8), } -# Tile sizes the sm_100a/sm_103a CUDA kernel runs natively; its CTAs own +# Tile sizes the ptx_blackwell kernel runs natively; its CTAs own # adjacent pairs of query tiles, so an odd count gets one zero-valid partner. -_SM100A_TILE_ELEMS = (64, 128) +_PTX_BLACKWELL_TILE_ELEMS = (64, 128) @torch.library.custom_op( - "fastvideo::h3_vsa_sm100a_from_mask_compat", + "fastvideo::h3_vsa_ptx_blackwell_from_mask_compat", mutates_args=(), device_types="cuda", ) -def _h3_vsa_sm100a_from_mask_compat( +def _h3_vsa_ptx_blackwell_from_mask_compat( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -130,10 +130,10 @@ def _h3_vsa_sm100a_from_mask_compat( variable_block_sizes: torch.Tensor, ) -> torch.Tensor: """Compile-safe mask adapter for kernel wheels predating the mask API.""" - if _sm100a is None or map_to_index is None: - raise RuntimeError("The sm100a compatibility route requires the raw kernel and map_to_index") + if _ptx_blackwell is None or map_to_index is None: + raise RuntimeError("The ptx_blackwell compatibility route requires the raw kernel and map_to_index") q2k_idx, q2k_num = map_to_index(block_map) - out, _ = _sm100a.block_sparse_attn_sm100a( + out, _ = _ptx_blackwell.block_sparse_attn_ptx_blackwell( q, k, v, @@ -145,8 +145,8 @@ def _h3_vsa_sm100a_from_mask_compat( return out -@torch.library.register_fake("fastvideo::h3_vsa_sm100a_from_mask_compat") -def _h3_vsa_sm100a_from_mask_compat_fake( +@torch.library.register_fake("fastvideo::h3_vsa_ptx_blackwell_from_mask_compat") +def _h3_vsa_ptx_blackwell_from_mask_compat_fake( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -157,12 +157,13 @@ def _h3_vsa_sm100a_from_mask_compat_fake( return torch.empty_like(q) -def _sm100a_has_compile_safe_mask_route(sm100a_mod: Any) -> bool: - return (callable(getattr(sm100a_mod, "block_sparse_attn_sm100a_from_mask", None)) - or (callable(getattr(sm100a_mod, "block_sparse_attn_sm100a", None)) and map_to_index is not None)) +def _ptx_blackwell_has_compile_safe_mask_route(ptx_blackwell_mod: Any) -> bool: + return (callable(getattr(ptx_blackwell_mod, "block_sparse_attn_ptx_blackwell_from_mask", None)) + or (callable(getattr(ptx_blackwell_mod, "block_sparse_attn_ptx_blackwell", None)) + and map_to_index is not None)) -def _sm100a_from_mask( +def _ptx_blackwell_from_mask( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -170,10 +171,10 @@ def _sm100a_from_mask( variable_block_sizes: torch.Tensor, ) -> tuple[torch.Tensor, None]: """Use the native mask entry when installed, otherwise the local adapter.""" - native = getattr(_sm100a, "block_sparse_attn_sm100a_from_mask", None) + native = getattr(_ptx_blackwell, "block_sparse_attn_ptx_blackwell_from_mask", None) if callable(native): return native(q, k, v, block_map, variable_block_sizes) - return _h3_vsa_sm100a_from_mask_compat(q, k, v, block_map, variable_block_sizes), None + return _h3_vsa_ptx_blackwell_from_mask_compat(q, k, v, block_map, variable_block_sizes), None def token_tile_and_valid(variable_block_sizes: torch.Tensor, @@ -359,8 +360,8 @@ class MiniMaxH3VSAMetadata(AttentionMetadata): video_tile_spans: tuple[tuple[int, int], ...] span_sparsities: tuple[float, ...] # tokens per tile (256, 128 or 64); selects the tile geometry AND the - # kernel route in forward() (256 -> VSA-256 CuTe/Triton, 128 -> sm_100a - # CUDA, 64 -> native Triton or opt-in sm_100a CUDA) + # kernel route in forward() (256 -> VSA-256 CuTe/Triton, 128 -> + # ptx_blackwell, 64 -> native Triton or opt-in ptx_blackwell) tile_elems: int = _TILE_ELEMS # layers forced dense regardless of sparsity (probe-guided opt-outs) dense_layers: tuple[int, ...] = () @@ -520,21 +521,22 @@ def _build_block_mask( return mask -def _sm100a_unavailable_reason(sm100a_mod: Any, query_bhsd: torch.Tensor, variable_block_sizes: torch.Tensor, - grad_mode: bool) -> str | None: +def _ptx_blackwell_unavailable_reason(ptx_blackwell_mod: Any, query_bhsd: torch.Tensor, + variable_block_sizes: torch.Tensor, grad_mode: bool) -> str | None: """Why the opt-in data-center Blackwell route cannot run here, or None if it can. Pure decision logic, split out so the routing is unit-testable without a - GPU or the compiled extension (tests substitute ``sm100a_mod``). Order + GPU or the compiled extension (tests substitute ``ptx_blackwell_mod``). Order matters only for the message: the cheapest, most actionable reason first. """ - if sm100a_mod is None: - return "fastvideo_kernel.block_sparse_attn_sm100a is not installed" + if ptx_blackwell_mod is None: + return "fastvideo_kernel.block_sparse_attn_ptx_blackwell is not installed" if grad_mode: return "inputs require grad; this backend keeps grad paths on Triton" - if not sm100a_mod.is_supported(query_bhsd, variable_block_sizes): - return ("block_sparse_attn_sm100a.is_supported returned False (needs an sm_100 or sm_103 device, a built " - "extension, bf16, head_dim 128, an even tile count, and integer tile sizes)") + if not ptx_blackwell_mod.is_supported(query_bhsd, variable_block_sizes): + return ( + "block_sparse_attn_ptx_blackwell.is_supported returned False (needs an sm_100 or sm_103 device, a built " + "extension, bf16, head_dim 128, an even tile count, and integer tile sizes)") return None @@ -559,13 +561,13 @@ class MiniMaxH3VSAImpl(AttentionImpl): # Generic torch.compile must not specialize the shared VSA forward on # the Python ``layer_idx`` value of each of H3's 50 blocks. This # tensor is prepared after weights load and drives only the compiled - # dense-layer decision; it does not opt the module into sm_100a. + # dense-layer decision; it does not opt the module into ptx_blackwell. self._compile_layer_idx: torch.Tensor | None = None # None means the regional-compile preparation hook has not run. The # eager path deliberately ignores this cache and preserves its # request-time env/probe/fallback behavior; only Dynamo capture reads # the prepared, static route. - self._regional_compile_sm100a_enabled: bool | None = None + self._regional_compile_ptx_blackwell_enabled: bool | None = None # Tile-128 route per (device, dtype, head size): None when the CUDA # kernel can run it, else why not. The kernel's predicate otherwise # depends only on the tile-128 buffer contract (contiguous BHSD, @@ -578,7 +580,7 @@ class MiniMaxH3VSAImpl(AttentionImpl): self._compile_layer_idx = torch.tensor(self.layer_idx, device=device, dtype=torch.int64) def prepare_for_regional_compile(self, device: torch.device) -> str | None: - """Resolve the inference-only sm_100a route before fullgraph capture. + """Resolve the inference-only ptx_blackwell route before fullgraph capture. The ordinary eager route probes the environment, extension, device, and tensor contract at every call so it can warn and fall back. Those @@ -589,35 +591,35 @@ class MiniMaxH3VSAImpl(AttentionImpl): """ if self._compile_layer_idx is None: self.prepare_for_compile(device) - requested = envs.FASTVIDEO_VSA_SM100A.get() + requested = envs.FASTVIDEO_VSA_PTX_BLACKWELL.get() enabled = False - reason = None if requested else f"{VSA_SM100A_ENV}=1 is required for compile-safe VSA-H3 attention" + reason = None if requested else f"{VSA_PTX_BLACKWELL_ENV}=1 is required for compile-safe VSA-H3 attention" if requested: - if _sm100a is None: - reason = "fastvideo_kernel.block_sparse_attn_sm100a is not installed" - elif not _sm100a_has_compile_safe_mask_route(_sm100a): - reason = ("neither a native block_sparse_attn_sm100a_from_mask entry nor the raw sm100a " + if _ptx_blackwell is None: + reason = "fastvideo_kernel.block_sparse_attn_ptx_blackwell is not installed" + elif not _ptx_blackwell_has_compile_safe_mask_route(_ptx_blackwell): + reason = ("neither a native block_sparse_attn_ptx_blackwell_from_mask entry nor the raw ptx_blackwell " "kernel plus map_to_index compatibility route is installed") else: - # Two 64-token blocks exercise the exact sm_100a inference + # Two 64-token blocks exercise the exact ptx_blackwell inference # specialization while keeping the one-time probe tiny. The # kernel predicate checks extension presence, CUDA capability, # dtype/layout, head size, block size, and even block count # without reading metadata tensor contents. probe_query = torch.empty((1, 1, 128, self.head_size), device=device, dtype=torch.bfloat16) probe_block_sizes = torch.full((2, ), 64, device=device, dtype=torch.int32) - reason = _sm100a_unavailable_reason( - _sm100a, + reason = _ptx_blackwell_unavailable_reason( + _ptx_blackwell, probe_query, probe_block_sizes, grad_mode=False, ) enabled = reason is None - self._regional_compile_sm100a_enabled = enabled + self._regional_compile_ptx_blackwell_enabled = enabled if enabled: route = ("native fastvideo-kernel mask entry" if callable( - getattr(_sm100a, "block_sparse_attn_sm100a_from_mask", None)) else + getattr(_ptx_blackwell, "block_sparse_attn_ptx_blackwell_from_mask", None)) else "FastVideo compatibility mask adapter") logger.info_once(f"VSA-H3 regional compile mask route: {route}") if requested and reason is not None: @@ -632,7 +634,7 @@ class MiniMaxH3VSAImpl(AttentionImpl): sites in ``forward()`` read it immediately). A grad-tracking forward instead receives a fresh buffer and leaves the holder untouched, so the builder never retains autograd state across steps. Odd no-grad - sm100a requests (tile 128 always, tile 64 when opted in) carry one + ptx_blackwell requests (tile 128 always, tile 64 when opted in) carry one additional all-zero tile internally; metadata and all observable outputs retain the logical geometry. """ @@ -643,18 +645,19 @@ class MiniMaxH3VSAImpl(AttentionImpl): n_tiles = attn_metadata.variable_block_sizes.numel() grad_mode = torch.is_grad_enabled() and x.requires_grad compiling = torch.compiler.is_compiling() - regional_compiling = compiling and self._regional_compile_sm100a_enabled is True + regional_compiling = compiling and self._regional_compile_ptx_blackwell_enabled is True if regional_compiling: - sm100a_requested = bool(self._regional_compile_sm100a_enabled) + ptx_blackwell_requested = bool(self._regional_compile_ptx_blackwell_enabled) elif compiling: # Training/generic compile keeps the long-standing Triton route. - sm100a_requested = False + ptx_blackwell_requested = False else: - sm100a_requested = envs.FASTVIDEO_VSA_SM100A.get() - # Tile 128 has no route other than the sm100a CUDA kernel. - sm100a_route = attn_metadata.tile_elems == 128 or (attn_metadata.tile_elems == 64 and sm100a_requested) - needs_sm100a_pair = n_tiles % 2 != 0 and not grad_mode and sm100a_route - kernel_tiles = n_tiles + int(needs_sm100a_pair) + ptx_blackwell_requested = envs.FASTVIDEO_VSA_PTX_BLACKWELL.get() + # Tile 128 has no route other than the ptx_blackwell CUDA kernel. + ptx_blackwell_route = attn_metadata.tile_elems == 128 or (attn_metadata.tile_elems == 64 + and ptx_blackwell_requested) + needs_ptx_blackwell_pair = n_tiles % 2 != 0 and not grad_mode and ptx_blackwell_route + kernel_tiles = n_tiles + int(needs_ptx_blackwell_pair) target_shape = (x.shape[0], kernel_tiles * attn_metadata.tile_elems, x.shape[-2], x.shape[-1]) # A grad-tracking forward must not reuse the builder-owned buffer. The @@ -675,17 +678,17 @@ class MiniMaxH3VSAImpl(AttentionImpl): holder = attn_metadata.tile_buf_holder if holder is None: raise RuntimeError("VSA-H3 metadata has no builder-owned tile buffer holder") - if (not compiling and attn_metadata.tile_elems in _SM100A_TILE_ELEMS + if (not compiling and attn_metadata.tile_elems in _PTX_BLACKWELL_TILE_ELEMS and envs.FASTVIDEO_H3_VSA_HEADS_FIRST_TILE.get() and supports_heads_first_scatter(x, attn_metadata.untile_combined_index)): - return self._tile_heads_first(x, attn_metadata, holder, kernel_tiles, needs_sm100a_pair) + return self._tile_heads_first(x, attn_metadata, holder, kernel_tiles, needs_ptx_blackwell_pair) buffer_matches = (holder.buffer is not None and holder.buffer.shape == target_shape and holder.buffer.dtype == x.dtype and holder.buffer.device == x.device) if buffer_matches and holder.untile_geometry is not attn_metadata.untile_combined_index: holder.buffer.zero_() holder.buffer = scatter_into_tile_buf(x, target_shape, attn_metadata.untile_combined_index, holder.buffer) holder.untile_geometry = attn_metadata.untile_combined_index - if needs_sm100a_pair: + if needs_ptx_blackwell_pair: # A prior even geometry can reuse this allocation and may have # written the last tile as logical data. holder.buffer[:, n_tiles * attn_metadata.tile_elems:].zero_() @@ -693,7 +696,7 @@ class MiniMaxH3VSAImpl(AttentionImpl): @staticmethod def _tile_heads_first(x: torch.Tensor, attn_metadata: MiniMaxH3VSAMetadata, holder: _MiniMaxH3VSATileBufferHolder, - kernel_tiles: int, needs_sm100a_pair: bool) -> torch.Tensor: + kernel_tiles: int, needs_ptx_blackwell_pair: bool) -> torch.Tensor: """``tile()`` into a heads-first buffer; returns its ``[B, S_pad, H, D]`` view. The 64/128-token kernels transpose the tiled tensor to BHSD and make it @@ -710,7 +713,7 @@ class MiniMaxH3VSAImpl(AttentionImpl): scatter_rows_heads_first(x, attn_metadata.untile_combined_index, buffer) holder.heads_first_buffer = buffer holder.heads_first_geometry = attn_metadata.untile_combined_index - if needs_sm100a_pair: + if needs_ptx_blackwell_pair: buffer[:, :, n_tiles * attn_metadata.tile_elems:].zero_() return buffer.transpose(1, 2) @@ -729,7 +732,7 @@ class MiniMaxH3VSAImpl(AttentionImpl): attn_metadata: MiniMaxH3VSAMetadata, ) -> torch.Tensor: compiling = torch.compiler.is_compiling() - regional_compiling = compiling and self._regional_compile_sm100a_enabled is True + regional_compiling = compiling and self._regional_compile_ptx_blackwell_enabled is True tile_elems = attn_metadata.tile_elems if regional_compiling and tile_elems != 64: @@ -746,18 +749,18 @@ class MiniMaxH3VSAImpl(AttentionImpl): # environment switch while Dynamo captures a regional full graph. # The metadata always describes the trained logical geometry. # ``tile()`` may append exactly one transport-only partner for an odd - # tile-64/128 sm100a call. Keep score selection and the gate branch on + # tile-64/128 ptx_blackwell call. Keep score selection and the gate branch on # the logical prefix, and reject every other shape before a kernel sees it. n_tiles = attn_metadata.variable_block_sizes.numel() logical_seq_len = n_tiles * tile_elems pair_pad_seq_len = logical_seq_len + tile_elems - pair_pad_is_valid = tile_elems in _SM100A_TILE_ELEMS and n_tiles % 2 != 0 + pair_pad_is_valid = tile_elems in _PTX_BLACKWELL_TILE_ELEMS and n_tiles % 2 != 0 allowed_seq_lengths = (logical_seq_len, pair_pad_seq_len) if pair_pad_is_valid else (logical_seq_len, ) if query.shape[1] not in allowed_seq_lengths: - expected = (f"the logical length {logical_seq_len} or one sm100a partner tile " + expected = (f"the logical length {logical_seq_len} or one ptx_blackwell partner tile " f"({pair_pad_seq_len})" if pair_pad_is_valid else f"the logical length {logical_seq_len}") raise ValueError(f"VSA-H3 tiled query has length {query.shape[1]}, expected {expected}.") - has_sm100a_pair = query.shape[1] == pair_pad_seq_len + has_ptx_blackwell_pair = query.shape[1] == pair_pad_seq_len for name, tensor in (("key", key), ("value", value)): if tensor.shape[1] != query.shape[1]: raise ValueError(f"VSA-H3 tiled {name} length {tensor.shape[1]} does not match query " @@ -809,7 +812,7 @@ class MiniMaxH3VSAImpl(AttentionImpl): # preserves the eager dense-layer contract without a Python branch. mask = mask | force_dense - if tile_elems in _SM100A_TILE_ELEMS: + if tile_elems in _PTX_BLACKWELL_TILE_ELEMS: # Native 64/128-token path: the block map is already at the # kernels' granularity. These entries take BHSD ([B, H, S_pad, D]); # mirror block_sparse_attn_256_bshd's Triton branch and transpose @@ -828,26 +831,26 @@ class MiniMaxH3VSAImpl(AttentionImpl): if not sm89_strided: q_bhsd, k_bhsd, v_bhsd = (t.contiguous() for t in (q_bhsd, k_bhsd, v_bhsd)) - sm100a_mask = mask - sm100a_variable_block_sizes = attn_metadata.variable_block_sizes - if has_sm100a_pair: + ptx_blackwell_mask = mask + ptx_blackwell_variable_block_sizes = attn_metadata.variable_block_sizes + if has_ptx_blackwell_pair: # The synthetic tile is neither a logical query nor key. Its # all-False row yields q2k_num=0, the all-False column keeps it # out of real rows, and vbs=0 masks all of its key slots. - sm100a_mask = torch.nn.functional.pad(mask, (0, 1, 0, 1), value=False) - sm100a_variable_block_sizes = torch.nn.functional.pad( + ptx_blackwell_mask = torch.nn.functional.pad(mask, (0, 1, 0, 1), value=False) + ptx_blackwell_variable_block_sizes = torch.nn.functional.pad( attn_metadata.variable_block_sizes, (0, 1), value=0, ) - # sm_100a CUDA forward, opt-in at tile 64 and the only route at + # ptx_blackwell forward, opt-in at tile 64 and the only route at # tile 128. Forward-only: tile-64 grad-tracking calls stay on Triton # so autograd keeps the Triton fwd+bwd pairing untouched. The # kernel does return an LSE in Triton's M format, so a future # fwd/bwd pairing is possible, but it is not built here. grad_mode = torch.is_grad_enabled() and (query.requires_grad or key.requires_grad or value.requires_grad) - use_sm100a = False + use_ptx_blackwell = False if tile_elems == 128: # The CUDA kernel's 128-token instantiation is the only tile-128 # implementation, so it needs no opt-in and never falls back. @@ -855,74 +858,78 @@ class MiniMaxH3VSAImpl(AttentionImpl): raise NotImplementedError("VSA-H3 tile 128 supports no-grad inference forwards only.") route_key = (q_bhsd.device, q_bhsd.dtype, int(q_bhsd.shape[-1])) if route_key not in self._tile128_route: - reason = _sm100a_unavailable_reason(_sm100a, q_bhsd, sm100a_variable_block_sizes, grad_mode) + reason = _ptx_blackwell_unavailable_reason(_ptx_blackwell, q_bhsd, + ptx_blackwell_variable_block_sizes, grad_mode) if reason is None and map_to_index is None: reason = "fastvideo_kernel.triton_kernels.index (map_to_index) is not importable" self._tile128_route[route_key] = reason reason = self._tile128_route[route_key] if reason is not None: - raise RuntimeError(f"VSA-H3 tile 128 requires the sm_100a/sm_103a CUDA block-sparse kernel: " + raise RuntimeError(f"VSA-H3 tile 128 requires the ptx_blackwell block-sparse kernel: " f"{reason}. Use VSA_tile_size=64 or 256 elsewhere.") - use_sm100a = True + use_ptx_blackwell = True elif regional_compiling: - if self._regional_compile_sm100a_enabled is None: + if self._regional_compile_ptx_blackwell_enabled is None: raise RuntimeError( - "VSA-H3 sm_100a routing was not resolved before torch.compile; " + "VSA-H3 ptx_blackwell routing was not resolved before torch.compile; " "call prepare_for_regional_compile(device) on every MiniMaxH3VSAImpl after loading weights.") # The preparation probe established module/device/kernel # support. Keep only static tensor/geometry facts here; no # env access, device-capability query, or is_supported call may # enter the Dynamo graph. - if not (self._regional_compile_sm100a_enabled and not grad_mode and q_bhsd.dtype == torch.bfloat16 - and q_bhsd.shape[-1] == 128 and sm100a_variable_block_sizes.numel() % 2 == 0): + if not (self._regional_compile_ptx_blackwell_enabled and not grad_mode + and q_bhsd.dtype == torch.bfloat16 and q_bhsd.shape[-1] == 128 + and ptx_blackwell_variable_block_sizes.numel() % 2 == 0): raise RuntimeError( - "VSA-H3 regional fullgraph compile requires the prepared sm_100a BF16/head-128 route " + "VSA-H3 regional fullgraph compile requires the prepared ptx_blackwell BF16/head-128 route " "on a supported device; disable inference_torch_compile for this request.") - use_sm100a = True - elif not compiling and envs.FASTVIDEO_VSA_SM100A.get(): - reason = _sm100a_unavailable_reason(_sm100a, q_bhsd, sm100a_variable_block_sizes, grad_mode) + use_ptx_blackwell = True + elif not compiling and envs.FASTVIDEO_VSA_PTX_BLACKWELL.get(): + reason = _ptx_blackwell_unavailable_reason(_ptx_blackwell, q_bhsd, ptx_blackwell_variable_block_sizes, + grad_mode) if reason is None and map_to_index is None: reason = "fastvideo_kernel.triton_kernels.index (map_to_index) is not importable" if reason is None: - use_sm100a = True + use_ptx_blackwell = True else: - logger.warning_once(f"{VSA_SM100A_ENV}=1 but falling back to the Triton-64 kernels: {reason}") + logger.warning_once( + f"{VSA_PTX_BLACKWELL_ENV}=1 but falling back to the Triton-64 kernels: {reason}") - if use_sm100a: + if use_ptx_blackwell: # prepare_for_regional_compile logs the selected mask route before # capture. Logging from this branch would itself break a # ``fullgraph=True`` forward. if not compiling: logger.info_once( - f"MiniMax-H3 VSA tile-{tile_elems} forward: using the sm100a/sm103a CUDA block-sparse kernel") + f"MiniMax-H3 VSA tile-{tile_elems} forward: using the ptx_blackwell block-sparse kernel") if regional_compiling: # The compile-safe wrapper keeps both Triton mask # compaction and the raw pybind launch behind one # fake-backed custom-op boundary. - out_bhsd, _ = _sm100a_from_mask( + out_bhsd, _ = _ptx_blackwell_from_mask( q_bhsd, k_bhsd, v_bhsd, - sm100a_mask, - sm100a_variable_block_sizes, + ptx_blackwell_mask, + ptx_blackwell_variable_block_sizes, ) else: # Preserve the established eager/index-native route and # compatibility with older kernel wheels. Per-row counts # are non-uniform (prefix queries are dense; video queries # run prefix+top-k); the kernel reads each query tile's own count. - q2k_idx, q2k_num = map_to_index(sm100a_mask) - out_bhsd, _ = _sm100a.block_sparse_attn_sm100a( + q2k_idx, q2k_num = map_to_index(ptx_blackwell_mask) + out_bhsd, _ = _ptx_blackwell.block_sparse_attn_ptx_blackwell( q_bhsd, k_bhsd, v_bhsd, q2k_idx, q2k_num, - sm100a_variable_block_sizes.to(torch.int32), + ptx_blackwell_variable_block_sizes.to(torch.int32), need_lse=False, ) else: - if has_sm100a_pair: + if has_ptx_blackwell_pair: q_bhsd = q_bhsd[:, :, :logical_seq_len] k_bhsd = k_bhsd[:, :, :logical_seq_len] v_bhsd = v_bhsd[:, :, :logical_seq_len] @@ -945,7 +952,7 @@ class MiniMaxH3VSAImpl(AttentionImpl): mask, attn_metadata.variable_block_sizes, ) - if has_sm100a_pair and use_sm100a: + if has_ptx_blackwell_pair and use_ptx_blackwell: out_bhsd = out_bhsd[:, :, :logical_seq_len] out = out_bhsd.transpose(1, 2).contiguous() # Fine attention is complete. Release its layout copies before the diff --git a/fastvideo/envs.py b/fastvideo/envs.py index f14d63b04..78c9dba60 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -346,11 +346,12 @@ FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN = EnvBool( category="attention", doc="MiniMax-H3 dense DiT self-attention uses the FlashAttention-4 packed-varlen entry point. This changes the " "floating-point reduction order, so it is an inference-only opt-in.") -FASTVIDEO_VSA_SM100A = EnvBool( +FASTVIDEO_VSA_PTX_BLACKWELL = EnvBool( False, category="attention", - doc="VIDEO_SPARSE_ATTN_H3 sends no-grad tile-64 forwards to the data-center Blackwell (sm_100a) kernel. " - "fastvideo-kernel reads the same variable with the same rule.") + doc="VIDEO_SPARSE_ATTN_H3 sends no-grad tile-64 forwards to the data-center Blackwell (sm_100a / sm_103a) kernel. " + "fastvideo-kernel reads the same variable with the same rule.", + deprecated_names=("FASTVIDEO_VSA_SM100A", )) FASTVIDEO_VSA_TRITON = EnvBool( False, category="attention", diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index f6f574021..082e39abb 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -283,7 +283,7 @@ class FastVideoArgs: # compile ported from hao-ai-lab/FastVideo#1718. Applied by the loader # right after the transformer loads, with fullgraph=True and inductor # options {emulate_precision_casts: True} injected (no user kwargs - # needed). MiniMax-H3 VSA is supported only by its compile-safe sm_100a + # needed). MiniMax-H3 VSA is supported only by its compile-safe ptx_blackwell # tile-64 inference route; other VSA routes degrade the transformer to # eager with one warning. Dense FA2/FA3/FA4 inference uses compile-visible # custom-op boundaries. Opt-in via FASTVIDEO_INFERENCE_TORCH_COMPILE=1 (folded in @@ -297,7 +297,7 @@ class FastVideoArgs: # VSA parameters. None means unset: __post_init__ fills the checkpoint's # trained value when the checkpoint fixes one, else 0.0 and 256. VSA_sparsity: float | None = None # inference/validation sparsity - VSA_tile_size: int | None = None # VSA-H3 tokens per tile (256, 128, 64); 128 = sm_100a CUDA only + VSA_tile_size: int | None = None # VSA-H3 tokens per tile (256, 128, 64); 128 = ptx_blackwell only # V-MoBA parameters moba_config_path: str | None = None @@ -896,7 +896,7 @@ class FastVideoArgs: type=int, default=FastVideoArgs.VSA_tile_size, help="VSA-H3 tile size in tokens (256, 128 or 64; default: the checkpoint's trained value, else 256); " - "64 runs the native Triton block-sparse path, 128 requires the sm_100a/sm_103a CUDA kernel", + "64 runs the native Triton block-sparse path, 128 requires the ptx_blackwell kernel", ) # Master port for distributed training/inference diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index 070d13637..4eb579d9f 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -984,7 +984,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT): layer identity so repeated blocks share one Dynamo graph. Generic and training compile retain their established attention dispatch; only the inference loader's separate ``prepare_for_regional_compile`` hook - may preselect the inference-only sm_100a path. + may preselect the inference-only ptx_blackwell path. The inference-only Triton fusions expose fake-backed custom operators, so Dynamo can keep them active as opaque nodes inside each fullgraph diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 974bbdb92..afd76ba5e 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -639,7 +639,7 @@ def _regional_compile_unsupported_reason( The legacy VSA backend remains outside the fullgraph support envelope. MiniMax H3's VSA backend is supported only through the inference-only - sm_100a tile-64 route; its regional hook resolves loaded compression + ptx_blackwell tile-64 route; its regional hook resolves loaded compression gates and probes the kernel before block capture. """ try: @@ -698,9 +698,9 @@ def _regional_compile_unsupported_reason( if envs.FASTVIDEO_H3_VSA_PROBE.get(): return ("FASTVIDEO_H3_VSA_PROBE records tensors and files from the VSA-H3 attention body, which " "regional fullgraph compile cannot capture; this model stays eager") - if not envs.FASTVIDEO_VSA_SM100A.get(): - return ("VIDEO_SPARSE_ATTN_H3 regional compile requires the compile-safe sm_100a route " - "(FASTVIDEO_VSA_SM100A=1); Triton/CuTe VSA stays eager") + if not envs.FASTVIDEO_VSA_PTX_BLACKWELL.get(): + return ("VIDEO_SPARSE_ATTN_H3 regional compile requires the compile-safe ptx_blackwell route " + "(FASTVIDEO_VSA_PTX_BLACKWELL=1); Triton/CuTe VSA stays eager") if vsa_tile_size != 64: return ("VIDEO_SPARSE_ATTN_H3 regional compile requires VSA_tile_size=64; " f"got {vsa_tile_size!r}, so tile-256/CuTe VSA stays eager") diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py index 94e050be1..e04ac8d33 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py @@ -270,7 +270,7 @@ class MiniMaxH3DenoisingStage(PipelineStage): vsa_dense_layers = tuple(batch.extra.get("vsa_dense_layers", ())) vsa_dense_first_n = int(batch.extra.get("vsa_dense_first_n_steps", 0)) # Run-level tile geometry (256 default, 64 = native Triton path, - # 128 = sm_100a CUDA), plumbed like the run-level sparsity; the + # 128 = ptx_blackwell), plumbed like the run-level sparsity; the # builder validates the value against VSA_H3_TILE_SHAPES. vsa_tile_size = int(fastvideo_args.VSA_tile_size) if vsa_ref_keep_rate is not None: diff --git a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py b/fastvideo/tests/attention/test_vsa_h3_blackwell_route.py similarity index 76% rename from fastvideo/tests/attention/test_vsa_h3_sm100a_route.py rename to fastvideo/tests/attention/test_vsa_h3_blackwell_route.py index 0eaa79130..cccfbcdd1 100644 --- a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py +++ b/fastvideo/tests/attention/test_vsa_h3_blackwell_route.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -"""Checks for the VSA-H3 tile-64 sm_100a/sm_103a route and odd-tile transport. +"""Checks for the VSA-H3 tile-64 ptx_blackwell route and odd-tile transport. -The opt-in third kernel route (``FASTVIDEO_VSA_SM100A=1``) must (a) stay off by +The opt-in third kernel route (``FASTVIDEO_VSA_PTX_BLACKWELL=1``) must (a) stay off by default, (b) engage only when the extension is present, the device qualifies, and the forward carries no grad, (c) add one zero-valid partner for an odd logical tile count, and (d) strip that partner before any fallback or returned @@ -16,8 +16,9 @@ import torch import fastvideo.attention.backends.video_sparse_attn_h3 as vsa_h3 import fastvideo.envs as envs import fastvideo.logger as fastvideo_logger -from fastvideo.attention.backends.video_sparse_attn_h3 import (VSA_SM100A_ENV, MiniMaxH3VSAImpl, - MiniMaxH3VSAMetadataBuilder, _sm100a_unavailable_reason) +from fastvideo.attention.backends.video_sparse_attn_h3 import (VSA_PTX_BLACKWELL_ENV, MiniMaxH3VSAImpl, + MiniMaxH3VSAMetadataBuilder, + _ptx_blackwell_unavailable_reason) # Small tile-64 geometry: 2 prefix segments + a (4,4,8)-token video grid. _SPEC = dict(raw_latent_shape=(4, 8, 16), patch_size=(1, 2, 2), prefix_segments=(70, 30)) @@ -52,8 +53,8 @@ def _tiled_qkv(meta, requires_grad=False): torch.randn(1, s_pad, _HEADS, _DIM, dtype=torch.bfloat16, requires_grad=requires_grad) for _ in range(3)) -class _FakeSm100a: - """Stands in for fastvideo_kernel.block_sparse_attn_sm100a.""" +class _FakePtxBlackwell: + """Stands in for fastvideo_kernel.block_sparse_attn_ptx_blackwell.""" def __init__(self, supported=True): self.supported = supported @@ -65,12 +66,12 @@ class _FakeSm100a: self.support_calls.append((q, variable_block_sizes)) return self.supported - def block_sparse_attn_sm100a(self, q, k, v, q2k_idx, q2k_num, variable_block_sizes, need_lse=True): + def block_sparse_attn_ptx_blackwell(self, q, k, v, q2k_idx, q2k_num, variable_block_sizes, need_lse=True): self.calls.append(dict(q=q, q2k_idx=q2k_idx, q2k_num=q2k_num, vbs=variable_block_sizes, need_lse=need_lse)) return q.clone(), None - def block_sparse_attn_sm100a_from_mask(self, q, k, v, block_map, variable_block_sizes): + def block_sparse_attn_ptx_blackwell_from_mask(self, q, k, v, block_map, variable_block_sizes): self.mask_calls.append(dict(q=q, block_map=block_map, vbs=variable_block_sizes)) return q.clone(), None @@ -100,9 +101,9 @@ class _FakeTriton: @pytest.fixture() def routed(monkeypatch): """Backend with both kernel entries faked; returns (fakes, run).""" - fake_sm = _FakeSm100a() + fake_sm = _FakePtxBlackwell() fake_triton = _FakeTriton() - monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) + monkeypatch.setattr(vsa_h3, "_ptx_blackwell", fake_sm) monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", fake_triton) monkeypatch.setattr(vsa_h3, "map_to_index", _fake_map_to_index) meta = _build_meta() @@ -118,24 +119,24 @@ def routed(monkeypatch): def test_reason_covers_every_precondition(): q = torch.randn(1, _HEADS, 128, _DIM) vbs = torch.full((2, ), 64, dtype=torch.long) - assert "not installed" in _sm100a_unavailable_reason(None, q, vbs, grad_mode=False) - ok = _FakeSm100a(supported=True) - assert "require grad" in _sm100a_unavailable_reason(ok, q, vbs, grad_mode=True) - bad = _FakeSm100a(supported=False) - assert "is_supported" in _sm100a_unavailable_reason(bad, q, vbs, grad_mode=False) - assert _sm100a_unavailable_reason(ok, q, vbs, grad_mode=False) is None + assert "not installed" in _ptx_blackwell_unavailable_reason(None, q, vbs, grad_mode=False) + ok = _FakePtxBlackwell(supported=True) + assert "require grad" in _ptx_blackwell_unavailable_reason(ok, q, vbs, grad_mode=True) + bad = _FakePtxBlackwell(supported=False) + assert "is_supported" in _ptx_blackwell_unavailable_reason(bad, q, vbs, grad_mode=False) + assert _ptx_blackwell_unavailable_reason(ok, q, vbs, grad_mode=False) is None def test_prepare_for_regional_compile_resolves_supported_route(monkeypatch, env_overrides): - fake_sm = _FakeSm100a(supported=True) - monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + fake_sm = _FakePtxBlackwell(supported=True) + monkeypatch.setattr(vsa_h3, "_ptx_blackwell", fake_sm) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) unsupported = impl.prepare_for_regional_compile(torch.device("cpu")) assert unsupported is None - assert impl._regional_compile_sm100a_enabled is True + assert impl._regional_compile_ptx_blackwell_enabled is True assert len(fake_sm.support_calls) == 1 probe_q, probe_vbs = fake_sm.support_calls[0] assert probe_q.shape == (1, 1, 128, _DIM) @@ -147,29 +148,29 @@ def test_prepare_for_regional_compile_resolves_supported_route(monkeypatch, env_ def test_prepare_for_regional_compile_env_off_skips_probe(monkeypatch, env_overrides): - fake_sm = _FakeSm100a(supported=True) - monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(None)) + fake_sm = _FakePtxBlackwell(supported=True) + monkeypatch.setattr(vsa_h3, "_ptx_blackwell", fake_sm) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(None)) impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) unsupported = impl.prepare_for_regional_compile(torch.device("cpu")) assert unsupported is not None - assert VSA_SM100A_ENV in unsupported + assert VSA_PTX_BLACKWELL_ENV in unsupported assert impl._compile_layer_idx is not None - assert impl._regional_compile_sm100a_enabled is False + assert impl._regional_compile_ptx_blackwell_enabled is False assert fake_sm.support_calls == [] def test_prepare_for_regional_compile_requires_mask_entry(monkeypatch, env_overrides): - class _IndexOnlySm100a: + class _IndexOnlyPtxBlackwell: def is_supported(self, q, variable_block_sizes): raise AssertionError("missing mask entry must be rejected before the support probe") - monkeypatch.setattr(vsa_h3, "_sm100a", _IndexOnlySm100a()) - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + monkeypatch.setattr(vsa_h3, "_ptx_blackwell", _IndexOnlyPtxBlackwell()) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) warnings = [] monkeypatch.setattr(vsa_h3.logger, "warning_once", warnings.append) impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) @@ -178,19 +179,19 @@ def test_prepare_for_regional_compile_requires_mask_entry(monkeypatch, env_overr assert unsupported is not None assert impl._compile_layer_idx is not None - assert impl._regional_compile_sm100a_enabled is False + assert impl._regional_compile_ptx_blackwell_enabled is False assert len(warnings) == 1 assert "compatibility route" in warnings[0] def test_mask_dispatch_uses_local_compatibility_route_for_older_kernel_wheel(monkeypatch): - class _RawOnlySm100a: + class _RawOnlyPtxBlackwell: - def block_sparse_attn_sm100a(self, *args, **kwargs): + def block_sparse_attn_ptx_blackwell(self, *args, **kwargs): raise AssertionError("raw entry must remain behind the compatibility custom op") - fake_sm = _RawOnlySm100a() - monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) + fake_sm = _RawOnlyPtxBlackwell() + monkeypatch.setattr(vsa_h3, "_ptx_blackwell", fake_sm) monkeypatch.setattr(vsa_h3, "map_to_index", lambda mask: (mask, mask)) calls = [] @@ -198,23 +199,23 @@ def test_mask_dispatch_uses_local_compatibility_route_for_older_kernel_wheel(mon calls.append((block_map, variable_block_sizes)) return q + 1 - monkeypatch.setattr(vsa_h3, "_h3_vsa_sm100a_from_mask_compat", compat) + monkeypatch.setattr(vsa_h3, "_h3_vsa_ptx_blackwell_from_mask_compat", compat) q = torch.zeros(1, 1, 128, _DIM) mask = torch.ones(1, 1, 2, 2, dtype=torch.bool) vbs = torch.full((2, ), 64, dtype=torch.int32) - out, lse = vsa_h3._sm100a_from_mask(q, q, q, mask, vbs) + out, lse = vsa_h3._ptx_blackwell_from_mask(q, q, q, mask, vbs) - assert vsa_h3._sm100a_has_compile_safe_mask_route(fake_sm) + assert vsa_h3._ptx_blackwell_has_compile_safe_mask_route(fake_sm) torch.testing.assert_close(out, q + 1) assert lse is None and calls == [(mask, vbs)] def test_prepared_route_fullgraph_avoids_eager_dispatch(monkeypatch, env_overrides): """Capture must not revisit env/capability checks or raw map_to_index.""" - fake_sm = _FakeSm100a(supported=True) - monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + fake_sm = _FakePtxBlackwell(supported=True) + monkeypatch.setattr(vsa_h3, "_ptx_blackwell", fake_sm) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) monkeypatch.setattr(vsa_h3, "probe_enabled", lambda: None) impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) impl.prepare_for_regional_compile(torch.device("cpu")) @@ -239,10 +240,10 @@ def test_prepared_route_fullgraph_avoids_eager_dispatch(monkeypatch, env_overrid return q.clone(), None monkeypatch.setattr(fake_sm, "is_supported", fail_eager_dispatch) - monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a", fail_eager_dispatch) - monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a_from_mask", compile_safe_from_mask) + monkeypatch.setattr(fake_sm, "block_sparse_attn_ptx_blackwell", fail_eager_dispatch) + monkeypatch.setattr(fake_sm, "block_sparse_attn_ptx_blackwell_from_mask", compile_safe_from_mask) monkeypatch.setattr(vsa_h3, "map_to_index", fail_eager_dispatch) - monkeypatch.setattr(vsa_h3, "_sm100a_unavailable_reason", fail_eager_dispatch) + monkeypatch.setattr(vsa_h3, "_ptx_blackwell_unavailable_reason", fail_eager_dispatch) monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", fail_eager_dispatch) def run(q, k, v): @@ -259,9 +260,9 @@ def test_prepared_route_fullgraph_avoids_eager_dispatch(monkeypatch, env_overrid def test_prepared_route_reuses_graph_across_layer_indices_and_preserves_dense_overrides(monkeypatch, env_overrides): """Fifty H3 blocks must not specialize the shared graph on layer_idx.""" - fake_sm = _FakeSm100a(supported=True) - monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + fake_sm = _FakePtxBlackwell(supported=True) + monkeypatch.setattr(vsa_h3, "_ptx_blackwell", fake_sm) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) monkeypatch.setattr(vsa_h3, "probe_enabled", lambda: None) meta = _build_meta(sparsity=0.5, dense_layers=(0, 17), prefix_segments=(64, 64)) assert meta.dense_layers_tensor.tolist() == [0, 17] @@ -289,10 +290,10 @@ def test_prepared_route_reuses_graph_across_layer_indices_and_preserves_dense_ov implementations.append(impl) monkeypatch.setattr(fake_sm, "is_supported", fail_eager_dispatch) - monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a", fail_eager_dispatch) - monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a_from_mask", compile_safe_from_mask) + monkeypatch.setattr(fake_sm, "block_sparse_attn_ptx_blackwell", fail_eager_dispatch) + monkeypatch.setattr(fake_sm, "block_sparse_attn_ptx_blackwell_from_mask", compile_safe_from_mask) monkeypatch.setattr(vsa_h3, "map_to_index", fail_eager_dispatch) - monkeypatch.setattr(vsa_h3, "_sm100a_unavailable_reason", fail_eager_dispatch) + monkeypatch.setattr(vsa_h3, "_ptx_blackwell_unavailable_reason", fail_eager_dispatch) monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", fail_eager_dispatch) torch._dynamo.reset() @@ -308,10 +309,10 @@ def test_prepared_route_reuses_graph_across_layer_indices_and_preserves_dense_ov def test_generic_compile_reuses_graph_across_layer_indices_and_stays_on_triton(monkeypatch, env_overrides): - """Pipeline compile must share one graph without selecting sm_100a.""" - fake_sm = _FakeSm100a(supported=True) - monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + """Pipeline compile must share one graph without selecting ptx_blackwell.""" + fake_sm = _FakePtxBlackwell(supported=True) + monkeypatch.setattr(vsa_h3, "_ptx_blackwell", fake_sm) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) monkeypatch.setattr(vsa_h3, "probe_enabled", lambda: None) meta = _build_meta(sparsity=0.5, dense_layers=(0, 17), prefix_segments=(64, 64)) q, k, v = _tiled_qkv(meta) @@ -320,14 +321,14 @@ def test_generic_compile_reuses_graph_across_layer_indices_and_stays_on_triton(m del k, v, variable_block_sizes return q + block_map.all().to(q.dtype), None - def fail_sm100a(*args, **kwargs): - raise AssertionError("generic torch.compile unexpectedly selected sm_100a") + def fail_ptx_blackwell(*args, **kwargs): + raise AssertionError("generic torch.compile unexpectedly selected ptx_blackwell") monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", fake_triton) - monkeypatch.setattr(fake_sm, "is_supported", fail_sm100a) - monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a", fail_sm100a) - monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a_from_mask", fail_sm100a) - monkeypatch.setattr(vsa_h3, "_sm100a_unavailable_reason", fail_sm100a) + monkeypatch.setattr(fake_sm, "is_supported", fail_ptx_blackwell) + monkeypatch.setattr(fake_sm, "block_sparse_attn_ptx_blackwell", fail_ptx_blackwell) + monkeypatch.setattr(fake_sm, "block_sparse_attn_ptx_blackwell_from_mask", fail_ptx_blackwell) + monkeypatch.setattr(vsa_h3, "_ptx_blackwell_unavailable_reason", fail_ptx_blackwell) implementations = [] for layer_idx in range(20): @@ -364,7 +365,7 @@ def test_generic_compile_reuses_graph_across_layer_indices_and_stays_on_triton(m def test_default_off_routes_triton(routed, env_overrides): fake_sm, fake_triton, run, _ = routed - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(None)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(None)) run() assert fake_triton.calls == 1 assert fake_sm.calls == [] @@ -372,7 +373,7 @@ def test_default_off_routes_triton(routed, env_overrides): def test_unprepared_training_compile_keeps_existing_triton_route(routed, monkeypatch, env_overrides): fake_sm, fake_triton, run, _ = routed - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) run(requires_grad=True) @@ -381,9 +382,9 @@ def test_unprepared_training_compile_keeps_existing_triton_route(routed, monkeyp assert fake_sm.calls == [] -def test_env_on_routes_sm100a_with_index_metadata(routed, monkeypatch, env_overrides): +def test_env_on_routes_ptx_blackwell_with_index_metadata(routed, monkeypatch, env_overrides): fake_sm, fake_triton, run, meta = routed - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) messages = [] monkeypatch.setattr(vsa_h3.logger, "info_once", messages.append) out = run() @@ -396,14 +397,14 @@ def test_env_on_routes_sm100a_with_index_metadata(routed, monkeypatch, env_overr assert call["q2k_idx"].shape[-1] == n_tiles and call["q2k_idx"].dtype == torch.int32 assert call["vbs"].dtype == torch.int32 assert call["need_lse"] is False - assert messages == ["MiniMax-H3 VSA tile-64 forward: using the sm100a/sm103a CUDA block-sparse kernel"] + assert messages == ["MiniMax-H3 VSA tile-64 forward: using the ptx_blackwell block-sparse kernel"] # BHSD kernel result comes back in the backend's BSHD layout assert out.shape == (1, n_tiles * 64, _HEADS, _DIM) -def test_sm100a_engagement_receipt_logs_once_without_stacklevel_conflict(routed, monkeypatch, env_overrides): +def test_ptx_blackwell_engagement_receipt_logs_once_without_stacklevel_conflict(routed, monkeypatch, env_overrides): fake_sm, fake_triton, run, _ = routed - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) records = [] def capture_log(level, msg, *args, **kwargs): @@ -420,17 +421,17 @@ def test_sm100a_engagement_receipt_logs_once_without_stacklevel_conflict(routed, assert fake_triton.calls == 0 assert len(fake_sm.calls) == 2 assert records == [(logging.INFO, - "MiniMax-H3 VSA tile-64 forward: using the sm100a/sm103a CUDA block-sparse kernel", (), + "MiniMax-H3 VSA tile-64 forward: using the ptx_blackwell block-sparse kernel", (), {"stacklevel": 2})] def test_env_on_grad_inputs_fall_back_to_triton(routed, env_overrides): fake_sm, fake_triton, run, _ = routed - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) run(requires_grad=True) assert fake_triton.calls == 1 assert fake_sm.calls == [] - # ...but the same process still routes no-grad forwards to sm_100a + # ...but the same process still routes no-grad forwards to ptx_blackwell run(requires_grad=False) assert len(fake_sm.calls) == 1 @@ -438,7 +439,7 @@ def test_env_on_grad_inputs_fall_back_to_triton(routed, env_overrides): def test_env_on_unsupported_warns_once_and_falls_back(routed, monkeypatch, env_overrides): fake_sm, fake_triton, run, _ = routed fake_sm.supported = False - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) warnings = [] monkeypatch.setattr(vsa_h3.logger, "warning_once", warnings.append) run() @@ -447,13 +448,13 @@ def test_env_on_unsupported_warns_once_and_falls_back(routed, monkeypatch, env_o assert fake_sm.calls == [] assert len(warnings) == 2 # warning_once dedups by message; both carry the same one line assert warnings[0] == warnings[1] - assert VSA_SM100A_ENV in warnings[0] and "is_supported" in warnings[0] + assert VSA_PTX_BLACKWELL_ENV in warnings[0] and "is_supported" in warnings[0] def test_env_on_missing_module_warns_and_falls_back(routed, monkeypatch, env_overrides): fake_sm, fake_triton, run, _ = routed - monkeypatch.setattr(vsa_h3, "_sm100a", None) - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + monkeypatch.setattr(vsa_h3, "_ptx_blackwell", None) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) warnings = [] monkeypatch.setattr(vsa_h3.logger, "warning_once", warnings.append) run() @@ -464,7 +465,7 @@ def test_env_on_missing_module_warns_and_falls_back(routed, monkeypatch, env_ove def test_env_on_no_grad_context_detaches_route_from_leaf_flags(routed, env_overrides): """A requires_grad leaf under torch.no_grad() is still a no-grad forward.""" fake_sm, fake_triton, run, meta = routed - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) q, k, v = _tiled_qkv(meta, requires_grad=True) with torch.no_grad(): @@ -481,16 +482,16 @@ def test_env_on_no_grad_context_detaches_route_from_leaf_flags(routed, env_overr (True, True, 0), ], ) -def test_preprocess_adds_partner_only_for_odd_no_grad_sm100a( +def test_preprocess_adds_partner_only_for_odd_no_grad_ptx_blackwell( env_overrides, env_enabled, requires_grad, extra_tiles, ): if env_enabled: - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) else: - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(None)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(None)) meta = _build_meta() n_tiles = meta.variable_block_sizes.numel() assert n_tiles % 2 == 1 @@ -512,9 +513,9 @@ def test_preprocess_adds_partner_only_for_odd_no_grad_sm100a( assert torch.count_nonzero(tiled[:, n_tiles * 64:]) == 0 -def test_odd_partner_is_visible_only_to_sm100a_transport(routed, env_overrides): +def test_odd_partner_is_visible_only_to_ptx_blackwell_transport(routed, env_overrides): fake_sm, fake_triton, _, meta = routed - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) raw_qkv = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, dtype=torch.bfloat16) tiled_qkv = impl.preprocess_qkv(raw_qkv, meta) @@ -542,7 +543,7 @@ def test_odd_partner_is_visible_only_to_sm100a_transport(routed, env_overrides): def test_unsupported_odd_route_strips_partner_before_triton(routed, monkeypatch, env_overrides): fake_sm, _, _, meta = routed fake_sm.supported = False - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) monkeypatch.setattr(vsa_h3.logger, "warning_once", lambda *_args, **_kwargs: None) calls = [] @@ -568,7 +569,7 @@ def test_unsupported_odd_route_strips_partner_before_triton(routed, monkeypatch, def test_geometry_change_clears_reused_partner_tile(env_overrides): - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) builder = MiniMaxH3VSAMetadataBuilder() even_meta = _build_meta(prefix_segments=(70, 70), builder=builder) odd_meta = _build_meta(prefix_segments=(70, 30), builder=builder) @@ -604,11 +605,11 @@ def test_forward_rejects_non_contract_transport_shapes(routed): impl.forward(logical, logical, logical, partner, meta) -def test_real_sm100a_odd_route_matches_triton_oracle(monkeypatch, env_overrides): +def test_real_ptx_blackwell_odd_route_matches_triton_oracle(monkeypatch, env_overrides): """Exercise the padded route through the actual data-center Blackwell extension.""" if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}: pytest.skip("requires a GB200/B300 (sm_100a/sm_103a) compute node") - if vsa_h3._sm100a is None or not vsa_h3._sm100a._HAS_VSA_SM100A: + if vsa_h3._ptx_blackwell is None or not vsa_h3._ptx_blackwell._HAS_VSA_PTX_BLACKWELL: pytest.skip("requires a fastvideo_kernel build containing data-center Blackwell VSA") device = torch.device("cuda") @@ -619,7 +620,7 @@ def test_real_sm100a_odd_route_matches_triton_oracle(monkeypatch, env_overrides) raw_qkv = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, device=device, dtype=torch.bfloat16) with torch.inference_mode(): - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(None)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(None)) triton_tiled = impl.preprocess_qkv(raw_qkv, meta) triton_query, triton_key, triton_value = triton_tiled.chunk(3, dim=0) triton_output = impl.postprocess_output( @@ -629,27 +630,27 @@ def test_real_sm100a_odd_route_matches_triton_oracle(monkeypatch, env_overrides) raise AssertionError("odd no-grad VSA-H3 unexpectedly fell back to Triton-64") monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", reject_triton) - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) - sm100a_tiled = impl.preprocess_qkv(raw_qkv, meta) - sm100a_query, sm100a_key, sm100a_value = sm100a_tiled.chunk(3, dim=0) - sm100a_output = impl.postprocess_output( - impl.forward(sm100a_query, sm100a_key, sm100a_value, None, meta), meta) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) + ptx_blackwell_tiled = impl.preprocess_qkv(raw_qkv, meta) + ptx_blackwell_query, ptx_blackwell_key, ptx_blackwell_value = ptx_blackwell_tiled.chunk(3, dim=0) + ptx_blackwell_output = impl.postprocess_output( + impl.forward(ptx_blackwell_query, ptx_blackwell_key, ptx_blackwell_value, None, meta), meta) torch.cuda.synchronize() - assert sm100a_tiled.shape[1] == triton_tiled.shape[1] + 64 - assert sm100a_output.shape == triton_output.shape == (1, meta.total_seq_length, _HEADS, _DIM) - torch.testing.assert_close(sm100a_output.float(), triton_output.float(), atol=0.04, rtol=0.02) + assert ptx_blackwell_tiled.shape[1] == triton_tiled.shape[1] + 64 + assert ptx_blackwell_output.shape == triton_output.shape == (1, meta.total_seq_length, _HEADS, _DIM) + torch.testing.assert_close(ptx_blackwell_output.float(), triton_output.float(), atol=0.04, rtol=0.02) -def test_real_sm100a_odd_preprocess_forward_postprocess_fullgraph(env_overrides): +def test_real_ptx_blackwell_odd_preprocess_forward_postprocess_fullgraph(env_overrides): """Capture #1745's odd transport buffer mutation with real Inductor.""" if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}: pytest.skip("requires a GB200/B300 (sm_100a/sm_103a) compute node") - if vsa_h3._sm100a is None or not vsa_h3._sm100a._HAS_VSA_SM100A: + if vsa_h3._ptx_blackwell is None or not vsa_h3._ptx_blackwell._HAS_VSA_PTX_BLACKWELL: pytest.skip("requires a fastvideo_kernel build containing data-center Blackwell VSA") device = torch.device("cuda") - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) meta = _build_meta(device=device, sparsity=0.5) assert meta.variable_block_sizes.numel() % 2 == 1 impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) diff --git a/fastvideo/tests/attention/test_vsa_h3_heads_first_tile.py b/fastvideo/tests/attention/test_vsa_h3_heads_first_tile.py index 8b320c8a1..8f99a945a 100644 --- a/fastvideo/tests/attention/test_vsa_h3_heads_first_tile.py +++ b/fastvideo/tests/attention/test_vsa_h3_heads_first_tile.py @@ -37,13 +37,13 @@ def _metadata(seq: int, n_tiles: int, tile_elems: int) -> SimpleNamespace: def _impl() -> vsa.MiniMaxH3VSAImpl: impl = object.__new__(vsa.MiniMaxH3VSAImpl) - impl._regional_compile_sm100a_enabled = None + impl._regional_compile_ptx_blackwell_enabled = None return impl @pytest.fixture -def sm100a(): - with envs.FASTVIDEO_VSA_SM100A.override(True): +def ptx_blackwell(): + with envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True): yield @@ -64,7 +64,7 @@ def test_scatter_rows_heads_first_matches_indexed_assignment() -> None: @pytest.mark.gpu @pytest.mark.parametrize(("tile_elems", "n_tiles"), [(128, 37), (128, 40), (64, 41)]) -def test_heads_first_tile_equals_the_default_tile(sm100a, tile_elems: int, n_tiles: int) -> None: +def test_heads_first_tile_equals_the_default_tile(ptx_blackwell, tile_elems: int, n_tiles: int) -> None: _require_cuda_triton() torch.manual_seed(1) meta = _metadata(seq=n_tiles * tile_elems - 300, n_tiles=n_tiles, tile_elems=tile_elems) @@ -79,7 +79,7 @@ def test_heads_first_tile_equals_the_default_tile(sm100a, tile_elems: int, n_til assert heads_first.transpose(1, 2).is_contiguous() assert _same_bits(default, heads_first) # The tile-score pooling reads the logical prefix of the view directly - # (forward() drops the sm100a pairing tile); its FP32 sums must not move. + # (forward() drops the ptx_blackwell pairing tile); its FP32 sums must not move. logical = n_tiles * tile_elems assert _same_bits(vsa._pool_tiles(default[:, :logical], meta.variable_block_sizes, tile_elems), vsa._pool_tiles(heads_first[:, :logical], meta.variable_block_sizes, tile_elems)) diff --git a/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py b/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py index ec48bbc84..cfe7728dd 100644 --- a/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py +++ b/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py @@ -5,7 +5,7 @@ With vsa_ref_policy="p2_multi_region" every Ref2VA reference video is its own sparse region, tiled in place at its packed offset, while text, audio, and image references stay dense. Each video query keeps its own top-k of every region: the reference keep rate may differ from the target's. The kernel -route for 128-token tiles is exercised through a stand-in for the sm_100a +route for 128-token tiles is exercised through a stand-in for the ptx_blackwell extension; real kernels are not needed. """ @@ -403,14 +403,14 @@ def test_build_multi_region_missing_ref_keep_rate(): # --------------------------------------------------------------------------- -# 128-token tiles: sm_100a/sm_103a CUDA only +# 128-token tiles: ptx_blackwell only # --------------------------------------------------------------------------- _HEADS, _DIM = 2, 128 -class _FakeSm100a: - """Stands in for fastvideo_kernel.block_sparse_attn_sm100a.""" +class _FakePtxBlackwell: + """Stands in for fastvideo_kernel.block_sparse_attn_ptx_blackwell.""" def __init__(self, supported=True): self.supported = supported @@ -419,7 +419,7 @@ class _FakeSm100a: def is_supported(self, q, variable_block_sizes): return self.supported - def block_sparse_attn_sm100a(self, q, k, v, q2k_idx, q2k_num, variable_block_sizes, need_lse=True): + def block_sparse_attn_ptx_blackwell(self, q, k, v, q2k_idx, q2k_num, variable_block_sizes, need_lse=True): self.calls.append(dict(q=q, q2k_num=q2k_num, vbs=variable_block_sizes, need_lse=need_lse)) return q.clone(), None @@ -435,16 +435,16 @@ def _fake_map_to_index(block_map): @pytest.fixture() def tile128(monkeypatch, env_overrides): - fake = _FakeSm100a() + fake = _FakePtxBlackwell() def forbid_triton(*args, **kwargs): raise AssertionError("tile 128 must never fall back to the Triton-64 kernel") - monkeypatch.setattr(vsa_h3, "_sm100a", fake) + monkeypatch.setattr(vsa_h3, "_ptx_blackwell", fake) monkeypatch.setattr(vsa_h3, "map_to_index", _fake_map_to_index) monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", forbid_triton) # Tile 128 has no opt-in: it routes to the CUDA kernel with the switch off. - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(None)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(None)) return fake @@ -482,9 +482,9 @@ def test_forward_tile128_even_tile_count_no_partner_tile(tile128): def test_forward_tile128_unavailable_cuda_kernel(tile128, monkeypatch): meta = _build_r2v(sparsity=0.5, tile_size=128) tile128.supported = False - with pytest.raises(RuntimeError, match="tile 128 requires the sm_100a/sm_103a CUDA block-sparse kernel"): + with pytest.raises(RuntimeError, match="tile 128 requires the ptx_blackwell block-sparse kernel"): _run_tile128(meta) - monkeypatch.setattr(vsa_h3, "_sm100a", None) + monkeypatch.setattr(vsa_h3, "_ptx_blackwell", None) with pytest.raises(RuntimeError, match="not installed"): _run_tile128(meta) @@ -505,7 +505,7 @@ def test_forward_tile128_route_resolved_once_per_device(tile128): raw = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, dtype=torch.bfloat16) query, key, value = impl.preprocess_qkv(raw, meta).chunk(3, dim=0) for _ in range(2): - with pytest.raises(RuntimeError, match="tile 128 requires the sm_100a/sm_103a CUDA block-sparse kernel"): + with pytest.raises(RuntimeError, match="tile 128 requires the ptx_blackwell block-sparse kernel"): impl.forward(query, key, value, None, meta) assert probes == [_CPU, _CPU] @@ -520,7 +520,7 @@ def test_forward_tile128_requires_no_grad(tile128): def test_forward_tile128_rejects_regional_fullgraph_compile(tile128, monkeypatch): meta = _build_r2v(sparsity=0.5, tile_size=128) impl = _impl(_DIM) - impl._regional_compile_sm100a_enabled = True + impl._regional_compile_ptx_blackwell_enabled = True monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) q = torch.zeros(1, meta.variable_block_sizes.numel() * 128, _HEADS, _DIM, dtype=torch.bfloat16) with pytest.raises(RuntimeError, match="requires 64-token tiles"): @@ -538,14 +538,14 @@ def test_build_tile128_uses_4x4x8_tiles(): @pytest.mark.parametrize("sparsity,ref_keep_rate", [(0.5, 0.25), (0.9, 0.1)]) -def test_forward_tile128_real_sm100a_kernel_matches_oracle(env_overrides, sparsity, ref_keep_rate): +def test_forward_tile128_real_ptx_blackwell_kernel_matches_oracle(env_overrides, sparsity, ref_keep_rate): """The real 128-token CUDA forward on a multi-region layout against an FP32 masked-SDPA oracle.""" if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}: pytest.skip("requires an sm_100a/sm_103a GPU (B200, B300, GB200, GB300)") - forwards = getattr(vsa_h3._sm100a, "_FWD_BY_BLOCK", {}) if vsa_h3._sm100a is not None else {} + forwards = getattr(vsa_h3._ptx_blackwell, "_FWD_BY_BLOCK", {}) if vsa_h3._ptx_blackwell is not None else {} if forwards.get(128) is None or vsa_h3.map_to_index is None: - pytest.skip("requires a fastvideo_kernel build with the 128-token sm_100a forward") - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(None)) + pytest.skip("requires a fastvideo_kernel build with the 128-token ptx_blackwell forward") + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(None)) device = torch.device("cuda") meta = _build_r2v(sparsity=sparsity, tile_size=128, ref_keep_rate=ref_keep_rate, device=device) n_tiles = meta.variable_block_sizes.numel() diff --git a/fastvideo/tests/inference/test_basic_fasth3_profile.py b/fastvideo/tests/inference/test_basic_fasth3_profile.py index bdac5df89..b47994902 100644 --- a/fastvideo/tests/inference/test_basic_fasth3_profile.py +++ b/fastvideo/tests/inference/test_basic_fasth3_profile.py @@ -42,7 +42,7 @@ def test_default_all_profile_matches_fastest_contract(tmp_path): assert args.num_gpus == 4 assert args.vsa_sparsity == 0.9 assert args.vsa_tile_size == 64 - assert args.vsa_kernel == "sm100a" + assert args.vsa_kernel == "ptx_blackwell" assert args.seed == 1000 assert args.warmup_seed == 999 assert args.warmup is True @@ -69,7 +69,7 @@ def test_default_all_profile_matches_fastest_contract(tmp_path): "vae_parallel_decode_strategy": "gather", } - assert environment["FASTVIDEO_VSA_SM100A"] == "1" + assert environment["FASTVIDEO_VSA_PTX_BLACKWELL"] == "1" assert environment["FASTVIDEO_VSA_CUTEDSL"] == "0" assert environment["FASTVIDEO_DISABLE_ATTENTION_COMPILE"] == "0" assert environment["FASTVIDEO_FA4"] == "1" @@ -175,7 +175,7 @@ def test_opt_outs_override_inherited_environment(monkeypatch): config = fasth3.build_generator_config(args) assert "FASTVIDEO_H3_VSA_PROBE" not in fasth3.os.environ - assert fasth3.os.environ["FASTVIDEO_VSA_SM100A"] == "0" + assert fasth3.os.environ["FASTVIDEO_VSA_PTX_BLACKWELL"] == "0" assert fasth3.os.environ["FASTVIDEO_VSA_CUTEDSL"] == "0" assert fasth3.os.environ["FASTVIDEO_DISABLE_ATTENTION_COMPILE"] == "0" assert fasth3.os.environ["FASTVIDEO_FA4"] == "0" @@ -208,7 +208,7 @@ def test_hopper_route_selects_tk_kernel_and_resident_fp8_text_encoder(): config = fasth3.build_generator_config(args) assert environment["FASTVIDEO_VSA_TK"] == "1" - assert environment["FASTVIDEO_VSA_SM100A"] == "0" + assert environment["FASTVIDEO_VSA_PTX_BLACKWELL"] == "0" assert environment["FASTVIDEO_FA4"] == "0" assert config.engine.use_fsdp_inference is True assert config.engine.offload.text_encoder is False @@ -230,7 +230,7 @@ def test_tk_profile_requires_the_sm90_extension(monkeypatch): def test_selected_fast_profile_requires_its_optional_routes(monkeypatch): monkeypatch.setattr(fasth3, "_fa4_is_installed", lambda: False) - monkeypatch.setattr(fasth3, "_sm100a_kernel_is_installed", lambda: False) + monkeypatch.setattr(fasth3, "_ptx_blackwell_kernel_is_installed", lambda: False) with pytest.raises(RuntimeError, match="flash-attn-4"): fasth3.validate_profile_dependencies(_args()) @@ -240,7 +240,7 @@ def test_selected_fast_profile_requires_its_optional_routes(monkeypatch): fasth3.validate_profile_dependencies(_args("--no-fa4", "--vsa-kernel", "triton")) monkeypatch.setattr(fasth3, "_fa4_is_installed", lambda: True) - monkeypatch.setattr(fasth3, "_sm100a_kernel_is_installed", lambda: True) + monkeypatch.setattr(fasth3, "_ptx_blackwell_kernel_is_installed", lambda: True) fasth3.validate_profile_dependencies(_args()) diff --git a/fastvideo/tests/inference/test_inference_regional_compile.py b/fastvideo/tests/inference/test_inference_regional_compile.py index 7a5fd1626..1ef2265f7 100644 --- a/fastvideo/tests/inference/test_inference_regional_compile.py +++ b/fastvideo/tests/inference/test_inference_regional_compile.py @@ -72,10 +72,10 @@ def test_supported_backends_allow_compile(backend_name, env_overrides) -> None: assert _regional_compile_unsupported_reason(_init_params_for(backend_name)) is None -def test_h3_vsa_sm100a_tile64_allows_compile(env_overrides) -> None: +def test_h3_vsa_ptx_blackwell_tile64_allows_compile(env_overrides) -> None: env_overrides.enter_context(envs.FASTVIDEO_DISABLE_ATTENTION_COMPILE.override(None)) env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_PROBE.override(None)) - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) reason = _regional_compile_unsupported_reason( _init_params_for("VIDEO_SPARSE_ATTN_H3"), @@ -85,14 +85,14 @@ def test_h3_vsa_sm100a_tile64_allows_compile(env_overrides) -> None: assert reason is None -@pytest.mark.parametrize(("sm100a", "tile_size"), [(False, 64), (True, 256), (True, None)]) -def test_h3_vsa_unsupported_compile_route_degrades_to_eager(sm100a, tile_size, env_overrides) -> None: +@pytest.mark.parametrize(("ptx_blackwell", "tile_size"), [(False, 64), (True, 256), (True, None)]) +def test_h3_vsa_unsupported_compile_route_degrades_to_eager(ptx_blackwell, tile_size, env_overrides) -> None: env_overrides.enter_context(envs.FASTVIDEO_DISABLE_ATTENTION_COMPILE.override(None)) env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_PROBE.override(None)) - if sm100a: - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) + if ptx_blackwell: + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True)) else: - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(None)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(None)) reason = _regional_compile_unsupported_reason( _init_params_for("VIDEO_SPARSE_ATTN_H3"), @@ -219,11 +219,11 @@ def test_minimax_h3_prepare_for_regional_compile_does_not_require_quantized_q_we def test_minimax_h3_prepare_for_regional_compile_propagates_backend_rejection() -> None: model = _gate_model(2.0) impl = model.transformer_blocks[0].attn.distributed_attention.attn_impl - impl.unsupported = "sm100a probe failed" + impl.unsupported = "ptx_blackwell probe failed" reason = model.prepare_for_regional_compile() - assert reason == "sm100a probe failed" + assert reason == "ptx_blackwell probe failed" def test_unprepared_vsa_gate_cannot_mutate_cache_during_compile(monkeypatch) -> None: diff --git a/fastvideo/tests/train/utils/test_activation_checkpoint_ops_retains_attention.py b/fastvideo/tests/train/utils/test_activation_checkpoint_ops_retains_attention.py index 80b1d4f28..279b59511 100644 --- a/fastvideo/tests/train/utils/test_activation_checkpoint_ops_retains_attention.py +++ b/fastvideo/tests/train/utils/test_activation_checkpoint_ops_retains_attention.py @@ -27,8 +27,8 @@ _RETAINED_OP_MODULES = { "fastvideo::_flash_attn_cute_fp4_forward": "fastvideo.attention.utils.flash_attn_cute", "fastvideo::_flash_attn_no_pad_forward": "fastvideo.attention.utils.flash_attn_no_pad", "fastvideo::_flash_attn_varlen_qk_no_pad_forward": "fastvideo.attention.utils.flash_attn_no_pad", + "fastvideo_kernel::block_sparse_attn_ptx_blackwell": "fastvideo_kernel.block_sparse_attn", "fastvideo_kernel::block_sparse_attn_sm90": "fastvideo_kernel.block_sparse_attn", - "fastvideo_kernel::block_sparse_attn_sm100a": "fastvideo_kernel.block_sparse_attn", "fastvideo_kernel::block_sparse_attn_triton": "fastvideo_kernel.block_sparse_attn", } diff --git a/fastvideo/tests/transformers/test_minimax_h3_tile_first.py b/fastvideo/tests/transformers/test_minimax_h3_tile_first.py index fa2ac25fc..130316f2f 100644 --- a/fastvideo/tests/transformers/test_minimax_h3_tile_first.py +++ b/fastvideo/tests/transformers/test_minimax_h3_tile_first.py @@ -66,7 +66,7 @@ def test_tile_first_matches_generic_vsa_with_partial_tiles(env_overrides, distri env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("VIDEO_SPARSE_ATTN_H3")) env_overrides.enter_context(envs.override_external("FASTVIDEO_VSA_TRITON", "1")) - env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(False)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(False)) env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_FP4.override(False)) env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_TILE_FIRST.override(False)) env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_SM89_KERNEL.override("original")) diff --git a/fastvideo/train/utils/activation_checkpoint.py b/fastvideo/train/utils/activation_checkpoint.py index c815ba6ae..80b07dcca 100644 --- a/fastvideo/train/utils/activation_checkpoint.py +++ b/fastvideo/train/utils/activation_checkpoint.py @@ -53,8 +53,8 @@ _SELECTIVE_ACTIVATION_CHECKPOINTING_OP_NAMES = { "fastvideo::_flash_attn_varlen_qk_no_pad_forward", # VSA dispatches block_sparse_attn; video_sparse_attn is its Python entry # point, not an op, and naming that here would match nothing. + "fastvideo_kernel::block_sparse_attn_ptx_blackwell", "fastvideo_kernel::block_sparse_attn_sm90", - "fastvideo_kernel::block_sparse_attn_sm100a", "fastvideo_kernel::block_sparse_attn_triton", } diff --git a/fastvideo/training/activation_checkpoint.py b/fastvideo/training/activation_checkpoint.py index 2cddede54..cf7084c31 100644 --- a/fastvideo/training/activation_checkpoint.py +++ b/fastvideo/training/activation_checkpoint.py @@ -35,8 +35,8 @@ _SELECTIVE_ACTIVATION_CHECKPOINTING_OP_NAMES = { "fastvideo::_flash_attn_cute_fp4_forward", "fastvideo::_flash_attn_no_pad_forward", "fastvideo::_flash_attn_varlen_qk_no_pad_forward", + "fastvideo_kernel::block_sparse_attn_ptx_blackwell", "fastvideo_kernel::block_sparse_attn_sm90", - "fastvideo_kernel::block_sparse_attn_sm100a", "fastvideo_kernel::block_sparse_attn_triton", } diff --git a/pyproject.toml b/pyproject.toml index 27831d34f..11bb254d4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -119,7 +119,7 @@ members = ["apps/dreamverse", "apps/infinite_livestream"] # PyPI's aarch64 fastvideo-kernel wheels do not cover the NVIDIA DGX Spark / # GB10 (sm_121), so on ARM build from the in-tree source instead. The FastH3 # extra also selects the in-tree source on x86_64 so a PR checkout can build -# its matching sm_100a route before that version reaches PyPI. Other x86_64 +# its matching ptx_blackwell route before that version reaches PyPI. Other x86_64 # installs keep the published wheel. This table is uv-only: plain pip ignores # it. Pass --no-sources to bypass it (e.g. when the kernel is already # installed, as in the Spark install guide). diff --git a/scripts/benchmarks/minimax_h3_4090/bench_pod.py b/scripts/benchmarks/minimax_h3_4090/bench_pod.py index b1f4dde62..83d8645f0 100644 --- a/scripts/benchmarks/minimax_h3_4090/bench_pod.py +++ b/scripts/benchmarks/minimax_h3_4090/bench_pod.py @@ -97,7 +97,7 @@ def main(): os.environ.setdefault("FASTVIDEO_MINIMAX_H3_FUSIONS", "all") os.environ.setdefault("FASTVIDEO_H3_VAE_TILE_BATCH", "28") os.environ.setdefault("FASTVIDEO_VSA_TRITON", "1") - os.environ.setdefault("FASTVIDEO_VSA_SM100A", "0") + os.environ.setdefault("FASTVIDEO_VSA_PTX_BLACKWELL", "0") os.environ.setdefault("FASTVIDEO_FA4", "0") os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") ap_table = os.environ.get("FASTVIDEO_H3_ADALN_TABLE") diff --git a/scripts/benchmarks/minimax_h3_pro6000/app.py b/scripts/benchmarks/minimax_h3_pro6000/app.py index 40e087a63..5b734d986 100644 --- a/scripts/benchmarks/minimax_h3_pro6000/app.py +++ b/scripts/benchmarks/minimax_h3_pro6000/app.py @@ -28,7 +28,7 @@ image = ( f"git -C /cutlass checkout {CUTLASS_COMMIT}") .env({"FLASHINFER_CUDA_ARCH_LIST": "12.0a", "FLASHINFER_WORKSPACE_BASE": "/vol/cache/flashinfer", "TORCHINDUCTOR_CACHE_DIR": "/vol/cache/inductor", "TRITON_CACHE_DIR": "/vol/cache/triton", - "FASTVIDEO_VSA_SM100A": "0", "FASTVIDEO_FA4": "0", "FASTVIDEO_STAGE_LOGGING": "1", + "FASTVIDEO_VSA_PTX_BLACKWELL": "0", "FASTVIDEO_FA4": "0", "FASTVIDEO_STAGE_LOGGING": "1", "HF_HUB_ENABLE_HF_TRANSFER": "1"}) ) app = modal.App("h3-pro6000-fastest", image=image) diff --git a/tests/bench_block_sparse_sm100a.py b/tests/bench_block_sparse_blackwell.py similarity index 92% rename from tests/bench_block_sparse_sm100a.py rename to tests/bench_block_sparse_blackwell.py index 8e75f1b15..e6cf0d779 100644 --- a/tests/bench_block_sparse_sm100a.py +++ b/tests/bench_block_sparse_blackwell.py @@ -1,19 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 -"""In-framework benchmark of the sm_100a/sm_103a VSA backend at real Wan shapes. +"""In-framework benchmark of the ptx_blackwell VSA backend at real Wan shapes. Not the powers-of-two bench grid: this builds the metadata FastVideo actually builds, at the deployed latent shapes, and drives the same block_sparse_attn_from_indices entry point the model calls -- so it measures the path that will run, including the dispatch and the index tensors as constructed rather than synthesised. - PYTHONPATH=fastvideo-kernel/python python tests/bench_block_sparse_sm100a.py + PYTHONPATH=fastvideo-kernel/python python tests/bench_block_sparse_blackwell.py """ import time import torch -from fastvideo_kernel import block_sparse_attn_sm100a as vsa +from fastvideo_kernel import block_sparse_attn_ptx_blackwell as vsa from fastvideo_kernel.block_sparse_attn import block_sparse_attn_triton from fastvideo_kernel.triton_kernels.index import map_to_index from fastvideo_kernel.vsa_utils import build_vsa_metadata @@ -63,7 +63,7 @@ def main(): idx, num = idx.to(torch.int32).contiguous(), num.to(torch.int32).contiguous() ok = vsa.is_supported(q, vbs) - ours = timed(lambda: vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs)) if ok else float("nan") + ours = timed(lambda: vsa.block_sparse_attn_ptx_blackwell(q, k, v, idx, num, vbs)) if ok else float("nan") # Triton needs 64-token blocks: expand each 128-block into its two halves. if block == 128: # Triton is 64-granular; expand only when we run 128 diff --git a/tests/local_tests/test_cookbook_serving.py b/tests/local_tests/test_cookbook_serving.py index e214ce1a3..f0a2e11e8 100644 --- a/tests/local_tests/test_cookbook_serving.py +++ b/tests/local_tests/test_cookbook_serving.py @@ -65,7 +65,7 @@ def test_spark_preview_is_a_runtime_with_one_or_two_devices(): "gpu_count": 1, "evidence": "source-configured", } - assert profile["command"].startswith("FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0") + assert profile["command"].startswith("FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0") assert "openai_fasth3_spark.yaml" in profile["command"] assert "--server.host 127.0.0.1" in profile["command"] diff --git a/tests/test_block_sparse_sm100a.py b/tests/test_block_sparse_blackwell.py similarity index 92% rename from tests/test_block_sparse_sm100a.py rename to tests/test_block_sparse_blackwell.py index b449a2da8..8b43084a2 100644 --- a/tests/test_block_sparse_sm100a.py +++ b/tests/test_block_sparse_blackwell.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Correctness tests for the sm_100a/sm_103a CUDA block-sparse VSA forward. +"""Correctness tests for the ptx_blackwell block-sparse VSA forward. Compared against an explicit PyTorch reference rather than the Triton kernel: Triton's block-sparse forward is hardcoded to 64-token blocks (BLOCK_M = BLOCK_N = 64) while this @@ -8,7 +8,7 @@ different sparsity granularities. The reference below applies exactly the semant kernel is supposed to implement -- selected blocks only, keys past variable_block_sizes masked. Every case runs at both block sizes. -Run with: python -m pytest tests/test_block_sparse_sm100a.py -v +Run with: python -m pytest tests/test_block_sparse_blackwell.py -v """ import os @@ -18,14 +18,14 @@ import sys import pytest import torch -from fastvideo_kernel import block_sparse_attn_sm100a as vsa +from fastvideo_kernel import block_sparse_attn_ptx_blackwell as vsa HEAD_DIM = 128 BLOCK_SIZES = [64, 128] pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)} - or not vsa._HAS_VSA_SM100A, + or not vsa._HAS_VSA_PTX_BLACKWELL, reason="requires data-center Blackwell (sm_100a/sm_103a) and a built fastvideo_kernel extension", ) @@ -89,7 +89,7 @@ def run_and_compare(block, ragged, num_blocks=8, topk=4, heads=4, atol=0.02): q, k, v, idx, num, vbs = make_case(block, num_blocks=num_blocks, topk=topk, heads=heads, ragged=ragged) assert vsa.is_supported(q, vbs) - got, got_lse = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs) + got, got_lse = vsa.block_sparse_attn_ptx_blackwell(q, k, v, idx, num, vbs) ref, ref_lse = reference(q, k, v, idx, num, vbs, block) got_o = got if vsa.BHSD else got.transpose(1, 2) @@ -127,7 +127,7 @@ def test_sequence_lengths(block, num_blocks): def test_lse_is_not_vacuous(block): """Guards the lse assertion: a wrong lse must actually fail the comparison.""" q, k, v, idx, num, vbs = make_case(block, ragged=True) - _, got_lse = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs) + _, got_lse = vsa.block_sparse_attn_ptx_blackwell(q, k, v, idx, num, vbs) _, ref_lse = reference(q, k, v, idx, num, vbs, block) assert (got_lse.float() - (ref_lse + 1.0)).abs().max().item() > 0.5 @@ -138,7 +138,7 @@ def test_inference_no_lse_matches_eager_under_fullgraph_compile(block): q, k, v, idx, num, vbs = make_case(block, ragged=True) def inference_path(q, k, v, idx, num, vbs): - out, lse = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs, need_lse=False) + out, lse = vsa.block_sparse_attn_ptx_blackwell(q, k, v, idx, num, vbs, need_lse=False) assert lse is None return out @@ -154,7 +154,7 @@ def test_inference_custom_op_fake_matches_real(block): """Pin the custom op schema and fake output metadata used by Dynamo.""" q, k, v, idx, num, vbs = make_case(block, ragged=True) torch.library.opcheck( - torch.ops.fastvideo_kernel.block_sparse_attn_sm100a_inference.default, + torch.ops.fastvideo_kernel.block_sparse_attn_ptx_blackwell_inference.default, (q, k, v, idx, num, vbs), test_utils=("test_schema", "test_faketensor", "test_aot_dispatch_dynamic"), ) @@ -162,18 +162,18 @@ def test_inference_custom_op_fake_matches_real(block): @pytest.mark.parametrize("block", BLOCK_SIZES) def test_mask_inference_matches_index_route_under_fullgraph_compile(block): - """H3's bool mask compaction and sm_100a launch stay in one opaque op.""" + """H3's bool mask compaction and ptx_blackwell launch stay in one opaque op.""" batch, heads, num_blocks = 1, 4, 8 q, k, v, idx, num, vbs = make_case(block, batch=batch, heads=heads, num_blocks=num_blocks, ragged=True) block_map = make_block_map(idx, num, batch, heads, num_blocks) def inference_path(q, k, v, block_map, vbs): - out, lse = vsa.block_sparse_attn_sm100a_from_mask(q, k, v, block_map, vbs) + out, lse = vsa.block_sparse_attn_ptx_blackwell_from_mask(q, k, v, block_map, vbs) assert lse is None return out with torch.inference_mode(): - expected, _ = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs, need_lse=False) + expected, _ = vsa.block_sparse_attn_ptx_blackwell(q, k, v, idx, num, vbs, need_lse=False) compiled = torch.compile(inference_path, backend="eager", fullgraph=True) actual = compiled(q, k, v, block_map, vbs) torch.testing.assert_close(actual, expected, atol=0, rtol=0) @@ -186,7 +186,7 @@ def test_mask_inference_custom_op_fake_matches_real(block): q, k, v, idx, num, vbs = make_case(block, batch=batch, heads=heads, num_blocks=num_blocks, ragged=True) block_map = make_block_map(idx, num, batch, heads, num_blocks) torch.library.opcheck( - torch.ops.fastvideo_kernel.block_sparse_attn_sm100a_from_mask_inference.default, + torch.ops.fastvideo_kernel.block_sparse_attn_ptx_blackwell_from_mask_inference.default, (q, k, v, block_map, vbs), test_utils=("test_schema", "test_faketensor", "test_aot_dispatch_dynamic"), ) @@ -251,7 +251,7 @@ def test_nonuniform_counts_straddling_a_pair(block): |out diff| ~0.6 on this exact case), leaving every other tile correct.""" q, k, v, idx, num, vbs = make_two_class_case(block) assert vsa.is_supported(q, vbs) - got, got_lse = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs) + got, got_lse = vsa.block_sparse_attn_ptx_blackwell(q, k, v, idx, num, vbs) ref, ref_lse = reference(q, k, v, idx, num, vbs, block) got_o = got if vsa.BHSD else got.transpose(1, 2) diff = (got_o.float() - ref).abs().max().item() @@ -269,7 +269,7 @@ def test_same_seed_determinism(block): runs = [] for _ in range(2): q, k, v, idx, num, vbs = make_case(block, ragged=True, seed=3) - o, m = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs) + o, m = vsa.block_sparse_attn_ptx_blackwell(q, k, v, idx, num, vbs) torch.cuda.synchronize() runs.append((o, m)) assert torch.equal(runs[0][0], runs[1][0]), "out differs across same-seed runs" @@ -281,7 +281,7 @@ def _zero_count_case_main(block): zero_rows = (5, 6, 8, 9) # odd member; even member with nonzero sibling 7; a whole pair q, k, v, idx, num, vbs = make_two_class_case(block, zero_rows=zero_rows) assert vsa.is_supported(q, vbs) - got, got_lse = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs) + got, got_lse = vsa.block_sparse_attn_ptx_blackwell(q, k, v, idx, num, vbs) torch.cuda.synchronize() ref, ref_lse = reference(q, k, v, idx, num, vbs, block) got_o = (got if vsa.BHSD else got.transpose(1, 2)).float()