diff --git a/.github/workflows/publish-kernel.yml b/.github/workflows/publish-kernel.yml index e73956c700..1d3a9aa32c 100644 --- a/.github/workflows/publish-kernel.yml +++ b/.github/workflows/publish-kernel.yml @@ -68,7 +68,7 @@ jobs: - os: ubuntu-22.04 arch: x86_64 wheel-plat: manylinux_2_35_x86_64 - # aarch64 is Blackwell (GB200 sm_100a + DGX Spark / consumer sm_120a), not + # aarch64 is Blackwell (GB200 sm_100a + sm_120a + DGX Spark sm_121a), not # Hopper, and Blackwell needs CUDA >= 12.8 — so only the cu130 leg applies. # Added via include so x86 keeps cu126 + cu130 while aarch64 stays cu130-only. include: @@ -164,19 +164,20 @@ jobs: cd fastvideo-kernel git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized # Release builds run on GPU-less runners, so set kernels + arch explicitly: - # * aarch64 = Blackwell (GB200 sm_100a + DGX Spark/consumer sm_120a), NOT - # Hopper, so TK (sm_90a wgmma) is OFF. The C++ FP4 (attn_qat_infer, SM120) - # covers sm_120a; turbodiffusion covers sm_100a+sm_120a. The sm_100 FP4 - # forward is the FA4 CuTe DSL path in the fastvideo package (PR #1221), + # * aarch64 = Blackwell (GB200 sm_100a + sm_120a + DGX Spark sm_121a), NOT + # Hopper, so TK (sm_90a wgmma) is OFF. The C++ FP4 (attn_qat_infer) + # covers sm_120a+sm_121a; turbodiffusion covers every listed arch. The + # sm_100 FP4 forward is the FA4 CuTe DSL path in the fastvideo package (PR #1221), # JIT-compiled at runtime — not built into this wheel. # * x86_64 cu130 = Hopper TK + data-center Blackwell sm_100a VSA # + consumer Blackwell sm_120a FP4. # * x86_64 cu126 = Hopper TK only (older drivers; CUDA < 12.8 has no FP4). - # The per-arch split in CMakeLists pins the FP4 targets to sm_120a and builds - # the main extension for the full arch list. CMAKE_BUILD_PARALLEL_LEVEL caps + # The per-arch split in CMakeLists pins the FP4 targets to requested + # sm_120a/sm_121a and builds the main extension for the full arch list. + # CMAKE_BUILD_PARALLEL_LEVEL caps # Ninja so heavy CUTLASS/TK template TUs don't OOM the 16 GB runner (exit 143). if [ "${{ matrix.platform.arch }}" = "aarch64" ]; then - export TORCH_CUDA_ARCH_LIST="10.0a;12.0a" + export TORCH_CUDA_ARCH_LIST="10.0a;12.0a;12.1a" export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=OFF -DFASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON" export CMAKE_BUILD_PARALLEL_LEVEL=1 elif [ "${{ matrix.torch-cuda.torch-cuda-short }}" = "cu130" ]; then @@ -250,7 +251,8 @@ jobs: - name: Download PyPI wheels # Publish the cu130 (CUDA 13) wheels to PyPI for both architectures: # x86_64 — Hopper sm_90a TK + consumer Blackwell sm_120a FP4 - # aarch64 — Blackwell: turbodiffusion (sm_100a/sm_120a) + C++ FP4 (sm_120a); + # aarch64 — Blackwell: turbodiffusion (sm_100a/sm_120a/sm_121a) + # + C++ FP4 (sm_120a/sm_121a); # no TK (Hopper). sm_100 FP4 forward is the FA4 CuTe DSL path in the # fastvideo package (#1221), shipped/JIT separately. # The x86_64 cu126 wheel stays available as a build artifact / GitHub-release asset. diff --git a/docs/getting_started/installation/spark.md b/docs/getting_started/installation/spark.md index 326598826a..0cfaaa03d6 100644 --- a/docs/getting_started/installation/spark.md +++ b/docs/getting_started/installation/spark.md @@ -58,7 +58,7 @@ With no GPU visible the kernel build can't probe the arch and `auto` can't detec the driver — name both explicitly: ```bash -UV_TORCH_BACKEND=cu130 TORCH_CUDA_ARCH_LIST=12.1 uv pip install -e . +UV_TORCH_BACKEND=cu130 TORCH_CUDA_ARCH_LIST=12.1a uv pip install -e . ``` ## Verify the install @@ -128,8 +128,8 @@ uv pip install "https://github.com/mjun0812/flash-attention-prebuild-wheels/rele |---|---| | `Could NOT find Python (missing: ... Development.Module)` | venv built from system Python without headers. Recreate with `--python-preference only-managed` (add `--clear` to reuse the path), or `sudo apt install python3.12-dev`. | | kernel build can't find cutlass headers | Submodules not initialised — run the `git submodule update` step. | -| `fastvideo-kernel: could not determine the target CUDA architecture` | The build couldn't see a GPU and no arch was given. Build on the Spark itself, or pass `TORCH_CUDA_ARCH_LIST=12.1` (see [Building without a visible GPU](#building-without-a-visible-gpu-ci--docker)). | -| `nvcc fatal: Unsupported gpu architecture 'compute_121'` | `nvcc` older than CUDA 12.9/13. Confirm `nvcc --version` is 13.x and `CUDACXX=/usr/local/cuda/bin/nvcc`. | +| `fastvideo-kernel: could not determine the target CUDA architecture` | The build couldn't see a GPU and no arch was given. Build on the Spark itself, or pass `TORCH_CUDA_ARCH_LIST=12.1a` (see [Building without a visible GPU](#building-without-a-visible-gpu-ci--docker)). | +| `sm_121a requires CUDA Toolkit 13.0+` | Confirm `nvcc --version` is 13.x and `CUDACXX=/usr/local/cuda/bin/nvcc`; architecture-specific sm_120a code cannot run on sm_121. | | `ninja: command not found` (manual build only) | `uv pip install scikit-build-core cmake ninja setuptools wheel`. | If you hit other issues, please open an issue on our diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index 753fef2835..06bba0200b 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -33,7 +33,7 @@ This page describes the various options for speeding up generation times in Fast - Video Sparse Attention: `FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN` - Sage Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN` - Sage Attention 3: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN_THREE` -- Attn-QAT inference (modified SageAttention3 FP4, sm_120/RTX 5090): `FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER` +- Attn-QAT inference (modified SageAttention3 FP4, sm_120a/sm_121a — RTX 5090 / DGX Spark): `FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER` - Video MoBA Attention: `FASTVIDEO_ATTENTION_BACKEND=VMOBA_ATTN` - Sparse Linear Attention: `FASTVIDEO_ATTENTION_BACKEND=SLA_ATTN` - SageSLA Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_SLA_ATTN` @@ -184,39 +184,45 @@ gen.generate_video(prompt="A raccoon in sunflowers", save_video=True) - Per-call cosine similarity vs BF16: ~0.99 (slight quantization error accumulates over denoising steps) - Only supports `headdim >= 128` -### NVFP4 + Attn-QAT (modified SageAttention3, Blackwell sm_120) +### NVFP4 Attn-QAT attention (modified SageAttention3, Blackwell sm_120a/sm_121a) -**`ATTN_QAT_INFER`** with **`transformer_quant=nvfp4_qat`** +**`ATTN_QAT_INFER`** -Runs the DiT fully in 4-bit: NVFP4 linear layers (activations quantized on the -fly) plus the modified SageAttention3 FP4 attention backend. This is the -inference half of the Quantization-Aware Distillation (QAD) recipe and the path -used for the RTX 5090 release. +Quantizes Q/K/V inside the modified SageAttention3 FP4 attention backend. On +DGX Spark, keep the transformer linear layers in BF16: the separate +`transformer_quant=nvfp4_qat` linear path is not supported on sm_121 yet. -The `attn_qat_infer` kernel hard-gates on **sm_120 (consumer Blackwell / RTX -5090)**; on other GPUs the backend logs a notice and falls back to Flash -Attention. See the [Attn-QAT paper](https://arxiv.org/abs/2603.00040). +The `attn_qat_infer` kernel targets consumer/workstation Blackwell — **sm_120a +(RTX 5090 / PRO 6000)** and **sm_121a (DGX Spark GB10)**; the block-scaled FP4 +MMA is numerically correct on both (GB10-verified: cos ~0.98 vs bf16 SDPA). +sm_121a support needs a CUDA 13 build (`TORCH_CUDA_ARCH_LIST=12.1a`). A runtime +capability gate falls back to Flash Attention on other GPUs. -Enable both halves — attention via the env var, linear via `transformer_quant`: +> **Quality note (stock weights):** on stock Wan-2.1 the FP4-attention output is +> below bf16 — QAD expects a QAT-distilled checkpoint the model was *trained* to +> tolerate FP4 attention with. Post-hoc on non-QAT weights, expect a quality +> cost; use the QAT checkpoint for release-grade output. + +See the [Attn-QAT paper](https://arxiv.org/abs/2603.00040). + +Enable FP4 attention with BF16 transformer linears on DGX Spark: ```python import os os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "ATTN_QAT_INFER" from fastvideo import VideoGenerator -from fastvideo.layers.quantization import get_quantization_config gen = VideoGenerator.from_pretrained( "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", num_gpus=1, - # Wan-2.1 uses the nvfp4_qat config (NVFP4 is LTX2-specific). Pass an - # instance — the bare string is not resolved on the from_pretrained path. - transformer_quant=get_quantization_config("nvfp4_qat")(), - use_fsdp_inference=False, # FSDP shards invalidate the FP4 tensor pointers + use_fsdp_inference=False, ) gen.generate(request={"prompt": "A raccoon in sunflowers", "output": {"save_video": True}}) ``` -Or run the example script: +The full QAD configuration, pairing this attention backend with +`transformer_quant=nvfp4_qat`, is currently limited to sm_120. Its RTX 5090 +example is: ```bash python examples/inference/optimizations/nvfp4_qat_wan2_1_1_3b.py diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index fdb0284260..90a72d0888 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -44,6 +44,7 @@ find_package(Python COMPONENTS Interpreter Development.Module REQUIRED) # --------------------------------------------------------------------------- if(NOT GPU_BACKEND STREQUAL "ROCM") if(DEFINED ENV{TORCH_CUDA_ARCH_LIST}) + set(TORCH_CUDA_ARCH_LIST "$ENV{TORCH_CUDA_ARCH_LIST}") message(STATUS "CUDA arch: TORCH_CUDA_ARCH_LIST=$ENV{TORCH_CUDA_ARCH_LIST} (from environment)") elseif(TORCH_CUDA_ARCH_LIST) set(ENV{TORCH_CUDA_ARCH_LIST} "${TORCH_CUDA_ARCH_LIST}") @@ -52,14 +53,14 @@ if(NOT GPU_BACKEND STREQUAL "ROCM") set(_FV_ARCH_LIST "") if(_FASTVIDEO_USER_CUDA_ARCH) # Caller pinned -DCMAKE_CUDA_ARCHITECTURES (which torch ignores); translate it - # to the TORCH_CUDA_ARCH_LIST spelling: "121" -> "12.1", "90a" -> "9.0a". + # to the TORCH_CUDA_ARCH_LIST spelling: "121a" -> "12.1a", "90a" -> "9.0a". # Only numeric spellings translate; keywords like "native"/"all" would # otherwise be mangled into nonsense ("nativ.e"). foreach(_fv_arch IN LISTS _FASTVIDEO_USER_CUDA_ARCH) if(NOT _fv_arch MATCHES "^[0-9]+[af]?$") message(FATAL_ERROR "fastvideo-kernel: CMAKE_CUDA_ARCHITECTURES='${_fv_arch}' is not " - "supported. Use a numeric arch (e.g. 90a, 121), set " + "supported. Use a numeric arch (e.g. 90a, 121a), set " "TORCH_CUDA_ARCH_LIST directly (e.g. 9.0a), or unset both to " "auto-detect from the visible GPU.") endif() @@ -72,7 +73,7 @@ if(NOT GPU_BACKEND STREQUAL "ROCM") else() # Best-effort probe of the visible GPU (mirrors build.sh detect_with_torch). execute_process( - COMMAND "${Python_EXECUTABLE}" -c "import torch; assert torch.cuda.is_available(); mj, mn = torch.cuda.get_device_capability(0); print(f'{mj}.{mn}a' if (mj, mn) in ((9, 0), (12, 0)) else f'{mj}.{mn}')" + COMMAND "${Python_EXECUTABLE}" -c "import torch; assert torch.cuda.is_available(); mj, mn = torch.cuda.get_device_capability(0); print(f'{mj}.{mn}a' if (mj, mn) in ((9, 0), (12, 0), (12, 1)) else f'{mj}.{mn}')" OUTPUT_VARIABLE _FV_ARCH_LIST OUTPUT_STRIP_TRAILING_WHITESPACE RESULT_VARIABLE _fv_detect_rc @@ -93,14 +94,26 @@ if(NOT GPU_BACKEND STREQUAL "ROCM") "fastvideo-kernel: could not determine the target CUDA architecture.\n" "Refusing to let torch auto-detect an arch that may not run on this GPU. " "Fix with one of:\n" - " - set TORCH_CUDA_ARCH_LIST (e.g. 12.1, or 9.0a for Hopper), or\n" - " - pass -DCMAKE_CUDA_ARCHITECTURES= (e.g. 121), or\n" + " - set TORCH_CUDA_ARCH_LIST (e.g. 12.1a, or 9.0a for Hopper), or\n" + " - pass -DCMAKE_CUDA_ARCHITECTURES= (e.g. 121a), or\n" " - build where the target GPU is visible to torch.\n" "Note: 'pip/uv pip install' builds under build isolation, which hides the " "GPU; set TORCH_CUDA_ARCH_LIST or add --no-build-isolation. " "fastvideo-kernel/build.sh sets all of this for you.") endif() endif() + + string(REGEX MATCH "(^|[; ,])((12\\.0a)|(120a)|(sm_120a))([; ,]|$)" + _FV_HAS_120A "${TORCH_CUDA_ARCH_LIST}") + string(REGEX MATCH "(^|[; ,])((12\\.1a)|(121a)|(sm_121a))([; ,]|$)" + _FV_HAS_121A "${TORCH_CUDA_ARCH_LIST}") + if(_FV_HAS_121A AND CUDAToolkit_VERSION VERSION_LESS 13.0) + message(FATAL_ERROR + "fastvideo-kernel: sm_121a requires CUDA Toolkit 13.0+, but found " + "${CUDAToolkit_VERSION}. Architecture-specific sm_120a code is not " + "compatible with sm_121; use CUDA 13 or remove 12.1a from " + "TORCH_CUDA_ARCH_LIST.") + endif() endif() # Robustly find Torch include paths using Python @@ -268,10 +281,7 @@ else() set(_WANTS_ATTN_QAT_INFER ON) elseif(FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER STREQUAL "AUTO") if(TORCH_CUDA_ARCH_LIST) - string(REGEX MATCH - "(^|[; ,])((12\\.0a)|(120a)|(sm_120a))([; ,]|$)" - _HAS_120A "${TORCH_CUDA_ARCH_LIST}") - if(_HAS_120A) + if(_FV_HAS_120A OR _FV_HAS_121A) set(_WANTS_ATTN_QAT_INFER ON) endif() else() @@ -299,11 +309,23 @@ else() endif() if(ENABLE_ATTN_QAT_INFER) + set(FASTVIDEO_FP4_CUDA_ARCHS "") + if(_FV_HAS_120A) + list(APPEND FASTVIDEO_FP4_CUDA_ARCHS "120a") + endif() + if(_FV_HAS_121A) + list(APPEND FASTVIDEO_FP4_CUDA_ARCHS "121a") + endif() + if(NOT FASTVIDEO_FP4_CUDA_ARCHS) + message(FATAL_ERROR + "FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON requires 12.0a or " + "12.1a in TORCH_CUDA_ARCH_LIST; got '${TORCH_CUDA_ARCH_LIST}'.") + endif() message(STATUS "attn_qat_infer kernels: ENABLED") else() message(STATUS "attn_qat_infer kernels: DISABLED " - "(requires CUDA 12.8+ and Blackwell sm_120a)") + "(requires CUDA 12.8+ and Blackwell sm_120a/sm_121a)") endif() endif() @@ -313,12 +335,12 @@ set(BUILD_CXX_KERNELS ON) # --------------------------------------------------------------------------- # Per-arch split for the Blackwell FP4 (attn_qat_infer) build # --------------------------------------------------------------------------- -# The FP4 kernels are sm_120a-only (they emit `cvt.e2m1x2` etc.), while the main +# The FP4 kernels are Blackwell-only (sm_120a/sm_121a; they emit `cvt.e2m1x2` etc.), while the main # extension (Hopper-only TK + generic turbodiffusion) targets the full arch list. # find_package(Torch) injects ONE global -gencode list into CMAKE_CUDA_FLAGS that # forces every target onto every arch, so the FP4 sources also get the sm_90a pass # and ptxas rejects their Blackwell instructions. Strip that global list and drive -# arch per target via CUDA_ARCHITECTURES instead (the fp4* targets pin 120a below; +# arch per target via CUDA_ARCHITECTURES instead (the fp4* targets pin the Blackwell archs below; # the main extension gets the full list). Only do this for the FP4 build with an # explicit arch list, so the cu126 / local autodetect paths stay untouched. if(ENABLE_ATTN_QAT_INFER AND TORCH_CUDA_ARCH_LIST) @@ -330,7 +352,7 @@ if(ENABLE_ATTN_QAT_INFER AND TORCH_CUDA_ARCH_LIST) set(FASTVIDEO_MAIN_CUDA_ARCHS "${TORCH_CUDA_ARCH_LIST}") string(REPLACE "sm_" "" FASTVIDEO_MAIN_CUDA_ARCHS "${FASTVIDEO_MAIN_CUDA_ARCHS}") string(REPLACE "." "" FASTVIDEO_MAIN_CUDA_ARCHS "${FASTVIDEO_MAIN_CUDA_ARCHS}") - message(STATUS "[per-arch] main extension archs=${FASTVIDEO_MAIN_CUDA_ARCHS}, fp4* archs=120a") + message(STATUS "[per-arch] main extension archs=${FASTVIDEO_MAIN_CUDA_ARCHS}, fp4* archs set per-target below") endif() # Compiler flags @@ -515,8 +537,11 @@ if(ENABLE_ATTN_QAT_INFER) $<$:-O3 -std=c++17> $<$:${ATTN_QAT_INFER_CUDA_FLAGS}> ) + # Compile only the FP4 architectures requested by the caller. The large + # template translation units use 8-12 GiB each, so an extra target matters. + message(STATUS "fp4* CUDA architectures: ${FASTVIDEO_FP4_CUDA_ARCHS}") set_target_properties(fp4attn_cuda PROPERTIES - CUDA_ARCHITECTURES "120a" + CUDA_ARCHITECTURES "${FASTVIDEO_FP4_CUDA_ARCHS}" CXX_STANDARD 17 CUDA_STANDARD 17 ) @@ -532,7 +557,7 @@ if(ENABLE_ATTN_QAT_INFER) $<$:${ATTN_QAT_INFER_CUDA_FLAGS}> ) set_target_properties(fp4quant_cuda PROPERTIES - CUDA_ARCHITECTURES "120a" + CUDA_ARCHITECTURES "${FASTVIDEO_FP4_CUDA_ARCHS}" CXX_STANDARD 17 CUDA_STANDARD 17 ) @@ -555,6 +580,6 @@ message(STATUS "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST}") message(STATUS "fastvideo_kernel_ops: ON (turbodiffusion int8-gemm/quant/rmsnorm/layernorm, all listed archs)") message(STATUS " + TK sta/block_sparse (sm_90a only): ${ENABLE_TK_KERNELS}") message(STATUS " + Ulysses NCCL-device all-to-all: ${ENABLE_ULYSSES_A2A}") -message(STATUS "fp4attn/fp4quant (sm_120a only, CUDA >= 12.8): ${ENABLE_ATTN_QAT_INFER}") +message(STATUS "fp4attn/fp4quant (sm_120a CUDA >= 12.8; sm_121a CUDA >= 13): ${ENABLE_ATTN_QAT_INFER}") message(STATUS "Triton fallbacks ship in python/fastvideo_kernel/triton_kernels regardless.") message(STATUS "============================================================") diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index 2e14b77495..e7d08ff8e8 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -12,7 +12,7 @@ Compiled CUDA extensions (CMake, see the build summary printed at the end of eve | same extension, optional part | ThunderKittens sliding-tile attention (`sta_fwd`) and VSA block-sparse (`block_sparse_fwd/bwd`) | `csrc/attention/*_h100.cu` | Hopper `sm_90a` only | `FASTVIDEO_KERNEL_BUILD_TK` (AUTO = ON iff `9.0a` is in the arch list; always OFF on aarch64 hosts — TK headers don't compile there) | | same extension, optional part | MiniMax-H3 block-sparse VSA forward (64- and 128-token blocks) | `csrc/attention/block_sparse*_sm100a.cu` | Blackwell `sm_100a` only | ON iff `10.0a` is in `TORCH_CUDA_ARCH_LIST` | | same extension, optional part | fused NVLink Ulysses all-to-all | `csrc/comm/ulysses_all_to_all.cu` | CUDA | `FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A` (AUTO = ON with NCCL 2.29+ device headers and library; always OFF on ROCm) | -| `fp4attn_cuda`, `fp4quant_cuda` | FP4 attention + quantization ("attn_qat_infer", modified SageAttention3) | `attn_qat_infer/` | consumer Blackwell `sm_120a` only, CUDA ≥ 12.8 | `FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER` (AUTO = ON iff `12.0a` is in the arch list) | +| `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) | Runtime-JIT kernels (no build step, ship in every wheel/image): @@ -28,11 +28,11 @@ Runtime-JIT kernels (no build step, ship in every wheel/image): |---|---|---|---|---|---|---| | PyPI wheels (`.github/workflows/publish-kernel.yml`) | version bump in `fastvideo-kernel/pyproject.toml` on main, or manual dispatch | x86_64 cu126 | `9.0a` | ON | AUTO | — (CUDA < 12.8) | | | | x86_64 cu130 | `9.0a;10.0a;12.0a` | ON | AUTO | ON | -| | | aarch64 cu130 | `10.0a;12.0a` | — | AUTO | ON | +| | | aarch64 cu130 | `10.0a;12.0a;12.1a` | — | AUTO | ON | | Docker images `ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev` (`.github/workflows/infra-build-image.yml`) | `docker/Dockerfile` changes on main, or manual dispatch | amd64 cuda12.6.3 + cuda13.0.0 | `9.0a` | ON | AUTO | — | | | | arm64 cuda12.6.3 (GH200) | `9.0a` | — (aarch64) | AUTO | — | | | | arm64 cuda13.0.0 (GB10 / DGX Spark) | `12.1` | — | AUTO | — | -| Local `./build.sh` | manual | probes the visible GPU via torch | detected | ON iff sm_90 (non-aarch64 host) | AUTO | ON iff sm_120 | +| Local `./build.sh` | manual | probes the visible GPU via torch | detected | ON iff sm_90 (non-aarch64 host) | AUTO | ON iff sm_120 or sm_121 | Notes: diff --git a/fastvideo-kernel/build.sh b/fastvideo-kernel/build.sh index c7ed847c36..1e430e1b70 100755 --- a/fastvideo-kernel/build.sh +++ b/fastvideo-kernel/build.sh @@ -137,15 +137,34 @@ if [ "${GPU_BACKEND}" = "CUDA" ]; then # Respect explicit overrides. if [ -z "${TORCH_CUDA_ARCH_LIST:-}" ]; then - if [ "${cc_major}" = "9" ] && [ "${cc_minor}" = "0" ]; then - export TORCH_CUDA_ARCH_LIST="9.0a" - elif [ "${cc_major}" = "12" ] && [ "${cc_minor}" = "0" ]; then - # Blackwell sm_120 needs the arch-conditional 'a' suffix so CMake's - # AUTO gate (matches 12.0a/120a/sm_120a) builds the attn_qat_infer - # (modified SageAttention3 FP4) kernels instead of silently skipping. - export TORCH_CUDA_ARCH_LIST="12.0a" - else - export TORCH_CUDA_ARCH_LIST="${cc_major}.${cc_minor}" + case "${cc_major}.${cc_minor}" in + 9.0|12.0|12.1) + # Architecture-specific instructions require the `a` target. + export TORCH_CUDA_ARCH_LIST="${cc_major}.${cc_minor}a" + ;; + *) + export TORCH_CUDA_ARCH_LIST="${cc_major}.${cc_minor}" + ;; + esac + fi + + if [[ "${TORCH_CUDA_ARCH_LIST}" =~ (^|[\;,[:space:]])(12\.1a|121a|sm_121a)([\;,[:space:]]|$) ]]; then + cuda_compiler="${CUDACXX:-}" + if [ -z "${cuda_compiler}" ] && [ -n "${CUDA_HOME:-}" ] && [ -x "${CUDA_HOME}/bin/nvcc" ]; then + cuda_compiler="${CUDA_HOME}/bin/nvcc" + fi + if [ -z "${cuda_compiler}" ]; then + cuda_compiler="$(command -v nvcc || true)" + fi + if [ -z "${cuda_compiler}" ]; then + echo "ERROR: sm_121a requires CUDA Toolkit 13.0+, but nvcc was not found." >&2 + echo " Set CUDACXX or CUDA_HOME to a CUDA 13 toolkit." >&2 + exit 1 + fi + cuda_major="$("${cuda_compiler}" --version 2>/dev/null | sed -n 's/.*release \([0-9][0-9]*\)\..*/\1/p' | tail -n 1)" + if [[ ! "${cuda_major}" =~ ^[0-9]+$ ]] || (( cuda_major < 13 )); then + echo "ERROR: sm_121a requires CUDA Toolkit 13.0+; ${cuda_compiler} reports CUDA ${cuda_major:-unknown}." >&2 + exit 1 fi fi diff --git a/fastvideo-kernel/tests/test_attn_qat_infer.py b/fastvideo-kernel/tests/test_attn_qat_infer.py index b231cbc1c4..ada6677ccd 100644 --- a/fastvideo-kernel/tests/test_attn_qat_infer.py +++ b/fastvideo-kernel/tests/test_attn_qat_infer.py @@ -5,7 +5,7 @@ Compares causal and non-causal outputs against a naive float32 reference to verify that the V-row permutation in scaled_fp4_quant_trans_kernel is correct (or absent). -Requires a Blackwell GPU (sm_120a) and fp4attn_cuda / fp4quant_cuda extensions built +Requires a Blackwell GPU (sm_120a or sm_121a) and fp4attn_cuda / fp4quant_cuda extensions built via `cd fastvideo-kernel && ./build.sh`. Run from the fastvideo-kernel directory: @@ -24,9 +24,9 @@ import torch.nn.functional as F from torch.nn.attention import SDPBackend, sdpa_kernel -# The FP4 extensions are only compiled under the sm_120a (Blackwell) arch +# The FP4 extensions are only compiled under the sm_120a/sm_121a arch # gate; on other GPUs the api import below would die at collection time. -pytest.importorskip("fp4attn_cuda", reason="ATTN_QAT_INFER FP4 kernels require a sm_120a build") +pytest.importorskip("fp4attn_cuda", reason="ATTN_QAT_INFER FP4 kernels require a sm_120a/sm_121a build") from attn_qat_infer.api import sageattn_blackwell @@ -129,6 +129,23 @@ def test_accuracy_sdpa(causal: bool, B: int, H: int, L: int, D: int): assert cos >= 0.97, f"({label}) B={B} H={H} L={L} D={D} cos_sim={cos:.4f} < 0.97" +@pytest.mark.parametrize( + "q_len,kv_len", + [(384, 512), (384, 257)], + ids=["wan_t2v_cross_attention", "wan_i2v_image_cross_attention"], +) +def test_cross_attention_unequal_sequence_lengths(q_len: int, kv_len: int): + """Wan cross-attention keeps video queries separate from text/image keys.""" + torch.manual_seed(42) + q = torch.randn(1, 4, q_len, 128, dtype=torch.bfloat16, device=DEVICE) + k = torch.randn(1, 4, kv_len, 128, dtype=torch.bfloat16, device=DEVICE) + v = torch.randn(1, 4, kv_len, 128, dtype=torch.bfloat16, device=DEVICE) + ref = reference_sdpa(q, k, v, is_causal=False) + out = sageattn_blackwell(q.clone(), k.clone(), v.clone(), is_causal=False) + cos = cosine_similarity(out, ref) + assert cos >= 0.97, f"q_len={q_len} kv_len={kv_len} cos_sim={cos:.4f} < 0.97" + + if __name__ == "__main__": print("=" * 65) print("sageattn_blackwell (ATTN_QAT_INFER) inference correctness tests") diff --git a/fastvideo-kernel/tests/test_build_arch.py b/fastvideo-kernel/tests/test_build_arch.py new file mode 100644 index 0000000000..8c5cb768f7 --- /dev/null +++ b/fastvideo-kernel/tests/test_build_arch.py @@ -0,0 +1,78 @@ +import os +import shutil +import subprocess +from pathlib import Path + +import pytest + + +BUILD_SCRIPT = Path(__file__).parents[1] / "build.sh" + + +def _write_executable(path: Path, contents: str) -> None: + path.write_text(contents) + path.chmod(0o755) + + +def _run_build(tmp_path: Path, capability: str, cuda_release: str) -> tuple[subprocess.CompletedProcess[str], Path]: + kernel_dir = tmp_path / "fastvideo-kernel" + kernel_dir.mkdir() + shutil.copy2(BUILD_SCRIPT, kernel_dir / "build.sh") + (kernel_dir / "include/cutlass/include").mkdir(parents=True) + (kernel_dir / "include/tk/include").mkdir(parents=True) + + result_dir = tmp_path / "result" + result_dir.mkdir() + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + _write_executable(fake_bin / "git", "#!/bin/sh\nexit 1\n") + _write_executable(fake_bin / "python3", "#!/bin/sh\nprintf '%s\\n' \"$FAKE_CUDA_CAPABILITY\"\n") + _write_executable( + fake_bin / "nvcc", + "#!/bin/sh\nprintf 'Cuda compilation tools, release %s, V%s.0\\n' \"$FAKE_CUDA_RELEASE\" \"$FAKE_CUDA_RELEASE\"\n", + ) + _write_executable( + fake_bin / "uv", + "#!/bin/sh\nprintf '%s' \"${TORCH_CUDA_ARCH_LIST:-}\" > \"$FV_TEST_RESULT/arch\"\n" + "printf '%s' \"${CMAKE_ARGS:-}\" > \"$FV_TEST_RESULT/cmake_args\"\n", + ) + + env = os.environ.copy() + for name in ("CMAKE_ARGS", "CONDA_PREFIX", "CUDA_HOME", "CUDACXX", "TORCH_CUDA_ARCH_LIST", "VIRTUAL_ENV"): + env.pop(name, None) + env.update({ + "FAKE_CUDA_CAPABILITY": capability, + "FAKE_CUDA_RELEASE": cuda_release, + "FV_TEST_RESULT": str(result_dir), + "PATH": f"{fake_bin}{os.pathsep}{env['PATH']}", + }) + result = subprocess.run( + ["bash", "build.sh"], + cwd=kernel_dir, + env=env, + capture_output=True, + text=True, + check=False, + ) + return result, result_dir + + +@pytest.mark.parametrize( + "capability,torch_arch,cmake_arch", + [("9.0", "9.0a", "90a"), ("12.0", "12.0a", "120"), ("12.1", "12.1a", "121")], +) +def test_build_script_maps_architecture_specific_targets( + tmp_path: Path, capability: str, torch_arch: str, cmake_arch: str +) -> None: + result, result_dir = _run_build(tmp_path, capability, "13.0") + + assert result.returncode == 0, result.stderr + assert (result_dir / "arch").read_text() == torch_arch + assert f"-DCMAKE_CUDA_ARCHITECTURES={cmake_arch}" in (result_dir / "cmake_args").read_text() + + +def test_build_script_rejects_sm121_before_cmake_on_cuda_12(tmp_path: Path) -> None: + result, _ = _run_build(tmp_path, "12.1", "12.9") + + assert result.returncode != 0 + assert "sm_121a requires CUDA Toolkit 13.0+" in result.stderr diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index 4c971ca9e4..ee4efc08d6 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -172,7 +172,8 @@ def __init__(self, softmax_scale=None, causal=False, supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA)) + AttentionBackendEnum.TORCH_SDPA, + AttentionBackendEnum.ATTN_QAT_INFER)) def forward(self, x: torch.Tensor, context: torch.Tensor, context_lens: int): r""" diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 771bc3c323..6205262b98 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -67,8 +67,8 @@ def __init__(self, transformer, scheduler, pipeline=None, transformer_2=None, va dtype=torch.float16, # TODO(will): hack supported_attention_backends=(AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.BSA_ATTN, AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA, - AttentionBackendEnum.SAGE_ATTN_THREE), # hack + AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN_THREE, + AttentionBackendEnum.ATTN_QAT_INFER), # hack # Build metadata for the backend this transformer actually resolved # instead of re-deriving it from the environment. The two agreed # only when the request arrived via the env var: a request passed as