From 7376b4db9033ad330786eb8345dc9c54d8da3cea Mon Sep 17 00:00:00 2001 From: Daniel Guo Date: Sun, 31 May 2026 16:43:22 -0700 Subject: [PATCH] Run weighted_rigid_align's Kabsch SVD in float32 (fix bf16 crash) weighted_rigid_align runs torch.linalg.svd directly on the (possibly bf16) covariance matrix. PyTorch has no SVD kernel for bfloat16/half (CUDA gesvdjBatched or CPU), so any caller running in bf16 -- e.g. rfd3 inference with motif realignment (allow_realignment / center_option=motif) -- crashes with: RuntimeError: "svd_cuda_gesvdjBatched" not implemented for 'BFloat16' SVD is numerically sensitive anyway, so compute the rotation in float32 and cast it back to the input dtype. Verified: a bf16 input goes from this crash to a correct alignment (RMSD ~0 vs a known rigid transform), output dtype preserved. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/foundry/utils/alignment.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/foundry/utils/alignment.py b/src/foundry/utils/alignment.py index c5d066dd..285090a1 100644 --- a/src/foundry/utils/alignment.py +++ b/src/foundry/utils/alignment.py @@ -57,7 +57,9 @@ def weighted_rigid_align( # Computation of the covariance matrix C = torch.einsum("bji,bjk->bik", w_resolved[..., None] * X_gt_resolved, X_resolved) - U, S, V = torch.linalg.svd(C) + # SVD has no bf16/half CUDA kernel and is numerically sensitive; compute the + # rotation in float32 and cast it back to the input dtype below. + U, S, V = torch.linalg.svd(C.float()) R = U @ V B, _, _ = X_L.shape @@ -71,7 +73,7 @@ def weighted_rigid_align( det = torch.linalg.det(R) F[..., -1, -1] = torch.sign(det) - R = U @ F @ V + R = (U @ F @ V).to(X_gt_L.dtype) X_gt_L = X_gt_L - u_X_gt.unsqueeze(-2) X_align_L = X_gt_L @ R + u_X.unsqueeze(-2)