diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index 70b7df159..9697275b9 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -237,9 +237,11 @@ sglang serve --model-path meituan-longcat/LongCat-Image \ --port 30010 ``` -SANA-Video supports the same path for fixed-resolution serving. Its default -text stage emits a fixed 300-token prompt shape, so the runtime reuses one -graph across prompt lengths without padding it to a generic text bucket: +SANA-Video supports the same path for fixed-shape serving. Its default text +stage emits a fixed 300-token prompt shape, so the runtime reuses one graph +across prompt lengths without padding it to a generic text bucket. BCG also +requires the serving frame count to match warmup; the example captures the +model's default 81-frame signature: ```bash Command sglang serve \ diff --git a/docs/docs/sglang-diffusion/fused_kernels.mdx b/docs/docs/sglang-diffusion/fused_kernels.mdx index bee260769..e2ef68bf3 100644 --- a/docs/docs/sglang-diffusion/fused_kernels.mdx +++ b/docs/docs/sglang-diffusion/fused_kernels.mdx @@ -147,9 +147,8 @@ Kernels are written against a specific eager chain in a specific model, so cover | LTX-2 | QK-norm + split RoPE, ada-values split, RMSNorm+modulate, modulate, residual-gate add, linear+GELU | | LTX-2.5 decoder | paired 3D RoPE with shared axis-table cache | | HunyuanVideo | QKV+RoPE pack, strided QK RMSNorm, linear+GELU | -| SANA-Video | paired fp64 interleaved RoPE | | Sana | LN+modulate, GLUMB bias+SiLU / bias+GLU, residual-gate add | -| SANA-Video | Packed QKV/KV; BF16-input linear attention at `quality=high` | +| SANA-Video | Packed QKV/KV; paired fp64 interleaved RoPE; LN+modulate, GLUMB bias+SiLU / bias+GLU, and residual-gate add during BCG; BF16-input linear attention at `quality=high` | | Sana-WM | bidirectional gated delta-net, fused QK inverse-RMS | | Wan | temb table slices; VAE cat+pad and DupUp3D add, `channels_last_3d` RMSNorm+SiLU | | Cosmos3 / Krea2 / MiniMax-H3 | QK-norm + RoPE (Krea2 also CuTe-DSL norm+scale/shift; MiniMax-H3 also indexed modulation) | diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana.py b/python/sglang/multimodal_gen/runtime/models/dits/sana.py index 470fbdcae..2af9923b2 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/sana.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/sana.py @@ -47,7 +47,7 @@ def _eager_ln_modulate( return norm(x) * (1 + scale) + shift -def _sana_ln_modulate( +def sana_ln_modulate( norm: nn.LayerNorm, x: torch.Tensor, scale: torch.Tensor, @@ -176,7 +176,7 @@ def _conv2d_without_bias(conv: nn.Conv2d, x: torch.Tensor) -> torch.Tensor: ) -def _sana_conv_bias_silu(conv: nn.Conv2d, x: torch.Tensor) -> torch.Tensor: +def sana_conv_bias_silu(conv: nn.Conv2d, x: torch.Tensor) -> torch.Tensor: if conv.bias is None or not _use_sana_bcg_fast_path(x): return F.silu(_mps_safe_conv2d(conv, x)) @@ -204,7 +204,7 @@ def _sana_conv_bias_silu(conv: nn.Conv2d, x: torch.Tensor) -> torch.Tensor: ) -def _sana_conv_bias_glu(conv: nn.Conv2d, x: torch.Tensor) -> torch.Tensor: +def sana_conv_bias_glu(conv: nn.Conv2d, x: torch.Tensor) -> torch.Tensor: if conv.bias is None or not _use_sana_bcg_fast_path(x): hidden_states = _mps_safe_conv2d(conv, x) hidden_states, gate = torch.chunk(hidden_states, 2, dim=1) @@ -245,7 +245,7 @@ def _sana_conv_bias_glu(conv: nn.Conv2d, x: torch.Tensor) -> torch.Tensor: ) -def _sana_residual_gate_add( +def sana_residual_gate_add( residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor ) -> torch.Tensor: if torch.compiler.is_compiling(): @@ -307,7 +307,7 @@ class SanaModulatedNorm(nn.Module): def forward(self, x, temb, scale_shift_table): scale_shift_table = _mps_match_dtype(scale_shift_table, temb) shift, scale = (scale_shift_table[None] + temb[:, None]).chunk(2, dim=1) - return _sana_ln_modulate(self.norm, x, scale, shift) + return sana_ln_modulate(self.norm, x, scale, shift) class GLUMBConv(nn.Module): @@ -329,8 +329,8 @@ class GLUMBConv(nn.Module): self.conv_point = nn.Conv2d(hidden_channels, out_channels, 1, 1, 0, bias=False) def forward(self, hidden_states): - hidden_states = _sana_conv_bias_silu(self.conv_inverted, hidden_states) - hidden_states = _sana_conv_bias_glu(self.conv_depth, hidden_states) + hidden_states = sana_conv_bias_silu(self.conv_inverted, hidden_states) + hidden_states = sana_conv_bias_glu(self.conv_depth, hidden_states) hidden_states = _mps_safe_conv2d(self.conv_point, hidden_states) return hidden_states @@ -475,20 +475,20 @@ class SanaTransformerBlock(nn.Module): scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) ).chunk(6, dim=1) - norm_hidden = _sana_ln_modulate(self.norm1, hidden_states, scale_msa, shift_msa) + norm_hidden = sana_ln_modulate(self.norm1, hidden_states, scale_msa, shift_msa) attn_output = self.attn1(norm_hidden) - hidden_states = _sana_residual_gate_add(hidden_states, attn_output, gate_msa) + hidden_states = sana_residual_gate_add(hidden_states, attn_output, gate_msa) attn_output = self.attn2( hidden_states, encoder_hidden_states, encoder_attention_mask ) hidden_states = hidden_states + attn_output - norm_hidden = _sana_ln_modulate(self.norm2, hidden_states, scale_mlp, shift_mlp) + norm_hidden = sana_ln_modulate(self.norm2, hidden_states, scale_mlp, shift_mlp) norm_hidden = norm_hidden.unflatten(1, (height, width)).permute(0, 3, 1, 2) ff_output = self.ff(norm_hidden) ff_output = ff_output.flatten(2, 3).permute(0, 2, 1) - hidden_states = _sana_residual_gate_add(hidden_states, ff_output, gate_mlp) + hidden_states = sana_residual_gate_add(hidden_states, ff_output, gate_mlp) return hidden_states diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana_video.py b/python/sglang/multimodal_gen/runtime/models/dits/sana_video.py index ac4235ccc..4fa3acfab 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/sana_video.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/sana_video.py @@ -27,7 +27,13 @@ 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.models.dits.sana import SanaAdaLayerNormSingle +from sglang.multimodal_gen.runtime.models.dits.sana import ( + SanaAdaLayerNormSingle, + sana_conv_bias_glu, + sana_conv_bias_silu, + sana_ln_modulate, + sana_residual_gate_add, +) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @@ -180,7 +186,6 @@ class GLUMBTempConv(nn.Module): def __init__(self, channels: int, expand_ratio: float) -> None: super().__init__() hidden_channels = int(expand_ratio * channels) - self.nonlinearity = nn.SiLU() self.conv_inverted = nn.Conv2d(channels, hidden_channels * 2, 1) self.conv_depth = nn.Conv2d( hidden_channels * 2, @@ -203,10 +208,8 @@ class GLUMBTempConv(nn.Module): hidden_states = hidden_states.reshape( batch_size * num_frames, height, width, channels ).permute(0, 3, 1, 2) - hidden_states = self.nonlinearity(self.conv_inverted(hidden_states)) - hidden_states = self.conv_depth(hidden_states) - hidden_states, gate = hidden_states.chunk(2, dim=1) - hidden_states = hidden_states * self.nonlinearity(gate) + hidden_states = sana_conv_bias_silu(self.conv_inverted, hidden_states) + hidden_states = sana_conv_bias_glu(self.conv_depth, hidden_states) hidden_states = self.conv_point(hidden_states) temporal = hidden_states.reshape( @@ -397,20 +400,24 @@ class SanaVideoTransformerBlock(nn.Module): + timestep.reshape(batch_size, timestep.shape[1], 6, -1) ).unbind(dim=2) - norm_hidden_states = self.norm1(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa - hidden_states = hidden_states + gate_msa * self.attn1( - norm_hidden_states.to(hidden_states.dtype), rotary_emb + norm_hidden_states = sana_ln_modulate( + self.norm1, hidden_states, scale_msa, shift_msa + ) + hidden_states = sana_residual_gate_add( + hidden_states, + self.attn1(norm_hidden_states.to(hidden_states.dtype), rotary_emb), + gate_msa, ) hidden_states = hidden_states + self.attn2( hidden_states, encoder_hidden_states, encoder_attention_mask ) - norm_hidden_states = self.norm2(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp + norm_hidden_states = sana_ln_modulate( + self.norm2, hidden_states, scale_mlp, shift_mlp + ) norm_hidden_states = norm_hidden_states.unflatten(1, (frames, height, width)) ff_output = self.ff(norm_hidden_states).flatten(1, 3) - return hidden_states + gate_mlp * ff_output + return sana_residual_gate_add(hidden_states, ff_output, gate_mlp) class SanaVideoModulatedNorm(nn.Module): @@ -427,7 +434,7 @@ class SanaVideoModulatedNorm(nn.Module): shift, scale = ( scale_shift_table[None, None] + embedded_timestep[:, :, None] ).unbind(dim=2) - return self.norm(hidden_states) * (1 + scale) + shift + return sana_ln_modulate(self.norm, hidden_states, scale, shift) class SanaVideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin): diff --git a/python/sglang/multimodal_gen/test/unit/test_sana_video.py b/python/sglang/multimodal_gen/test/unit/test_sana_video.py index 5b99165ef..6a14bbc91 100644 --- a/python/sglang/multimodal_gen/test/unit/test_sana_video.py +++ b/python/sglang/multimodal_gen/test/unit/test_sana_video.py @@ -1,6 +1,7 @@ from types import SimpleNamespace import torch +import torch.nn.functional as F from sglang.multimodal_gen.configs.pipeline_configs.sana_video import ( SanaVideoPipelineConfig, @@ -8,6 +9,8 @@ from sglang.multimodal_gen.configs.pipeline_configs.sana_video import ( from sglang.multimodal_gen.configs.sample.sana_video import SanaVideoSamplingParams from sglang.multimodal_gen.registry import get_model_info from sglang.multimodal_gen.runtime.models.dits.sana_video import ( + GLUMBTempConv, + SanaVideoModulatedNorm, SanaVideoRotaryPosEmbed, apply_interleaved_rotary_emb, apply_interleaved_rotary_emb_pair, @@ -99,3 +102,47 @@ def test_sana_video_paired_rope_falls_back_to_eager_on_cpu(): assert torch.equal(query_out, apply_interleaved_rotary_emb(query, cos, sin)) assert torch.equal(key_out, apply_interleaved_rotary_emb(key, cos, sin)) + + +def test_sana_video_glumb_shared_fast_path_preserves_reference(): + torch.manual_seed(0) + module = GLUMBTempConv(channels=4, expand_ratio=2.0).eval() + hidden_states = torch.randn(1, 2, 3, 5, 4) + + with torch.no_grad(): + actual = module(hidden_states) + + batch_size, num_frames, height, width, channels = hidden_states.shape + expected = hidden_states.reshape( + batch_size * num_frames, height, width, channels + ).permute(0, 3, 1, 2) + expected = F.silu(module.conv_inverted(expected)) + expected = module.conv_depth(expected) + expected, gate = expected.chunk(2, dim=1) + expected = expected * F.silu(gate) + expected = module.conv_point(expected) + temporal = expected.reshape( + batch_size, num_frames, channels, height * width + ).permute(0, 2, 1, 3) + expected = temporal + module.conv_temp(temporal) + expected = expected.permute(0, 2, 3, 1).reshape( + batch_size, num_frames, height, width, channels + ) + + assert torch.equal(actual, expected) + + +def test_sana_video_modulated_norm_shared_fast_path_preserves_reference(): + torch.manual_seed(1) + module = SanaVideoModulatedNorm(dim=8, eps=1e-6).eval() + hidden_states = torch.randn(2, 5, 8) + embedded_timestep = torch.randn(2, 1, 8) + scale_shift_table = torch.randn(2, 8) + + actual = module(hidden_states, embedded_timestep, scale_shift_table) + shift, scale = ( + scale_shift_table[None, None] + embedded_timestep[:, :, None] + ).unbind(dim=2) + expected = module.norm(hidden_states) * (1 + scale) + shift + + assert torch.equal(actual, expected)