diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py b/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py index e12bd874d..def48ee48 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py @@ -19,6 +19,10 @@ import torch.nn as nn import torch.nn.functional as F from diffusers.models.embeddings import TimestepEmbedding, Timesteps +from sglang.kernels.ops.diffusion.residual_gate_add import ( + can_use_residual_gate_add_cuda, + residual_gate_add_cuda, +) from sglang.multimodal_gen.configs.models.dits.ernie_image import ( ErnieImageDitConfig, ) @@ -40,6 +44,40 @@ 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.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +_ERNIE_RESIDUAL_GATE_CUDA_DISABLED = False + + +def _ernie_residual_gate_add( + residual: torch.Tensor, + update: torch.Tensor, + gate: torch.Tensor, +) -> torch.Tensor: + """Single-kernel ``residual + gate * update``, bit-exact vs the eager pair. + + Restricted to half dtypes: there the kernel reproduces the eager pair's + two-step rounding exactly (verified by ``torch.equal``), while for fp32 it + would contract to an fma (one rounding) and stop being bit-exact. + """ + global _ERNIE_RESIDUAL_GATE_CUDA_DISABLED + + if ( + not _ERNIE_RESIDUAL_GATE_CUDA_DISABLED + and residual.dtype in (torch.float16, torch.bfloat16) + and can_use_residual_gate_add_cuda(residual, update, gate) + ): + try: + return residual_gate_add_cuda(residual, update, gate) + except Exception as exc: + if torch.compiler.is_compiling(): + raise + logger.warning_once(f"Disabling ERNIE residual-gate CUDA fast path: {exc}") + _ERNIE_RESIDUAL_GATE_CUDA_DISABLED = True + + return residual + gate * update def _rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor: @@ -256,13 +294,13 @@ class ErnieImageSharedAdaLNBlock(nn.Module): ) -> torch.Tensor: residual = x x = self.adaLN_sa_ln(x) * (1 + scale_msa) + shift_msa - x = residual + gate_msa * self.self_attention( + attn_out = self.self_attention( x, rotary_pos_emb, attn_mask=attn_mask, attn_mask_meta=attn_mask_meta ) + residual = _ernie_residual_gate_add(residual, attn_out, gate_msa) - residual = x - x = self.adaLN_mlp_ln(x) * (1 + scale_mlp) + shift_mlp - x = residual + gate_mlp * self.mlp(x) + x = self.adaLN_mlp_ln(residual) * (1 + scale_mlp) + shift_mlp + x = _ernie_residual_gate_add(residual, self.mlp(x), gate_mlp) return x @@ -473,8 +511,10 @@ class ErnieImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin) c = self.time_embedding(t_emb.to(dtype=dtype)) mod_params = self.adaLN_modulation(c) + # .contiguous() is a bit-exact copy of the tiny (B, 1, D) modulation + # params; the fused residual-gate kernel requires dense inputs. shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( - t.unsqueeze(1) for t in mod_params.chunk(6, dim=-1) + t.unsqueeze(1).contiguous() for t in mod_params.chunk(6, dim=-1) ) for layer in self.layers: diff --git a/test/registered/kernels/ops/diffusion/test_ernie_residual_gate_add.py b/test/registered/kernels/ops/diffusion/test_ernie_residual_gate_add.py new file mode 100644 index 000000000..221a762ef --- /dev/null +++ b/test/registered/kernels/ops/diffusion/test_ernie_residual_gate_add.py @@ -0,0 +1,35 @@ +"""ERNIE residual-gate fast path must stay bit-exact vs the eager pair.""" + +import sys + +import pytest +import torch + +from sglang.multimodal_gen.runtime.models.dits.ernie_image import ( + _ernie_residual_gate_add, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large") +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +def test_residual_gate_add_is_bit_exact(dtype): + # Real ERNIE-Image shapes: hidden 4096, 1024^2 image tokens + text tokens. + # fp32 exercises the eager fallback (fast path is half-dtype only). + torch.manual_seed(0) + residual = torch.randn(1, 4216, 4096, device="cuda", dtype=dtype) + update = torch.randn_like(residual) + gate = torch.randn(1, 1, 4096, device="cuda", dtype=dtype) + out = _ernie_residual_gate_add(residual, update, gate) + assert torch.equal(out, residual + gate * update) + + # Full-shape gate takes the same kernel path and must stay exact too. + gate_full = gate.expand_as(residual).contiguous() + out_full = _ernie_residual_gate_add(residual, update, gate_full) + assert torch.equal(out_full, residual + gate_full * update) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__]))