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

[Diffusion] Fuse LongCat Image normalization and modulation (#38530)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-09-09 11:12:17 +08:00
committed by GitHub
parent 76eea36e38
commit 2952c8d5ea
2 changed files with 209 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py
+78 -10
View File
@@ -32,6 +32,7 @@ from diffusers.models.normalization import (
AdaLayerNormZeroSingle,
)
from sglang.kernels.ops import diffusion as diffusion_ops
from sglang.kernels.ops.diffusion import (
BitExactFusionGate,
can_use_fused_inplace_qknorm_rope,
@@ -62,6 +63,76 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
_LONGCAT_QKNORM_ROPE = BitExactFusionGate("LongCat fused QKNorm+RoPE")
_LONGCAT_LN_MOD = BitExactFusionGate("LongCat fused LN+modulate", per_signature=True)
def _longcat_norm_modulate(
norm: nn.Module,
x: torch.Tensor,
scale: torch.Tensor,
shift: torch.Tensor,
) -> torch.Tensor:
if torch.is_grad_enabled() or torch.compiler.is_compiling():
return norm(x) * (1 + scale[:, None]) + shift[:, None]
# LayerNorm's reduction depends on the live aten dispatch. Verify each
# shape/stride before using the bit-exact fused kernel, outside capture.
if (
not _LONGCAT_LN_MOD.disabled
and x.is_cuda
and diffusion_ops.is_plain_layer_norm(norm, x.shape[-1])
and diffusion_ops.can_use_fused_layernorm_modulate(x, scale, shift)
):
sig = (
x.shape,
x.stride(),
scale.shape,
scale.stride(),
shift.shape,
shift.stride(),
norm.eps,
x.dtype,
x.device,
)
verified = _LONGCAT_LN_MOD.is_verified(sig)
if verified or not torch.cuda.is_current_stream_capturing():
try:
out = diffusion_ops.fused_layernorm_modulate(x, scale, shift, norm.eps)
except Exception as exc:
_LONGCAT_LN_MOD.on_exception(exc, logger=logger)
else:
if verified:
return out
reference = norm(x) * (1 + scale[:, None]) + shift[:, None]
return _LONGCAT_LN_MOD.accept_or_fallback(
out, reference, sig=sig, logger=logger
)
return norm(x) * (1 + scale[:, None]) + shift[:, None]
class _LongCatAdaLayerNormZero(AdaLayerNormZero):
def forward(
self,
x: torch.Tensor,
timestep: Optional[torch.Tensor] = None,
class_labels: Optional[torch.LongTensor] = None,
hidden_dtype: Optional[torch.dtype] = None,
emb: Optional[torch.Tensor] = None,
):
if self.emb is not None:
emb = self.emb(timestep, class_labels, hidden_dtype=hidden_dtype)
emb = self.linear(self.silu(emb))
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(
6, dim=1
)
x = _longcat_norm_modulate(self.norm, x, scale_msa, shift_msa)
return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
class _LongCatAdaLayerNormZeroSingle(AdaLayerNormZeroSingle):
def forward(self, x: torch.Tensor, emb: Optional[torch.Tensor] = None):
emb = self.linear(self.silu(emb))
shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1)
return _longcat_norm_modulate(self.norm, x, scale_msa, shift_msa), gate_msa
def _longcat_qknorm_rope_reference(
@@ -500,7 +571,7 @@ class _SingleTransformerBlock(nn.Module):
):
super().__init__()
self.mlp_hidden_dim = int(dim * mlp_ratio)
self.norm = AdaLayerNormZeroSingle(dim)
self.norm = _LongCatAdaLayerNormZeroSingle(dim)
# proj_mlp: ColumnParallelLinear with gather_output=False keeps output
# head-sharded, consistent with attn_output from _LongCatSingleAttention.
self.proj_mlp = ColumnParallelLinear(
@@ -619,8 +690,8 @@ class _TransformerBlock(nn.Module):
prefix: str = "",
):
super().__init__()
self.norm1 = AdaLayerNormZero(dim)
self.norm1_context = AdaLayerNormZero(dim)
self.norm1 = _LongCatAdaLayerNormZero(dim)
self.norm1_context = _LongCatAdaLayerNormZero(dim)
self.attn = _LongCatJointAttention(
dim=dim,
num_attention_heads=num_attention_heads,
@@ -666,9 +737,8 @@ class _TransformerBlock(nn.Module):
hidden_states, attn_output, gate_msa.unsqueeze(1)
)
norm_hidden_states = self.norm2(hidden_states)
norm_hidden_states = (
norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
norm_hidden_states = _longcat_norm_modulate(
self.norm2, hidden_states, scale_mlp, shift_mlp
)
ff_output = self.ff(norm_hidden_states)
hidden_states = residual_gate_add(
@@ -679,10 +749,8 @@ class _TransformerBlock(nn.Module):
encoder_hidden_states, context_attn_output, c_gate_msa.unsqueeze(1)
)
norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
norm_encoder_hidden_states = (
norm_encoder_hidden_states * (1 + c_scale_mlp[:, None])
+ c_shift_mlp[:, None]
norm_encoder_hidden_states = _longcat_norm_modulate(
self.norm2_context, encoder_hidden_states, c_scale_mlp, c_shift_mlp
)
context_ff_output = self.ff_context(norm_encoder_hidden_states)
encoder_hidden_states = residual_gate_add(
test/registered/kernel/diffusion/test_longcat_image_norm_modulate.py
+131
View File
@@ -0,0 +1,131 @@
"""LongCat normalization parity and graph-safe fusion dispatch."""
import unittest
from unittest.mock import patch
import torch
from diffusers.models.normalization import AdaLayerNormZero, AdaLayerNormZeroSingle
import sglang.multimodal_gen.runtime.models.dits.longcat_image as longcat
from sglang.kernels.ops.diffusion import BitExactFusionGate, modulate_scale_shift
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")
class TestLongCatNormModulation(CustomTestCase):
def setUp(self):
super().setUp()
self.original_gate = longcat._LONGCAT_LN_MOD
longcat._LONGCAT_LN_MOD = BitExactFusionGate("test", per_signature=True)
torch.manual_seed(42)
def tearDown(self):
longcat._LONGCAT_LN_MOD = self.original_gate
super().tearDown()
def require_cuda(self):
if not torch.cuda.is_available():
self.skipTest("CUDA required")
@torch.inference_mode()
def test_adaln_checkpoint_and_output_parity(self):
for device, dtype, dim, seq in [
("cpu", torch.float32, 64, 17),
("cuda", torch.bfloat16, 3072, 512),
("cuda", torch.bfloat16, 3072, 4608),
]:
if device == "cuda" and not torch.cuda.is_available():
continue
for reference_cls, candidate_cls in [
(AdaLayerNormZero, longcat._LongCatAdaLayerNormZero),
(AdaLayerNormZeroSingle, longcat._LongCatAdaLayerNormZeroSingle),
]:
with self.subTest(device=device, seq=seq, cls=reference_cls.__name__):
reference = reference_cls(dim).to(device=device, dtype=dtype)
candidate = candidate_cls(dim).to(device=device, dtype=dtype)
candidate.load_state_dict(reference.state_dict(), strict=True)
x = torch.randn(1, seq, dim, device=device, dtype=dtype)
emb = torch.randn(1, dim, device=device, dtype=dtype)
expected, actual = reference(x, emb=emb), candidate(x, emb=emb)
for a, b in zip(expected, actual, strict=True):
self.assertTrue(torch.equal(a, b))
if device == "cuda":
self.assertTrue(longcat._LONGCAT_LN_MOD.verified)
self.assertFalse(longcat._LONGCAT_LN_MOD.disabled)
def inputs(self, seq=4096):
self.require_cuda()
x = torch.randn(1, seq, 3072, device="cuda", dtype=torch.bfloat16)
modulation = torch.randn(1, 6 * 3072, device="cuda", dtype=torch.bfloat16)
shift, scale, *_ = modulation.chunk(6, dim=-1)
norm = torch.nn.LayerNorm(3072, elementwise_affine=False, eps=1e-6).cuda()
return norm, x, scale, shift
def test_grad_enabled_uses_differentiable_reference(self):
norm, x, scale, shift = self.inputs(seq=17)
with torch.inference_mode():
longcat._longcat_norm_modulate(norm, x, scale, shift)
self.assertTrue(longcat._LONGCAT_LN_MOD.verified)
leaves = [t.detach().clone().requires_grad_() for t in (x, scale, shift)]
refs = [t.detach().clone().requires_grad_() for t in leaves]
with patch.object(longcat.diffusion_ops, "fused_layernorm_modulate") as fused:
actual = longcat._longcat_norm_modulate(norm, *leaves)
actual.float().sum().backward()
fused.assert_not_called()
expected = norm(refs[0]) * (1 + refs[1][:, None]) + refs[2][:, None]
expected.float().sum().backward()
self.assertTrue(torch.equal(actual, expected))
for a, b in zip(leaves, refs, strict=True):
self.assertIsNotNone(a.grad)
self.assertTrue(torch.equal(a.grad, b.grad))
@torch.inference_mode()
def test_changed_inputs_are_used_by_graph_replay(self):
norm, x, scale, shift = self.inputs()
longcat._longcat_norm_modulate(norm, x, scale, shift)
self.assertTrue(longcat._LONGCAT_LN_MOD.verified)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
actual = longcat._longcat_norm_modulate(norm, x, scale, shift)
x.add_(0.25)
scale.neg_()
shift.mul_(0.5)
graph.replay()
expected = norm(x) * (1 + scale[:, None]) + shift[:, None]
self.assertTrue(torch.equal(actual, expected))
@torch.inference_mode()
def test_unverified_capture_uses_eager_reference(self):
norm, x, scale, shift = self.inputs(seq=17)
expected = modulate_scale_shift(norm(x), scale, shift)
with patch.object(longcat.diffusion_ops, "fused_layernorm_modulate") as fused:
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
actual = longcat._longcat_norm_modulate(norm, x, scale, shift)
graph.replay()
fused.assert_not_called()
self.assertFalse(longcat._LONGCAT_LN_MOD.verified)
self.assertTrue(torch.equal(actual, expected))
@torch.inference_mode()
def test_mismatch_disables_fusion_and_returns_reference(self):
norm, x, scale, shift = self.inputs(seq=17)
expected = norm(x) * (1 + scale[:, None]) + shift[:, None]
with patch.object(
longcat.diffusion_ops,
"fused_layernorm_modulate",
return_value=torch.zeros_like(x),
):
actual = longcat._longcat_norm_modulate(norm, x, scale, shift)
self.assertTrue(longcat._LONGCAT_LN_MOD.disabled)
self.assertTrue(torch.equal(actual, expected))
with patch.object(longcat.diffusion_ops, "fused_layernorm_modulate") as fused:
actual = longcat._longcat_norm_modulate(norm, x, scale, shift)
fused.assert_not_called()
self.assertTrue(torch.equal(actual, expected))
if __name__ == "__main__":
unittest.main()