[refactor]: linear/mlp FP4 path additions for Wan-2.1 (Attn-QAT 6/12)

[SolitaryThinker]

Slice 6 of 12 in the PR #1225 decomposition. Tier 2 — backward-compat additions, gated paths only. No FP4 activation in this slice.

Summary

Plumbing-only slice that wires an optional quant_config: QuantizationConfig | None = None through the Wan-2.1 DiT constructor chain so a future NVFP4QAT-configured Wan-2.1 build (later slice in this stack) can quantize its attention QKV/out projections and FFN without further refactor. Also adds opt-in shape-tracking instrumentation on ReplicatedLinear for QAT-aware backends to discover which GEMM shapes need quantized kernels.

Every new code path is gated. Defaults (quant_config=None, enable_shape_tracking=False) preserve current behavior bit-identically.

Files

File	Δ	What it adds
fastvideo/layers/linear.py	+48	Adds class attrs enable_shape_tracking = False and _shape_to_layer_types plus helper methods get_shape_mapping, reset_shape_tracking, _track_shape, print_shape_summary to ReplicatedLinear. Inserts a single if self.enable_shape_tracking: self._track_shape(...) call in forward after the existing GEMM; diagnostic output goes through the module logger. No constructor params changed; no new dependencies.
fastvideo/layers/mlp.py	+12	Adds quant_config: QuantizationConfig | None = None to MLP.__init__ and threads it (plus an explicit prefix) into both ReplicatedLinear sub-layers. Imports QuantizationConfig from fastvideo.layers.quantization.
fastvideo/models/dits/wanvideo.py	+67	Adds quant_config (and prefix where missing) to WanSelfAttention, WanI2VCrossAttention, WanTransformerBlock, WanTransformerBlock_VSA, and reads config.quant_config in WanTransformer3DModel.__init__. Threads the kwarg through every ReplicatedLinear and MLP construction call. config.quant_config already exists on DiTBaseConfig; no new config field is introduced in this slice.
fastvideo/tests/layers/test_linear_mlp_quant_gating.py	+49	New CPU-only regression test for the gated surfaces (see Test plan).

Files in PR #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 [feat] Dreamverse 13/14: Activate LTX2 integration #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 #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 [refactor]: shared attention infra additions for QAT-compat (Attn-QAT 5/12) #1383 merged at ba75ad82dbe4a7069412494c051c1c69155fdc9d).
• No stack dependency — this is a clean linear PR off main.

Provenance

Files extracted from PR #1225 (#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.

Two cleanups were applied beyond the source extraction:

• The dormant process_weights_after_loading calls the source staged in MLP.__init__ were dropped. Every quant method currently in the tree resolves to the base-class no-op for that hook (FP4 weight conversion is loader-driven, not done per-layer at construction), so the calls were inert. Removing them also keeps MLP consistent with the attention layers, which thread quant_config without any post-load hook.
• The shape-tracking registry on ReplicatedLinear was tightened: the redundant _unique_shapes set was folded into _shape_to_layer_types (whose keys already encode uniqueness; entries are now set-valued), and per-shape diagnostics moved from print to the module logger.

Test plan

Adds fastvideo/tests/layers/test_linear_mlp_quant_gating.py — a CPU-only regression test (no GPU required) for the gated surfaces this slice introduces:

• ReplicatedLinear shape tracking is off by default (no shapes recorded), and records distinct (input, output) shapes / resets cleanly when enable_shape_tracking is explicitly enabled.
• MLP with quant_config=None (the default) keeps both sub-layers on UnquantizedLinearMethod and produces correctly-shaped output.

The quant_config plumbing stays 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.

Pre-commit (yapf + ruff + codespell + mypy) covers linear.py / mlp.py / wanvideo.py; fastvideo/tests/ is excluded from those hooks by design.

Sequence

Attn-QAT-Stack: 6/12. Earlier merged slices: 4/12 (PR #1358), 5/12 (PR #1383). Out of scope for this slice: the actual FP4 activation switch, weight-loading conversion, and any cross-cutting config registration (later slices).

[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

• #approved-reviews-by>=1

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 introduces shape tracking capabilities for GEMM operations in the ReplicatedLinear layer and integrates quantization configuration support throughout the MLP and WanTransformer3DModel components. The changes enable the recording of unique input and output shape pairs to assist with future quantization-aware training. Feedback identifies opportunities to simplify the shape tracking logic by removing redundant data structures and suggests refining the quantization initialization in the MLP layer to avoid unnecessary method calls and premature weight processing.

[gemini-code-assist]

The _unique_shapes set is redundant because the keys of the _shape_to_layer_types dictionary already track the unique (input_shape, output_shape) pairs.

Suggested change

	_unique_shapes: set[tuple[torch.Size, torch.Size]] = set()
	_shape_to_layer_types: dict[tuple[torch.Size, torch.Size], list[str]] = {}

[gemini-code-assist]

Since _unique_shapes is redundant, it can be removed from the reset logic.

Suggested change

	def reset_shape_tracking(cls) -> None:
	"""Clear tracked shapes and layer type mappings."""
	cls._unique_shapes.clear()
	cls._shape_to_layer_types.clear()
	@classmethod
	def reset_shape_tracking(cls) -> None:
	"""Clear tracked shapes and layer type mappings."""
	cls._shape_to_layer_types.clear()

[gemini-code-assist]

The existence check can be performed directly against the _shape_to_layer_types dictionary, eliminating the need for the separate _unique_shapes set.

Suggested change

	if shape_key not in self._unique_shapes:
	self._unique_shapes.add(shape_key)
	self._shape_to_layer_types[shape_key] = []
	if shape_key not in self._shape_to_layer_types:
	self._shape_to_layer_types[shape_key] = []

[gemini-code-assist]

The call to get_quant_method is redundant because ReplicatedLinear.__init__ (via LinearBase) already performs this lookup and stores the result in self.fc_in.quant_method. Additionally, calling process_weights_after_loading inside __init__ is premature as the weights have only just been initialized (via torch.empty) and not yet loaded from a checkpoint. This processing should typically occur after the state dict has been loaded into the model.

Suggested change

	if quant_config is not None:
	quant_method = self.fc_in.quant_config.get_quant_method(self.fc_in, f"{prefix}.fc_in")
	if quant_method is not None:
	quant_method.process_weights_after_loading(self.fc_in)
	if quant_config is not None:
	self.fc_in.quant_method.process_weights_after_loading(self.fc_in)

[gemini-code-assist]

Similar to fc_in, the explicit call to get_quant_method is redundant here as the method is already resolved and available via self.fc_out.quant_method. Also, calling process_weights_after_loading in the constructor is likely premature.

Suggested change

	if quant_config is not None:
	quant_method = self.fc_out.quant_config.get_quant_method(self.fc_out, f"{prefix}.fc_out")
	if quant_method is not None:
	quant_method.process_weights_after_loading(self.fc_out)
	if quant_config is not None:
	self.fc_out.quant_method.process_weights_after_loading(self.fc_out)

[mergify]

Pre-commit checks failed

Hi @SolitaryThinker, 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.

[SolitaryThinker]

/merge