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

[diffusion] performance: refactor diffusion fuse qkv and apply to qwen-image (#14793)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2025-12-10 18:55:41 +08:00
committed by GitHub
parent 02f1e81e2d
commit 12b7a4fab0
4 changed files with 155 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/models/dits/flux.py
+17 -48
View File
@@ -43,6 +43,10 @@ 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.models.dits.utils import (
delete_projection_layers,
fuse_linear_projections,
)
from sglang.multimodal_gen.runtime.platforms import (
AttentionBackendEnum,
current_platform,
@@ -169,51 +173,20 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
if self.fused_projections:
return
device = self.to_q.weight.data.device
dtype = self.to_q.weight.data.dtype
concatenated_weights = torch.cat(
[self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data]
self.to_qkv = fuse_linear_projections(
self.to_q, self.to_k, self.to_v, self.use_bias, ReplicatedLinear
)
in_features = concatenated_weights.shape[1]
out_features = concatenated_weights.shape[0]
self.to_qkv = ReplicatedLinear(in_features, out_features, bias=self.use_bias)
self.to_qkv.weight.data = concatenated_weights.to(device=device, dtype=dtype)
if self.use_bias:
concatenated_bias = torch.cat(
[self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data]
)
self.to_qkv.bias.data = concatenated_bias.to(device=device, dtype=dtype)
delete_projection_layers(self, ["to_q", "to_k", "to_v"])
if self.added_kv_proj_dim is not None:
concatenated_weights = torch.cat(
[
self.add_q_proj.weight.data,
self.add_k_proj.weight.data,
self.add_v_proj.weight.data,
]
self.to_added_qkv = fuse_linear_projections(
self.add_q_proj,
self.add_k_proj,
self.add_v_proj,
self.added_proj_bias,
ReplicatedLinear,
)
in_features = concatenated_weights.shape[1]
out_features = concatenated_weights.shape[0]
self.to_added_qkv = ReplicatedLinear(
in_features, out_features, bias=self.added_proj_bias
)
self.to_added_qkv.weight.data = concatenated_weights.to(
device=device, dtype=dtype
)
if self.added_proj_bias:
concatenated_bias = torch.cat(
[
self.add_q_proj.bias.data,
self.add_k_proj.bias.data,
self.add_v_proj.bias.data,
]
)
self.to_added_qkv.bias.data = concatenated_bias.to(
device=device, dtype=dtype
)
delete_projection_layers(self, ["add_q_proj", "add_k_proj", "add_v_proj"])
self.fused_projections = True
@@ -530,13 +503,9 @@ class FluxTransformer2DModel(CachableDiT):
)
def fuse_qkv_projections(self):
for block in self.transformer_blocks:
if hasattr(block.attn, "fuse_projections") and getattr(
block.attn, "_supports_qkv_fusion", True
):
block.attn.fuse_projections()
for block in self.single_transformer_blocks:
for block in list(self.transformer_blocks) + list(
self.single_transformer_blocks
):
if hasattr(block.attn, "fuse_projections") and getattr(
block.attn, "_supports_qkv_fusion", True
):