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

python/sglang/multimodal_gen/configs/models/dits/zimage.py

@@ -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

python/sglang/multimodal_gen/runtime/models/dits/zimage.py

@@ -8,7 +8,10 @@ import torch.nn.functional as F
 from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig
 from sglang.multimodal_gen.runtime.layers.attention import USPAttention
 from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
-from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
+from sglang.multimodal_gen.runtime.layers.linear import (
+    QKVParallelLinear,
+    ReplicatedLinear,
+)
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import _apply_rotary_emb
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
@@ -104,9 +107,13 @@ class ZImageAttention(nn.Module):
         self.head_dim = dim // num_heads
         self.qk_norm = qk_norm
 
-        self.to_q = ReplicatedLinear(dim, dim, bias=False)
+        self.to_qkv = QKVParallelLinear(
-        self.to_k = ReplicatedLinear(dim, self.head_dim * num_kv_heads, bias=False)
+            hidden_size=dim,
-        self.to_v = ReplicatedLinear(dim, self.head_dim * num_kv_heads, bias=False)
+            head_size=self.head_dim,
+            total_num_heads=num_heads,
+            total_num_kv_heads=num_kv_heads,
+            bias=False,
+        )
 
         if self.qk_norm:
             self.norm_q = RMSNorm(self.head_dim, eps=eps)
@@ -135,9 +142,9 @@ class ZImageAttention(nn.Module):
         hidden_states: torch.Tensor,
         freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
     ):
-        q, _ = self.to_q(hidden_states)
+        qkv, _ = self.to_qkv(hidden_states)
-        k, _ = self.to_k(hidden_states)
+        kv_dim = self.head_dim * self.num_kv_heads
-        v, _ = self.to_v(hidden_states)
+        q, k, v = torch.split(qkv, [self.dim, kv_dim, kv_dim], dim=-1)
 
         q = q.view(*q.shape[:-1], self.num_heads, self.head_dim)
         k = k.view(*k.shape[:-1], self.num_kv_heads, self.head_dim)
@@ -344,6 +351,7 @@ class RopeEmbedder:
 class ZImageTransformer2DModel(CachableDiT):
     _supports_gradient_checkpointing = True
     _no_split_modules = ["ZImageTransformerBlock"]
+    param_names_mapping = ZImageDitConfig().arch_config.param_names_mapping
 
     def __init__(
         self,