[feat] Upstream Attn-QAT Video Diffusion Code - #1225
Closed
RandNMR73 wants to merge 203 commits into
Closed
Conversation
SolitaryThinker
pushed a commit
that referenced
this pull request
May 13, 2026
Per @SolitaryThinker review: the .codex empty marker file was extracted from PR #1225 (sync-branch) but is not wanted on origin/main. Dropping it from Slice 1 of the Attn-QAT decomposition. Attn-QAT-Stack: 1.1/12
SolitaryThinker
added a commit
that referenced
this pull request
May 14, 2026
Adds NVFP4QATConfig — the Attn-QAT flavor of FP4 quantization. Distinct from main's NVFP4Config (NVIDIA Blackwell hardware FP4 inference); this config is for the QAT training side of the Attn-QAT stack. Pure deadcode: no existing code path imports NVFP4QATConfig yet. The `nvfp4_qat` method literal in QuantizationMethods makes it selectable, but the caller lands in slice 3+. Renamed from PR-1225's original `Fp4Config` / `fp4_config.py` / `"fp4"` literal to avoid collision with PR #1334's NVFP4 plumbing already merged to main. Extracted from PR #1225 (#1225) by @RandNMR73. Source SHA: 3f818d0 Attn-QAT-Stack: 2/12 Co-Authored-By: jzhang38 <42993249+jzhang38@users.noreply.github.com> Co-Authored-By: RandNMR73 <99706358+RandNMR73@users.noreply.github.com>
SolitaryThinker
added a commit
that referenced
this pull request
May 15, 2026
Adds the FastVideo-native FP4 linear forward helper module extracted from PR-1225. The source module does not define a public Fp4Linear class, so this slice keeps the original fp4linear.py filename and helper names. Pure deadcode: no existing code path imports this helper yet. Activation lands in Slice-12 of the decomposition. Extracted from PR #1225 (#1225) by @RandNMR73. Source SHA: 3f818d0. Attn-QAT-Stack: 3/12 Co-Authored-By: jzhang38 <42993249+jzhang38@users.noreply.github.com> Co-Authored-By: RandNMR73 <99706358+RandNMR73@users.noreply.github.com>
SolitaryThinker
added a commit
to SolitaryThinker/FastVideo
that referenced
this pull request
May 16, 2026
…/12)
Add fastvideo/attention/backends/attn_qat_{infer,train}.py extracted from PR hao-ai-lab#1225.
Backends not yet registered in the selector — full deadcode until activation (slice 12).
Attn-QAT-Stack: 4/12
Co-Authored-By: jzhang38 <42993249+jzhang38@users.noreply.github.com>
Co-Authored-By: RandNMR73 <99706358+RandNMR73@users.noreply.github.com>
This was referenced May 16, 2026
SolitaryThinker
added a commit
to SolitaryThinker/FastVideo
that referenced
this pull request
May 23, 2026
… 5/12) Slice 5/12 of the PR hao-ai-lab#1225 decomposition (Attn-QAT stack). Tier-2: backward-compatible additions to the shared attention infra that several backends consume. What this adds (per file): * `fastvideo/attention/backends/abstract.py` (+6/-2) - Moves `VSA_sparsity: float = 0.0` (kw-only) up to the base `AttentionMetadata` dataclass so non-VSA backends can satisfy a uniform metadata interface without re-declaring the field. - Adds `AttentionMetadata.__getattr__` raising `AttributeError` so callers that typo a field get a stable, matchable error message instead of relying on default object behavior. - Loosens `AttentionMetadataBuilder.build` signature from `**kwargs: dict[str, Any]` to `**kwargs: Any` to match how subclass builders are typed throughout the codebase. * `fastvideo/attention/backends/bsa_attn.py` (+3/-5) - Collapses the duplicated double-fallback flash-attn import block to a single import from `fastvideo.attention.utils.flash_attn_no_pad` (which now centralizes the fallback chain — see below). When the centralized impl isn't available, sets the symbol to None and `FLASH_ATTN_AVAILABLE=False`. * `fastvideo/attention/backends/video_sparse_attn.py` (+2/-3) - Removes the redundant `VSA_sparsity: float` field from the `VideoSparseAttentionMetadata` subclass (it is now inherited from the base with a default of 0.0). VSA continues to pass `VSA_sparsity=...` as a kwarg into the metadata constructor. - Adds `-> None` return annotations to `__init__` / `prepare` on the builder. * `fastvideo/attention/backends/vmoba.py` (+8/-4) - Adds `-> None` return annotations to `__init__` / `prepare`. - Tightens the `device` parameter to `torch.device | None = None` (was an unannotated default-`None`). - Narrows the `attn_metadata` parameter on `VMOBAAttentionImpl.forward` from the base `AttentionMetadata` to the concrete `VideoMobaAttentionMetadata`. - Adds an explicit `assert self.layer_idx is not None` guard and an `else: raise ValueError` for the moba_layer chunk-selection switch so a misconfigured config can't silently use an uninitialized `moba_chunk_size`. - Pre-declares `moba_chunk_size` with its union type for mypy. * `fastvideo/attention/layer.py` (+2) - Adds an `attention_mask: torch.Tensor | None = None` parameter to `DistributedAttention_VSA.forward` (downstream backends in the stack consume it). * `fastvideo/attention/utils/flash_attn_no_pad.py` (+47/-33) - Centralizes the varlen-flash-attn fallback chain into a single `_resolve_flash_attn_varlen_func()` helper. The fallback order is: `fastvideo.attention.utils.flash_attn_cute` → `flash_attn_interface` → `flash_attn`. Module-level `flash_attn_varlen_func_impl` is now this helper's output, so consuming backends (bsa, vsa) import one stable symbol. - Adds type annotations to every public function and to the helper. Files in PR hao-ai-lab#1225 considered but NOT applied: * `fastvideo/attention/backends/sage_attn3.py` — the source-SHA diff shrinks `get_supported_head_sizes()` from `[64, 128, 256]` to `[64, 128]`. That removal is unrelated to QAT-compat (head_size=256 has been supported since the original SAGE3 backend landed in hao-ai-lab#815) and is treated as an obsolete edit from the source branch's history. Not applied; current main's behavior is preserved. Test coverage added (`fastvideo/tests/attention/`): * `test_attention_metadata_base.py` — 8 CPU-only smoke tests for the new base-class semantics: `VSA_sparsity` default 0.0, kw-only-ness, subclass inheritance and override, `__getattr__` `AttributeError` message contract, and `asdict_zerocopy` field handling. * `test_flash_attn_no_pad_resolver.py` — 3 tests for the fallback chain (cute → interface → flash_attn). Skips at module level when `flash_attn` is not installed (the file has an unconditional top-level import that pre-exists this slice). Pre-commit gate (yapf + ruff + codespell + mypy) passes on all changed files. Local pytest of the new tests: 8 passed, 1 skipped (resolver suite, no flash_attn installed in this env). Source: extracted from PR hao-ai-lab#1225 (hao-ai-lab#1225, SHA 3f818d0) and 3-way-merged onto current main (slice 4 / PR hao-ai-lab#1358 merged as fda0203). The `bsa_attn.py` and `video_sparse_attn.py` files required 3-way merge because main has diverged from the source SHA's merge-base; both merged cleanly with no manual conflict resolution. Co-Authored-By: jzhang38 <42993249+jzhang38@users.noreply.github.com> Co-Authored-By: RandNMR73 <99706358+RandNMR73@users.noreply.github.com> Attn-QAT-Stack: 5/12
SolitaryThinker
added a commit
to SolitaryThinker/FastVideo
that referenced
this pull request
May 24, 2026
Slice 6 of 12 in the PR hao-ai-lab#1225 decomposition. Tier 2 — backward- compat additions, gated paths only. No activation in this slice. What this adds -------------- * fastvideo/layers/linear.py (+54): adds opt-in shape-tracking instrumentation to ``ReplicatedLinear`` so upcoming QAT-aware backends can discover which GEMM shapes need quantized kernels. Gated by the class attr ``enable_shape_tracking = False``; the default forward path is bit-identical to pre-slice behavior. Adds ``get_shape_mapping``, ``reset_shape_tracking``, ``_track_shape``, ``print_shape_summary``. No new constructor params. * fastvideo/layers/mlp.py (+22): adds an optional ``quant_config: QuantizationConfig | None = None`` kwarg to ``MLP.__init__`` and threads it (plus an explicit ``prefix``) into the two underlying ``ReplicatedLinear`` instances. When ``quant_config is not None``, runs ``process_weights_after_loading`` on each sub-layer's resolved quant method. When ``quant_config is None`` (default), behavior is unchanged: ``ReplicatedLinear`` falls back to ``UnquantizedLinearMethod`` exactly as before. * fastvideo/models/dits/wanvideo.py (+67): wires ``quant_config`` through ``WanSelfAttention``, ``WanI2VCrossAttention``, ``WanTransformerBlock``, ``WanTransformerBlock_VSA``, and ``WanTransformer3DModel`` constructors so a future ``NVFP4QAT``- configured Wan2.1 build can quantize its attention QKV/out projections and FFN. Reads ``config.quant_config`` from ``WanVideoConfig`` (the field is already present on the shared ``DiTBaseConfig``). All new kwargs default to ``None``; default Wan2.1 path stays bit-identical. Files in PR hao-ai-lab#1225 considered but NOT applied -------------------------------------------- The source-SHA ``fastvideo/layers/linear.py`` also contains several edits that pre-date current ``main`` and would silently regress it: * Removal of the ``NVFP4Config``-only-quantizes-a-curated-subset explanatory comments in ``LinearBase.__init__`` and ``ReplicatedLinear.__init__`` (added on main as part of slice 3 / PR hao-ai-lab#1336). * Removal of the ``if self.quant_method is None: self.quant_method = UnquantizedLinearMethod()`` fallback inside ``LinearBase.__init__`` (also part of the slice 3 hardening). * A constructor / ``create_weights`` reformat from multi-line to compact one-line style — pure style noise. * ``assert self.quant_method is not None`` → ``if self.quant_method is None: self.quant_method = UnquantizedLinearMethod()`` in ``ColumnParallelLinear.__init__/forward`` and ``RowParallelLinear.__init__/forward``. ``LinearBase.__init__`` on current ``main`` already guarantees ``quant_method`` is non-None, so the source PR's defensive checks would be no-ops; they pre-date the slice 3 base-class hardening. * The same ``if quant_method is None`` defensive insert in ``ReplicatedLinear.forward`` — also a no-op against current ``main`` for the same reason. None of the skipped edits affect the FP4 path; current ``main``'s behavior on those lines is strictly stronger than the source SHA's. This mirrors slice 5's intentional skip of the ``sage_attn3.py`` head_size removal (see PR hao-ai-lab#1383). Also dropped: an unused ``from contextlib import nullcontext`` import that the source PR staged in ``wanvideo.py`` for a deeper-stack slice (ruff would reject it as unused). Stacking -------- Base: ``main`` (slice 5 / PR hao-ai-lab#1383 merged at ``ba75ad82dbe4a7069412494c051c1c69155fdc9d``). No stack dependency — this is a clean linear PR off ``main``. Provenance ---------- Files extracted from PR hao-ai-lab#1225 (hao-ai-lab#1225) at source SHA ``3f818d0fc532ec6494b465967d5f485150917d0c`` and audited against current ``main``. ``mlp.py`` and ``wanvideo.py`` (modulo the dropped unused import) were applied directly — ``main`` had not diverged from the source's merge-base for those files. ``linear.py`` was hand-merged to preserve current ``main``'s slice-3 hardening (see the ``NOT applied`` list above); only the additive shape-tracking surface was carried over. Pre-commit gate (yapf + ruff + codespell + mypy) passes on all three changed files. Test plan --------- No new tests this slice. The shape-tracking surface is opt-in instrumentation (default disabled) and the ``quant_config`` plumbing is dormant until a future slice sets ``config.quant_config`` to a non-None value. The activation slice (12/12) will carry the contract test for the full FP4 Wan-2.1 path. Sequence -------- Attn-QAT-Stack: 6/12. Earlier merged slices: 4/12 (PR hao-ai-lab#1358), 5/12 (PR hao-ai-lab#1383). Out of scope for this slice: the actual FP4 activation switch, weight-loading conversion, and any cross-cutting config registration (later slices). Co-Authored-By: Peiyuan Zhang <a1286225768@gmail.com> Co-Authored-By: Matthew Noto <notomatthew31@gmail.com>
Contributor
|
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
added a commit
to SolitaryThinker/FastVideo
that referenced
this pull request
Jun 8, 2026
Slice 6 of 12 in the PR hao-ai-lab#1225 decomposition. Tier 2 — backward- compat additions, gated paths only. No activation in this slice. What this adds -------------- * fastvideo/layers/linear.py (+54): adds opt-in shape-tracking instrumentation to ``ReplicatedLinear`` so upcoming QAT-aware backends can discover which GEMM shapes need quantized kernels. Gated by the class attr ``enable_shape_tracking = False``; the default forward path is bit-identical to pre-slice behavior. Adds ``get_shape_mapping``, ``reset_shape_tracking``, ``_track_shape``, ``print_shape_summary``. No new constructor params. * fastvideo/layers/mlp.py (+22): adds an optional ``quant_config: QuantizationConfig | None = None`` kwarg to ``MLP.__init__`` and threads it (plus an explicit ``prefix``) into the two underlying ``ReplicatedLinear`` instances. When ``quant_config is not None``, runs ``process_weights_after_loading`` on each sub-layer's resolved quant method. When ``quant_config is None`` (default), behavior is unchanged: ``ReplicatedLinear`` falls back to ``UnquantizedLinearMethod`` exactly as before. * fastvideo/models/dits/wanvideo.py (+67): wires ``quant_config`` through ``WanSelfAttention``, ``WanI2VCrossAttention``, ``WanTransformerBlock``, ``WanTransformerBlock_VSA``, and ``WanTransformer3DModel`` constructors so a future ``NVFP4QAT``- configured Wan2.1 build can quantize its attention QKV/out projections and FFN. Reads ``config.quant_config`` from ``WanVideoConfig`` (the field is already present on the shared ``DiTBaseConfig``). All new kwargs default to ``None``; default Wan2.1 path stays bit-identical. Files in PR hao-ai-lab#1225 considered but NOT applied -------------------------------------------- The source-SHA ``fastvideo/layers/linear.py`` also contains several edits that pre-date current ``main`` and would silently regress it: * Removal of the ``NVFP4Config``-only-quantizes-a-curated-subset explanatory comments in ``LinearBase.__init__`` and ``ReplicatedLinear.__init__`` (added on main as part of slice 3 / PR hao-ai-lab#1336). * Removal of the ``if self.quant_method is None: self.quant_method = UnquantizedLinearMethod()`` fallback inside ``LinearBase.__init__`` (also part of the slice 3 hardening). * A constructor / ``create_weights`` reformat from multi-line to compact one-line style — pure style noise. * ``assert self.quant_method is not None`` → ``if self.quant_method is None: self.quant_method = UnquantizedLinearMethod()`` in ``ColumnParallelLinear.__init__/forward`` and ``RowParallelLinear.__init__/forward``. ``LinearBase.__init__`` on current ``main`` already guarantees ``quant_method`` is non-None, so the source PR's defensive checks would be no-ops; they pre-date the slice 3 base-class hardening. * The same ``if quant_method is None`` defensive insert in ``ReplicatedLinear.forward`` — also a no-op against current ``main`` for the same reason. None of the skipped edits affect the FP4 path; current ``main``'s behavior on those lines is strictly stronger than the source SHA's. This mirrors slice 5's intentional skip of the ``sage_attn3.py`` head_size removal (see PR hao-ai-lab#1383). Also dropped: an unused ``from contextlib import nullcontext`` import that the source PR staged in ``wanvideo.py`` for a deeper-stack slice (ruff would reject it as unused). Stacking -------- Base: ``main`` (slice 5 / PR hao-ai-lab#1383 merged at ``ba75ad82dbe4a7069412494c051c1c69155fdc9d``). No stack dependency — this is a clean linear PR off ``main``. Provenance ---------- Files extracted from PR hao-ai-lab#1225 (hao-ai-lab#1225) at source SHA ``3f818d0fc532ec6494b465967d5f485150917d0c`` and audited against current ``main``. ``mlp.py`` and ``wanvideo.py`` (modulo the dropped unused import) were applied directly — ``main`` had not diverged from the source's merge-base for those files. ``linear.py`` was hand-merged to preserve current ``main``'s slice-3 hardening (see the ``NOT applied`` list above); only the additive shape-tracking surface was carried over. Pre-commit gate (yapf + ruff + codespell + mypy) passes on all three changed files. Test plan --------- No new tests this slice. The shape-tracking surface is opt-in instrumentation (default disabled) and the ``quant_config`` plumbing is dormant until a future slice sets ``config.quant_config`` to a non-None value. The activation slice (12/12) will carry the contract test for the full FP4 Wan-2.1 path. Sequence -------- Attn-QAT-Stack: 6/12. Earlier merged slices: 4/12 (PR hao-ai-lab#1358), 5/12 (PR hao-ai-lab#1383). Out of scope for this slice: the actual FP4 activation switch, weight-loading conversion, and any cross-cutting config registration (later slices). Co-Authored-By: Peiyuan Zhang <a1286225768@gmail.com> Co-Authored-By: Matthew Noto <notomatthew31@gmail.com>
alexzms
added a commit
that referenced
this pull request
Jun 12, 2026
Add the attn_qat_infer Blackwell FP4 attention + quantization CUDA kernels (modified SageAttention3) into fastvideo-kernel, gated behind a build flag and landed as deadcode (not yet wired into a backend). - fastvideo-kernel/attn_qat_infer/: Blackwell FP4 attention (blackwell/api.cu) and 4D FP4 quantization (quantization/fp4_quantization_4d.cu) CUDA kernels, plus their Python wrappers (api.py) and microbenchmarks. - CMake: new FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER option (AUTO/ON/OFF). AUTO only builds on CUDA Toolkit 12.8+ with Blackwell sm_120a; otherwise the kernels are skipped, so the change is inert on existing CI/GPUs. - pyproject/MANIFEST: ship the attn_qat_infer module in the wheel. This is the modified-SageAttention3-kernel item from the #1225 tracker. The attn_qat_infer / attn_qat_train attention backends already on main import these modules lazily, so they stay dormant until a follow-up PR wires the backend in. Part of #1225.
This was referenced Jun 12, 2026
alexzms
added a commit
that referenced
this pull request
Jun 12, 2026
Add the attn_qat_infer Blackwell FP4 attention + quantization CUDA kernels (modified SageAttention3) into fastvideo-kernel, gated behind a build flag and landed as deadcode (not yet wired into a backend). - fastvideo-kernel/attn_qat_infer/: Blackwell FP4 attention (blackwell/api.cu) and 4D FP4 quantization (quantization/fp4_quantization_4d.cu) CUDA kernels, plus their Python wrappers (api.py) and microbenchmarks. - CMake: new FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER option (AUTO/ON/OFF). AUTO only builds on CUDA Toolkit 12.8+ with Blackwell sm_120a; otherwise the kernels are skipped, so the change is inert on existing CI/GPUs. - pyproject/MANIFEST: ship the attn_qat_infer module in the wheel. This is the modified-SageAttention3-kernel item from the #1225 tracker. The attn_qat_infer / attn_qat_train attention backends already on main import these modules lazily, so they stay dormant until a follow-up PR wires the backend in. Part of #1225. Co-authored-by: William Lin <SolitaryThinker@users.noreply.github.com> Co-authored-by: Loay Rashid <42599591+loaydatrain@users.noreply.github.com> Co-authored-by: Kaiqin Kong <k1kong@ucsd.edu>
alexzms
added a commit
that referenced
this pull request
Jun 16, 2026
…ipe finale (12/12) Completes the QAD training recipe: quantization-aware DMD distillation of Wan2.1-T2V-1.3B down to 3 sampling steps, with generator-only Attn-QAT. - component_loader.py: generator-only QAT for DMD distillation. The teacher (real_score) and critic (fake_score) transformers load with the _loading_teacher_critic_model flag; mask the nvfp4_qat quant and the global ATTN_QAT_TRAIN attention env for them so only the generator runs fake-quant attention. Config-driven, no monkey-patching, reuses the existing flag. - distill_dmd_qat.sh: stage-2 DMD distillation script (3-step, generator init from the stage-1 finetune checkpoint). - README: the full two-stage recipe (QAT finetune -> QAT DMD distill to 3 steps). Verified end-to-end on Blackwell (GB200/sm_100): the generator loads with ATTN_QAT_TRAIN while teacher/critic load full precision; the DMD double loop runs (generator updates every generator_update_interval, critic every step, healthy loss), 3-step validation generates videos, checkpoint saved. Part of #1225. Co-authored-by: William Lin <SolitaryThinker@users.noreply.github.com> Co-authored-by: Loay Rashid <42599591+loaydatrain@users.noreply.github.com> Co-authored-by: Kaiqin Kong <k1kong@ucsd.edu>
This was referenced Jun 16, 2026
SolitaryThinker
added a commit
that referenced
this pull request
Jun 19, 2026
Add FP8 (e4m3) quantization-aware training for the DiT linear layers, mirroring the FP4 linear STE path (#1463). Small follow-on for FP8-class GPUs. - layers/fp8linear.py: _LinearFWD8BWD16Fn — FP8 forward (torch._scaled_mm on sm89+, bf16 fake-quant fallback on older GPUs) + full-precision backward (STE). Absmax tensor/row-wise scaling, FP8_MAX=448. From the fork's FP8 training work (commit 96011d29), adapted to the config-driven path (no monkey-patch flags). - layers/quantization/fp8_qat_train_config.py: a training quant method bridging the STE into quant_config, registered "fp8_qat_train" (Wan to_q/k/v/out + ffn). - register it in layers/quantization/__init__.py. No flashinfer needed; runs on any sm89+ GPU (and older via the bf16 fallback), not just Blackwell. Enable with --transformer-quant fp8_qat_train (the training CLI arg + string->config resolve added in #1463). Verified on Blackwell (GB200/sm_100): unit test = FP8 forward + nonzero weight grad (STE); a finetune smoke with --transformer-quant fp8_qat_train trains 10 steps with healthy, decreasing loss / grad. Part of #1225. Co-authored-by: William Lin <SolitaryThinker@users.noreply.github.com> Co-authored-by: Loay Rashid <42599591+loaydatrain@users.noreply.github.com> Co-authored-by: Kaiqin Kong <k1kong@ucsd.edu>
Collaborator
|
closing, as this was upstreamed in other PRs |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary