[diffusion] Reuse SANA fast paths in SANA-Video BCG (#35961)

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 \

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

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
 

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

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)