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

[Mister-Raggs]

Summary

The attn_qat_infer FP4 attention kernels (modified SageAttention3) were pinned
to sm_120a, with the build/docs treating them as "sm_120a-only." On real
hardware they're also numerically correct on sm_121a (DGX Spark / GB10) — the
"sm_120a-only" was an untested release-scope assumption (the kernels shipped on
RTX 5090), not an intrinsic limit. This makes them buildable + opt-in-selectable
on sm_121, with the correctness evidence below. Default behavior is unchanged
on every existing GPU — this is opt-in (FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER).

Update (validated on the QAD checkpoint): on the checkpoint this backend is
meant for (FastVideo/FastWan-QAD-1.3B), FP4 attention is quality-neutral vs
bf16, and the full 4-bit path (FP4 linear + FP4 attention) runs end-to-end
on sm_121. See the validation comment for the
quality read, the determinism-controlled reason SSIM is inapplicable, and the
A/B/C/D timing. The two notes about "stock weights look worse" and "FP4 linear
crashes" below are superseded by that comment and annotated inline.

Evidence — numerical correctness on sm_121 (GB10)

Built for sm_121a (TORCH_CUDA_ARCH_LIST=12.1a, CUDA 13; cuobjdump confirms
sm_121 SASS), then sageattn_blackwell vs bf16 SDPA on a GB10:

config	cos vs bf16
non-causal, D=128, L=256/512	0.981–0.982
non-causal, D=64, L=256/512	0.981–0.982
causal, D=128, L=256/512	0.985
causal, D=64, L=256/512	0.985

All ≥ 0.981 (min over 3 seeds each), no NaNs. For reference, an idealized
correct-math NVFP4 attention on the same random input caps ~0.975 — so this is at
the FP4 correctness ceiling. Robust across cross-attention (unequal q/kv), the
full Wan seqlen (32760), and sm_scale=None.

Note: an earlier internal measurement had this kernel wrong on sm_121 (cos
~0.52). That was on a pre-#1493/#1506 build; those two FP4-attn correctness
fixes (authored for sm_120) also repaired sm_121. This is on current main.

End-to-end on Wan-2.1-T2V-1.3B (sm_121)

FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER runs end-to-end through a real Wan
forward (self + cross attention via the kernel) with no crash. The full 4-bit
path (FP4 linear + FP4 attention) also runs end-to-end on sm_121 — all four
attention×linear precision combinations generate cleanly (see the validation
comment; QAD-checkpoint 3-step timing there too). Timing at 50 steps on stock
weights was ~10 min vs ~11 min bf16 (single run, ~9%) — modest as expected:
attention is a minority of Wan's denoise (linears dominate), and the win grows
with sequence length (attention is O(L²)) and on sm_120-tuned paths.

Quality

On the intended (QAD) checkpoint, FP4 attention is quality-neutral vs bf16.
Validated on FastVideo/FastWan-QAD-1.3B (Wan2.1-1.3B, 3-step, fixed seed) across
three prompts — visually indistinguishable from bf16; both share the 3-step
distillation's quality ceiling. (SSIM-vs-bf16 ~0.5 is 3-step trajectory
divergence, not loss — determinism-controlled; full proof + full-4-bit results +
timing in the validation comment.)

On stock (non-QAT) Wan-2.1, by contrast, FP4 attention is visibly below bf16
(frames below). That's expected and is not a kernel defect: QAD is
Quantization-Aware Distillation, so the model must be trained to tolerate
FP4 attention — which is exactly why the neutral result needs the QAT checkpoint,
and why release-grade absolute quality on a full-step path is a training
follow-up (below).

BF16 (stock, 50-step):

FP4 (stock, 50-step):

Changes

• fastvideo-kernel/CMakeLists.txt — build the FP4 targets for sm_120a and
  sm_121a (added only on CUDA ≥ 13.0, since sm_121a needs CUDA 13 ptxas;
  pre-13 toolchains keep the exact sm_120a-only build). AUTO-gate + local-GPU
  detect recognize sm_121a.
• wanvideo.py / denoising.py — add ATTN_QAT_INFER to the attention +
  denoising-stage allowlists so it's opt-in-selectable on Wan (default unchanged).
• docs/inference/optimizations.md — correct "sm_120a-only" → sm_120a/sm_121a,
  with the quality caveat.

Not in this PR (follow-ups)

• Full-step QAT checkpoint for sharper absolute quality. The only public QAT
  weights are the 3-step distill, so today's demo is capped by 3-step. A full-step
  QAT-finetune (stage 1 of the QAD recipe, before the DMD step-distillation) would
  give full-sharpness FP4-attention output — a training-side follow-up.
• FP4 linear crashes on Wan/sm_121 — resolved / was stale: FP4 linear
  runs correctly on sm_121 on current main (isolated + end-to-end; see the
  validation comment). Full 4-bit works.

Validation

• Kernel build + cuobjdump sm_121 + cos sweep: on a GB10 (aarch64, CUDA 13,
  torch 2.11+cu130). The sm_121a build path is validated on-silicon; maintainers
  should confirm it on the CI toolchain/arch matrix (it's CUDA-version-guarded so
  it can't change the existing sm_120 build).

🤖 Generated with Claude Code

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed

This rule is failing.

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[gemini-code-assist]

Code Review

This pull request extends support for Blackwell FP4 attention inference (ATTN_QAT_INFER) to include the sm_121a architecture (DGX Spark GB10) alongside sm_120a. It updates the CMake build configuration to dynamically enable sm_121a compilation when CUDA 12.9+ is available, updates documentation, and registers the backend in the model and pipeline configurations. The review feedback highlights a potential build failure when TORCH_CUDA_ARCH_LIST contains 12.1a but the CUDA Toolkit version is below 12.9, suggesting a fallback mechanism to downgrade 121a to 120a in those environments.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If TORCH_CUDA_ARCH_LIST contains 12.1a (either via auto-detection on a GB10 GPU or manual configuration) but the CUDA Toolkit version is less than 12.9 (e.g., CUDA 12.8), the build of fastvideo_kernel_ops will fail because CUDA 12.8 does not support the sm_121a architecture. Since sm_120a binaries are compatible with sm_121a GPUs, we should sanitize FASTVIDEO_MAIN_CUDA_ARCHS to downgrade 121a to 120a when compiling with CUDA < 12.9.

    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}")
    if(CUDAToolkit_VERSION VERSION_LESS 12.9)
        string(REPLACE "121a" "120a" FASTVIDEO_MAIN_CUDA_ARCHS "${FASTVIDEO_MAIN_CUDA_ARCHS}")
        list(REMOVE_DUPLICATES FASTVIDEO_MAIN_CUDA_ARCHS)
    endif()
    set(FASTVIDEO_FP4_ARCHS_MSG "sm_120a")
    if(NOT CUDAToolkit_VERSION VERSION_LESS 12.9)
        set(FASTVIDEO_FP4_ARCHS_MSG "sm_120a/sm_121a")
    endif()
    message(STATUS "[per-arch] main extension archs=${FASTVIDEO_MAIN_CUDA_ARCHS}, fp4* archs=${FASTVIDEO_FP4_ARCHS_MSG}")

[Copilot]

Pull request overview

Enables the ATTN_QAT_INFER (FP4 attention) path to be buildable/opt-in selectable on Blackwell sm_121a (DGX Spark / GB10) in addition to existing sm_120a, while keeping default behavior unchanged unless users explicitly select the backend.

Changes:

• Extend fastvideo-kernel CMake arch detection/gating to recognize 12.1a and build FP4 targets for 120a plus 121a (when the CUDA toolkit supports it).
• Allow ATTN_QAT_INFER to be selected in Wan attention and the denoising stage allowlist (opt-in via FASTVIDEO_ATTENTION_BACKEND).
• Update inference optimization docs to reflect sm_120a/sm_121a support and add a quality caveat for stock (non-QAT) checkpoints.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

File	Description
fastvideo-kernel/CMakeLists.txt	Detects/permits Blackwell 12.1a and adjusts FP4 build arch targeting and status messaging.
fastvideo/models/dits/wanvideo.py	Adds ATTN_QAT_INFER to Wan attention backend allowlist so it can be selected via env var.
fastvideo/pipelines/stages/denoising.py	Adds ATTN_QAT_INFER to the stage allowlist used during attention backend selection.
docs/inference/optimizations.md	Updates documentation to state sm_120a/sm_121a support and documents the stock-weights quality caveat.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[SolitaryThinker]

/merge

[Mister-Raggs]

Update — validated on the QAD checkpoint, and full 4-bit runs on sm_121

Follow-up on a GB10 (aarch64, CUDA 13, torch 2.11+cu130, flashinfer 0.6.14) using
the public quantization-aware-distilled Wan2.1-1.3B checkpoint
FastVideo/FastWan-QAD-1.3B
— the checkpoint this backend is actually meant for. Two results that update the
description above.

1. FP4 attention is quality-neutral on the intended checkpoint

The description's "stock weights look below bf16" caveat was purely a
stock-weights artifact. On the QAD checkpoint (Wan2.1-1.3B, 480×832×77, 3-step,
no CFG, fixed seed), FP4 attention (ATTN_QAT_INFER) is visually
indistinguishable from bf16 across three different prompts. Both arms are equally
limited by the 3-step distillation — that ceiling is the checkpoint's, not FP4's.

Why the eye and not SSIM: SSIM-vs-bf16 lands ~0.5 here, but that is 3-step
trajectory divergence, not quality loss — shown three ways on the box:

• a video vs itself → SSIM 1.0000 (metric is sane);
• rerunning the bf16 reference at fixed seed → 1.0000 (pipeline is deterministic);
• FP4 linear — independently quality-neutral ([perf] LTX2 FP4-linear on DGX Spark (GB10/sm_121): ~24% faster denoise at 1080p #1594) — drops SSIM the same
  ~0.5 as FP4 attention.

So any FP4 numerical nudge reshuffles high-frequency texture through the 3-step
sampler and craters frame-aligned SSIM without touching perceptual quality (the
same reason #1594 rejects SSIM for few-step samplers). On the eye, FP4 ≈ bf16.

2. Full 4-bit (FP4 linear + FP4 attention) runs end-to-end on sm_121

The description lists "FP4 linear crashes on Wan/sm_121" as a follow-up. On
current main it does not crash — that June segfault was stale. Isolated:
flashinfer nvfp4_quantize + mm_fp4 (cutlass) are correct at M=256/1559/32760,
K=N=1536 (cutlass/cudnn/auto all fine; only trtllm is an unsupported-at-cc121
guarded refusal, not a crash). End-to-end, all four attention×linear precision
combinations generate cleanly, including full 4-bit.

Timing (QAD-1.3B, 480×832×77, 3-step, eager, denoise mean over 3 runs)

attention	linear	denoise
bf16	bf16	40.59 s
FP4	bf16	38.24 s
bf16	FP4	41.16 s
FP4	FP4	38.94 s

FP4 attention ≈ 6% denoise, consistent across linear precision. FP4 linear is
break-even at this size/res in eager — it pays off at higher resolution (cf.
#1594's ~24% at 1080p on LTX2, where the linear matmuls are large). So at
1.3B/480p the value is the attention path.

Remaining follow-up

Absolute quality is capped by the 3-step distillation, and there is no public
full-step QAT checkpoint — the only QAT-trained weights that exist are this
3-step distill. A sharper full-step FP4-attention demo needs a full-step
QAT-finetuned checkpoint (the stage-1 output of the QAD recipe, before the DMD
step-distillation) — a training-side follow-up.

Repro harness (A/B/C/D quality, FP4-linear isolator, SSIM/PSNR compare) on
spark/qad-fp4-quality.

[SolitaryThinker]

Maintainer repair wave is pushed at exact head ea1a5220fa97e783e93bfe61b14eccbab0a06eca (no merge action).

Canonical rebase:

• live main: 6e25d94ebca1f46c70332dbd21193793296ec8a0
• original head: 79858cd8a1fde2a648ab8f5e2bbaccd9e1b1c89f
• rebase-only tip: 07de6f040a7ec4b3ee038eae3050c366a8051719
• repair commit: ea1a5220fa97e783e93bfe61b14eccbab0a06eca ([bugfix]: harden SM121 Attn-QAT support)
• range-diff preserved both contributor commits exactly:

1: 73f75dd52 = 1: 833d73521
2: 79858cd8a = 2: 07de6f040

The repair makes the standard local build emit 12.1a on SM121, fails clearly when SM121 is requested with CUDA <13, removes the invalid 121a -> 120a rewrite, and compiles FP4 only for requested SM120/SM121 targets. Runtime availability now gates on actual SM120/SM121 capability so unsupported devices retain the documented Flash/SDPA fallback. Spark docs use FP4 attention with BF16 linears and keep the separate nvfp4_qat linear path SM120-only. The branch also adds build-arch regressions, Fastcheck collection for the runtime capability test, and unequal-length Wan T2V/I2V accuracy cases.

Local evidence before push:

• build/capability regressions: 8 passed
• attention suite: 22 passed, 37 skipped
• combined targeted run: 26 passed, 37 skipped
• bash -n, git diff --check, and project hooks: passed

Exact-head DGX Spark proof:

• NVIDIA GB10 / SM121, aarch64, driver 580.78, PyTorch 2.12.0+cu130, CUDA toolkit/runtime 13.0
• override-free fastvideo-kernel/./build.sh detected 12.1, exported TORCH_CUDA_ARCH_LIST=12.1a, enabled Attn-QAT, selected FP4 target 121a only, and built/installed both extensions in 4m24s
• exact worktree import asserted at /mnt/fv-pr1598-ea1a/fastvideo; fp4attn_cuda and fp4quant_cuda imported and is_attn_qat_infer_available() returned True
• full on-GPU correctness file: 32 passed in 6.09s
• T2V cross-attention (q=384, kv=512) vs BF16 SDPA: cosine 0.9819354415, max abs 0.09716797, mean abs 0.01087134, RMSE 0.01383005
• I2V image cross-attention (q=384, kv=257) vs BF16 SDPA: cosine 0.9828431010, max abs 0.13183594, mean abs 0.01507255, RMSE 0.01917533

The fresh Spark venv initially lacked the repository's declared uv build prerequisite; after installing uv into that existing venv, the same override-free build command passed. This was environment setup, before architecture detection, not a source failure.

Fresh CI may evaluate this new head. No approval, retry, label, or merge action was taken.

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease

[mergify]

Pre-commit checks failed

Hi @Mister-Raggs, the pre-commit checks have failed. To fix them locally:

# Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install

# Run all checks and auto-fix what's possible
pre-commit run --all-files

Common fixes:

• yapf: yapf -i <file> (formatting)
• ruff: ruff check --fix <file> (linting)
• codespell: codespell --write-changes <file> (spelling)

After fixing, commit and push the changes. The checks will re-run automatically.

For future commits, pre-commit will run automatically on changed files before each commit.

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease