[PyT] Fix PyTorch BF16 reduction numerics tests

tests/pytorch/test_numerics.py

@@ -4,6 +4,7 @@
 
 import math
 import os
+from contextlib import contextmanager
 from typing import Dict, List, Tuple, Optional
 import pytest
 
@@ -254,6 +255,27 @@ def reset_global_fp8_state():
     FP8GlobalStateManager.reset()
 
 
+@contextmanager
+def _disable_bf16_reduced_precision_reduction():
+    """TE disables cuBLASLt heuristics that store partial GEMM results in BF16.
+    This is to ensure precision is kept. This affects older archs like L40.
+    PyTorch does not do this by default, so we need to adjust PyTorch's
+    settings here to match TE's increased precision here.
+    """
+    original_value = torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction
+    torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
+    try:
+        yield
+    finally:
+        torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = original_value
+
+
+@pytest.fixture(autouse=True)
+def disable_bf16_reduced_precision_reduction():
+    with _disable_bf16_reduced_precision_reduction():
+        yield
+
+
 class TorchScaledMaskedSoftmax(nn.Module):
     def __init__(self) -> None:
         super().__init__()