diff --git a/python/sglang/kernels/ops/diffusion/triton/silu_mul_bitexact.py b/python/sglang/kernels/ops/diffusion/triton/silu_mul_bitexact.py new file mode 100644 index 000000000..f2df199f6 --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py b/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py index afd212558..b1e83ec0e 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py +++ b/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( diff --git a/python/sglang/multimodal_gen/test/unit/test_ideogram_rope_swiglu_fusion.py b/python/sglang/multimodal_gen/test/unit/test_ideogram_rope_swiglu_fusion.py new file mode 100644 index 000000000..8ad507f9b --- /dev/null +++ b/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()