diff --git a/python/sglang/multimodal_gen/runtime/layers/elementwise.py b/python/sglang/multimodal_gen/runtime/layers/elementwise.py new file mode 100644 index 000000000..4c5e50857 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/layers/elementwise.py @@ -0,0 +1,35 @@ +import torch + +from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp +from sglang.multimodal_gen.runtime.layers.triton_ops import fuse_scale_shift_kernel + + +class MulAdd(CustomOp): + """ + Fuse elementwise mul and add + Input: a, b, c, OptionalInt[k] + Output: a * (k + b) + c + """ + + def __init__(self, prefix: str = ""): + super().__init__() + + def forward_native( + self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0 + ) -> torch.Tensor: + # a.shape: [batch_size, seq_len, inner_dim] + if b.dim() == 4: + # b.shape: [batch_size, num_frames, 1, inner_dim] + num_frames = b.shape[1] + frame_seqlen = a.shape[1] // num_frames + return c + ( + a.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * b + ).flatten(1, 2) + else: + # b.shape: [batch_size, 1, inner_dim] + return c + a * (k + b) + + def forward_cuda( + self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0 + ): + return fuse_scale_shift_kernel(a, b, c, scale_constant=k) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 6b42cafd1..0fb825261 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -238,31 +238,6 @@ class LayerNorm(CustomOp): return s -class ScaleResidual(nn.Module): - """ - Applies gated residual connection. - """ - - def __init__(self, prefix: str = ""): - super().__init__() - - def forward( - self, residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor - ) -> torch.Tensor: - """Apply gated residual connection.""" - # x.shape: [batch_size, seq_len, inner_dim] - if gate.dim() == 4: - # gate.shape: [batch_size, num_frames, 1, inner_dim] - num_frames = gate.shape[1] - frame_seqlen = x.shape[1] // num_frames - return residual + ( - x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate - ).flatten(1, 2) - else: - # gate.shape: [batch_size, 1, inner_dim] - return residual + x * gate - - # adapted from Diffusers: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/normalization.py # NOTE(will): Needed to match behavior of diffusers and wan2.1 even while using # FSDP's MixedPrecisionPolicy diff --git a/python/sglang/multimodal_gen/runtime/layers/triton_ops.py b/python/sglang/multimodal_gen/runtime/layers/triton_ops.py index aa0fc17b9..931e9a473 100644 --- a/python/sglang/multimodal_gen/runtime/layers/triton_ops.py +++ b/python/sglang/multimodal_gen/runtime/layers/triton_ops.py @@ -25,6 +25,7 @@ def _fused_scale_shift_4d_kernel( normalized_ptr, scale_ptr, shift_ptr, + scale_constant: tl.constexpr, # scale_constant is either 0 or 1. rows, inner_dim, seq_len, @@ -56,8 +57,8 @@ def _fused_scale_shift_4d_kernel( scale = tl.load(scale_ptrs, mask=mask, other=0.0) shift = tl.load(shift_ptrs, mask=mask, other=0.0) - one = tl.full([BLOCK_N], 1.0, dtype=scale.dtype) - output = normalized * (one + scale) + shift + scale_const_tensor = tl.full([BLOCK_N], scale_constant, dtype=scale.dtype) + output = normalized * (scale_const_tensor + scale) + shift tl.store(out_ptrs, output, mask=mask) @@ -67,6 +68,7 @@ def fuse_scale_shift_kernel_blc_opt( x_ptr, shift_ptr, scale_ptr, + scale_constant: tl.constexpr, # scale_constant is either 0 or 1., y_ptr, B, L, @@ -125,7 +127,7 @@ def fuse_scale_shift_kernel_blc_opt( ) scale = tl.load(scale_ptr + sc_off, mask=mask, other=0) - y = x * (1 + scale) + shift + y = x * (scale_constant + scale) + shift tl.store(y_ptr + x_off, y, mask=mask) @@ -221,6 +223,7 @@ def fuse_scale_shift_kernel( x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, + scale_constant: float = 1.0, block_l: int = 128, block_c: int = 128, ): @@ -251,6 +254,7 @@ def fuse_scale_shift_kernel( x_2d, scale_reshaped, shift_reshaped, + scale_constant, rows, C, L, @@ -306,6 +310,7 @@ def fuse_scale_shift_kernel( x, shift_blc if need_shift_scalar else shift_exp, scale_blc if need_scale_scalar else scale_exp, + scale_constant, output, B, L, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py index 6244f7390..0095a0591 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py @@ -26,11 +26,11 @@ import torch.distributed as dist from sglang.multimodal_gen.configs.models.dits import WanVideoConfig from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_world_size from sglang.multimodal_gen.runtime.layers.attention import LocalAttention +from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd from sglang.multimodal_gen.runtime.layers.layernorm import ( FP32LayerNorm, LayerNormScaleShift, RMSNorm, - ScaleResidual, ScaleResidualLayerNormScaleShift, ) from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear @@ -318,7 +318,7 @@ class CausalWanTransformerBlock(nn.Module): # 3. Feed-forward self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") - self.mlp_residual = ScaleResidual() + self.mlp_residual = MulAdd() self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) @@ -417,7 +417,7 @@ class CausalWanTransformerBlock(nn.Module): # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) - hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) + hidden_states = self.mlp_residual(ff_output, c_gate_msa, hidden_states) hidden_states = hidden_states.to(orig_dtype) return hidden_states diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py index e8dc063cc..80539de66 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py @@ -15,10 +15,10 @@ from sglang.multimodal_gen.runtime.layers.attention import ( LocalAttention, UlyssesAttention, ) +from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd from sglang.multimodal_gen.runtime.layers.layernorm import ( LayerNormScaleShift, RMSNorm, - ScaleResidual, ScaleResidualLayerNormScaleShift, ) from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear @@ -81,7 +81,7 @@ class MMDoubleStreamBlock(nn.Module): self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift( hidden_size, norm_type="layer", elementwise_affine=False, dtype=dtype ) - self.img_mlp_residual = ScaleResidual() + self.img_mlp_residual = MulAdd() # Image attention components self.img_attn_qkv = ReplicatedLinear( @@ -127,7 +127,7 @@ class MMDoubleStreamBlock(nn.Module): self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift( hidden_size, norm_type="layer", elementwise_affine=False, dtype=dtype ) - self.txt_mlp_residual = ScaleResidual() + self.txt_mlp_residual = MulAdd() # Text attention components self.txt_attn_qkv = ReplicatedLinear( @@ -231,7 +231,7 @@ class MMDoubleStreamBlock(nn.Module): # Process image MLP img_mlp_out = self.img_mlp(img_mlp_input) - img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate) + img = self.img_mlp_residual(img_mlp_out, img_mlp_gate, img_residual) # Process text attention output txt_attn_out, _ = self.txt_attn_proj( @@ -245,7 +245,7 @@ class MMDoubleStreamBlock(nn.Module): # Process text MLP txt_mlp_out = self.txt_mlp(txt_mlp_input) - txt = self.txt_mlp_residual(txt_residual, txt_mlp_out, txt_mlp_gate) + txt = self.txt_mlp_residual(txt_mlp_out, txt_mlp_gate, txt_residual) return img, txt @@ -304,7 +304,7 @@ class MMSingleStreamBlock(nn.Module): elementwise_affine=False, dtype=dtype, ) - self.output_residual = ScaleResidual() + self.output_residual = MulAdd() # Activation function self.mlp_act = nn.GELU(approximate="tanh") @@ -384,7 +384,7 @@ class MMSingleStreamBlock(nn.Module): output, _ = self.linear2(combined) # Apply residual connection with gating using fused operation - return self.output_residual(x, output, mod_gate) + return self.output_residual(output, mod_gate, x) class HunyuanVideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): 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 5d5a795d2..ad5c96229 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -17,6 +17,7 @@ from diffusers.models.normalization import AdaLayerNormContinuous from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.layers.attention import USPAttention +from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd from sglang.multimodal_gen.runtime.layers.layernorm import ( LayerNorm, RMSNorm, @@ -28,7 +29,6 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( ) from sglang.multimodal_gen.runtime.layers.triton_ops import ( fuse_scale_shift_gate_select01_kernel, - fuse_scale_shift_kernel, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import ( @@ -350,7 +350,7 @@ class QwenEmbedLayer3DRope(nn.Module): if idx != layer_num: video_freq = self._compute_video_freqs(frame, height, width, idx) else: - ### For the condition image, we set the layer index to -1 + # For the condition image, we set the layer index to -1 video_freq = self._compute_condition_freqs(frame, height, width) video_freq = video_freq.to(device) vid_freqs.append(video_freq) @@ -673,6 +673,8 @@ class QwenImageTransformerBlock(nn.Module): self.txt_mlp = FeedForward( dim=dim, dim_out=dim, activation_fn="gelu-approximate" ) + # Utils + self.fuse_mul_add = MulAdd() def _modulate(self, x, mod_params, index=None): shift, scale, gate = mod_params.chunk(3, dim=-1) @@ -714,14 +716,14 @@ class QwenImageTransformerBlock(nn.Module): ) gate_result = torch.where(mask, gate0.unsqueeze(1), gate1.unsqueeze(1)) return ( - fuse_scale_shift_kernel(x, scale_result, shift_result), + self.fuse_mul_add(x, scale_result, shift_result, k=1.0), gate_result, ) else: shift_result = shift.unsqueeze(1) scale_result = scale.unsqueeze(1) gate_result = gate.unsqueeze(1) - return fuse_scale_shift_kernel(x, scale_result, shift_result), gate_result + return self.fuse_mul_add(x, scale_result, shift_result, k=1.0), gate_result def forward( self, @@ -759,8 +761,10 @@ class QwenImageTransformerBlock(nn.Module): # 4. Splits results back to separate streams joint_attention_kwargs = joint_attention_kwargs or {} attn_output = self.attn( - hidden_states=img_modulated, # Image stream (will be processed as "sample") - encoder_hidden_states=txt_modulated, # Text stream (will be processed as "context") + # Image stream (will be processed as "sample") + hidden_states=img_modulated, + # Text stream (will be processed as "context") + encoder_hidden_states=txt_modulated, encoder_hidden_states_mask=encoder_hidden_states_mask, image_rotary_emb=image_rotary_emb, **joint_attention_kwargs, @@ -780,13 +784,15 @@ class QwenImageTransformerBlock(nn.Module): img_normed2, img_mod2, modulate_index ) img_mlp_output = self.img_mlp(img_modulated2) - hidden_states = hidden_states + img_gate2 * img_mlp_output + hidden_states = self.fuse_mul_add(img_mlp_output, img_gate2, hidden_states) # Process text stream - norm2 + MLP txt_normed2 = self.txt_norm2(encoder_hidden_states) txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2) txt_mlp_output = self.txt_mlp(txt_modulated2) - encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output + encoder_hidden_states = self.fuse_mul_add( + txt_mlp_output, txt_gate2, encoder_hidden_states + ) # Clip to prevent overflow for fp16 if encoder_hidden_states.dtype == torch.float16: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index d90d96b1c..a7dd4879a 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -20,11 +20,11 @@ from sglang.multimodal_gen.runtime.layers.attention import ( UlyssesAttention_VSA, USPAttention, ) +from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd from sglang.multimodal_gen.runtime.layers.layernorm import ( FP32LayerNorm, LayerNormScaleShift, RMSNorm, - ScaleResidual, ScaleResidualLayerNormScaleShift, tensor_parallel_rms_norm, ) @@ -382,7 +382,7 @@ class WanTransformerBlock(nn.Module): # 3. Feed-forward self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") - self.mlp_residual = ScaleResidual() + self.mlp_residual = MulAdd() self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) @@ -488,7 +488,7 @@ class WanTransformerBlock(nn.Module): # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) - hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) + hidden_states = self.mlp_residual(ff_output, c_gate_msa, hidden_states) hidden_states = hidden_states.to(orig_dtype) return hidden_states @@ -582,7 +582,7 @@ class WanTransformerBlock_VSA(nn.Module): # 3. Feed-forward self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") - self.mlp_residual = ScaleResidual() + self.mlp_residual = MulAdd() self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) @@ -669,7 +669,7 @@ class WanTransformerBlock_VSA(nn.Module): # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) - hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) + hidden_states = self.mlp_residual(ff_output, c_gate_msa, hidden_states) hidden_states = hidden_states.to(orig_dtype) return hidden_states