reduce the kernel test coverage for simulator (#36)

reduce the kernel test coverage for simulator

Signed-off-by: chzhang <[email protected]>

Author: chaojun-zhang

tests/conftest.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+
+def pytest_generate_tests(metafunc):
+    use_mini_pytest_profiler = os.getenv("XPU_KERNEL_PYTEST_PROFILER",
+                                         "") == "MINI"
+    if not use_mini_pytest_profiler:
+        return
+
+    module = metafunc.module
+
+    func_pytest_params = getattr(module, "MINI_PYTEST_PARAMS", {})
+    profile = func_pytest_params.get(metafunc.function.__name__, None)
+
+    if not profile:
+        profile = func_pytest_params.get('default', None)
+
+    if not profile:
+        return
+
+    for param_name, values in profile.items():
+        if param_name in metafunc.fixturenames:
+            new_markers = []
+            for mark in metafunc.definition.own_markers:
+                if mark.name == "parametrize" and mark.args[0] != param_name:
+                    new_markers.append(mark)
+                metafunc.definition.own_markers = new_markers
+            metafunc.parametrize(param_name, values)

tests/test_activation.py

@@ -15,6 +15,14 @@
     f"xpu:{i}" for i in range(1 if torch.xpu.device_count() == 1 else 2)
 ]
 
+#override pytest parameters when enable mini pytest
+MINI_PYTEST_PARAMS = {
+    "default": {
+        "num_tokens": [1],
+        "d": [128],
+    },
+}
+
 
 @pytest.mark.parametrize("activation",
                          ["silu_and_mul", "mul_and_silu", "gelu", "gelu_tanh"])

tests/test_cache.py

@@ -36,6 +36,22 @@
 
 KV_CACHE_DTYPE = ["auto"]  # FIXME: will add "fp8" when accuracy is improved
 
+#override pytest parameters when enable mini pytest
+MINI_PYTEST_PARAMS = {
+    "default": {
+        "num_tokens": [1],
+        "head_size": [64, 80],
+    },
+    "test_concat_and_cache_mla": {
+        "num_tokens": [1],
+        "num_blocks": [32],
+    },
+    "test_gather_cache_mla": {
+        "num_blocks": [32],
+        "max_seq_len": [64],
+    },
+}
+
 
 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
 @pytest.mark.parametrize("num_heads", NUM_HEADS)

tests/test_deepseek_scaling_rope.py

@@ -6,6 +6,13 @@
 
 DEVICE = torch.device("xpu")
 
+#override pytest parameters when enable mini pytest
+MINI_PYTEST_PARAMS = {
+    "default": {
+        "batch": [1],
+    },
+}
+
 
 def _rotate_neox(x):
     x1 = x[..., :x.shape[-1] // 2]

tests/test_fp8_quant.py

@@ -129,6 +129,14 @@ def seed_everything(seed):
 SEEDS = [0]
 FP8_DTYPES = [torch.float8_e5m2, torch.float8_e4m3fn]
 
+#override pytest parameters when enable mini pytest
+MINI_PYTEST_PARAMS = {
+    "default": {
+        "num_tokens": [1, 7, 83],
+        "hidden_size": [1, 2, 3, 4, 16],
+    },
+}
+
 
 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
 @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)

tests/test_grouped_topk.py

@@ -6,6 +6,13 @@
                                        fused_grouped_topk_sycl, grouped_topk)
 from tests.utils import seed_everything
 
+#override pytest parameters when enable mini pytest
+MINI_PYTEST_PARAMS = {
+    "default": {
+        "n_hidden": [128, 256],
+    },
+}
+
 
 @pytest.mark.parametrize("n_token", [1, 33, 64])
 @pytest.mark.parametrize("n_hidden", [1024, 2048])

tests/test_layernorm.py

@@ -17,6 +17,14 @@
     f"xpu:{i}" for i in range(1 if torch.xpu.device_count() == 1 else 2)
 ]
 
+#override pytest parameters when enable mini pytest
+MINI_PYTEST_PARAMS = {
+    "default": {
+        "num_tokens": [7],
+        "hidden_size": [8],
+    },
+}
+
 
 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
 @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)

tests/test_moe_sum.py

@@ -12,6 +12,14 @@
 
 TOP_KS = [2, 6]
 
+#override pytest parameters when enable mini pytest
+MINI_PYTEST_PARAMS = {
+    "default": {
+        "m": [1, 33],
+        "k": [128, 256],
+    },
+}
+
 
 @pytest.mark.parametrize("m", [1, 33, 64, 222])
 @pytest.mark.parametrize("topk", TOP_KS)

tests/test_rotary_embedding.py

@@ -31,6 +31,16 @@ def rotary_embedding_opcheck(rot,
                  rot.is_neox_style))
 
 
+#override pytest parameters when enable mini pytest
+MINI_PYTEST_PARAMS = {
+    "default": {
+        "max_position": [11, 256],
+        "head_size": [32],
+        "seq_len": [11, 128],
+    },
+}
+
+
 @pytest.mark.parametrize("device", ["xpu"])
 @pytest.mark.parametrize("max_position", [11, 4096, 32768])
 @pytest.mark.parametrize("is_neox_style", [True, False])

tests/test_swigluoai_and_mul.py

@@ -13,6 +13,14 @@
     f"xpu:{i}" for i in range(1 if torch.xpu.device_count() == 1 else 2)
 ]
 
+#override pytest parameters when enable mini pytest
+MINI_PYTEST_PARAMS = {
+    "default": {
+        "num_tokens": [1, 7],
+        "d": [32, 64],
+    },
+}
+
 default_atol = {torch.float16: 1e-3, torch.bfloat16: 1e-3, torch.float: 1e-5}
 default_rtol = {
     torch.float16: 1e-3,