diff --git a/python/sglang/kernels/ops/diffusion/triton/silu_mul_bitexact.py b/python/sglang/kernels/ops/diffusion/triton/silu_mul_bitexact.py index f2df199f6..9163707b3 100644 --- a/python/sglang/kernels/ops/diffusion/triton/silu_mul_bitexact.py +++ b/python/sglang/kernels/ops/diffusion/triton/silu_mul_bitexact.py @@ -40,6 +40,26 @@ def _silu_mul_kernel( tl.store(out_ptr + offs, s * b, mask=mask) # store rounds the multiply +@triton.jit +def _packed_silu_mul_kernel( + out_ptr, + x_ptr, + num_rows, + row_stride, + D: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + block = tl.program_id(1).to(tl.int64) + cols = block * BLOCK + tl.arange(0, BLOCK) + mask = (row < num_rows) & (cols < D) + row_base = row * row_stride + a = tl.load(x_ptr + row_base + cols, mask=mask, other=0.0).to(tl.float32) + b = tl.load(x_ptr + row_base + D + cols, mask=mask, other=0.0).to(tl.float32) + s = round_bf16_to_fp32(a * tl.sigmoid(a)) + tl.store(out_ptr + row * D + cols, s * b, mask=mask) + + def can_use_fused_silu_mul(a: torch.Tensor, b: torch.Tensor) -> bool: return ( a.dtype is torch.bfloat16 @@ -76,3 +96,32 @@ def fused_silu_mul_bitexact(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: BLOCK=1024, ) return out + + +def fused_packed_silu_mul_bitexact(x: torch.Tensor) -> torch.Tensor: + """Bit-exact SwiGLU over a contiguous packed ``[..., 2 * D]`` input.""" + if not ( + x.is_cuda + and x.dtype is torch.bfloat16 + and x.dim() == 3 + and x.stride(-1) == 1 + and x.stride(-2) >= x.shape[-1] + and x.stride(0) == x.shape[1] * x.stride(1) + and x.shape[-1] % 2 == 0 + and x.numel() > 0 + ): + raise RuntimeError("unsupported input for packed fused SiLU-mul") + hidden = x.shape[-1] // 2 + rows = x.numel() // x.shape[-1] + row_stride = x.stride(-2) + out = torch.empty((*x.shape[:-1], hidden), dtype=x.dtype, device=x.device) + with torch.cuda.device(x.device): + _packed_silu_mul_kernel[(rows, triton.cdiv(hidden, 1024))]( + out, + x, + rows, + row_stride, + D=hidden, + BLOCK=1024, + ) + return out diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index d3d6a7432..ee64cdfe6 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -16,11 +16,21 @@ from typing import Any, Dict, List, Optional, Tuple import torch import torch.nn as nn +import torch.nn.functional as F from diffusers.models.attention import AttentionModuleMixin from diffusers.models.embeddings import TimestepEmbedding, Timesteps from diffusers.models.normalization import AdaLayerNormContinuous +from sglang.kernels.ops.diffusion.bitexact_gate import BitExactFusionGate from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add +from sglang.kernels.ops.diffusion.triton.layernorm_modulate import ( + can_use_fused_layernorm_modulate, + fused_layernorm_modulate_raw, + is_plain_layer_norm, +) +from sglang.kernels.ops.diffusion.triton.silu_mul_bitexact import ( + fused_packed_silu_mul_bitexact, +) from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig from sglang.multimodal_gen.runtime.distributed import ( divide, @@ -69,6 +79,118 @@ logger = init_logger(__name__) # pylint: disable=invalid-name _get_qkv_projections = get_qkv_projections +_FLUX2_LN_MOD = BitExactFusionGate("FLUX.2 fused LN+modulate", per_signature=True) +_FLUX2_LN_MOD_SIGS = _FLUX2_LN_MOD.verified_sigs +assert _FLUX2_LN_MOD_SIGS is not None +_FLUX2_SWIGLU = BitExactFusionGate("FLUX.2 fused SwiGLU", per_signature=True) +_FLUX2_SWIGLU_SIGS = _FLUX2_SWIGLU.verified_sigs +assert _FLUX2_SWIGLU_SIGS is not None + + +def _flux2_norm_modulate( + norm: nn.Module, + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, +) -> torch.Tensor: + """Bit-exact single-kernel ``LN(x) * (1 + scale) + shift``.""" + # Preserve the original expression for Dynamo/Inductor. This direct + # Triton dispatch is intentionally an eager fast path. + if torch.compiler.is_compiling(): + return norm(x) * (1 + scale) + shift + + scale_row = scale.squeeze(1) if scale.dim() == 3 and scale.shape[1] == 1 else scale + shift_row = shift.squeeze(1) if shift.dim() == 3 and shift.shape[1] == 1 else shift + if ( + _FLUX2_LN_MOD.disabled + or not is_plain_layer_norm(norm, x.shape[-1]) + or not can_use_fused_layernorm_modulate(x, scale_row, shift_row) + ): + return norm(x) * (1 + scale) + shift + + # The bit-exact contract is set by dtype/reduction width/affine-row + # layout, not by the number of independent rows. Excluding sequence + # length lets the representative warmup verify the real prompt path too. + sig = ( + x.dtype, + x.device, + x.shape[0], + x.shape[-1], + x.stride(-1), + scale_row.stride(0) if scale_row.shape[0] > 1 else x.shape[-1], + shift_row.stride(0) if shift_row.shape[0] > 1 else x.shape[-1], + norm.eps, + ) + verified = sig in _FLUX2_LN_MOD_SIGS + if not verified and torch.cuda.is_current_stream_capturing(): + return norm(x) * (1 + scale) + shift + try: + # Direct dispatch avoids custom-op overhead on this eager-only path. + out = fused_layernorm_modulate_raw(x, scale_row, shift_row, norm.eps) + except Exception as exc: + _FLUX2_LN_MOD.on_exception(exc, logger=logger) + return norm(x) * (1 + scale) + shift + if verified: + return out + ref = norm(x) * (1 + scale) + shift + return _FLUX2_LN_MOD.accept_or_fallback( + out, + ref, + sig=sig, + logger=logger, + mismatch_msg=( + "FLUX.2 fused LN+modulate fast path is not bit-exact on this " + "platform; falling back to eager" + ), + ) + + +def _flux2_swiglu(x: torch.Tensor) -> torch.Tensor: + """Bit-exact fused SwiGLU for the packed FLUX.2 FFN projection.""" + half = x.shape[-1] // 2 + # Let Inductor fuse the reference expression in torch.compile mode. + if torch.compiler.is_compiling(): + return F.silu(x[..., :half]) * x[..., half:] + + # Sequence length only changes the launch grid; D and row stride define + # how the two packed halves are addressed and therefore need verification. + sig = (x.dtype, x.device, x.shape[0], x.shape[-1], x.stride(-2), x.stride(-1)) + verified = sig in _FLUX2_SWIGLU_SIGS + can_fuse = ( + not _FLUX2_SWIGLU.disabled + and x.is_cuda + and x.dtype is torch.bfloat16 + and x.dim() == 3 + and x.stride(-1) == 1 + and x.stride(-2) >= x.shape[-1] + and x.stride(0) == x.shape[1] * x.stride(1) + and x.shape[-1] % 2 == 0 + and x.numel() > 0 + ) + # Per-signature verification may compare tensors and synchronize. Never + # verify a new layout while a CUDA graph is being captured. + if can_fuse and not verified and torch.cuda.is_current_stream_capturing(): + return F.silu(x[..., :half]) * x[..., half:] + if can_fuse: + try: + out = fused_packed_silu_mul_bitexact(x) + except Exception as exc: + _FLUX2_SWIGLU.on_exception(exc, logger=logger) + else: + if verified: + return out + return _FLUX2_SWIGLU.accept_or_fallback( + out, + F.silu(x[..., :half]) * x[..., half:], + sig=sig, + logger=logger, + mismatch_msg=( + "FLUX.2 fused SwiGLU fast path is not bit-exact on this " + "platform; falling back to eager" + ), + ) + return F.silu(x[..., :half]) * x[..., half:] + class Flux2SwiGLU(nn.Module): """ @@ -76,14 +198,8 @@ class Flux2SwiGLU(nn.Module): layer fused into the first linear layer of the FF sub-block. Thus, this module has no trainable parameters. """ - def __init__(self): - super().__init__() - self.gate_fn = nn.SiLU() - def forward(self, x: torch.Tensor) -> torch.Tensor: - x1, x2 = x.chunk(2, dim=-1) - x = self.gate_fn(x1) * x2 - return x + return _flux2_swiglu(x) class Flux2FeedForward(nn.Module): @@ -623,8 +739,9 @@ class Flux2SingleTransformerBlock(nn.Module): mod_shift, mod_scale, mod_gate = temb_mod_params - norm_hidden_states = self.norm(hidden_states) - norm_hidden_states = (1 + mod_scale) * norm_hidden_states + mod_shift + norm_hidden_states = _flux2_norm_modulate( + self.norm, hidden_states, mod_scale, mod_shift + ) joint_attention_kwargs = joint_attention_kwargs or {} attn_output = self.attn( @@ -737,14 +854,17 @@ class Flux2TransformerBlock(nn.Module): ) = temb_mod_params_txt # Img stream - norm_hidden_states = self.norm1(hidden_states) - norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa + norm_hidden_states = _flux2_norm_modulate( + self.norm1, hidden_states, scale_msa, shift_msa + ) # Conditioning txt stream - norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states) - norm_encoder_hidden_states = ( - 1 + c_scale_msa - ) * norm_encoder_hidden_states + c_shift_msa + norm_encoder_hidden_states = _flux2_norm_modulate( + self.norm1_context, + encoder_hidden_states, + c_scale_msa, + c_shift_msa, + ) # Attention on concatenated img + txt stream attention_outputs = self.attn( @@ -760,8 +880,9 @@ class Flux2TransformerBlock(nn.Module): # Process attention outputs for the image stream (`hidden_states`). hidden_states = residual_gate_add(hidden_states, attn_output, gate_msa) - norm_hidden_states = self.norm2(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp + norm_hidden_states = _flux2_norm_modulate( + self.norm2, hidden_states, scale_mlp, shift_mlp + ) ff_output = self.ff(norm_hidden_states) hidden_states = residual_gate_add(hidden_states, ff_output, gate_mlp) @@ -771,9 +892,11 @@ class Flux2TransformerBlock(nn.Module): encoder_hidden_states, context_attn_output, c_gate_msa ) - norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) - norm_encoder_hidden_states = ( - norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp + norm_encoder_hidden_states = _flux2_norm_modulate( + self.norm2_context, + encoder_hidden_states, + c_scale_mlp, + c_shift_mlp, ) context_ff_output = self.ff_context(norm_encoder_hidden_states) diff --git a/test/registered/kernels/ops/diffusion/test_flux2_eager_fusions.py b/test/registered/kernels/ops/diffusion/test_flux2_eager_fusions.py new file mode 100644 index 000000000..5301c9ddd --- /dev/null +++ b/test/registered/kernels/ops/diffusion/test_flux2_eager_fusions.py @@ -0,0 +1,101 @@ +"""FLUX.2 eager fusions must be bit-exact for real packed/view layouts.""" + +import unittest +from unittest.mock import patch + +import torch +import torch.nn.functional as F + +import sglang.multimodal_gen.runtime.models.dits.flux_2 as flux2 +from sglang.multimodal_gen.runtime.models.dits.flux_2 import ( + _flux2_norm_modulate, + _flux2_swiglu, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=12, stage="base-b-kernel-unit", runner_config="1-gpu-large") + + +@unittest.skipUnless(torch.cuda.is_available(), "CUDA required") +class TestFlux2EagerFusions(CustomTestCase): + def setUp(self): + flux2._FLUX2_LN_MOD.disabled = False + flux2._FLUX2_LN_MOD.verified = False + flux2._FLUX2_LN_MOD_SIGS.clear() + flux2._FLUX2_SWIGLU.disabled = False + flux2._FLUX2_SWIGLU.verified = False + flux2._FLUX2_SWIGLU_SIGS.clear() + + def test_norm_modulate_is_bit_exact_across_sequence_lengths(self): + torch.manual_seed(0) + hidden = 256 + norm = torch.nn.LayerNorm( + hidden, eps=1e-6, elementwise_affine=False, device="cuda" + ) + # FLUX.2 modulation values are views of one packed projection. + params = torch.randn(1, 1, 6 * hidden, device="cuda").bfloat16() + shift, scale = params.chunk(6, dim=-1)[:2] + + for seq in (17, 65): + x = torch.randn(1, seq, hidden, device="cuda").bfloat16() + expected = norm(x) * (1 + scale) + shift + actual = _flux2_norm_modulate(norm, x, scale, shift) + self.assertTrue(torch.equal(actual, expected)) + + self.assertFalse(flux2._FLUX2_LN_MOD.disabled) + self.assertEqual(len(flux2._FLUX2_LN_MOD_SIGS), 1) + + def test_packed_swiglu_is_bit_exact_for_contiguous_and_strided_views(self): + torch.manual_seed(1) + hidden = 384 + inputs = [ + torch.randn(1, 19, 2 * hidden, device="cuda").bfloat16(), + torch.randn(1, 19, 3 * hidden, device="cuda").bfloat16()[..., : 2 * hidden], + ] + for x in inputs: + expected = F.silu(x[..., :hidden]) * x[..., hidden:] + actual = _flux2_swiglu(x) + self.assertTrue(torch.equal(actual, expected)) + + self.assertFalse(flux2._FLUX2_SWIGLU.disabled) + self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 2) + + def test_fp16_preserves_reference_path(self): + x = torch.randn(1, 17, 512, device="cuda", dtype=torch.float16) + expected = F.silu(x[..., :256]) * x[..., 256:] + actual = _flux2_swiglu(x) + self.assertTrue(torch.equal(actual, expected)) + self.assertFalse(flux2._FLUX2_SWIGLU.disabled) + + def test_packed_swiglu_rejects_non_dense_outer_stride(self): + base = torch.randn(2, 23, 512, device="cuda", dtype=torch.bfloat16) + x = base[:, :19] + self.assertNotEqual(x.stride(0), x.shape[1] * x.stride(1)) + + expected = F.silu(x[..., :256]) * x[..., 256:] + actual = _flux2_swiglu(x) + self.assertTrue(torch.equal(actual, expected)) + self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 0) + + def test_new_swiglu_signature_is_not_verified_during_graph_capture(self): + first = torch.randn(1, 17, 512, device="cuda", dtype=torch.bfloat16) + self.assertTrue( + torch.equal( + _flux2_swiglu(first), + F.silu(first[..., :256]) * first[..., 256:], + ) + ) + self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 1) + + second = torch.randn(1, 19, 768, device="cuda", dtype=torch.bfloat16) + with patch("torch.cuda.is_current_stream_capturing", return_value=True): + actual = _flux2_swiglu(second) + + expected = F.silu(second[..., :384]) * second[..., 384:] + self.assertTrue(torch.equal(actual, expected)) + self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 1) + + +if __name__ == "__main__": + unittest.main()