per-device capability probe: pre-Ada GPUs auto-fall back to upcast path

libs/ltx_core/quantization/fp8_torch_mm.py

@@ -17,6 +17,22 @@ from ltx_core.quantization.fp8_cast import _replace_fwd_with_upcast
 
 _F8_MAX = 448.0  # float8_e4m3fn finite max
 
+_SM_OK: dict = {}  # device -> bool: native fp8 mm available (sm_89+)
+
+
+def _scaled_mm_supported(device) -> bool:
+    ok = _SM_OK.get(device)
+    if ok is None:
+        try:
+            ok = torch.cuda.get_device_capability(device) >= (8, 9)
+        except Exception:
+            ok = False
+        _SM_OK[device] = ok
+        if not ok:
+            print(f"[ltx_core] fp8 scaled-mm: {device} lacks sm_89+ fp8 cores; "
+                  "falling back to per-layer upcast on this device.", flush=True)
+    return ok
+
 
 def _replace_fwd_with_scaled_mm(layer: torch.nn.Linear) -> None:
     """Swap layer.forward for a runtime-branching fp8 path.
@@ -41,8 +57,9 @@ def _replace_fwd_with_scaled_mm(layer: torch.nn.Linear) -> None:
             return layer.original_forward(*args, **_kwargs)
         if layer.bias is not None and layer.bias.dtype == torch.float8_e4m3fn:
             layer.bias.data = layer.bias.data.to(torch.bfloat16)  # one-time restore
-        if (not (x.is_cuda and w.is_cuda)) or w.shape[0] % 16 or w.shape[1] % 16:
-            # CPU pass or non-16-aligned dims: per-layer upcast math.
+        if (not (x.is_cuda and w.is_cuda)) or w.shape[0] % 16 or w.shape[1] % 16 \
+                or not _scaled_mm_supported(x.device):
+            # CPU pass, non-16-aligned dims, or pre-Ada GPU: per-layer upcast math.
             w_up = w.to(x.dtype)
             b = layer.bias.to(x.dtype) if layer.bias is not None else None
             return torch.nn.functional.linear(x, w_up, b)