From 591cfb088127c4ff86bd189dd633423f2d6bbeda Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Thu, 6 Aug 2026 22:54:40 +0800 Subject: [PATCH] [diffusion] FLUX.2 bit-exact residual-gate fast path (H200 klein-4B 50-step denoise -1.2%) (#33823) Co-authored-by: Claude Fable 5 --- .../runtime/models/dits/flux_2.py | 54 ++++++++++++++++--- .../ops/diffusion/test_residual_gate_add.py | 5 +- 2 files changed, 51 insertions(+), 8 deletions(-) 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 300a65c56..e875b0598 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -20,6 +20,10 @@ from diffusers.models.attention import AttentionModuleMixin from diffusers.models.embeddings import TimestepEmbedding, Timesteps from diffusers.models.normalization import AdaLayerNormContinuous +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.flux import FluxConfig from sglang.multimodal_gen.runtime.distributed import ( divide, @@ -65,6 +69,40 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name +_FLUX2_RESIDUAL_GATE_CUDA_DISABLED = False + + +def _flux2_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. The + kernel's row-broadcast gate only covers ``[1, ..., 1, D]``; batched + ``[B>1, 1, D]`` gates fail ``can_use_residual_gate_add_cuda`` and take the + eager fallback below. + """ + global _FLUX2_RESIDUAL_GATE_CUDA_DISABLED + + if ( + not _FLUX2_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 FLUX.2 residual-gate CUDA fast path: {exc}") + _FLUX2_RESIDUAL_GATE_CUDA_DISABLED = True + + return residual + gate * update + def _get_qkv_projections( attn: "Flux2Attention", hidden_states, encoder_hidden_states=None @@ -656,7 +694,7 @@ class Flux2SingleTransformerBlock(nn.Module): **joint_attention_kwargs, ) - hidden_states = hidden_states + mod_gate * attn_output + hidden_states = _flux2_residual_gate_add(hidden_states, attn_output, mod_gate) if hidden_states.dtype == torch.float16: hidden_states = hidden_states.clip(-65504, 65504) @@ -774,18 +812,18 @@ class Flux2TransformerBlock(nn.Module): attn_output, context_attn_output = attention_outputs # Process attention outputs for the image stream (`hidden_states`). - attn_output = gate_msa * attn_output - hidden_states = hidden_states + attn_output + hidden_states = _flux2_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 ff_output = self.ff(norm_hidden_states) - hidden_states = hidden_states + gate_mlp * ff_output + hidden_states = _flux2_residual_gate_add(hidden_states, ff_output, gate_mlp) # Process attention outputs for the text stream (`encoder_hidden_states`). - context_attn_output = c_gate_msa * context_attn_output - encoder_hidden_states = encoder_hidden_states + context_attn_output + encoder_hidden_states = _flux2_residual_gate_add( + encoder_hidden_states, context_attn_output, c_gate_msa + ) norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) norm_encoder_hidden_states = ( @@ -793,7 +831,9 @@ class Flux2TransformerBlock(nn.Module): ) context_ff_output = self.ff_context(norm_encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output + encoder_hidden_states = _flux2_residual_gate_add( + encoder_hidden_states, context_ff_output, c_gate_mlp + ) if encoder_hidden_states.dtype == torch.float16: encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) diff --git a/test/registered/kernels/ops/diffusion/test_residual_gate_add.py b/test/registered/kernels/ops/diffusion/test_residual_gate_add.py index 81480a31e..b45cddb0a 100644 --- a/test/registered/kernels/ops/diffusion/test_residual_gate_add.py +++ b/test/registered/kernels/ops/diffusion/test_residual_gate_add.py @@ -18,10 +18,13 @@ CASES = [ ((1, 512, 4096), (1, 512, 4096)), ((1, 17, 65), (1, 1, 65)), ((1, 17, 65), (1, 17, 65)), - # FLUX.1 1024^2 shapes: dual-stream image/text and single-stream joint. + # FLUX.1 / FLUX.2-klein 1024^2 shapes (D=3072): dual-stream image/text + # and single-stream/joint concat; gates are [1, 1, D] modulation rows. ((1, 4096, 3072), (1, 1, 3072)), ((1, 512, 3072), (1, 1, 3072)), ((1, 4608, 3072), (1, 1, 3072)), + # FLUX.2-dev (D=6144) joint sequence. + ((1, 4608, 6144), (1, 1, 6144)), ]