[diffusion] ERNIE-Image bit-exact residual-gate fast path (H200 1024^2 e2e 16.17 -> 15.75 s) (#33734)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

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(residual) * (1 + scale_mlp) + shift_mlp
-        x = self.adaLN_mlp_ln(x) * (1 + scale_mlp) + shift_mlp
+        x = _ernie_residual_gate_add(residual, self.mlp(x), gate_mlp)
-        x = residual + gate_mlp * self.mlp(x)
 
         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:

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