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
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

19
python/sglang/multimodal_gen/configs/models/dits/flux.py
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",
}
)

20
python/sglang/multimodal_gen/configs/models/dits/qwenimage.py
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",
}
)

4
python/sglang/multimodal_gen/docs/profiling.md
View File

@@ -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

4
python/sglang/multimodal_gen/runtime/loader/component_loader.py
View File

@@ -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

98
python/sglang/multimodal_gen/runtime/models/dits/flux.py
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,

106
python/sglang/multimodal_gen/runtime/models/dits/flux_2.py
View File

@@ -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,

103
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
View File

@@ -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,

43
python/sglang/multimodal_gen/runtime/models/dits/utils.py
View File

@@ -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)

6
python/sglang/multimodal_gen/test/server/perf_baselines.json
View File

@@ -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": {