[codex] Optimize Z-Image packed QKV

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

@@ -49,6 +49,9 @@ class ZImageArchConfig(DiTArchConfig):
 
     param_names_mapping: dict = field(
         default_factory=lambda: {
+            r"(.*)\.attention\.to_q\.weight$": (r"\1.attention.to_qkv.weight", 0, 3),
+            r"(.*)\.attention\.to_k\.weight$": (r"\1.attention.to_qkv.weight", 1, 3),
+            r"(.*)\.attention\.to_v\.weight$": (r"\1.attention.to_qkv.weight", 2, 3),
             r"(.*)\.feed_forward\.w1\.weight$": (r"\1.feed_forward.w13.weight", 0, 2),
             r"(.*)\.feed_forward\.w3\.weight$": (r"\1.feed_forward.w13.weight", 1, 2),
             r"(.*)\.feed_forward\.w1\.(lora_A|lora_B)$": (

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

@@ -159,7 +159,7 @@ def __init__(
         self.local_num_kv_heads = num_kv_heads // tp_size
 
         kv_dim = self.head_dim * num_kv_heads
-        self.use_fused_qkv = isinstance(quant_config, NunchakuConfig)
+        self.use_fused_qkv = True
 
         if self.use_fused_qkv:
             self.to_qkv = MergedColumnParallelLinear(