layers only run for hip devices

Author: PhaneeshB

sharktank/tests/layers/ffn_with_iree_test.py

@@ -7,6 +7,7 @@
 import pytest
 from sharktank.utils._helpers import run_iree_vs_torch_fx
 from sharktank.utils._iree_compile_flags_config import LLM_HIP_COMPILE_FLAGS
+from sharktank.utils.testing import is_hip_condition
 
 
 class FFN(torch.nn.Module):
@@ -24,8 +25,9 @@ def forward(self, x):
             return self.w_down(torch.nn.functional.gelu(self.w_up(x)))
 
 
+@pytest.mark.skipif(not is_hip_condition, reason="Test requires HIP device")
 @pytest.mark.parametrize("dtype,atol", [(torch.float32, 1e-4), (torch.float16, 1e-4)])
-def test_ffn_iree_vs_eager(dtype, atol):
+def test_ffn_mock_iree_vs_eager(dtype, atol):
     torch.manual_seed(42)
     m = FFN(hidden=64, inter=128, dtype=dtype, activation="silu")
     x = torch.randn(2, 8, 64, dtype=dtype)

sharktank/tests/layers/linear_with_iree_test.py

@@ -8,6 +8,7 @@
 import pytest
 from sharktank.utils._helpers import run_iree_vs_torch_fx
 from sharktank.utils._iree_compile_flags_config import LLM_HIP_COMPILE_FLAGS
+from sharktank.utils.testing import is_hip_condition
 
 
 class Linear(torch.nn.Module):
@@ -18,9 +19,9 @@ def __init__(self, in_f, out_f, bias=False, dtype=torch.float32):
     def forward(self, x):
         return self.lin(x)
 
-
+@pytest.mark.skipif(not is_hip_condition, reason="Test requires HIP device")
 @pytest.mark.parametrize("dtype,atol", [(torch.float32, 1e-4), (torch.float16, 1e-4)])
-def test_linear_iree_vs_eager(dtype, atol):
+def test_linear_mock_iree_vs_eager(dtype, atol):
     torch.manual_seed(42)
     m = Linear(64, 64, bias=False, dtype=dtype)
     x = torch.randn(2, 8, 64, dtype=dtype)

sharktank/tests/layers/output_lm_test_with_iree.py

@@ -12,7 +12,7 @@
 from sharktank.types import Dataset, Theta
 from sharktank.layers.configs import LlamaModelConfig
 from sharktank.utils._iree_compile_flags_config import LLM_HIP_COMPILE_FLAGS
-
+from sharktank.utils.testing import is_hip_condition
 
 class OutputLMHead(torch.nn.Module):
     """Standalone output_lm_head block extracted from PagedLlmModelV1"""
@@ -86,6 +86,7 @@ def create_output_lm_head_from_irpa(
 
 
 # Test cases
+@pytest.mark.skipif(not is_hip_condition, reason="Test requires HIP device")
 @pytest.mark.parametrize("dtype,atol", [(torch.float16, 1e-4)])
 def test_output_lm_head_iree_vs_eager(request, dtype, atol):
     """
@@ -116,7 +117,7 @@ def test_output_lm_head_iree_vs_eager(request, dtype, atol):
             parameters_path=irpa_path,
         )
 
-
+@pytest.mark.skipif(not is_hip_condition, reason="Test requires HIP device")
 def test_output_lm_head_mock():
     """
     Mock test with synthetic weights for OutputLMHead functionality.

sharktank/tests/layers/rms_norm_with_iree_test.py

@@ -8,7 +8,7 @@
 import pytest
 from sharktank.utils._helpers import run_iree_vs_torch_fx
 from sharktank.utils._iree_compile_flags_config import LLM_HIP_COMPILE_FLAGS
-
+from sharktank.utils.testing import is_hip_condition
 
 class RMSNorm(torch.nn.Module):
     def __init__(self, hidden=64, eps=1e-5, dtype=torch.float32):
@@ -25,6 +25,7 @@ def forward(self, x):
         return y * self.weight  # broadcast over last dim
 
 
+@pytest.mark.skipif(not is_hip_condition, reason="Test requires HIP device")
 @pytest.mark.parametrize("dtype,atol", [(torch.float32, 1e-4), (torch.bfloat16, 1e-2)])
 def test_rms_norm_iree_vs_eager(dtype, atol):
     torch.manual_seed(42)

sharktank/tests/layers/token_embedding_with_iree_test.py

@@ -11,6 +11,7 @@
 from sharktank.types.theta import Dataset
 from sharktank.utils._helpers import run_iree_vs_torch_fx, validate_and_get_irpa_path
 from sharktank.utils._iree_compile_flags_config import LLM_HIP_COMPILE_FLAGS
+from sharktank.utils.testing import is_hip_condition
 
 
 class TokenEmbeddingSmall(torch.nn.Module):
@@ -22,7 +23,7 @@ def __init__(self, vocab_size=128, hidden=64, dtype=torch.float32):
     def forward(self, ids: torch.Tensor):
         return self.weight[ids]
 
-
+@pytest.mark.skipif(not is_hip_condition, reason="Test requires HIP device")
 @pytest.mark.parametrize("dtype,atol", [(torch.float16, 1e-4)])
 def test_token_embedding_iree_vs_eager(request, dtype, atol):
     torch.manual_seed(42)
@@ -43,6 +44,7 @@ def test_token_embedding_iree_vs_eager(request, dtype, atol):
     )
 
 
+@pytest.mark.skipif(not is_hip_condition, reason="Test requires HIP device")
 @pytest.mark.parametrize("dtype,atol", [(torch.float16, 1e-4)])
 def test_token_embedding_mock_iree_vs_eager(dtype, atol):
     torch.manual_seed(42)