pytorch - 💡(How to fix) Fix [sm_120] Non-deterministic segfault with sustained complex linear-algebra workload + allocator churn on Blackwell [3 comments, 2 participants]

"""
Minimal repro: Blackwell (sm_120) segfault with sustained complex linalg + allocator churn.
Crashes after ~2-11 samples on sm_120. Runs all 500 on sm_86.
"""
import torch, gc

class WorkloadState:
    """Pure allocator churn — allocates GPU tensors on init, freed on del.
    Does NOT participate in computation."""
    def __init__(self, F, T, M, N, K, device):
        self.W = torch.rand(N, F, K, device=device) + 1e-10
        self.H = torch.rand(N, K, T, device=device) + 1e-10
        self.G = torch.rand(N, M, device=device) + 1e-10
        self.X = torch.randn(F, T, M, dtype=torch.complex64, device=device)
        self.PSD = torch.einsum('nfk, nkt -> ftn', self.W, self.H)
        self.Y = torch.einsum('ftn, nm -> ftm', self.PSD, self.G) + 1e-10
        self.X_pow = (self.X * self.X.conj()).real + 1e-10
        self.Q = torch.eye(M, dtype=torch.complex64, device=device).unsqueeze(0).expand(F, -1, -1).contiguous()
        self.XX = torch.einsum('ftm, ftn -> ftmn', self.X, self.X.conj())
        self.eye = torch.eye(M, dtype=torch.complex64, device=device).unsqueeze(0).expand(F, -1, -1)

def iteration(F, T, M, N, K, device):
    W = torch.rand(N, F, K, device=device) + 1e-10
    H = torch.rand(N, K, T, device=device) + 1e-10
    G = torch.rand(N, M, device=device) + 1e-10
    X = torch.randn(F, T, M, dtype=torch.complex64, device=device)
    PSD = torch.einsum('nfk, nkt -> ftn', W, H)
    Y = torch.einsum('ftn, nm -> ftm', PSD, G) + 1e-10
    X_pow = (X * X.conj()).real + 1e-10

    for n in range(N):
        num = torch.einsum('kt, m, ftm, ftm -> fk', H[n], G[n], X_pow, Y.pow(-2))
        den = torch.einsum('kt, m, ftm -> fk', H[n], G[n], Y.pow(-1))
        W[n] = W[n] * torch.sqrt(num / (den + 1e-10)) + 1e-10
        PSD[:, :, n] = torch.einsum('fk, kt -> ft', W[n], H[n])
        Y = torch.einsum('ftn, nm -> ftm', PSD, G) + 1e-10
        num = torch.einsum('fk, m, ftm, ftm -> kt', W[n], G[n], X_pow, Y.pow(-2))
        den = torch.einsum('fk, m, ftm -> kt', W[n], G[n], Y.pow(-1))
        H[n] = H[n] * torch.sqrt(num / (den + 1e-10)) + 1e-10
        PSD[:, :, n] = torch.einsum('fk, kt -> ft', W[n], H[n])
        Y = torch.einsum('ftn, nm -> ftm', PSD, G) + 1e-10

    Q = torch.eye(M, dtype=torch.complex64, device=device).unsqueeze(0).expand(F, -1, -1).contiguous()
    XX = torch.einsum('ftm, ftn -> ftmn', X, X.conj())
    eye = torch.eye(M, dtype=torch.complex64, device=device).unsqueeze(0).expand(F, -1, -1)
    for m in range(M):
        V = torch.einsum('ftmn, ft -> fmn', XX, (1.0 / Y[..., m]).to(torch.complex64)) / T
        tmp = torch.linalg.solve(Q @ V, eye)
        t_m = tmp[..., m]
        d = torch.sqrt(torch.einsum('fi, fij, fj -> f', t_m.conj(), V, t_m))
        Q[:, m] = (t_m / d[:, None]).conj()

device = 'cuda'
print(f'PyTorch {torch.__version__}, CUDA {torch.version.cuda}')
print(f'GPU: {torch.cuda.get_device_name()} (sm_{torch.cuda.get_device_capability()[0]}{torch.cuda.get_device_capability()[1]})')
F, T, M, N, K = 256, 500, 6, 3, 64
for i in range(500):
    state = WorkloadState(F, T, M, N, K, device)
    for _ in range(200):
        iteration(F, T, M, N, K, device)
    del state; gc.collect(); torch.cuda.empty_cache()
    print(f'Sample {i+1}: OK', flush=True)
print('ALL PASSED')

---

PyTorch 2.7.1+cu128, CUDA 12.8
GPU: NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition (sm_120)
Sample 1: OK
Sample 2: OK

---

PyTorch 2.7.1+cu128, CUDA 12.8
GPU: NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition (sm_120)
Sample 1: OK
...
Sample 11: OK

---

PyTorch 2.5.1, CUDA 12.4
GPU: NVIDIA GeForce RTX 3090 (sm_86)
Sample 1: OK
...
Sample 500: OK
ALL 500 PASSED — no segfault