only run tests on compute capability 9.0+

Author: danielvegamyhre

test/prototype/scaled_grouped_mm/test_kernels.py

@@ -11,9 +11,14 @@
 
 # We need to skip before doing any imports which would use triton, since
 # triton won't be available on CPU builds and torch < 2.5
-if not (TORCH_VERSION_AT_LEAST_2_5 and torch.cuda.is_available()):
+if not (
+    TORCH_VERSION_AT_LEAST_2_5
+    and torch.cuda.is_available()
+    and torch.cuda.get_device_capability()[0] >= 9
+):
     pytest.skip("Unsupported PyTorch version", allow_module_level=True)
 
+
 from torchao.prototype.scaled_grouped_mm.kernels.jagged_float8_scales import (
     triton_fp8_col_major_jagged_colwise_scales,
     triton_fp8_row_major_jagged_rowwise_scales,

test/prototype/scaled_grouped_mm/test_scaled_grouped_mm.py

@@ -11,7 +11,11 @@
 
 # We need to skip before doing any imports which would use triton, since
 # triton won't be available on CPU builds and torch < 2.5
-if not (TORCH_VERSION_AT_LEAST_2_5 and torch.cuda.is_available()):
+if not (
+    TORCH_VERSION_AT_LEAST_2_5
+    and torch.cuda.is_available()
+    and torch.cuda.get_device_capability()[0] >= 9
+):
     pytest.skip("Unsupported PyTorch version", allow_module_level=True)