From 079afaffb19abbe42c09f0ebb8322d6e9f6f2905 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Tue, 1 Sep 2026 08:38:33 +0800 Subject: [PATCH] [Diffusion] Fuse FLUX.2 gated residual normalization on Blackwell (#37112) Co-authored-by: Cursor --- .../csrc/diffusion/flux2_gated_resnorm.cuh | 231 ++++++++++++++++++ .../sglang/kernels/ops/diffusion/__init__.py | 3 + .../diffusion/norm/flux2_gated_resnorm_jit.py | 119 +++++++++ .../runtime/models/dits/flux_2.py | 138 ++++++++--- .../jit/test_flux2_gated_resnorm.py | 68 ++++++ 5 files changed, 528 insertions(+), 31 deletions(-) create mode 100644 python/sglang/kernels/jit/csrc/diffusion/flux2_gated_resnorm.cuh create mode 100644 python/sglang/kernels/ops/diffusion/norm/flux2_gated_resnorm_jit.py create mode 100644 test/registered/jit/test_flux2_gated_resnorm.py diff --git a/python/sglang/kernels/jit/csrc/diffusion/flux2_gated_resnorm.cuh b/python/sglang/kernels/jit/csrc/diffusion/flux2_gated_resnorm.cuh new file mode 100644 index 000000000..51c6f8332 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/diffusion/flux2_gated_resnorm.cuh @@ -0,0 +1,231 @@ +// Bit-exact FLUX.2 gated residual + LayerNorm + adaLN modulation. +// +// The D=6144 kernel reproduces both the eager bf16 residual update and +// PyTorch's 128-thread vectorized LayerNorm Welford tree. Unsupported shapes +// stay on the existing unfused model path. + +#pragma once + +#include + +#include +#include +#include + +#include + +namespace sglang { + +namespace flux2_gated_resnorm { + +constexpr int kHidden = 6144; +constexpr int kThreads = 128; +constexpr int kWarps = kThreads / device::kWarpThreads; +constexpr int kVecElems = 4; +constexpr int kIterations = kHidden / (kThreads * kVecElems); + +static_assert(kWarps == 4); +static_assert(kIterations == 12); + +struct Params { + void* output; + void* residual_out; + const void* update; + const void* residual; + const void* gate; + const void* scale; + const void* shift; + float eps; +}; + +struct WelfordState { + float mean; + float m2; + float count; +}; + +SGL_DEVICE float reciprocal_nr(float x) { + float out; + asm volatile( + "{\n\t" + ".reg .f32 e, e2;\n\t" + "rcp.approx.f32 %0, %1;\n\t" + "fma.rn.f32 e, %1, %0, 0fBF800000;\n\t" + "sub.ftz.f32 e2, 0f80000000, e;\n\t" + "fma.rn.f32 %0, %0, e2, %0;\n\t" + "}" + : "=&f"(out) + : "f"(x)); + return out; +} + +SGL_DEVICE float div_rn(float numerator, float denominator) { + float out; + asm volatile("div.rn.f32 %0, %1, %2;" : "=f"(out) : "f"(numerator), "f"(denominator)); + return out; +} + +SGL_DEVICE WelfordState welford_push(WelfordState state, float value) { + const float delta = __fsub_rn(value, state.mean); + const float count = __fadd_rn(state.count, 1.0f); + const float mean = __fmaf_rn(delta, reciprocal_nr(count), state.mean); + const float centered = __fsub_rn(value, mean); + return WelfordState{mean, __fmaf_rn(delta, centered, state.m2), count}; +} + +SGL_DEVICE WelfordState welford_combine(WelfordState lower, WelfordState upper) { + const float count = __fadd_rn(upper.count, lower.count); + if (count <= 0.0f) { + return WelfordState{0.0f, 0.0f, count}; + } + const float coefficient = reciprocal_nr(count); + const float delta = __fsub_rn(lower.mean, upper.mean); + const float lower_fraction = __fmul_rn(coefficient, lower.count); + const float delta_squared = __fmul_rn(delta, delta); + const float upper_fraction = __fmul_rn(upper.count, coefficient); + const float m2_sum = __fadd_rn(upper.m2, lower.m2); + const float lower_weighted_mean = __fmul_rn(lower_fraction, lower.mean); + const float mean = __fmaf_rn(upper.mean, upper_fraction, lower_weighted_mean); + const float upper_delta_squared = __fmul_rn(upper.count, delta_squared); + const float m2 = __fmaf_rn(lower_fraction, upper_delta_squared, m2_sum); + return WelfordState{mean, m2, count}; +} + +SGL_DEVICE WelfordState warp_welford(WelfordState state, int lane) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + const WelfordState upper{ + __shfl_down_sync(0xffffffffu, state.mean, offset), + __shfl_down_sync(0xffffffffu, state.m2, offset), + __shfl_down_sync(0xffffffffu, state.count, offset), + }; + if (lane < offset) { + state = welford_combine(state, upper); + } + } + return state; +} + +__global__ void kernel(const Params __grid_constant__ params) { + using namespace device; + using Vec = AlignedVector; + + const int row = blockIdx.x; + const int tid = threadIdx.x; + const int lane = tid & int(kWarpThreads - 1); + const int warp = tid >> 5; + const int row_offset = row * kHidden; + + WelfordState state{0.0f, 0.0f, 0.0f}; + +#pragma unroll + for (int iteration = 0; iteration < kIterations; ++iteration) { + const int offset = row_offset + iteration * 512 + tid * kVecElems; + const int gate_offset = iteration * 512 + tid * kVecElems; + Vec update; + Vec residual; + Vec gate; + Vec residual_out; + update.load(static_cast(params.update) + offset); + residual.load(static_cast(params.residual) + offset); + gate.load(static_cast(params.gate) + gate_offset); +#pragma unroll + for (int element = 0; element < kVecElems; ++element) { + // Match residual_gate_add: bf16-round the product before the add. + const bf16_t product = + static_cast(__fmul_rn(static_cast(update[element]), static_cast(gate[element]))); + const bf16_t updated = + static_cast(__fadd_rn(static_cast(residual[element]), static_cast(product))); + residual_out[element] = updated; + state = welford_push(state, static_cast(updated)); + } + residual_out.store(static_cast(params.residual_out) + offset); + } + + // Match aten LayerNorm's four-warp tree: (0,2), (1,3), then (0,1). + state = warp_welford(state, lane); + __shared__ WelfordState warp_states[kWarps]; + __shared__ float shared_mean; + __shared__ float shared_rstd; + if (lane == 0) { + warp_states[warp] = state; + } + __syncthreads(); + if (tid == 0) { + const WelfordState pair02 = welford_combine(warp_states[0], warp_states[2]); + const WelfordState pair13 = welford_combine(warp_states[1], warp_states[3]); + const WelfordState total = welford_combine(pair02, pair13); + shared_mean = total.mean; + shared_rstd = rsqrtf(__fadd_rn(div_rn(total.m2, float(kHidden)), params.eps)); + } + __syncthreads(); + +#pragma unroll + for (int iteration = 0; iteration < kIterations; ++iteration) { + const int offset = row_offset + iteration * 512 + tid * kVecElems; + const int modulation_offset = iteration * 512 + tid * kVecElems; + Vec scale; + Vec shift; + Vec residual_out; + Vec output; + scale.load(static_cast(params.scale) + modulation_offset); + shift.load(static_cast(params.shift) + modulation_offset); + residual_out.load(static_cast(params.residual_out) + offset); +#pragma unroll + for (int element = 0; element < kVecElems; ++element) { + // Match aten's bf16 LayerNorm output, then eager adaLN's bf16 rounding + // after 1+scale, multiply, and shift addition. + const bf16_t normalized = static_cast( + __fmul_rn(__fsub_rn(static_cast(residual_out[element]), shared_mean), shared_rstd)); + const bf16_t one_plus_scale = static_cast(__fadd_rn(1.0f, static_cast(scale[element]))); + const bf16_t scaled = + static_cast(__fmul_rn(static_cast(normalized), static_cast(one_plus_scale))); + output[element] = static_cast(__fadd_rn(static_cast(scaled), static_cast(shift[element]))); + } + output.store(static_cast(params.output) + offset); + } +} + +struct Kernel { + static void + run(tvm::ffi::TensorView output, + tvm::ffi::TensorView residual_out, + tvm::ffi::TensorView residual, + tvm::ffi::TensorView update, + tvm::ffi::TensorView gate, + tvm::ffi::TensorView scale, + tvm::ffi::TensorView shift, + double eps) { + using namespace host; + auto rows = SymbolicSize{"rows"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({rows, kHidden}) + .with_dtype() + .with_device(device) + .verify(output) + .verify(residual_out) + .verify(residual) + .verify(update); + TensorMatcher({kHidden}).with_dtype().with_device(device).verify(gate).verify(scale).verify(shift); + RuntimeCheck(rows.unwrap() > 0, "rows must be positive"); + RuntimeCheck(rows.unwrap() <= int64_t(UINT32_MAX), "rows out of range"); + + const auto params = Params{ + .output = output.data_ptr(), + .residual_out = residual_out.data_ptr(), + .update = update.data_ptr(), + .residual = residual.data_ptr(), + .gate = gate.data_ptr(), + .scale = scale.data_ptr(), + .shift = shift.data_ptr(), + .eps = static_cast(eps), + }; + LaunchKernel(static_cast(rows.unwrap()), kThreads, device.unwrap())(kernel, params); + } +}; + +} // namespace flux2_gated_resnorm + +} // namespace sglang diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 5284a16af..411ec5131 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/python/sglang/kernels/ops/diffusion/__init__.py @@ -369,6 +369,9 @@ for _op, _backend, _target, _caps, _description in _SPECS: _EXPORTS: dict[str, str] = { "load_extension_with_recovery": "ext.loader", # Normalization: RMSNorm / LayerNorm / GroupNorm and their fused epilogues + "can_defer_flux2_gated_residual": "norm.flux2_gated_resnorm_jit", + "can_use_flux2_gated_resnorm": "norm.flux2_gated_resnorm_jit", + "flux2_gated_resnorm_raw": "norm.flux2_gated_resnorm_jit", "FLYDSL_NORM_MIN_ALIGNED_DIM": "norm.fused_residual_norm_flydsl", "flydsl_fused_residual_norm_scale_shift": "norm.fused_residual_norm_flydsl", "flydsl_norm_scale_shift": "norm.fused_residual_norm_flydsl", diff --git a/python/sglang/kernels/ops/diffusion/norm/flux2_gated_resnorm_jit.py b/python/sglang/kernels/ops/diffusion/norm/flux2_gated_resnorm_jit.py new file mode 100644 index 000000000..428fb838b --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/norm/flux2_gated_resnorm_jit.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.kernels.jit.utils import cache_once, load_jit + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +_HIDDEN = 6144 +_ALIGNMENT = 32 + + +def _blackwell_or_newer(device: torch.device) -> bool: + return ( + torch.cuda.is_available() and torch.cuda.get_device_capability(device)[0] >= 10 + ) + + +def _aligned(tensor: torch.Tensor) -> bool: + return tensor.data_ptr() % _ALIGNMENT == 0 + + +def _row_bf16(tensor: torch.Tensor, device: torch.device) -> torch.Tensor | None: + if not ( + isinstance(tensor, torch.Tensor) + and tensor.dtype == torch.bfloat16 + and tensor.is_cuda + and tensor.device == device + and tensor.stride(-1) == 1 + ): + return None + if tensor.shape == (_HIDDEN,): + row = tensor + elif tensor.shape in ((1, _HIDDEN), (1, 1, _HIDDEN)): + row = tensor.reshape(_HIDDEN) + else: + return None + return row if _aligned(row) else None + + +def can_defer_flux2_gated_residual( + residual: torch.Tensor, + update: torch.Tensor, + gate: torch.Tensor, +) -> bool: + if not ( + not torch.compiler.is_compiling() + and residual.dtype == torch.bfloat16 + and residual.is_cuda + and residual.dim() == 3 + and residual.shape[0] == 1 + and residual.shape[-1] == _HIDDEN + and residual.numel() > 0 + and residual.is_contiguous() + and _aligned(residual) + and update.dtype == residual.dtype + and update.device == residual.device + and update.shape == residual.shape + and update.is_contiguous() + and _aligned(update) + and _blackwell_or_newer(residual.device) + ): + return False + return _row_bf16(gate, residual.device) is not None + + +def can_use_flux2_gated_resnorm( + residual: torch.Tensor, + update: torch.Tensor, + gate: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, +) -> bool: + return can_defer_flux2_gated_residual(residual, update, gate) and all( + _row_bf16(tensor, residual.device) is not None for tensor in (scale, shift) + ) + + +@cache_once +def _module() -> Module: + return load_jit( + "flux2_gated_resnorm", + cuda_files=["diffusion/flux2_gated_resnorm.cuh"], + cuda_wrappers=[("run", "flux2_gated_resnorm::Kernel::run")], + ) + + +def flux2_gated_resnorm_raw( + residual: torch.Tensor, + update: torch.Tensor, + gate: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + output = torch.empty_like(residual) + residual_out = torch.empty_like(residual) + _module().run( + output.view(-1, _HIDDEN), + residual_out.view(-1, _HIDDEN), + residual.view(-1, _HIDDEN), + update.view(-1, _HIDDEN), + _row_bf16(gate, residual.device), + _row_bf16(scale, residual.device), + _row_bf16(shift, residual.device), + float(eps), + ) + return output, residual_out + + +__all__ = [ + "can_defer_flux2_gated_residual", + "can_use_flux2_gated_resnorm", + "flux2_gated_resnorm_raw", +] 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 588194eb6..184fa8670 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -23,7 +23,10 @@ from diffusers.models.normalization import AdaLayerNormContinuous from sglang.kernels.ops.diffusion import ( BitExactFusionGate, + can_defer_flux2_gated_residual, + can_use_flux2_gated_resnorm, can_use_fused_layernorm_modulate, + flux2_gated_resnorm_raw, fused_layernorm_modulate_raw, fused_packed_silu_mul_bitexact, is_plain_layer_norm, @@ -87,6 +90,38 @@ _FLUX2_SWIGLU = BitExactFusionGate("FLUX.2 fused SwiGLU", per_signature=True) _FLUX2_SWIGLU_SIGS = _FLUX2_SWIGLU.verified_sigs assert _FLUX2_SWIGLU_SIGS is not None +PendingGatedResidual = Tuple[torch.Tensor, torch.Tensor, torch.Tensor] + + +def _materialize_gated_residual(pending: PendingGatedResidual) -> torch.Tensor: + residual, update, gate = pending + return residual_gate_add(residual, update, gate) + + +def _defer_gated_residual( + residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor +) -> torch.Tensor | PendingGatedResidual: + if can_defer_flux2_gated_residual(residual, update, gate): + return residual, update, gate + return residual_gate_add(residual, update, gate) + + +def _flux2_gated_resnorm( + norm: nn.Module, + residual: torch.Tensor, + update: torch.Tensor, + gate: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + if is_plain_layer_norm(norm, residual.shape[-1]) and can_use_flux2_gated_resnorm( + residual, update, gate, scale, shift + ): + return flux2_gated_resnorm_raw(residual, update, gate, scale, shift, norm.eps) + + residual = residual_gate_add(residual, update, gate) + return _flux2_norm_modulate(norm, residual, scale, shift), residual + def _flux2_norm_modulate( norm: nn.Module, @@ -748,7 +783,7 @@ class Flux2SingleTransformerBlock(nn.Module): def forward( self, - hidden_states: torch.Tensor, + hidden_states: torch.Tensor | PendingGatedResidual, encoder_hidden_states: Optional[torch.Tensor], temb_mod_params: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, @@ -756,18 +791,25 @@ class Flux2SingleTransformerBlock(nn.Module): split_hidden_states: bool = False, text_seq_len: Optional[int] = None, num_replicated_prefix: int = 0, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> torch.Tensor | PendingGatedResidual: # If encoder_hidden_states is None, hidden_states is assumed to have encoder_hidden_states already # concatenated if encoder_hidden_states is not None: + assert isinstance(hidden_states, torch.Tensor) text_seq_len = encoder_hidden_states.shape[1] hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) mod_shift, mod_scale, mod_gate = temb_mod_params - norm_hidden_states = _flux2_norm_modulate( - self.norm, hidden_states, mod_scale, mod_shift - ) + if isinstance(hidden_states, tuple): + residual, update, gate = hidden_states + norm_hidden_states, hidden_states = _flux2_gated_resnorm( + self.norm, residual, update, gate, mod_scale, mod_shift + ) + else: + 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( @@ -777,11 +819,16 @@ class Flux2SingleTransformerBlock(nn.Module): **joint_attention_kwargs, ) - hidden_states = residual_gate_add(hidden_states, attn_output, mod_gate) - if hidden_states.dtype == torch.float16: + hidden_states = _defer_gated_residual(hidden_states, attn_output, mod_gate) + if ( + isinstance(hidden_states, torch.Tensor) + and hidden_states.dtype == torch.float16 + ): hidden_states = hidden_states.clip(-65504, 65504) if split_hidden_states: + if isinstance(hidden_states, tuple): + hidden_states = _materialize_gated_residual(hidden_states) encoder_hidden_states, hidden_states = ( hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:], @@ -847,8 +894,8 @@ class Flux2TransformerBlock(nn.Module): def forward( self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, + hidden_states: torch.Tensor | PendingGatedResidual, + encoder_hidden_states: torch.Tensor | PendingGatedResidual, temb_mod_params_img: Tuple[ Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ... ], @@ -858,7 +905,9 @@ class Flux2TransformerBlock(nn.Module): freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, num_replicated_prefix: int = 0, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> Tuple[ + torch.Tensor | PendingGatedResidual, torch.Tensor | PendingGatedResidual + ]: joint_attention_kwargs = joint_attention_kwargs or {} # Modulation parameters shape: [1, 1, self.dim] @@ -880,17 +929,34 @@ class Flux2TransformerBlock(nn.Module): ) = temb_mod_params_txt # Img stream - norm_hidden_states = _flux2_norm_modulate( - self.norm1, hidden_states, scale_msa, shift_msa - ) + if isinstance(hidden_states, tuple): + residual, update, gate = hidden_states + norm_hidden_states, hidden_states = _flux2_gated_resnorm( + self.norm1, residual, update, gate, scale_msa, shift_msa + ) + else: + norm_hidden_states = _flux2_norm_modulate( + self.norm1, hidden_states, scale_msa, shift_msa + ) # Conditioning txt stream - norm_encoder_hidden_states = _flux2_norm_modulate( - self.norm1_context, - encoder_hidden_states, - c_scale_msa, - c_shift_msa, - ) + if isinstance(encoder_hidden_states, tuple): + residual, update, gate = encoder_hidden_states + norm_encoder_hidden_states, encoder_hidden_states = _flux2_gated_resnorm( + self.norm1_context, + residual, + update, + gate, + c_scale_msa, + c_shift_msa, + ) + else: + 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( @@ -904,32 +970,36 @@ class Flux2TransformerBlock(nn.Module): attn_output, context_attn_output = attention_outputs # Process attention outputs for the image stream (`hidden_states`). - hidden_states = residual_gate_add(hidden_states, attn_output, gate_msa) - - norm_hidden_states = _flux2_norm_modulate( - self.norm2, hidden_states, scale_mlp, shift_mlp + norm_hidden_states, hidden_states = _flux2_gated_resnorm( + self.norm2, + hidden_states, + attn_output, + gate_msa, + scale_mlp, + shift_mlp, ) ff_output = self.ff(norm_hidden_states) - hidden_states = residual_gate_add(hidden_states, ff_output, gate_mlp) + hidden_states = _defer_gated_residual(hidden_states, ff_output, gate_mlp) # Process attention outputs for the text stream (`encoder_hidden_states`). - encoder_hidden_states = residual_gate_add( - encoder_hidden_states, context_attn_output, c_gate_msa - ) - - norm_encoder_hidden_states = _flux2_norm_modulate( + norm_encoder_hidden_states, encoder_hidden_states = _flux2_gated_resnorm( self.norm2_context, encoder_hidden_states, + context_attn_output, + c_gate_msa, c_scale_mlp, c_shift_mlp, ) context_ff_output = self.ff_context(norm_encoder_hidden_states) - encoder_hidden_states = residual_gate_add( + encoder_hidden_states = _defer_gated_residual( encoder_hidden_states, context_ff_output, c_gate_mlp ) - if encoder_hidden_states.dtype == torch.float16: + if ( + isinstance(encoder_hidden_states, torch.Tensor) + and encoder_hidden_states.dtype == torch.float16 + ): encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) return encoder_hidden_states, hidden_states @@ -1284,6 +1354,10 @@ class Flux2Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): joint_attention_kwargs=joint_attention_kwargs, num_replicated_prefix=num_replicated_prefix, ) + if isinstance(encoder_hidden_states, tuple): + encoder_hidden_states = _materialize_gated_residual(encoder_hidden_states) + if isinstance(hidden_states, tuple): + hidden_states = _materialize_gated_residual(hidden_states) # Concatenate text and image streams for single-block inference; # join_seqs relocates any SP text tail-pad behind the image once for # the whole trunk (see sp_shard.join_seqs for why). @@ -1301,6 +1375,8 @@ class Flux2Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): text_seq_len=txt_real, num_replicated_prefix=num_replicated_prefix, ) + if isinstance(hidden_states, tuple): + hidden_states = _materialize_gated_residual(hidden_states) # Remove text (and any tail pad) from the concatenated stream img_end = hidden_states.shape[1] - sp_txt_pad hidden_states = hidden_states[:, txt_real:img_end, ...] diff --git a/test/registered/jit/test_flux2_gated_resnorm.py b/test/registered/jit/test_flux2_gated_resnorm.py new file mode 100644 index 000000000..6444eaff9 --- /dev/null +++ b/test/registered/jit/test_flux2_gated_resnorm.py @@ -0,0 +1,68 @@ +"""Blackwell coverage for the FLUX.2 gated residual normalization fast path.""" + +import unittest +from unittest.mock import patch + +import torch +import torch.nn as nn + +import sglang.multimodal_gen.runtime.models.dits.flux_2 as flux2 +from sglang.kernels.ops.diffusion import residual_gate_add +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + + +def _is_blackwell() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10 + + +@unittest.skipUnless(_is_blackwell(), "FLUX.2 gated residual fusion requires SM100+") +class TestFlux2GatedResnorm(CustomTestCase): + def test_gated_residual_norm_modulate_is_bit_exact(self): + torch.manual_seed(0) + hidden = 6144 + shape = (1, 64, hidden) + residual = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + update = torch.randn_like(residual) + params = torch.randn(1, 1, 3 * hidden, device="cuda").bfloat16() + gate, scale, shift = params.chunk(3, dim=-1) + norm = nn.LayerNorm(hidden, elementwise_affine=False, eps=1e-6, device="cuda") + + expected_residual = residual_gate_add(residual, update, gate) + expected = norm(expected_residual) * (1 + scale) + shift + actual, actual_residual = flux2._flux2_gated_resnorm( + norm, residual, update, gate, scale, shift + ) + + self.assertIsInstance( + flux2._defer_gated_residual(residual, update, gate), tuple + ) + self.assertTrue(torch.equal(actual_residual, expected_residual)) + self.assertTrue(torch.equal(actual, expected)) + + def test_gated_residual_defer_rejects_unsupported_inputs(self): + torch.manual_seed(0) + for batch, dtype in ((1, torch.float16), (2, torch.bfloat16)): + residual = torch.randn(batch, 17, 6144, device="cuda", dtype=dtype) + update = torch.randn_like(residual) + gate = torch.randn(batch, 1, 6144, device="cuda", dtype=dtype) + expected = residual_gate_add(residual, update, gate) + actual = flux2._defer_gated_residual(residual, update, gate) + + self.assertIsInstance(actual, torch.Tensor) + self.assertTrue(torch.equal(actual, expected)) + + residual = torch.randn(1, 17, 6144, device="cuda", dtype=torch.bfloat16) + update = torch.randn_like(residual) + gate = torch.randn(1, 1, 6144, device="cuda", dtype=torch.bfloat16) + with patch("torch.compiler.is_compiling", return_value=True): + actual = flux2._defer_gated_residual(residual, update, gate) + + self.assertIsInstance(actual, torch.Tensor) + self.assertTrue(torch.equal(actual, residual_gate_add(residual, update, gate))) + + +if __name__ == "__main__": + unittest.main()