[diffusion] perf: support pack qkv for Z-Image (#15191)

2025-12-16 00:22:24 +08:00
parent 1dedb63860
commit c0d94440b7
2 changed files with 32 additions and 7 deletions
@@ -26,6 +26,23 @@ class ZImageArchConfig(DiTArchConfig):
    axes_dims: Tuple[int, int, int] = (32, 48, 48)
    axes_lens: Tuple[int, int, int] = (1024, 512, 512)

+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            # (param_name, shard_name, shard_id)
+            (".to_qkv", ".to_q", "q"),
+            (".to_qkv", ".to_k", "k"),
+            (".to_qkv", ".to_v", "v"),
+        ]
+    )
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"(.*)\.to_q\.weight$": (r"\1.to_qkv.weight", 0, 3),
+            r"(.*)\.to_k\.weight$": (r"\1.to_qkv.weight", 1, 3),
+            r"(.*)\.to_v\.weight$": (r"\1.to_qkv.weight", 2, 3),
+        }
+    )
+
    def __post_init__(self):
        super().__post_init__()
        self.out_channels = self.out_channels or self.in_channels