chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 Packages Projects Releases Wiki Activity

[diffusion] refactor: refactor fuse qkv with QKVParallelLinear linear (#15090)

Browse Source 
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Xiaoyu Zhang
2025-12-15 00:33:29 +08:00
committed by GitHub
co-authored by Mick
parent d277a86dea
commit 64b5c3ab90
9 changed files with 118 additions and 285 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/models/dits/flux.py
+19
View File
@@ -21,8 +21,27 @@ class FluxArchConfig(DiTArchConfig):
guidance_embeds: bool = False
axes_dims_rope: Tuple[int, int, int] = (16, 56, 56)
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"),
(".to_added_qkv", ".add_q_proj", "q"),
(".to_added_qkv", ".add_k_proj", "k"),
(".to_added_qkv", ".add_v_proj", "v"),
]
)
param_names_mapping: dict = field(
default_factory=lambda: {
# QKV fusion mappings
r"(.*)\.to_q\.(weight|bias)$": (r"\1.to_qkv.\2", 0, 3),
r"(.*)\.to_k\.(weight|bias)$": (r"\1.to_qkv.\2", 1, 3),
r"(.*)\.to_v\.(weight|bias)$": (r"\1.to_qkv.\2", 2, 3),
r"(.*)\.add_q_proj\.(weight|bias)$": (r"\1.to_added_qkv.\2", 0, 3),
r"(.*)\.add_k_proj\.(weight|bias)$": (r"\1.to_added_qkv.\2", 1, 3),
r"(.*)\.add_v_proj\.(weight|bias)$": (r"\1.to_added_qkv.\2", 2, 3),
r"transformer\.(\w*)\.(.*)$": r"\1.\2",
}
)
python/sglang/multimodal_gen/configs/models/dits/qwenimage.py
+20
View File
@@ -21,8 +21,28 @@ class QwenImageArchConfig(DiTArchConfig):
guidance_embeds: bool = False
axes_dims_rope: Tuple[int, int, int] = (16, 56, 56)
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"),
(".to_added_qkv", ".add_q_proj", "q"),
(".to_added_qkv", ".add_k_proj", "k"),
(".to_added_qkv", ".add_v_proj", "v"),
]
)
param_names_mapping: dict = field(
default_factory=lambda: {
# QKV fusion mappings
r"(.*)\.to_q\.(weight|bias)$": (r"\1.to_qkv.\2", 0, 3),
r"(.*)\.to_k\.(weight|bias)$": (r"\1.to_qkv.\2", 1, 3),
r"(.*)\.to_v\.(weight|bias)$": (r"\1.to_qkv.\2", 2, 3),
r"(.*)\.add_q_proj\.(weight|bias)$": (r"\1.to_added_qkv.\2", 0, 3),
r"(.*)\.add_k_proj\.(weight|bias)$": (r"\1.to_added_qkv.\2", 1, 3),
r"(.*)\.add_v_proj\.(weight|bias)$": (r"\1.to_added_qkv.\2", 2, 3),
# LoRA mappings
r"^(transformer_blocks\.\d+\.attn\..*\.lora_[AB])\.default$": r"\1",
}
)