[feat] QAD 5090: Wire the Attn-QAT inference attention backend (8/12)

[alexzms]

Purpose

Continues the Attn-QAT upstreaming stack (#1225). Wires in the
AttnQatInferBackend that #1358 landed as deadcode, now that its modified
SageAttention3 FP4 kernels land in #1455. This makes the 4-bit attention path
selectable end-to-end for Wan-2.1 inference.

Changes

• platforms/interface.py: add ATTN_QAT_INFER to AttentionBackendEnum. The
  name matches AttnQatInferBackend.get_name(), so backend_name_to_enum()
  resolves the FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER override.
• platforms/cuda.py: dispatch ATTN_QAT_INFER → AttnQatInferBackend, guarded
  by is_attn_qat_infer_available(); falls back to Flash Attention when the
  kernel is not built (non-Blackwell, or before [kernel] QAD 5090: Add modified SageAttention3 FP4 inference kernels (7/12) #1455 merges).
• configs/models/dits/base.py: add ATTN_QAT_INFER to the default
  _supported_attention_backends so Wan can select it.

The training-side ATTN_QAT_TRAIN backend depends on a separate Triton kernel
and is deferred to the training upstream.

Test Plan / Test Results

• ✅ name ↔ enum closed: AttnQatInferBackend.get_name() == "ATTN_QAT_INFER"
  is a valid AttentionBackendEnum member.
• ✅ Dispatch verified on a Blackwell node (using the [kernel] QAD 5090: Add modified SageAttention3 FP4 inference kernels (7/12) #1455-built kernel):
  • kernel present → get_attn_backend_cls(ATTN_QAT_INFER) returns
    fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend
  • kernel absent → logs the fallback and returns FlashAttentionBackend
• ✅ pre-commit (yapf / ruff / mypy / codespell) passes.

Notes

• Depends on [kernel] QAD 5090: Add modified SageAttention3 FP4 inference kernels (7/12) #1455 (the kernel). Until that merges, is_attn_qat_infer_available()
  returns False on main, so this safely no-ops to Flash Attention.
• Full runtime/numeric parity is sm_120 (RTX 5090)–specific by kernel design.

Part of #1225.

[mergify]

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 PR merge requirements

Wonderful, this rule succeeded.

• #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 introduces support for a new attention backend, ATTN_QAT_INFER (modified SageAttention3 FP4), across the configuration, platform interface, and CUDA platform implementation. Feedback suggests wrapping the backend import and availability check in a try...except block to prevent potential runtime crashes and ensure a graceful fallback to Flash Attention, as well as removing an unused import.

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]

To improve robustness, wrap the import and availability check of the ATTN_QAT_INFER backend in a try...except block. This prevents potential runtime crashes (such as AttributeError if the kernel is missing expected attributes, or OSError due to CUDA driver/toolkit mismatches) and ensures a graceful fallback to Flash Attention. Additionally, we can remove the unused AttnQatInferBackend import and its associated # noqa: F401 suppression.

            try:
                from fastvideo.attention.backends.attn_qat_infer import is_attn_qat_infer_available
                if is_attn_qat_infer_available():
                    logger.info("Using Attn-QAT inference (modified SageAttention3 FP4) backend.")
                    return "fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend"
            except Exception as e:
                logger.info(e)
            logger.info("Attn-QAT inference kernel is not built. Fall back to Flash Attention.")

[SolitaryThinker]

Hi @alexzms — automated review from Gob, one of @SolitaryThinker's AI reviewers. Findings aren't all human-verified; ping @SolitaryThinker if anything looks off.

TL;DR

This activation-gate review found the slice-7 sm_scale carryover is still not fixed: AttnQatInferImpl.forward() passes sm_scale=self.softmax_scale, but sageattn_blackwell(... **kwargs) ignores it and blockscaled_fp4_attn hardcodes 1/sqrt(D). Because slice 8 now makes ATTN_QAT_INFER selectable through the DiT base config, that silent scale mismatch should be fixed before merge. The in-place k carryover remains latent S2 because the activated path does not pass enable_smoothing_k=True.

Verdict: request-changes

• S0 (blockers): 0
• S1 (must-fix): 1
• S2 (should-fix): 2
• S3 (discussion): not shown here; see review.md

---

Carryover status from #1455 (slice 7 → slice 8 activation)

Slice-7 finding (S2 there, gated by deadcode)	Status at HEAD 04bde902	Severity now
sm_scale swallow → kernel hardcodes scale	Now reachable: wrapper passes sm_scale, kernel still ignores it, and the backend is now selectable via the default DiT backend allow-list	S1
In-place k mutation via enable_smoothing_k=True	Persistent but not currently triggered: defaults remain false and the activated wrapper passes no smoothing flag	S2

Prior gemini-code-assist concerns (status at 04bde902)

Prior concern	Status at HEAD	Severity now
Wrap backend import in try/except for graceful fallback	Safe at HEAD: the selector import is inside the ATTN_QAT_INFER branch, and the backend lazily imports the kernel under try/except ImportError before returning unavailable	n/a
Remove unused import	Lint-clean at HEAD: the import is retained with the same # noqa: F401 validation pattern used by sibling backend branches	n/a

Findings (formatted for upload)

[S1] sm_scale is still swallowed now that the backend is selectable

What: AttnQatInferImpl.forward() passes sm_scale=self.softmax_scale into the kernel wrapper, but sageattn_blackwell(... **kwargs) documents those extra arguments as ignored, and blockscaled_fp4_attn() recomputes softmax_scale = (qlist[0].shape[-1] * 2) ** (-0.5) before calling fp4attn_cuda.fwd(...).

Why it matters: After this PR, FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER can select this backend for DiTs inheriting the default supported-backend list. Any model or attention site with a non-default scale will silently run with 1/sqrt(D) instead of the configured scale, producing a correctness or quality regression without an exception.

Suggested fix: Plumb the scale end-to-end: add an explicit sm_scale/softmax_scale parameter to sageattn_blackwell() and blockscaled_fp4_attn(), pass it to fp4attn_cuda.fwd(...), and keep the current 1/sqrt(D) computation only as the fallback when the caller passes None.

Evidence: fastvideo/attention/backends/attn_qat_infer.py:117-124, fastvideo-kernel/attn_qat_infer/api.py:140-144, fastvideo-kernel/attn_qat_infer/api.py:159, fastvideo/configs/models/dits/base.py:22-28

---

— Gob (@SolitaryThinker's AI reviewer). Full review (including S2/S3 items) is archived locally.

[SolitaryThinker]

/merge

[SolitaryThinker]

Hi @alexzms — automated re-review from Gob, one of @SolitaryThinker's AI reviewers. Findings aren't all human-verified; ping @SolitaryThinker if anything looks off.

TL;DR

F1 is resolved at 8996245d4f: sm_scale is now explicit in sageattn_blackwell, passed to blockscaled_fp4_attn, resolved into softmax_scale, and sent to fp4attn_cuda.fwd(...). This is a differential re-review against my prior request-changes review at 04bde9028f; the remaining items are unchanged latent/follow-up S2s plus S3 hygiene, so the verdict moves to approve-with-followup, not full approve.

Verdict: approve-with-followup

• S0 (blockers): 0
• S1 (must-fix): 0
• S2 (should-fix): 2
• S3 (discussion): not shown here; see review.md

---

Prior findings status (against rereview HEAD 8996245d4f)

Prior finding	Status at HEAD	New severity
F1 [S1] sm_scale swallow → kernel hardcodes scale	✅ RESOLVED in 8996245d	n/a
F2 [S2] In-place k mutation via enable_smoothing_k=True	⏸️ PERSISTENT, still latent because current caller does not pass enable_smoothing_k=True	S2 follow-up
F3 [S2] Base-level opt-in broader than Wan-2.1 target	⏸️ PERSISTENT, base default allow-list still includes ATTN_QAT_INFER	S2 follow-up
F4 [S3] Missing Attn-QAT-Stack: 8/12 trailer	⏸️ PERSISTENT	S3
F5 [S3] Dispatch branch lacks unit test	⏸️ PERSISTENT	S3

Address commit(s): 8996245d ("[bugfix]: plumb sm_scale through sageattn_blackwell + blockscaled_fp4_attn (#1457)").

The key fixed code in fastvideo-kernel/attn_qat_infer/api.py is:

def blockscaled_fp4_attn(qlist: Tuple,
                         klist: Tuple,
                         vlist: Tuple,
                         delta_s: torch.Tensor,
                         KL: int,
                         is_causal: bool = False,
                         per_block_mean: bool = True,
                         is_bf16: bool = True,
                         single_level_p_quant: bool = False,
                         sm_scale: float | None = None
                        ):
    softmax_scale = sm_scale if sm_scale is not None else (qlist[0].shape[-1] * 2) ** (-0.5)
    return fp4attn_cuda.fwd(qlist[0], klist[0], vlist[0], qlist[1], klist[1], vlist[1], delta_s, KL, None, softmax_scale, is_causal, per_block_mean, is_bf16, single_level_p_quant)

def sageattn_blackwell(q, k, v, attn_mask = None, is_causal = False, per_block_mean = True, single_level_p_quant = True, sm_scale: float | None = None, **kwargs):

    o_fp4 = blockscaled_fp4_attn(
    qlist_from_cuda,
    klist_from_cuda,
    vlist_from_cuda,
    delta_s,
    KL,
    is_causal,
    per_block_mean,
    is_bf16,
    single_level_p_quant,
    sm_scale
    )[0][:, :, :QL, :].contiguous()

Findings

No new S0/S1 findings surfaced in the differential re-review. I am moving the prior verdict up one tier because the prior S1 is fixed, but I am not marking this as full approve because F2 and F3 remain S2 follow-ups: preprocess_qkv() still contains the latent in-place k -= k.mean(...) path, and the backend remains enabled at the base DiT allow-list rather than scoped to Wan-only.

---

— Gob (@SolitaryThinker's AI reviewer). Full review (including S3 items + verification log) is archived locally.