[diffusion] Ideogram-4: fuse Qwen3-style RoPE and SwiGLU silu-mul (denoise -5.1% H100 / -4.7% H200, bit-exact) (#34314)

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

python/sglang/kernels/ops/diffusion/triton/silu_mul_bitexact.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Bit-exact fused ``silu(a) * b`` over two same-shape tensors.
+
+For SwiGLU MLPs whose gate/up projections are separate GEMMs (so the
+concatenated-input ``silu_and_mul`` kernels don't apply without an extra
+full-width ``cat`` pass), this fuses the eager pair
+
+    ``s = F.silu(a)`` (one kernel) ``out = s * b`` (another kernel)
+
+into one pass while reproducing both aten bf16 rounding boundaries:
+``silu`` is a single aten op (fp32 opmath, one round), the multiply rounds
+once more.  ``tl.sigmoid`` lowers to the same fp32 sigmoid aten uses, which
+makes the replication exact (verified ``torch.equal`` on 1M random bf16
+values); callers still verify the first call and fall back on mismatch.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+from sglang.kernels.ops.diffusion.triton.numerics import round_bf16_to_fp32
+from sglang.srt.utils.custom_op import register_custom_op
+
+
+@triton.jit
+def _silu_mul_kernel(
+    out_ptr,
+    a_ptr,
+    b_ptr,
+    numel,
+    BLOCK: tl.constexpr,
+):
+    offs = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < numel
+    a = tl.load(a_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+    b = tl.load(b_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+    s = round_bf16_to_fp32(a * tl.sigmoid(a))
+    tl.store(out_ptr + offs, s * b, mask=mask)  # store rounds the multiply
+
+
+def can_use_fused_silu_mul(a: torch.Tensor, b: torch.Tensor) -> bool:
+    return (
+        a.dtype is torch.bfloat16
+        and b.dtype is torch.bfloat16
+        and a.is_cuda
+        and b.is_cuda
+        and a.device == b.device
+        and a.shape == b.shape
+        and a.is_contiguous()
+        and b.is_contiguous()
+        and a.numel() > 0
+    )
+
+
+def _fake_silu_mul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    return torch.empty_like(a)
+
+
+@register_custom_op(
+    op_name="triton_fused_silu_mul_bitexact",
+    mutates_args=[],
+    fake_impl=_fake_silu_mul,
+)
+def fused_silu_mul_bitexact(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """``silu(a) * b``, bit-exact vs the eager two-kernel chain."""
+    out = torch.empty_like(a)
+    numel = a.numel()
+    with torch.cuda.device(a.device):
+        _silu_mul_kernel[(triton.cdiv(numel, 1024),)](
+            out,
+            a,
+            b,
+            numel,
+            BLOCK=1024,
+        )
+    return out

python/sglang/multimodal_gen/runtime/models/dits/ideogram.py

@@ -7,12 +7,23 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from sglang.kernels.ops.diffusion.bitexact_gate import (
+    BitExactFusionGate,
+    tensors_equal,
+)
 from sglang.kernels.ops.diffusion.fused_gate_rmsnorm import (
     fused_gate_rmsnorm_active,
     fused_rmsnorm_scale,
     fused_rmsnorm_tanh_residual,
     mark_fused_gate_rmsnorm_site,
 )
+from sglang.kernels.ops.diffusion.triton.rope_rotate_half_bitexact import (
+    fused_rope_rotate_half_bitexact,
+)
+from sglang.kernels.ops.diffusion.triton.silu_mul_bitexact import (
+    can_use_fused_silu_mul,
+    fused_silu_mul_bitexact,
+)
 from sglang.multimodal_gen.configs.models.dits.ideogram import Ideogram4DiTConfig
 from sglang.multimodal_gen.configs.models.fsdp import is_layer
 from sglang.multimodal_gen.runtime.distributed import (
@@ -48,10 +59,115 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload im
 )
 from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT
 from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
 
 OUTPUT_IMAGE_INDICATOR = 2
 LLM_TOKEN_INDICATOR = 3
 
+_IDEOGRAM_ROPE = BitExactFusionGate("Ideogram fused RoPE")
+_IDEOGRAM_SWIGLU = BitExactFusionGate("Ideogram fused SiLU-mul")
+
+
+def _can_use_fused_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> bool:
+    # cos/sin are full-span (B, S, 1, D) rows broadcast over heads.
+    expected = (q.shape[0], q.shape[1], 1, q.shape[-1])
+    return (
+        q.dtype is torch.bfloat16
+        and k.dtype is torch.bfloat16
+        and q.is_cuda
+        and q.dim() == 4
+        and q.is_contiguous()
+        and k.shape == q.shape
+        and k.is_contiguous()
+        and k.device == q.device
+        and cos.dtype is torch.bfloat16
+        and sin.dtype is torch.bfloat16
+        and cos.device == q.device
+        and sin.device == q.device
+        and cos.shape == expected
+        and sin.shape == expected
+        and cos.is_contiguous()
+        and sin.is_contiguous()
+        and q.shape[-1] % 2 == 0
+    )
+
+
+def _ideogram_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Single-kernel Qwen3-style RoPE per projection, bit-exact vs eager.
+
+    The eager chain is ~6 kernels per projection (four muls, two add/subs,
+    plus the sliced ``empty_like`` fills); the Triton kernel reproduces every
+    aten bf16 rounding boundary (``round(round(q1*cos1) + round(-q2*sin1))``
+    equals the eager ``round(round(q1*cos1) - round(q2*sin1))`` exactly), so
+    it mounts by default with first-call self-verification.
+    """
+    verified = _IDEOGRAM_ROPE.verified
+    if (
+        not _IDEOGRAM_ROPE.disabled
+        and _can_use_fused_rope(q, k, cos, sin)
+        and (verified or _IDEOGRAM_ROPE.can_attempt_once())
+    ):
+        try:
+            cos_rows = cos.reshape(-1, cos.shape[-1])
+            sin_rows = sin.reshape(-1, sin.shape[-1])
+            q_fused = fused_rope_rotate_half_bitexact(q, cos_rows, sin_rows)
+            k_fused = fused_rope_rotate_half_bitexact(k, cos_rows, sin_rows)
+        except Exception as exc:
+            _IDEOGRAM_ROPE.on_exception(exc, logger=logger)
+        else:
+            if verified:
+                return q_fused, k_fused
+            return _IDEOGRAM_ROPE.accept_or_fallback(
+                (q_fused, k_fused),
+                qwen3_apply_rotary_pos_emb(q, k, cos, sin),
+                equal=tensors_equal,
+                logger=logger,
+                mismatch_msg=(
+                    "Ideogram fused RoPE fast path is not bit-exact on this "
+                    "platform; falling back to eager"
+                ),
+            )
+    return qwen3_apply_rotary_pos_emb(q, k, cos, sin)
+
+
+def _ideogram_swiglu(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """``silu(a) * b`` in one kernel, bit-exact vs the eager pair."""
+    verified = _IDEOGRAM_SWIGLU.verified
+    if (
+        not _IDEOGRAM_SWIGLU.disabled
+        and can_use_fused_silu_mul(a, b)
+        and (verified or _IDEOGRAM_SWIGLU.can_attempt_once())
+    ):
+        try:
+            out = fused_silu_mul_bitexact(a, b)
+        except Exception as exc:
+            _IDEOGRAM_SWIGLU.on_exception(exc, logger=logger)
+        else:
+            if verified:
+                return out
+            return _IDEOGRAM_SWIGLU.accept_or_fallback(
+                out,
+                F.silu(a) * b,
+                logger=logger,
+                mismatch_msg=(
+                    "Ideogram fused SiLU-mul fast path is not bit-exact on "
+                    "this platform; falling back to eager"
+                ),
+            )
+    return F.silu(a) * b
+
 
 class Ideogram4RMSNorm(nn.Module):
     def __init__(self, dim: int, eps: float = 1e-6) -> None:
@@ -255,7 +371,7 @@ class Ideogram4Attention(nn.Module):
         q, k, v = qkv.unbind(dim=2)
         q = self.norm_q(q)
         k = self.norm_k(k)
-        q, k = qwen3_apply_rotary_pos_emb(q, k, cos, sin)
+        q, k = _ideogram_rope(q, k, cos, sin)
         out = self.attn(q, k, v, attn_mask=attn_mask, attn_mask_meta=attn_mask_meta)
         out = out.reshape(batch_size, seq_len, self.local_num_heads * self.head_dim)
         return self.o(out)
@@ -299,7 +415,7 @@ class Ideogram4MLP(nn.Module):
         )
 
     def forward(self, x):
-        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+        return self.w2(_ideogram_swiglu(self.w1(x), self.w3(x)))
 
 
 def _norm_scale(

python/sglang/multimodal_gen/test/unit/test_ideogram_rope_swiglu_fusion.py

@@ -0,0 +1,63 @@
+import unittest
+
+import torch
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
+    qwen3_apply_rotary_pos_emb,
+)
+from sglang.multimodal_gen.runtime.models.dits.ideogram import (
+    _ideogram_rope,
+    _ideogram_swiglu,
+)
+
+
+def _make_qk_cos_sin(batch, seq, heads, head_dim, device):
+    q = torch.randn(batch, seq, heads, head_dim, device=device, dtype=torch.bfloat16)
+    k = torch.randn(batch, seq, heads, head_dim, device=device, dtype=torch.bfloat16)
+    # Qwen3-style full-span duplicated-halves cos/sin, row-broadcast over heads.
+    half = head_dim // 2
+    freqs = torch.randn(batch, seq, 1, half, device=device) * 3.0
+    emb = torch.cat((freqs, freqs), dim=-1)
+    cos = emb.cos().to(torch.bfloat16)
+    sin = emb.sin().to(torch.bfloat16)
+    return q, k, cos, sin
+
+
+class TestIdeogramRopeFusion(unittest.TestCase):
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_fused_rope_matches_eager(self):
+        torch.manual_seed(0)
+        for batch, seq, heads, head_dim in [(2, 257, 16, 128), (1, 64, 3, 64)]:
+            q, k, cos, sin = _make_qk_cos_sin(batch, seq, heads, head_dim, "cuda")
+            q_ref, k_ref = qwen3_apply_rotary_pos_emb(q, k, cos, sin)
+            q_fused, k_fused = _ideogram_rope(q, k, cos, sin)
+            self.assertTrue(torch.equal(q_ref, q_fused))
+            self.assertTrue(torch.equal(k_ref, k_fused))
+
+    def test_eager_fallback_cpu(self):
+        torch.manual_seed(1)
+        q, k, cos, sin = _make_qk_cos_sin(1, 9, 2, 32, torch.device("cpu"))
+        q_ref, k_ref = qwen3_apply_rotary_pos_emb(q, k, cos, sin)
+        q_out, k_out = _ideogram_rope(q, k, cos, sin)
+        self.assertTrue(torch.equal(q_ref, q_out))
+        self.assertTrue(torch.equal(k_ref, k_out))
+
+
+class TestIdeogramSwigluFusion(unittest.TestCase):
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_fused_swiglu_matches_eager(self):
+        torch.manual_seed(0)
+        a = torch.randn(2, 513, 3584, device="cuda", dtype=torch.bfloat16) * 4
+        b = torch.randn(2, 513, 3584, device="cuda", dtype=torch.bfloat16)
+        self.assertTrue(torch.equal(F.silu(a) * b, _ideogram_swiglu(a, b)))
+
+    def test_eager_fallback_cpu(self):
+        torch.manual_seed(1)
+        a = torch.randn(3, 8, dtype=torch.float32)
+        b = torch.randn(3, 8, dtype=torch.float32)
+        self.assertTrue(torch.equal(F.silu(a) * b, _ideogram_swiglu(a, b)))
+
+
+if __name__ == "__main__":
+    unittest.main()