ops: keep FP8 LoRA biases in compute dtype

The dynamic VBAR cast path applied want_requant to both weights and biases.
For an FP8 checkpoint with a low-VRAM LoRA patch, this returned an FP8 bias
to the optimized linear operation even though the activation dtype was FP16.

Restrict the optional requantized return value to weights while preserving
resident-storage updates for both parameters. Make the legacy FP8 linear path
release its VBAR/offload resources in a finally block so backend failures do
not leave pages pinned before the normal fallback runs.

Add regressions for the weight/bias dtype contract and exception cleanup.

Fixes #14952

Signed-off-by: liminfei <liminfei@amd.com>

comfy/ops.py

@@ -238,6 +238,7 @@ def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, w
     def post_cast(s, param_key, x, dtype, resident, update_weight):
         lowvram_fn = getattr(s, param_key + "_lowvram_function", None)
         fns = getattr(s, param_key + "_function", [])
+        requant = want_requant and param_key == "weight"
 
         if x is None:
             return None
@@ -255,13 +256,13 @@ def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, w
         if not resident and lowvram_fn is not None:
             x = to_dequant(x, dtype if compute_dtype is None else compute_dtype)
             x = lowvram_fn(x)
-            if (want_requant and len(fns) == 0 or update_weight):
+            if (requant and len(fns) == 0 or update_weight):
                 seed = comfy.utils.string_to_seed(s.seed_key)
                 if isinstance(orig, QuantizedTensor):
                     y = orig.requantize_from_float(x, scale="recalculate", stochastic_rounding=seed)
                 else:
                     y = comfy.float.stochastic_rounding(x, orig.dtype, seed=seed)
-            if want_requant and len(fns) == 0:
+            if requant and len(fns) == 0:
                 x = y
             if update_weight:
                 orig.copy_(y)
@@ -831,21 +832,22 @@ def fp8_linear(self, input):
         return None
     lora_compute_dtype=comfy.model_management.lora_compute_dtype(input.device)
     w, bias, offload_stream = cast_bias_weight(self, input, dtype=dtype, bias_dtype=input_dtype, offloadable=True, compute_dtype=lora_compute_dtype, want_requant=True)
-    scale_weight = torch.ones((), device=input.device, dtype=torch.float32)
+    try:
+        scale_weight = torch.ones((), device=input.device, dtype=torch.float32)
 
-    scale_input = torch.ones((), device=input.device, dtype=torch.float32)
-    input = torch.clamp(input, min=-448, max=448, out=input)
-    input_fp8 = input.to(dtype).contiguous()
-    layout_params_input = TensorCoreFP8Layout.Params(scale=scale_input, orig_dtype=input_dtype, orig_shape=tuple(input_fp8.shape))
-    quantized_input = QuantizedTensor(input_fp8, "TensorCoreFP8Layout", layout_params_input)
+        scale_input = torch.ones((), device=input.device, dtype=torch.float32)
+        input = torch.clamp(input, min=-448, max=448, out=input)
+        input_fp8 = input.to(dtype).contiguous()
+        layout_params_input = TensorCoreFP8Layout.Params(scale=scale_input, orig_dtype=input_dtype, orig_shape=tuple(input_fp8.shape))
+        quantized_input = QuantizedTensor(input_fp8, "TensorCoreFP8Layout", layout_params_input)
 
-    # Wrap weight in QuantizedTensor - this enables unified dispatch
-    # Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py!
-    layout_params_weight = TensorCoreFP8Layout.Params(scale=scale_weight, orig_dtype=input_dtype, orig_shape=tuple(w.shape))
-    quantized_weight = QuantizedTensor(w, "TensorCoreFP8Layout", layout_params_weight)
-    o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)
-
-    uncast_bias_weight(self, w, bias, offload_stream)
+        # Wrap weight in QuantizedTensor - this enables unified dispatch
+        # Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py!
+        layout_params_weight = TensorCoreFP8Layout.Params(scale=scale_weight, orig_dtype=input_dtype, orig_shape=tuple(w.shape))
+        quantized_weight = QuantizedTensor(w, "TensorCoreFP8Layout", layout_params_weight)
+        o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)
+    finally:
+        uncast_bias_weight(self, w, bias, offload_stream)
     if tensor_3d:
         o = o.reshape((input_shape[0], input_shape[1], w.shape[0]))
 

tests-unit/comfy_quant/test_fp8_ops.py

@@ -0,0 +1,74 @@
+import os
+import sys
+import unittest
+from types import SimpleNamespace
+from unittest import mock
+
+import torch
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
+
+from comfy.cli_args import args
+
+if not torch.cuda.is_available():
+    args.cpu = True
+
+from comfy import ops
+
+
+class TestFP8Ops(unittest.TestCase):
+    def test_vbar_requantizes_weight_but_not_bias(self):
+        weight = torch.zeros((16, 16), dtype=torch.float8_e4m3fn)
+        bias = torch.zeros(16, dtype=torch.float8_e4m3fn)
+        patch = lambda tensor: tensor
+        layer = SimpleNamespace(
+            weight=weight,
+            bias=bias,
+            weight_function=[],
+            bias_function=[],
+            weight_lowvram_function=patch,
+            bias_lowvram_function=patch,
+            seed_key="layer",
+            _prefetch={
+                "resident": False,
+                "xfer_dest": torch.empty(1, dtype=torch.uint8),
+                "needs_cast": False,
+                "cast_geometry": None,
+                "signature": (1,),
+            },
+        )
+
+        with (
+            mock.patch.object(ops.comfy.memory_management, "interpret_gathered_like", return_value=[weight, bias]),
+            mock.patch.object(ops.comfy.float, "stochastic_rounding", side_effect=lambda tensor, dtype, seed: tensor.to(dtype)) as rounding,
+        ):
+            cast_weight, cast_bias = ops.resolve_cast_module_with_vbar(
+                layer,
+                torch.float8_e4m3fn,
+                torch.device("cpu"),
+                torch.float16,
+                torch.float16,
+                True,
+            )
+
+        self.assertEqual(cast_weight.dtype, torch.float8_e4m3fn)
+        self.assertEqual(cast_bias.dtype, torch.float16)
+        self.assertEqual(rounding.call_count, 2)
+
+    def test_fp8_linear_unpins_vbar_when_linear_fails(self):
+        weight = torch.zeros((16, 16), dtype=torch.float8_e4m3fn)
+        bias = torch.zeros(16, dtype=torch.float8_e4m3fn)
+        vbar = object()
+        layer = SimpleNamespace(weight=weight, _v=vbar)
+        input_tensor = torch.zeros((2, 16), dtype=torch.float16)
+        offload_info = (None, torch.device("cpu"), None)
+
+        with (
+            mock.patch.object(ops, "cast_bias_weight", return_value=(weight, bias, offload_info)),
+            mock.patch.object(ops.comfy_aimdo.model_vbar, "vbar_unpin") as unpin,
+            mock.patch.object(torch.nn.functional, "linear", side_effect=RuntimeError("FP8 linear failed")),
+        ):
+            with self.assertRaisesRegex(RuntimeError, "FP8 linear failed"):
+                ops.fp8_linear(layer, input_tensor)
+
+        unpin.assert_called_once_with(vbar)