minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[NPU] [Diffusion] Performance Optimization for LTX-2 Model (#22445)

Browse Source 
Co-authored-by: gengjinsong <gengjinsong@huawei.com>
Co-authored-by: gengjinsong <904939979@qq.com>
Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
gjsheu
2026-06-19 08:58:45 +03:00
committed by GitHub
co-authored by gengjinsong gengjinsong ronnie_zheng
parent 3b417d3999
commit 59eb142ec2
2 changed files with 47 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/layernorm.py
+15
View File
@@ -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):
python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py
+32 -15
View File
@@ -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 ())