[bugfix]: harden SM121 Attn-QAT support

Author: SolitaryThinker

.github/workflows/publish-kernel.yml

@@ -67,7 +67,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:
@@ -163,15 +163,16 @@ 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 + 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;12.1a"
@@ -243,7 +244,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.

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

docs/inference/optimizations.md

@@ -146,20 +146,19 @@ 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_120a/sm_121a)
+### 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 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`). On other
-GPUs the backend logs a notice and falls back to Flash Attention.
+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.
 
 > **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
@@ -168,26 +167,24 @@ GPUs the backend logs a notice and falls back to Flash Attention.
 
 See the [Attn-QAT paper](https://arxiv.org/abs/2603.00040).
 
-Enable both halves — attention via the env var, linear via `transformer_quant`:
+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

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()
@@ -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
@@ -208,13 +221,7 @@ else()
         set(_WANTS_ATTN_QAT_INFER ON)
     elseif(FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER STREQUAL "AUTO")
         if(TORCH_CUDA_ARCH_LIST)
-            # Consumer/workstation Blackwell: sm_120a (RTX 5090 / PRO 6000) and
-            # sm_121a (DGX Spark GB10). The block-scaled FP4 MMA is valid on both
-            # (GB10-verified: cos ~0.98 vs bf16 SDPA across shapes).
-            string(REGEX MATCH
-                "(^|[; ,])((12\\.0a)|(120a)|(sm_120a)|(12\\.1a)|(121a)|(sm_121a))([; ,]|$)"
-                _HAS_BLACKWELL_FP4 "${TORCH_CUDA_ARCH_LIST}")
-            if(_HAS_BLACKWELL_FP4)
+            if(_FV_HAS_120A OR _FV_HAS_121A)
                 set(_WANTS_ATTN_QAT_INFER ON)
             endif()
         else()
@@ -242,6 +249,18 @@ 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
@@ -273,14 +292,6 @@ 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}")
-    # sm_121a needs CUDA 13's ptxas. On older toolchains, downgrade any requested
-    # 121a to 120a for the *main* extension too (not just the fp4* targets below),
-    # so a CUDA-12.x build with 12.1a in the arch list doesn't fail to compile —
-    # sm_120a PTX forward-JITs to sm_121 at runtime.
-    if(CUDAToolkit_VERSION VERSION_LESS 13.0)
-        string(REPLACE "121a" "120a" FASTVIDEO_MAIN_CUDA_ARCHS "${FASTVIDEO_MAIN_CUDA_ARCHS}")
-        list(REMOVE_DUPLICATES FASTVIDEO_MAIN_CUDA_ARCHS)
-    endif()
     message(STATUS "[per-arch] main extension archs=${FASTVIDEO_MAIN_CUDA_ARCHS}, fp4* archs set per-target below")
 endif()
 
@@ -421,15 +432,8 @@ if(ENABLE_ATTN_QAT_INFER)
         $<$<COMPILE_LANGUAGE:CXX>:-O3 -std=c++17>
         $<$<COMPILE_LANGUAGE:CUDA>:${ATTN_QAT_INFER_CUDA_FLAGS}>
     )
-    # FP4 kernels target consumer/workstation Blackwell: sm_120a (RTX 5090 /
-    # PRO 6000) and sm_121a (DGX Spark GB10). sm_121a needs CUDA 13's ptxas
-    # (CUDA 12.x lacks it), so it is added only on CUDA >= 13.0 — older
-    # toolchains keep the original sm_120a-only build unchanged. sm_121a
-    # GB10-verified (cos ~0.98 vs bf16).
-    set(FASTVIDEO_FP4_CUDA_ARCHS "120a")
-    if(NOT CUDAToolkit_VERSION VERSION_LESS 13.0)
-        list(APPEND FASTVIDEO_FP4_CUDA_ARCHS "121a")
-    endif()
+    # 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 "${FASTVIDEO_FP4_CUDA_ARCHS}"
@@ -470,6 +474,6 @@ message(STATUS "host / backend:       ${CMAKE_SYSTEM_PROCESSOR} / ${GPU_BACKEND}
 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 "fp4attn/fp4quant (sm_120a/sm_121a, 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 "============================================================")

fastvideo-kernel/README.md

@@ -30,7 +30,7 @@ Runtime-JIT kernels (no build step, ship in every wheel/image):
 | 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 | — |
 | | | arm64 cuda12.6.3 (GH200) | `9.0a` | — (aarch64) | — |
 | | | arm64 cuda13.0.0 (GB10 / DGX Spark) | `12.1` | — | — |
-| Local `./build.sh` | manual | probes the visible GPU via torch | detected | ON iff sm_90 (non-aarch64 host) | ON iff sm_120 |
+| Local `./build.sh` | manual | probes the visible GPU via torch | detected | ON iff sm_90 (non-aarch64 host) | ON iff sm_120 or sm_121 |
 
 Notes:
 

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
 

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:
@@ -23,9 +23,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
 
@@ -126,6 +126,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")