Skip to content

[kernel] Build + allow attn_qat_infer FP4 attention on sm_121a (DGX Spark) #1598

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Mister-Raggs wants to merge 4 commits into
base: main
Choose a base branch
from
Open

[kernel] Build + allow attn_qat_infer FP4 attention on sm_121a (DGX Spark) #1598

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
09357e5
[kernel] Build + allow attn_qat_infer FP4 attention on sm_121a (DGX S…
Mister-Raggs Jul 13, 2026
3a3bbc9
[kernel] Address review: guard sm_121a on CUDA>=13, downgrade main ar…
Mister-Raggs Jul 14, 2026
eed2e2f
[bugfix]: harden SM121 Attn-QAT support
SolitaryThinker Jul 15, 2026
aea590c
[test]: drop redundant attn_qat_infer capability test (superseded by …
Mister-Raggs Jul 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 11 additions & 9 deletions .github/workflows/publish-kernel.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down
6 changes: 3 additions & 3 deletions docs/getting_started/installation/spark.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
40 changes: 23 additions & 17 deletions docs/inference/optimizations.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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`
Expand Down Expand Up @@ -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
Expand Down
57 changes: 41 additions & 16 deletions fastvideo-kernel/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Expand All @@ -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()
Expand All @@ -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
Expand All @@ -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=<arch> (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=<arch> (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
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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()

Expand All @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -515,8 +537,11 @@ if(ENABLE_ATTN_QAT_INFER)
$<$<COMPILE_LANGUAGE:CXX>:-O3 -std=c++17>
$<$<COMPILE_LANGUAGE:CUDA>:${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
)
Expand All @@ -532,7 +557,7 @@ if(ENABLE_ATTN_QAT_INFER)
$<$<COMPILE_LANGUAGE:CUDA>:${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
)
Expand All @@ -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 "============================================================")
Loading
Loading