[Core] Optimize diffusion LoRA expand accumulation (#6268)

Signed-off-by: 0z5a <0z5a@users.noreply.github.com>
Signed-off-by: 0z5a <dinadegttyarenko6@gmail.com>
Co-authored-by: 0z5a <0z5a@users.noreply.github.com>

tests/diffusion/lora/test_base_linear.py

@@ -29,7 +29,7 @@ class _DummyQuantMethod:
         return y
 
 
-def test_diffusion_base_linear_apply_multi_slice():
+def test_diffusion_base_linear_apply_multi_slice(monkeypatch):
     # Build a fake diffusion LoRA layer with 2 slices and rank=2.
     layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
     layer.tp_size = 1
@@ -72,8 +72,18 @@ def test_diffusion_base_linear_apply_multi_slice():
     layer.lora_b_stacked = (b0, b1)
     layer.output_slices = out_slices
 
+    addmm_calls = []
+    original_addmm = torch.addmm
+
+    def _record_addmm(input, mat1, mat2, *, out=None):
+        addmm_calls.append((input, out))
+        return original_addmm(input, mat1, mat2, out=out)
+
+    monkeypatch.setattr(torch, "addmm", _record_addmm)
+
     x = torch.tensor([[1.0, 2.0, 3.0]])
-    out = layer.apply(x)
+    with torch.inference_mode():
+        out = layer.apply(x)
 
     # Base output is identity: [1,2,3]
     base_out = x @ base_weight.t()
@@ -87,6 +97,66 @@ def test_diffusion_base_linear_apply_multi_slice():
     delta1 = torch.tensor([[6.0]])
     expected = torch.cat([base_out[:, :2] + delta0, base_out[:, 2:3] + delta1], dim=-1)
     assert torch.allclose(out, expected)
+    assert len(addmm_calls) == 2
+    assert all(input is out for input, out in addmm_calls)
+
+
+def test_diffusion_base_linear_apply_preserves_autograd(monkeypatch):
+    layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
+    layer.tp_size = 1
+    layer.lora_config = _DummyLoRAConfig()
+
+    layer.base_layer = type("Base", (), {})()
+    layer.base_layer.quant_method = _DummyQuantMethod(torch.eye(2))
+
+    a = torch.tensor([[[[1.0, 2.0]]]], requires_grad=True)
+    b = torch.tensor([[[[3.0], [4.0]]]], requires_grad=True)
+    layer.lora_a_stacked = (a,)
+    layer.lora_b_stacked = (b,)
+    layer.output_slices = (2,)
+    layer._diffusion_lora_active_slices = (True,)
+
+    def _unexpected_addmm(*args, **kwargs):
+        raise AssertionError("the out= fast path must not run with autograd enabled")
+
+    monkeypatch.setattr(torch, "addmm", _unexpected_addmm)
+
+    x = torch.tensor([[2.0, 1.0]], requires_grad=True)
+    out = layer.apply(x)
+    assert torch.allclose(out, torch.tensor([[14.0, 17.0]]))
+
+    out.sum().backward()
+    assert torch.allclose(x.grad, torch.tensor([[8.0, 15.0]]))
+    assert torch.allclose(a.grad, torch.tensor([[[[14.0, 7.0]]]]))
+    assert torch.allclose(b.grad, torch.tensor([[[[4.0], [4.0]]]]))
+
+
+def test_diffusion_base_linear_apply_falls_back_for_mixed_output_dtype(monkeypatch):
+    layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
+    layer.tp_size = 1
+    layer.lora_config = _DummyLoRAConfig()
+
+    class _Float64OutputQuantMethod(_DummyQuantMethod):
+        def apply(self, base_layer, x: torch.Tensor, bias: torch.Tensor | None):
+            return super().apply(base_layer, x, bias).to(torch.float64)
+
+    layer.base_layer = type("Base", (), {})()
+    layer.base_layer.quant_method = _Float64OutputQuantMethod(torch.eye(2))
+    layer.lora_a_stacked = (torch.ones((1, 1, 1, 2)),)
+    layer.lora_b_stacked = (torch.tensor([[[[1.0], [2.0]]]]),)
+    layer.output_slices = (2,)
+    layer._diffusion_lora_active_slices = (True,)
+
+    def _unexpected_addmm(*args, **kwargs):
+        raise AssertionError("mixed output and LoRA dtypes must use the functional fallback")
+
+    monkeypatch.setattr(torch, "addmm", _unexpected_addmm)
+
+    with torch.inference_mode():
+        out = layer.apply(torch.tensor([[1.0, 2.0]]))
+
+    assert out.dtype == torch.float64
+    assert torch.allclose(out, torch.tensor([[4.0, 8.0]], dtype=torch.float64))
 
 
 def test_diffusion_lora_buffer_device_persists_across_reallocation(monkeypatch):