diff --git a/python/sglang/multimodal_gen/configs/models/dits/flux.py b/python/sglang/multimodal_gen/configs/models/dits/flux.py index f50c9b1e6..ee5a9867a 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/flux.py +++ b/python/sglang/multimodal_gen/configs/models/dits/flux.py @@ -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", } ) diff --git a/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py b/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py index f9003264b..782c39973 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py +++ b/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py @@ -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", } ) diff --git a/python/sglang/multimodal_gen/docs/profiling.md b/python/sglang/multimodal_gen/docs/profiling.md index 58a72f5fc..b675816da 100644 --- a/python/sglang/multimodal_gen/docs/profiling.md +++ b/python/sglang/multimodal_gen/docs/profiling.md @@ -47,6 +47,10 @@ By default, trace files are saved in the ./logs/ directory. The exact output fil ```bash [mm-dd hh:mm:ss] Saving profiler traces to: /sgl-workspace/sglang/logs/mocked_fake_id_for_offline_generate-5_steps-global-rank0.trace.json.gz ``` +{request_id}-{num_steps}_steps-global-rank{rank}.trace.json.gz +``` + +Example: `mocked_fake_id_for_offline_generate-5_steps-global-rank0.trace.json.gz` ### View Traces diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py index 1b7888208..db71023a6 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loader.py @@ -692,10 +692,6 @@ class TransformerLoader(ComponentLoader): model = model.eval() - if hasattr(model, "fuse_qkv_projections"): - logger.info("Fusing QKV projections for better performance") - model.fuse_qkv_projections() - return model diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index 40089bc10..436e60667 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index 4e83e0cb0..c64e47fdd 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -23,15 +23,12 @@ from diffusers.models.normalization import AdaLayerNormContinuous from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig 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 QKVParallelLinear 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, @@ -41,40 +38,18 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name -def _get_projections(attn: "Flux2Attention", 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( - attn: "Flux2Attention", 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"): - 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: "Flux2Attention", 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) + 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 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 class Flux2SwiGLU(nn.Module): @@ -120,8 +95,6 @@ class Flux2FeedForward(nn.Module): class Flux2Attention(torch.nn.Module, AttentionModuleMixin): - _supports_qkv_fusion = True - def __init__( self, query_dim: int, @@ -150,9 +123,13 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin): self.added_kv_proj_dim = added_kv_proj_dim self.added_proj_bias = added_proj_bias - self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias) - self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=bias) - self.to_v = torch.nn.Linear(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, + ) # QK Norm self.norm_q = RMSNorm(dim_head, eps=eps) @@ -165,14 +142,12 @@ class Flux2Attention(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 = torch.nn.Linear( - added_kv_proj_dim, self.inner_dim, bias=added_proj_bias - ) - self.add_k_proj = torch.nn.Linear( - added_kv_proj_dim, self.inner_dim, bias=added_proj_bias - ) - self.add_v_proj = torch.nn.Linear( - 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 = torch.nn.Linear(self.inner_dim, query_dim, bias=out_bias) @@ -189,30 +164,6 @@ class Flux2Attention(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, torch.nn.Linear - ) - 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, - torch.nn.Linear, - ) - delete_projection_layers(self, ["add_q_proj", "add_k_proj", "add_v_proj"]) - - self.fused_projections = True - def forward( self, hidden_states: torch.Tensor, @@ -655,6 +606,8 @@ class Flux2Transformer2DModel(CachableDiT): """ + param_names_mapping = FluxConfig().arch_config.param_names_mapping + def __init__(self, config: FluxConfig, hf_config: dict[str, Any]): super().__init__(config=config, hf_config=hf_config) patch_size: int = config.patch_size @@ -746,15 +699,6 @@ class Flux2Transformer2DModel(CachableDiT): self.gradient_checkpointing = False - 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, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 6808460bc..f8acd57bd 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -16,60 +16,35 @@ from diffusers.models.normalization import AdaLayerNormContinuous from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.layers.layernorm import LayerNorm, 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.triton_ops import ( apply_rotary_embedding, fuse_scale_shift_kernel, ) 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 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name -def _get_projections( - attn: "QwenImageCrossAttention", hidden_states, encoder_hidden_states=None -): - img_query, _ = attn.to_q(hidden_states) - img_key, _ = attn.to_k(hidden_states) - img_value, _ = attn.to_v(hidden_states) - - txt_query = txt_key = txt_value = None - if encoder_hidden_states is not None and hasattr(attn, "add_q_proj"): - txt_query, _ = attn.add_q_proj(encoder_hidden_states) - txt_key, _ = attn.add_k_proj(encoder_hidden_states) - txt_value, _ = attn.add_v_proj(encoder_hidden_states) - - return img_query, img_key, img_value, txt_query, txt_key, txt_value - - -def _get_fused_projections( +def _get_qkv_projections( attn: "QwenImageCrossAttention", hidden_states, encoder_hidden_states=None ): img_qkv, _ = attn.to_qkv(hidden_states) img_query, img_key, img_value = img_qkv.chunk(3, dim=-1) txt_query = txt_key = txt_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: txt_qkv, _ = attn.to_added_qkv(encoder_hidden_states) txt_query, txt_key, txt_value = txt_qkv.chunk(3, dim=-1) return img_query, img_key, img_value, txt_query, txt_key, txt_value -def _get_qkv_projections( - attn: "QwenImageCrossAttention", 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 QwenTimestepProjEmbeddings(nn.Module): def __init__(self, embedding_dim): super().__init__() @@ -260,8 +235,6 @@ class QwenEmbedRope(nn.Module): class QwenImageCrossAttention(nn.Module): - _supports_qkv_fusion = True - def __init__( self, dim: int, # query_dim @@ -286,27 +259,31 @@ class QwenImageCrossAttention(nn.Module): self.qk_norm = qk_norm self.eps = eps self.parallel_attention = parallel_attention + self.added_kv_proj_dim = added_kv_proj_dim + + # Use QKVParallelLinear for fused QKV projections + self.to_qkv = QKVParallelLinear( + hidden_size=dim, + head_size=head_dim, + total_num_heads=num_heads, + bias=True, + ) - # layers - self.to_q = ReplicatedLinear(dim, dim) - self.to_k = ReplicatedLinear(dim, dim) - self.to_v = ReplicatedLinear(dim, dim) if self.qk_norm: self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() + self.inner_dim = out_dim if out_dim is not None else head_dim * num_heads self.inner_kv_dim = self.inner_dim + if added_kv_proj_dim is not None: - self.add_k_proj = ReplicatedLinear( - added_kv_proj_dim, self.inner_kv_dim, bias=True + # Use QKVParallelLinear for added (encoder) QKV projections + self.to_added_qkv = QKVParallelLinear( + hidden_size=added_kv_proj_dim, + head_size=head_dim, + total_num_heads=num_heads, + bias=True, ) - self.add_v_proj = ReplicatedLinear( - added_kv_proj_dim, self.inner_kv_dim, bias=True - ) - if context_pre_only is not None: - self.add_q_proj = ReplicatedLinear( - added_kv_proj_dim, self.inner_dim, bias=True - ) if context_pre_only is not None and not context_pre_only: self.to_add_out = ReplicatedLinear(self.inner_dim, self.dim, bias=out_bias) @@ -338,31 +315,6 @@ class QwenImageCrossAttention(nn.Module): }, ) - self.fused_projections = False - self.added_kv_proj_dim_val = added_kv_proj_dim - - @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, use_bias=False, linear_cls=ReplicatedLinear - ) - delete_projection_layers(self, ["to_q", "to_k", "to_v"]) - - if self.added_kv_proj_dim_val is not None and hasattr(self, "add_q_proj"): - self.to_added_qkv = fuse_linear_projections( - self.add_q_proj, - self.add_k_proj, - self.add_v_proj, - use_bias=True, - linear_cls=ReplicatedLinear, - ) - delete_projection_layers(self, ["add_q_proj", "add_k_proj", "add_v_proj"]) - - self.fused_projections = True - def forward( self, hidden_states: torch.Tensor, @@ -579,6 +531,8 @@ class QwenImageTransformer2DModel(CachableDiT): _skip_layerwise_casting_patterns = ["pos_embed", "norm"] _repeated_blocks = ["QwenImageTransformerBlock"] + param_names_mapping = QwenImageDitConfig().arch_config.param_names_mapping + def __init__( self, config: QwenImageDitConfig, @@ -625,13 +579,6 @@ class QwenImageTransformer2DModel(CachableDiT): self.inner_dim, patch_size * patch_size * self.out_channels, bias=True ) - 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() - def forward( self, hidden_states: torch.Tensor, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/utils.py b/python/sglang/multimodal_gen/runtime/models/dits/utils.py deleted file mode 100644 index 74b2d4d67..000000000 --- a/python/sglang/multimodal_gen/runtime/models/dits/utils.py +++ /dev/null @@ -1,43 +0,0 @@ -from typing import Union - -import torch -import torch.nn as nn - -from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear - - -def fuse_linear_projections( - q_proj: Union[nn.Linear, ReplicatedLinear], - k_proj: Union[nn.Linear, ReplicatedLinear], - v_proj: Union[nn.Linear, ReplicatedLinear], - use_bias: bool, - linear_cls: type = None, -) -> Union[nn.Linear, ReplicatedLinear]: - device = q_proj.weight.data.device - dtype = q_proj.weight.data.dtype - - concatenated_weights = torch.cat( - [q_proj.weight.data, k_proj.weight.data, v_proj.weight.data] - ) - in_features = concatenated_weights.shape[1] - out_features = concatenated_weights.shape[0] - - if linear_cls is None: - linear_cls = type(q_proj) - - fused_layer = linear_cls(in_features, out_features, bias=use_bias) - fused_layer.weight.data = concatenated_weights.to(device=device, dtype=dtype) - - if use_bias: - concatenated_bias = torch.cat( - [q_proj.bias.data, k_proj.bias.data, v_proj.bias.data] - ) - fused_layer.bias.data = concatenated_bias.to(device=device, dtype=dtype) - - return fused_layer - - -def delete_projection_layers(module: nn.Module, layer_names: list[str]) -> None: - for name in layer_names: - if hasattr(module, name): - delattr(module, name) diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index f2890284c..e93a9d1e4 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -584,7 +584,7 @@ "48": 247.48, "49": 247.54 }, - "expected_e2e_ms": 16563.83, + "expected_e2e_ms": 18382.19, "expected_avg_denoise_ms": 260.76, "expected_median_denoise_ms": 247.84 }, @@ -706,7 +706,7 @@ "TimestepPreparationStage": 2.9, "LatentPreparationStage": 1.25, "ImageVAEEncodingStage": 1655.89, - "DenoisingStage": 100544.98, + "DenoisingStage": 106972.82, "DecodingStage": 1355.52, "per_frame_generation": null }, @@ -753,7 +753,7 @@ "39": 1599.78 }, "expected_e2e_ms": 123182.9887, - "expected_avg_denoise_ms": 2513.52, + "expected_avg_denoise_ms": 2831.00, "expected_median_denoise_ms": 1600.09 }, "wan2_1_i2v_14b_480P_2gpu": {