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

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,15 @@ 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
+    _shape_to_layer_types: dict[tuple[torch.Size, torch.Size], set[str]] = {}
+
     def __init__(
         self,
         input_size: int,
@@ -285,6 +294,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 +305,41 @@ 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._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._shape_to_layer_types:
+            self._shape_to_layer_types[shape_key] = set()
+            logger.debug("Layer: %s | input shape: %s --> output shape: %s, Quant Method: %s", self.prefix, input_shape,
+                         output_shape, self.quant_method.__class__.__name__)
+        self._shape_to_layer_types[shape_key].add(self.__class__.__name__)
+
+    @classmethod
+    def print_shape_summary(cls) -> None:
+        """Log a summary of all unique shapes and their layer types."""
+        if not cls._shape_to_layer_types:
+            logger.info("No shapes have been processed yet.")
+            return
+
+        lines = [
+            "=== Matrix Multiplication Shape Summary ===",
+            f"Total unique shapes: {len(cls._shape_to_layer_types)}",
+        ]
+        for i, (shape_key, layer_types) in enumerate(cls._shape_to_layer_types.items(), 1):
+            input_shape, output_shape = shape_key
+            lines.append(f"{i}. Input: {input_shape} → Output: {output_shape}")
+            lines.append(f"   Layer types: {', '.join(sorted(layer_types))}")
+        logger.info("\n".join(lines))
+
 
 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,27 @@ 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",
+        )
 
         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")
 
     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)
         ])

fastvideo/tests/layers/test_linear_mlp_quant_gating.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Regression coverage for PR #1390's S2-1 plumbing finding: the dormant FP4 shape-tracking path must stay
+gated off unless explicitly enabled, and the MLP quant_config=None default path must keep using ReplicatedLinear's
+unquantized fallback.
+"""
+from __future__ import annotations
+
+import torch
+
+from fastvideo.layers.linear import ReplicatedLinear, UnquantizedLinearMethod
+from fastvideo.layers.mlp import MLP
+
+
+def test_replicated_linear_shape_tracking_default_off() -> None:
+    ReplicatedLinear.reset_shape_tracking()
+    assert ReplicatedLinear.enable_shape_tracking is False
+
+    linear = ReplicatedLinear(input_size=8, output_size=4)
+    linear(torch.randn(2, 8))
+
+    assert len(ReplicatedLinear._shape_to_layer_types) == 0
+
+
+def test_replicated_linear_shape_tracking_enabled_records_unique_shapes() -> None:
+    ReplicatedLinear.reset_shape_tracking()
+    ReplicatedLinear.enable_shape_tracking = True
+    try:
+        linear = ReplicatedLinear(input_size=8, output_size=4)
+        linear(torch.randn(2, 8))
+        linear(torch.randn(3, 8))
+
+        assert len(ReplicatedLinear._shape_to_layer_types) == 2
+        for layer_types in ReplicatedLinear._shape_to_layer_types.values():
+            assert "ReplicatedLinear" in layer_types
+
+        ReplicatedLinear.reset_shape_tracking()
+        assert len(ReplicatedLinear._shape_to_layer_types) == 0
+    finally:
+        ReplicatedLinear.enable_shape_tracking = False
+
+
+def test_mlp_quant_config_none_uses_unquantized_path() -> None:
+    mlp = MLP(input_dim=8, mlp_hidden_dim=16)
+
+    assert isinstance(mlp.fc_in.quant_method, UnquantizedLinearMethod)
+    assert isinstance(mlp.fc_out.quant_method, UnquantizedLinearMethod)
+
+    output = mlp.forward(torch.randn(2, 8))
+    assert output.shape == (2, 8)