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

[alexzms]

Purpose

Continues the Attn-QAT upstreaming stack (#1225). First training-side slice:
makes the AttnQatTrainBackend (landed as deadcode in #1358) selectable for
quantization-aware finetuning, mirroring the inference wiring of #1457. The
fake-quant-in-backward attention path becomes reachable via
FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN.

This is config-driven and intentionally avoids the research fork's
monkey-patch module-swapping (traverse_swap_module / swap_*): main already
selects the attention backend from config in both training and inference, so no
training-pipeline change is needed.

Changes

• platforms/interface.py: add ATTN_QAT_TRAIN to AttentionBackendEnum.
• platforms/cuda.py: dispatch ATTN_QAT_TRAIN → AttnQatTrainBackend, guarded
  by is_attn_qat_train_available().
• attention/backends/attn_qat_train.py: add is_attn_qat_train_available()
  helper (mirrors is_attn_qat_infer_available()).
• configs/models/dits/base.py: add ATTN_QAT_TRAIN to the default supported set.

Test Plan / Test Results

• ✅ name ↔ enum closed: AttnQatTrainBackend.get_name() == "ATTN_QAT_TRAIN" is
  a valid enum member.
• ✅ Dispatch verified: get_attn_backend_cls(ATTN_QAT_TRAIN) returns the backend
  when the kernel is available; otherwise warns and falls back to Flash Attention.
• ✅ pre-commit (yapf / ruff / mypy / codespell) passes.

Notes

• The training Triton kernel fastvideo_kernel.triton_kernels.attn_qat_train is
  not upstreamed yet, so this is deadcode for now (selecting it falls back to
  Flash with a warning). A follow-up will land the kernel.
• Overlaps [feat] QAD 5090: Wire the Attn-QAT inference attention backend (8/12) #1457 on interface.py / cuda.py / base.py (adjacent enum /
  dispatch / supported-list lines). Whichever merges second needs a trivial
  rebase.
• Open design question: for training, should a missing kernel hard-error
  instead of silently falling back (to avoid "fake QAT")? Currently warn+fallback
  to match inference.

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

Waiting for

• check-success=fastcheck-passed
• check-success=full-suite-passed

This rule is failing.

• check-success=fastcheck-passed
• check-success=full-suite-passed
• #approved-reviews-by>=1
• 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 the ATTN_QAT_TRAIN (fake-quantized attention) backend, updating the attention backend enum, configuration files, and CUDA platform backend loader. The reviewer suggested raising an ImportError instead of silently falling back to Flash Attention when the QAT training kernel is unavailable, to prevent training without fake-quantization.

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]

Hard-error on missing QAT training kernel

For training, falling back to standard Flash Attention when the QAT kernel is missing is highly risky. QAT (Quantization-Aware Training) introduces fake-quantization noise during the forward and backward passes to prepare the model for low-precision deployment. If it silently falls back, the model will be trained as a standard float16/bfloat16 model, wasting significant compute resources and resulting in a model that cannot be quantized properly.

We should raise an ImportError instead of falling back to Flash Attention when ATTN_QAT_TRAIN is explicitly requested but unavailable. Additionally, we can simplify the import statement by removing the unused AttnQatTrainBackend import.

        elif selected_backend == AttentionBackendEnum.ATTN_QAT_TRAIN:
            from fastvideo.attention.backends.attn_qat_train import is_attn_qat_train_available
            if is_attn_qat_train_available():
                logger.info("Using Attn-QAT training (fake-quantized attention) backend.")
                return "fastvideo.attention.backends.attn_qat_train.AttnQatTrainBackend"
            raise ImportError(
                "Attn-QAT training kernel is not built or available. "
                "To avoid training without fake-quantization (which defeats the purpose of QAT), "
                "please ensure the fastvideo-kernel package is installed and compiled with QAT support."
            )

[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

[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

I recommend changing the training-side missing-kernel path to hard-error rather than warning and silently falling back, because otherwise a user can think they are QAT-training while actually training the attention path without fake quantization. Separately, this PR currently needs a mechanical rebase against current main after #1457 merged; that is a pre-merge metadata gate, not a code-quality concern.

Verdict: request-changes

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

---

Pre-merge gate: rebase required

PR is currently CONFLICTING with main (the needs-rebase label is applied). After #1457 merged at 20:12Z, interface.py, cuda.py, and base.py all gained the ATTN_QAT_INFER lines this PR also touches. Rebase is mechanical (adjacent enum/list additions, no semantic conflict) but blocks merge.

Prior gemini-code-assist concerns (status at 4e8ee637)

Prior concern	Status at HEAD	Severity now
Hard-error vs warn+fallback for missing training kernel	⏸️ PERSISTENT (same question, my recommendation: hard-error)	S2

Findings (formatted for upload)

[S2] Training-side fallback should hard-error, not warn+silent-fallback

What: CudaPlatform.get_attn_backend_cls() accepts FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN, checks is_attn_qat_train_available(), and when the Triton kernel is missing logs a warning before dropping through to the normal Flash/SDPA fallback path.

Why it matters: Silent fallback in training means the user thinks they're doing QAT but is actually full-precision-training the attention path. The model weights get shaped by mismatched distributions; the resulting "QAT" checkpoint is a lie — at inference time it gets quantized into a state it was never trained for. Inference (#1457) can defensibly warn+fallback because the user gets a quality regression, not a fundamentally different training run.

Suggested fix: For the ATTN_QAT_TRAIN branch only, raise an ImportError or RuntimeError when is_attn_qat_train_available() is false. Keep the inference branch's warn+fallback behavior unchanged if that remains the intended inference UX.

Evidence: fastvideo/platforms/cuda.py:143-150

---

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

[SolitaryThinker]

/merge