From 59eb142ec2c1e16017d1aa22578568b9ff188882 Mon Sep 17 00:00:00 2001 From: gjsheu Date: Fri, 19 Jun 2026 13:58:45 +0800 Subject: [PATCH] [NPU] [Diffusion] Performance Optimization for LTX-2 Model (#22445) Co-authored-by: gengjinsong Co-authored-by: gengjinsong <904939979@qq.com> Co-authored-by: ronnie_zheng --- .../runtime/layers/layernorm.py | 15 ++++++ .../runtime/models/dits/ltx_2.py | 47 +++++++++++++------ 2 files changed, 47 insertions(+), 15 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 3b73421d4..45ce3fc70 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -41,6 +41,9 @@ if _is_cuda or _is_xpu: if _is_npu: import torch_npu + from sgl_kernel_npu.norm.rmsnorm_without_weight import ( + fused_rmsnorm_without_weight, + ) if _is_musa: from sgl_kernel import fused_add_rmsnorm @@ -292,6 +295,18 @@ class RMSNorm(CustomOp): return f"hidden_size={self.hidden_size}, eps={self.variance_epsilon}" +@CustomOp.register("rms_norm_no_weight") +class RMSNormNoWeight(CustomOp): + def forward_native(self, x: torch.Tensor, eps: float) -> torch.Tensor: + return F.rms_norm(x, normalized_shape=(x.shape[-1],), eps=eps) + + def forward_cuda(self, x: torch.Tensor, eps: float) -> torch.Tensor: + return self.forward_native(x, eps=eps) + + def forward_npu(self, x: torch.Tensor, eps: float) -> torch.Tensor: + return fused_rmsnorm_without_weight(x, eps) + + # Copied and adapted from sglang @CustomOp.register("layer_norm") class LayerNorm(CustomOp): diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index a4874179e..57a737078 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -23,6 +23,7 @@ from sglang.multimodal_gen.runtime.distributed.communication_op import ( tensor_model_parallel_all_reduce, ) from sglang.multimodal_gen.runtime.layers.attention import LocalAttention, USPAttention +from sglang.multimodal_gen.runtime.layers.layernorm import RMSNormNoWeight from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, RowParallelLinear, @@ -35,9 +36,14 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload im LayerwiseOffloadableModuleMixin, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT -from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum +from sglang.multimodal_gen.runtime.platforms import ( + AttentionBackendEnum, + current_platform, +) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +_is_npu = current_platform.is_npu() + logger = init_logger(__name__) ADALN_NUM_BASE_PARAMS = 6 @@ -393,10 +399,6 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module): return cos_freqs.to(dtype=out_dtype), sin_freqs.to(dtype=out_dtype) -def rms_norm(x: torch.Tensor, eps: float) -> torch.Tensor: - return F.rms_norm(x, normalized_shape=(x.shape[-1],), eps=eps) - - class LTX2TextProjection(nn.Module): def __init__( self, @@ -839,6 +841,7 @@ class LTX2TransformerBlock(nn.Module): super().__init__() self.idx = idx self.norm_eps = norm_eps + self.rms_norm = RMSNormNoWeight() # LTX2.3 self.cross_attention_adaln = cross_attention_adaln self.use_local_av_cross_attention = use_local_av_cross_attention @@ -1015,7 +1018,7 @@ class LTX2TransformerBlock(nn.Module): self.scale_shift_table, batch_size, temb, slice(0, 3) ) norm_hidden_states = ( - rms_norm(hidden_states, self.norm_eps) * (1 + vscale_msa) + vshift_msa + self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_msa) + vshift_msa ) attn_hidden_states = self.attn1( norm_hidden_states, @@ -1031,7 +1034,8 @@ class LTX2TransformerBlock(nn.Module): self.audio_scale_shift_table, batch_size, temb_audio, slice(0, 3) ) norm_audio_hidden_states = ( - rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_msa) + ashift_msa + self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_msa) + + ashift_msa ) attn_audio_hidden_states = self.audio_attn1( norm_audio_hidden_states, @@ -1056,7 +1060,7 @@ class LTX2TransformerBlock(nn.Module): self.prompt_scale_shift_table, batch_size, temb_prompt, slice(None) ) norm_hidden_states = ( - rms_norm(hidden_states, self.norm_eps) * (1 + vscale_q) + vshift_q + self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_q) + vshift_q ) mod_encoder_hidden_states = ( encoder_hidden_states * (1 + v_prompt_scale) + v_prompt_shift @@ -1078,7 +1082,8 @@ class LTX2TransformerBlock(nn.Module): slice(None), ) norm_audio_hidden_states = ( - rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_q) + ashift_q + self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_q) + + ashift_q ) mod_audio_encoder_hidden_states = ( audio_encoder_hidden_states * (1 + a_prompt_scale) + a_prompt_shift @@ -1092,7 +1097,7 @@ class LTX2TransformerBlock(nn.Module): audio_hidden_states + attn_audio_hidden_states * agate_q ) else: - norm_hidden_states = rms_norm(hidden_states, self.norm_eps) + norm_hidden_states = self.rms_norm(hidden_states, self.norm_eps) attn_hidden_states = self.attn2( norm_hidden_states, context=encoder_hidden_states, @@ -1100,7 +1105,7 @@ class LTX2TransformerBlock(nn.Module): ) hidden_states = hidden_states + attn_hidden_states - norm_audio_hidden_states = rms_norm(audio_hidden_states, self.norm_eps) + norm_audio_hidden_states = self.rms_norm(audio_hidden_states, self.norm_eps) attn_audio_hidden_states = self.audio_attn2( norm_audio_hidden_states, context=audio_encoder_hidden_states, @@ -1108,8 +1113,8 @@ class LTX2TransformerBlock(nn.Module): ) audio_hidden_states = audio_hidden_states + attn_audio_hidden_states # 3. Audio-to-Video and Video-to-Audio Cross-Attention - norm_hidden_states = rms_norm(hidden_states, self.norm_eps) - norm_audio_hidden_states = rms_norm(audio_hidden_states, self.norm_eps) + norm_hidden_states = self.rms_norm(hidden_states, self.norm_eps) + norm_audio_hidden_states = self.rms_norm(audio_hidden_states, self.norm_eps) # Compute combined ada params video_per_layer_ca_scale_shift = self.video_a2v_cross_attn_scale_shift_table[ @@ -1221,7 +1226,7 @@ class LTX2TransformerBlock(nn.Module): self.scale_shift_table, batch_size, temb, slice(3, 6) ) norm_hidden_states = ( - rms_norm(hidden_states, self.norm_eps) * (1 + vscale_mlp) + vshift_mlp + self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_mlp) + vshift_mlp ) ff_output = self.ff(norm_hidden_states) hidden_states = hidden_states + ff_output * vgate_mlp @@ -1230,7 +1235,8 @@ class LTX2TransformerBlock(nn.Module): self.audio_scale_shift_table, batch_size, temb_audio, slice(3, 6) ) norm_audio_hidden_states = ( - rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_mlp) + ashift_mlp + self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_mlp) + + ashift_mlp ) audio_ff_output = self.audio_ff(norm_audio_hidden_states) audio_hidden_states = audio_hidden_states + audio_ff_output * agate_mlp @@ -1813,6 +1819,17 @@ class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin): audio_encoder_hidden_states = self.audio_caption_projection( audio_encoder_hidden_states ) + + if _is_npu: + # If the 'encoder_attention_mask' is provided and it is all ones, + # it can be set to 'None' to avoid the degradation of performance on the NPU side, + # where the mask, even though it has no affect, + # can lead to the introduction of multiple small operators. + if encoder_attention_mask is not None and torch.all( + encoder_attention_mask == 1 + ): + encoder_attention_mask = None + # 5. Run blocks skip_video_self_attn_blocks = set(skip_video_self_attn_blocks or ()) skip_audio_self_attn_blocks = set(skip_audio_self_attn_blocks or ())