[bugfix]: restore TF32 flags in transformer test teardown on torch 2.12

torch 2.12 computes torch.backends.cuda.matmul.allow_tf32 on cuBLASModule
instead of storing it, so monkeypatch.setattr cannot undo it and pytest's
teardown raised AttributeError for every Qwen-Image-2.1 attention and Hunyuan
batch test. A shared tf32_disabled fixture saves and restores the legacy flags.
This commit is contained in:
Aryan Kumar
2026-11-01 22:34:28 -08:00
parent 166b9b8a07
commit ec94425dbd
3 changed files with 22 additions and 6 deletions
+18
View File
@@ -45,3 +45,21 @@ def env_overrides():
"""
with contextlib.ExitStack() as stack:
yield stack
@pytest.fixture
def tf32_disabled():
"""Disable TF32 for CUDA matmuls and cuDNN during the test, then restore the previous flags.
On torch 2.12 ``torch.backends.cuda.matmul.allow_tf32`` is computed by ``cuBLASModule`` rather
than stored on it, so ``monkeypatch.setattr`` cannot undo it. Restore through the same legacy
flags: mixing them with the per-backend ``fp32_precision`` API makes later ``allow_tf32`` reads raise.
"""
matmul, cudnn = torch.backends.cuda.matmul, torch.backends.cudnn
previous = (matmul.allow_tf32, cudnn.allow_tf32)
matmul.allow_tf32 = False
cudnn.allow_tf32 = False
try:
yield
finally:
matmul.allow_tf32, cudnn.allow_tf32 = previous
@@ -25,7 +25,7 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires
@pytest.fixture
def hunyuan_setup(monkeypatch, request):
def hunyuan_setup(request):
# Environment writes go through the registry helpers (docs/contributing/env_vars.md).
with contextlib.ExitStack() as stack:
# Respect the port reserved by the CI runner when sharing a GPU host.
@@ -37,8 +37,7 @@ def hunyuan_setup(monkeypatch, request):
if "MASTER_ADDR" not in os.environ:
stack.enter_context(envs.override_external("MASTER_ADDR", "127.0.0.1"))
stack.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA"))
monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False)
monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False)
request.getfixturevalue("tf32_disabled")
request.getfixturevalue("distributed_setup")
yield
@@ -24,12 +24,11 @@ from fastvideo.layers.linear import ReplicatedLinear
from fastvideo.platforms import AttentionBackendEnum
@pytest.fixture
def qwen_setup(env_overrides, monkeypatch, request):
def qwen_setup(env_overrides, request):
if not torch.cuda.is_available():
pytest.skip("Requires one CUDA GPU")
env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA"))
monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False)
monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False)
request.getfixturevalue("tf32_disabled")
request.getfixturevalue("distributed_setup")
with torch.no_grad(), sdpa_kernel(SDPBackend.MATH):
yield torch.device("cuda")