Synthyra/FastESMFold

1import torch
2import torch._inductor.config as inductor_config
3import torch._dynamo as dynamo
4
5# Enable TensorFloat32 tensor cores for float32 matmul (Ampere+ GPUs)
6# Provides significant speedup with minimal precision loss
7torch.set_float32_matmul_precision('high')
8
9# Enable TF32 for matrix multiplications and cuDNN operations
10torch.backends.cuda.matmul.allow_tf32 = True
11torch.backends.cudnn.allow_tf32 = True
12
13# Enable cuDNN autotuner - finds fastest algorithms for your hardware
14# Best when input sizes are consistent; may slow down first iterations
15torch.backends.cudnn.benchmark = True
16
17# Deterministic operations off for speed (set True if reproducibility needed)
18torch.backends.cudnn.deterministic = False
19inductor_config.max_autotune_gemm_backends = "ATEN,CUTLASS,FBGEMM"    
20
21dynamo.config.capture_scalar_outputs = True
22torch._dynamo.config.recompile_limit = 16
23