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

Slice 6 of 12 in the PR #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 #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 #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 #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.

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 #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).

Co-Authored-By: Peiyuan Zhang <a1286225768@gmail.com>
Co-Authored-By: Matthew Noto <notomatthew31@gmail.com>

Author: 3 people

fastvideo/layers/linear.py

@@ -219,6 +219,16 @@ class ReplicatedLinear(LinearBase):
                         (e.g. model.layers.0.qkv_proj)
     """
 
+    # Opt-in instrumentation: when ``enable_shape_tracking`` is set to True,
+    # ``forward`` records every unique ``(input_shape, output_shape)`` pair
+    # observed across all ``ReplicatedLinear`` instances, along with the
+    # subclass name that produced it. Used by upcoming QAT-aware backends
+    # to discover which GEMM shapes need quantized kernels. Defaults to
+    # False; default forward path is bit-identical to pre-slice behavior.
+    enable_shape_tracking = False
+    _unique_shapes: set[tuple[torch.Size, torch.Size]] = set()
+    _shape_to_layer_types: dict[tuple[torch.Size, torch.Size], list[str]] = {}
+
     def __init__(
         self,
         input_size: int,
@@ -285,6 +295,8 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]:
         bias = self.bias if not self.skip_bias_add else None
         assert self.quant_method is not None
         output = self.quant_method.apply(self, x, bias)
+        if self.enable_shape_tracking:
+            self._track_shape(x.shape, output.shape)
         output_bias = self.bias if self.skip_bias_add else None
         return output, output_bias
 
@@ -294,6 +306,48 @@ def extra_repr(self) -> str:
         s += f", bias={self.bias is not None}"
         return s
 
+    @classmethod
+    def get_shape_mapping(cls) -> dict:
+        """Get the mapping from (input_shape, output_shape) to layer types."""
+        return cls._shape_to_layer_types.copy()
+
+    @classmethod
+    def reset_shape_tracking(cls) -> None:
+        """Clear tracked shapes and layer type mappings."""
+        cls._unique_shapes.clear()
+        cls._shape_to_layer_types.clear()
+
+    def _track_shape(self, input_shape: torch.Size, output_shape: torch.Size) -> None:
+        shape_key = (input_shape, output_shape)
+
+        if shape_key not in self._unique_shapes:
+            self._unique_shapes.add(shape_key)
+            self._shape_to_layer_types[shape_key] = []
+            print(f"Layer: {self.prefix} | input shape: {input_shape} --> "
+                  f"output shape: {output_shape}, Quant Method: "
+                  f"{self.quant_method.__class__.__name__}")
+
+        layer_type = self.__class__.__name__
+        if layer_type not in self._shape_to_layer_types[shape_key]:
+            self._shape_to_layer_types[shape_key].append(layer_type)
+
+    @classmethod
+    def print_shape_summary(cls) -> None:
+        """Print a summary of all unique shapes and their layer types."""
+        if not cls._shape_to_layer_types:
+            print("No shapes have been processed yet.")
+            return
+
+        print("\n=== Matrix Multiplication Shape Summary ===")
+        print(f"Total unique shapes: {len(cls._shape_to_layer_types)}")
+        print()
+
+        for i, (shape_key, layer_types) in enumerate(cls._shape_to_layer_types.items(), 1):
+            input_shape, output_shape = shape_key
+            print(f"{i}. Input: {input_shape} → Output: {output_shape}")
+            print(f"   Layer types: {', '.join(layer_types)}")
+            print()
+
 
 class ColumnParallelLinear(LinearBase):
     """Linear layer with column parallelism.

fastvideo/layers/mlp.py

@@ -5,6 +5,7 @@
 
 from fastvideo.layers.activation import get_act_fn
 from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.quantization import QuantizationConfig
 
 
 class MLP(nn.Module):
@@ -21,18 +22,35 @@ def __init__(
         act_type: str = "gelu_pytorch_tanh",
         dtype: torch.dtype | None = None,
         prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
     ):
         super().__init__()
         self.fc_in = ReplicatedLinear(
             input_dim,
             mlp_hidden_dim,  # For activation func like SiLU that need 2x width
             bias=bias,
-            params_dtype=dtype)
+            params_dtype=dtype,
+            quant_config=quant_config,
+            prefix=f"{prefix}.fc_in",
+        )
+        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)
 
         self.act = get_act_fn(act_type)
         if output_dim is None:
             output_dim = input_dim
-        self.fc_out = ReplicatedLinear(mlp_hidden_dim, output_dim, bias=bias, params_dtype=dtype)
+        self.fc_out = ReplicatedLinear(mlp_hidden_dim,
+                                       output_dim,
+                                       bias=bias,
+                                       params_dtype=dtype,
+                                       quant_config=quant_config,
+                                       prefix=f"{prefix}.fc_out")
+        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)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x, _ = self.fc_in(x)

fastvideo/models/dits/wanvideo.py

@@ -26,6 +26,7 @@
 from fastvideo.logger import init_logger
 from fastvideo.models.dits.base import BaseDiT
 from fastvideo.platforms import AttentionBackendEnum, current_platform
+from fastvideo.layers.quantization import QuantizationConfig
 
 from fastvideo.distributed.parallel_state import get_sp_world_size
 
@@ -106,7 +107,9 @@ def __init__(self,
                  window_size=(-1, -1),
                  qk_norm=True,
                  eps=1e-6,
-                 parallel_attention=False) -> None:
+                 parallel_attention=False,
+                 quant_config: QuantizationConfig | None = None,
+                 prefix: str = "") -> None:
         assert dim % num_heads == 0
         super().__init__()
         self.dim = dim
@@ -118,10 +121,10 @@ def __init__(self,
         self.parallel_attention = parallel_attention
 
         # layers
-        self.to_q = ReplicatedLinear(dim, dim)
-        self.to_k = ReplicatedLinear(dim, dim)
-        self.to_v = ReplicatedLinear(dim, dim)
-        self.to_out = ReplicatedLinear(dim, dim)
+        self.to_q = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_q")
+        self.to_k = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_k")
+        self.to_v = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_v")
+        self.to_out = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_out")
         self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
         self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
 
@@ -194,13 +197,15 @@ def __init__(
         qk_norm=True,
         eps=1e-6,
         supported_attention_backends: tuple[AttentionBackendEnum, ...]
-        | None = None
+        | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ) -> None:
         super().__init__(dim, num_heads, window_size, qk_norm, eps,
-                         supported_attention_backends)
+                         supported_attention_backends, quant_config=quant_config, prefix=prefix)
 
-        self.add_k_proj = ReplicatedLinear(dim, dim)
-        self.add_v_proj = ReplicatedLinear(dim, dim)
+        self.add_k_proj = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.add_k_proj")
+        self.add_v_proj = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.add_v_proj")
         self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
         self.norm_added_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
 
@@ -246,16 +251,17 @@ def __init__(self,
                  added_kv_proj_dim: int | None = None,
                  supported_attention_backends: tuple[AttentionBackendEnum, ...]
                  | None = None,
+                 quant_config: QuantizationConfig | None = None,
                  prefix: str = ""):
         super().__init__()
 
         # 1. Self-attention
         self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
-        self.to_q = ReplicatedLinear(dim, dim, bias=True)
-        self.to_k = ReplicatedLinear(dim, dim, bias=True)
-        self.to_v = ReplicatedLinear(dim, dim, bias=True)
+        self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_q")
+        self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_k")
+        self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_v")
 
-        self.to_out = ReplicatedLinear(dim, dim, bias=True)
+        self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_out")
         self.attn1 = DistributedAttention(
             num_heads=num_heads,
             head_size=dim // num_heads,
@@ -290,13 +296,17 @@ def __init__(self,
             self.attn2 = WanI2VCrossAttention(dim,
                                               num_heads,
                                               qk_norm=qk_norm,
-                                              eps=eps)
+                                              eps=eps,
+                                              quant_config=quant_config,
+                                              prefix=f"{prefix}.attn2")
         else:
             # T2V
             self.attn2 = WanT2VCrossAttention(dim,
                                               num_heads,
                                               qk_norm=qk_norm,
-                                              eps=eps)
+                                              eps=eps,
+                                              quant_config=quant_config,
+                                              prefix=f"{prefix}.attn2")
         self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
             dim,
             norm_type="layer",
@@ -306,7 +316,7 @@ def __init__(self,
             compute_dtype=torch.float32)
 
         # 3. Feed-forward
-        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
+        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn")
         self.mlp_residual = ScaleResidual()
 
         self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
@@ -406,17 +416,17 @@ def __init__(self,
                  added_kv_proj_dim: int | None = None,
                  supported_attention_backends: tuple[AttentionBackendEnum, ...]
                  | None = None,
+                 quant_config: QuantizationConfig | None = None,
                  prefix: str = ""):
         super().__init__()
 
         # 1. Self-attention
         self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
-        self.to_q = ReplicatedLinear(dim, dim, bias=True)
-        self.to_k = ReplicatedLinear(dim, dim, bias=True)
-        self.to_v = ReplicatedLinear(dim, dim, bias=True)
-        self.to_gate_compress = ReplicatedLinear(dim, dim, bias=True)
-
-        self.to_out = ReplicatedLinear(dim, dim, bias=True)
+        self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_q")
+        self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_k")
+        self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_v")
+        self.to_gate_compress = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_gate_compress")
+        self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_out")
         self.attn1 = DistributedAttention_VSA(
             num_heads=num_heads,
             head_size=dim // num_heads,
@@ -451,13 +461,17 @@ def __init__(self,
             self.attn2 = WanI2VCrossAttention(dim,
                                               num_heads,
                                               qk_norm=qk_norm,
-                                              eps=eps)
+                                              eps=eps,
+                                              quant_config=quant_config,
+                                              prefix=f"{prefix}.attn2")
         else:
             # T2V
             self.attn2 = WanT2VCrossAttention(dim,
                                               num_heads,
                                               qk_norm=qk_norm,
-                                              eps=eps)
+                                              eps=eps,
+                                              quant_config=quant_config,
+                                              prefix=f"{prefix}.attn2")
         self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
             dim,
             norm_type="layer",
@@ -467,7 +481,7 @@ def __init__(self,
             compute_dtype=torch.float32)
 
         # 3. Feed-forward
-        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
+        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn")
         self.mlp_residual = ScaleResidual()
 
         self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
@@ -556,6 +570,7 @@ class WanTransformer3DModel(BaseDiT):
     def __init__(self, config: WanVideoConfig, hf_config: dict[str,
                                                                Any]) -> None:
         super().__init__(config=config, hf_config=hf_config)
+        self.quant_config = config.quant_config
 
         inner_dim = config.num_attention_heads * config.attention_head_dim
         self.hidden_size = config.hidden_size
@@ -594,6 +609,7 @@ def __init__(self, config: WanVideoConfig, hf_config: dict[str,
                               config.eps,
                               config.added_kv_proj_dim,
                               self._supported_attention_backends,
+                              quant_config=config.quant_config,
                               prefix=f"{config.prefix}.blocks.{i}")
             for i in range(config.num_layers)
         ])