chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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/runtime/models/dits/flux.py
+22 -76
View File
@@ -36,17 +36,16 @@ from sglang.multimodal_gen.runtime.layers.attention import USPAttention
# from sglang.multimodal_gen.runtime.layers.layernorm import LayerNorm as LayerNorm
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.mlp import MLP
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
NDRotaryEmbedding,
_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,
@@ -56,45 +55,21 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__) # pylint: disable=invalid-name
def _get_projections(attn: "FluxAttention", hidden_states, encoder_hidden_states=None):
query, _ = attn.to_q(hidden_states)
key, _ = attn.to_k(hidden_states)
value, _ = attn.to_v(hidden_states)
encoder_query = encoder_key = encoder_value = None
if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None:
encoder_query, _ = attn.add_q_proj(encoder_hidden_states)
encoder_key, _ = attn.add_k_proj(encoder_hidden_states)
encoder_value, _ = attn.add_v_proj(encoder_hidden_states)
return query, key, value, encoder_query, encoder_key, encoder_value
def _get_fused_projections(
def _get_qkv_projections(
attn: "FluxAttention", hidden_states, encoder_hidden_states=None
):
qkv, _ = attn.to_qkv(hidden_states)
query, key, value = qkv.chunk(3, dim=-1)
encoder_query = encoder_key = encoder_value = None
if encoder_hidden_states is not None and hasattr(attn, "to_added_qkv"):
if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None:
added_qkv, _ = attn.to_added_qkv(encoder_hidden_states)
encoder_query, encoder_key, encoder_value = added_qkv.chunk(3, dim=-1)
return query, key, value, encoder_query, encoder_key, encoder_value
def _get_qkv_projections(
attn: "FluxAttention", hidden_states, encoder_hidden_states=None
):
if attn.fused_projections:
return _get_fused_projections(attn, hidden_states, encoder_hidden_states)
return _get_projections(attn, hidden_states, encoder_hidden_states)
class FluxAttention(torch.nn.Module, AttentionModuleMixin):
_supports_qkv_fusion = True
def __init__(
self,
query_dim: int,
@@ -125,11 +100,15 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
self.added_proj_bias = added_proj_bias
self.norm_q = RMSNorm(dim_head, eps=eps)
self.norm_k = RMSNorm(dim_head, eps=eps)
self.to_q = ReplicatedLinear(query_dim, self.inner_dim, bias=bias)
self.to_k = ReplicatedLinear(query_dim, self.inner_dim, bias=bias)
self.to_v = ReplicatedLinear(query_dim, self.inner_dim, bias=bias)
# Use QKVParallelLinear for fused QKV projections
self.to_qkv = QKVParallelLinear(
hidden_size=query_dim,
head_size=dim_head,
total_num_heads=num_heads,
bias=bias,
)
if not self.pre_only:
self.to_out = torch.nn.ModuleList([])
@@ -142,14 +121,12 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
if added_kv_proj_dim is not None:
self.norm_added_q = RMSNorm(dim_head, eps=eps)
self.norm_added_k = RMSNorm(dim_head, eps=eps)
self.add_q_proj = ReplicatedLinear(
added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
)
self.add_k_proj = ReplicatedLinear(
added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
)
self.add_v_proj = ReplicatedLinear(
added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
# Use QKVParallelLinear for added (encoder) QKV projections
self.to_added_qkv = QKVParallelLinear(
hidden_size=added_kv_proj_dim,
head_size=dim_head,
total_num_heads=num_heads,
bias=added_proj_bias,
)
self.to_add_out = ReplicatedLinear(self.inner_dim, query_dim, bias=out_bias)
@@ -166,30 +143,6 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
},
)
self.fused_projections = False
@torch.no_grad()
def fuse_projections(self):
if self.fused_projections:
return
self.to_qkv = fuse_linear_projections(
self.to_q, self.to_k, self.to_v, self.use_bias, ReplicatedLinear
)
delete_projection_layers(self, ["to_q", "to_k", "to_v"])
if self.added_kv_proj_dim is not None:
self.to_added_qkv = fuse_linear_projections(
self.add_q_proj,
self.add_k_proj,
self.add_v_proj,
self.added_proj_bias,
ReplicatedLinear,
)
delete_projection_layers(self, ["add_q_proj", "add_k_proj", "add_v_proj"])
self.fused_projections = True
def forward(
self,
x: torch.Tensor,
@@ -445,6 +398,8 @@ class FluxTransformer2DModel(CachableDiT):
Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
"""
param_names_mapping = FluxConfig().arch_config.param_names_mapping
def __init__(self, config: FluxConfig, hf_config: dict[str, Any]) -> None:
super().__init__(config=config, hf_config=hf_config)
self.config = config.arch_config
@@ -503,15 +458,6 @@ class FluxTransformer2DModel(CachableDiT):
bias=True,
)
def fuse_qkv_projections(self):
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
):
block.attn.fuse_projections()
def forward(
self,
hidden_states: torch.Tensor,