[Diffusion] Fuse FLUX.2 gated residual normalization on Blackwell (#37112)

Co-authored-by: Cursor <cursoragent@cursor.com>

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 <sgl_kernel/tensor.h>
+
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <cstdint>
+
+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<bf16_t, kVecElems>;
+
+  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<const bf16_t*>(params.update) + offset);
+    residual.load(static_cast<const bf16_t*>(params.residual) + offset);
+    gate.load(static_cast<const bf16_t*>(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<bf16_t>(__fmul_rn(static_cast<float>(update[element]), static_cast<float>(gate[element])));
+      const bf16_t updated =
+          static_cast<bf16_t>(__fadd_rn(static_cast<float>(residual[element]), static_cast<float>(product)));
+      residual_out[element] = updated;
+      state = welford_push(state, static_cast<float>(updated));
+    }
+    residual_out.store(static_cast<bf16_t*>(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<const bf16_t*>(params.scale) + modulation_offset);
+    shift.load(static_cast<const bf16_t*>(params.shift) + modulation_offset);
+    residual_out.load(static_cast<const bf16_t*>(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<bf16_t>(
+          __fmul_rn(__fsub_rn(static_cast<float>(residual_out[element]), shared_mean), shared_rstd));
+      const bf16_t one_plus_scale = static_cast<bf16_t>(__fadd_rn(1.0f, static_cast<float>(scale[element])));
+      const bf16_t scaled =
+          static_cast<bf16_t>(__fmul_rn(static_cast<float>(normalized), static_cast<float>(one_plus_scale)));
+      output[element] = static_cast<bf16_t>(__fadd_rn(static_cast<float>(scaled), static_cast<float>(shift[element])));
+    }
+    output.store(static_cast<bf16_t*>(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<kDLCUDA>();
+
+    TensorMatcher({rows, kHidden})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(output)
+        .verify(residual_out)
+        .verify(residual)
+        .verify(update);
+    TensorMatcher({kHidden}).with_dtype<bf16_t>().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<float>(eps),
+    };
+    LaunchKernel(static_cast<uint32_t>(rows.unwrap()), kThreads, device.unwrap())(kernel, params);
+  }
+};
+
+}  // namespace flux2_gated_resnorm
+
+}  // namespace sglang

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",

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",
+]

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, ...]