unbreak CI by fixing MX tests (#2208)

Summary:

#2201 broke CI:
1. some MX tests for fp4 are running on A10G instances, with skipping
   not being properly applied
   (https://hud.pytorch.org/pytorch/ao/commit/4bfd7c09ef4592eacbbf990aea6d6bda608865c1#42164784332-box)
2. some SQNR thresholds were to tight for fp4
   (https://hud.pytorch.org/pytorch/ao/commit/4bfd7c09ef4592eacbbf990aea6d6bda608865c1#42164784332-box)

This PR fixes both of these to get CI back to green (I hope). Note that
I can't repro 1 locally, so we'll have to land and see if it works.

Test Plan: CI

Reviewers:

Subscribers:

Tasks:

Tags:

Author: vkuzo

test/prototype/mx_formats/test_mx_linear.py

@@ -396,7 +396,6 @@ def test_inference_print_str():
 @pytest.mark.skipif(
     not TORCH_VERSION_AT_LEAST_2_8, reason="torch.compile requires PyTorch 2.8+"
 )
-@pytest.mark.skipif(not is_sm_at_least_100, reason="Reqs sm100")
 @pytest.mark.parametrize("elem_dtype", [torch.float8_e4m3fn, torch.float4_e2m1fn_x2])
 @pytest.mark.parametrize("bias", [True, False])
 @pytest.mark.parametrize("compile", [True, False])
@@ -405,9 +404,14 @@ def test_inference_subclass(elem_dtype, bias: bool, compile: bool):
     """
     Smoke test for inference compile
     """
+    # TODO(future): figure out why these CUDA capability conditions are not properly
+    # applied when inside `pytest.mark.skipif` for this test
     if elem_dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
         if not is_sm_at_least_89():
             pytest.skip("CUDA capability >= 8.9 required for float8 in triton")
+    elif elem_dtype == torch.float4_e2m1fn_x2:
+        if not is_sm_at_least_100():
+            pytest.skip("CUDA capability >= 10.0 required for float4 gemm")
 
     m = nn.Linear(32, 128, bias=bias, dtype=torch.bfloat16, device="cuda")
     m_mx = copy.deepcopy(m)

test/prototype/mx_formats/test_mx_tensor.py

@@ -66,7 +66,7 @@ def assert_sqnr_gt_threshold(orig, new, threshold):
     if elem_dtype is torch.float8_e4m3fn:
         assert_sqnr_gt_threshold(data_hp, data_mx_dq, 18.0)
     else:
-        assert_sqnr_gt_threshold(data_hp, data_mx_dq, 14.0)
+        assert_sqnr_gt_threshold(data_hp, data_mx_dq, 13.0)
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")