[diffusion] Cache fp32 layernorm params (#25847)

Co-authored-by: BBuf <bbuf@example.com>

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -387,14 +387,42 @@ class LayerNorm(CustomOp):
 # NOTE(will): Needed to match behavior of diffusers and wan2.1 even while using
 # FSDP's MixedPrecisionPolicy
 class FP32LayerNorm(nn.LayerNorm):
+    def _cached_fp32_param(
+        self, attr: str, param: torch.Tensor | None, device: torch.device
+    ) -> torch.Tensor | None:
+        if param is None:
+            return None
+
+        # Keep autograd semantics identical to the old path. The diffusion
+        # runtime enters here for inference, where grad is disabled.
+        if torch.is_grad_enabled():
+            return param.float().to(device=device)
+
+        key = (
+            param.data_ptr(),
+            param._version,
+            param.device,
+            device,
+            param.dtype,
+        )
+        cache = self.__dict__.get(attr)
+        if cache is not None and cache[0] == key:
+            return cache[1]
+
+        fp32_param = param.detach().to(device=device, dtype=torch.float32)
+        self.__dict__[attr] = (key, fp32_param)
+        return fp32_param
+
     def forward(self, inputs: torch.Tensor) -> torch.Tensor:
         origin_dtype = inputs.dtype
         device = inputs.device
+        weight = self._cached_fp32_param("_weight_fp32_cache", self.weight, device)
+        bias = self._cached_fp32_param("_bias_fp32_cache", self.bias, device)
         return F.layer_norm(
             inputs.float(),
             self.normalized_shape,
-            self.weight.float().to(device=device) if self.weight is not None else None,
-            self.bias.float().to(device=device) if self.bias is not None else None,
+            weight,
+            bias,
             self.eps,
         ).to(origin_dtype)
 

python/sglang/multimodal_gen/test/unit/test_fp32_layernorm.py

@@ -0,0 +1,71 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.runtime.layers.layernorm import FP32LayerNorm
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_fp32_layernorm_cache_matches_reference():
+    norm = FP32LayerNorm(16, eps=1e-5).cuda().to(torch.bfloat16)
+    inputs = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16)
+
+    with torch.no_grad():
+        actual = norm(inputs)
+        expected = F.layer_norm(
+            inputs.float(),
+            norm.normalized_shape,
+            norm.weight.float().to(device=inputs.device),
+            norm.bias.float().to(device=inputs.device),
+            norm.eps,
+        ).to(inputs.dtype)
+
+    torch.testing.assert_close(actual, expected)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_fp32_layernorm_cache_reuses_converted_params():
+    norm = FP32LayerNorm(16, eps=1e-5).cuda().to(torch.bfloat16)
+    inputs = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16)
+
+    with torch.no_grad():
+        norm(inputs)
+        weight_cache = norm.__dict__["_weight_fp32_cache"]
+        bias_cache = norm.__dict__["_bias_fp32_cache"]
+
+        norm(inputs)
+
+    assert norm.__dict__["_weight_fp32_cache"][1] is weight_cache[1]
+    assert norm.__dict__["_bias_fp32_cache"][1] is bias_cache[1]
+    assert "_weight_fp32_cache" not in norm.state_dict()
+    assert "_bias_fp32_cache" not in norm.state_dict()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_fp32_layernorm_cache_invalidates_on_param_update():
+    norm = FP32LayerNorm(16, eps=1e-5).cuda().to(torch.bfloat16)
+    inputs = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16)
+
+    with torch.no_grad():
+        norm(inputs)
+        first_key, first_weight = norm.__dict__["_weight_fp32_cache"]
+
+        norm.weight.add_(1.0)
+        norm(inputs)
+        second_key, second_weight = norm.__dict__["_weight_fp32_cache"]
+
+    assert second_key != first_key
+    assert second_weight is not first_weight
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_fp32_layernorm_grad_mode_preserves_autograd_path():
+    norm = FP32LayerNorm(16, eps=1e-5).cuda().to(torch.bfloat16)
+    inputs = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+
+    output = norm(inputs).float().sum()
+    output.backward()
+
+    assert inputs.grad is not None
+    assert "_weight_fp32_cache" not in norm.__dict__
+    assert "_bias_fp32_cache" not in norm.__dict__