Add skipIfA10G decorator (#982)

Author: yf225

benchmarks/run.py

@@ -45,6 +45,7 @@
 from torch.utils._pytree import tree_leaves
 from torch.utils._pytree import tree_map
 
+from helion._testing import get_nvidia_gpu_model
 from helion._utils import counters
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -59,23 +60,6 @@ def is_cuda() -> bool:
     return torch.version.cuda is not None
 
 
-def get_nvidia_gpu_model() -> str:
-    """
-    Retrieves the model of the NVIDIA GPU being used.
-    Will return the name of the first GPU listed.
-    Returns:
-        str: The model of the NVIDIA GPU or empty str if not found.
-    """
-    try:
-        model = subprocess.check_output(
-            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader,nounits"]
-        )
-        return model.decode().strip().split("\n")[0]
-    except OSError:
-        logger.warning("nvidia-smi not found. Returning empty str.")
-        return ""
-
-
 IS_B200 = is_cuda() and get_nvidia_gpu_model() == "NVIDIA B200"
 
 

helion/_testing.py

@@ -47,6 +47,19 @@ def is_cuda() -> bool:
     )
 
 
+def get_nvidia_gpu_model() -> str:
+    """
+    Retrieves the model of the NVIDIA GPU being used.
+    Will return the name of the current device.
+    Returns:
+        str: The model of the NVIDIA GPU or empty str if not found.
+    """
+    if torch.cuda.is_available():
+        props = torch.cuda.get_device_properties(torch.cuda.current_device())
+        return getattr(props, "name", "")
+    return ""
+
+
 def skipIfRefEager(reason: str) -> Callable[[Callable], Callable]:
     """Skip test if running in ref eager mode (HELION_INTERPRET=1)."""
     return unittest.skipIf(os.environ.get("HELION_INTERPRET") == "1", reason)
@@ -67,6 +80,13 @@ def skipIfXPU(reason: str) -> Callable[[Callable], Callable]:
     return unittest.skipIf(torch.xpu.is_available(), reason)  # pyright: ignore[reportAttributeAccessIssue]
 
 
+def skipIfA10G(reason: str) -> Callable[[Callable], Callable]:
+    """Skip test if running on A10G GPU"""
+    gpu_model = get_nvidia_gpu_model()
+    is_a10g = "A10G" in gpu_model
+    return unittest.skipIf(is_a10g, reason)
+
+
 def skipIfNotCUDA() -> Callable[[Callable], Callable]:
     """Skip test if not running on CUDA (NVIDIA GPU)."""
     return unittest.skipIf(