Refactor all tests

Signed-off-by: Pratikkumar Prajapati <pprajapati@tenstorrent.com>

Author: pprajapatiTT

.github/workflows/galaxy-deepseek-tests-long-impl.yaml

@@ -65,9 +65,9 @@ jobs:
         timeout-minutes: 360
         run: |
           uv pip install -r models/demos/deepseek_v3/reference/deepseek/requirements.txt
-          DEEPSEEK_MAX_SEQ_LEN=16384 pytest models/demos/deepseek_v3/tests/test_mla.py --timeout 3600 --durations=0
-          DEEPSEEK_MAX_SEQ_LEN=16384 pytest models/demos/deepseek_v3/tests/test_decoder_block.py --timeout 7200 --durations=0
-          DEEPSEEK_MAX_SEQ_LEN=2048 pytest models/demos/deepseek_v3/tests/test_model.py --timeout 10800 --durations=0
+          DEEPSEEK_MAX_SEQ_LEN_OVERRIDE=16384 pytest models/demos/deepseek_v3/tests/test_mla.py --timeout 3600 --durations=0
+          DEEPSEEK_MAX_SEQ_LEN_OVERRIDE=16384 pytest models/demos/deepseek_v3/tests/test_decoder_block.py --timeout 7200 --durations=0
+          DEEPSEEK_MAX_SEQ_LEN_OVERRIDE=2048 pytest models/demos/deepseek_v3/tests/test_model.py --timeout 10800 --durations=0
       - uses: tenstorrent/tt-metal/.github/actions/slack-report@main
         if: ${{ failure() }}
         with:

models/demos/deepseek_v3/conftest.py

@@ -182,13 +182,13 @@ def hf_config_short(request, hf_config):
     Build a shortened DeepSeek config for tests.
 
     Environment variables:
-        DEEPSEEK_MAX_SEQ_LEN: Optional override for `hf_config_short.max_seq_len`.
+        DEEPSEEK_MAX_SEQ_LEN_OVERRIDE: Optional override for `hf_config_short.max_seq_len`.
             When set (e.g. "32768"), tests that read `hf_config_short.max_seq_len`
             can exercise longer sequence lengths without modifying code.
     """
     hf_config_out = deepcopy(hf_config)
     hf_config_out.num_hidden_layers = getattr(request, "param", 1)
-    max_seq_len_override = os.getenv("DEEPSEEK_MAX_SEQ_LEN")
+    max_seq_len_override = os.getenv("DEEPSEEK_MAX_SEQ_LEN_OVERRIDE")
     if max_seq_len_override is not None:
         hf_config_out.max_seq_len = int(max_seq_len_override)
     else:

models/demos/deepseek_v3/tests/pytest_utils.py

@@ -6,6 +6,8 @@
 
 import pytest
 
+DEFAULT_PREFILL_SEQ_LEN = 128
+
 
 def expand_test_cases_with_position_ids_ranges(base_cases):
     """
@@ -121,46 +123,45 @@ def build_expanded_test_ids(expanded_cases):
     return expanded_ids
 
 
-def get_base_test_cases(users_per_row, prefill_seq_lens, include_decode_random_pos_ids=True):
+def get_base_test_cases(users_per_row, prefill_seq_len, include_decode_random_pos_ids=True):
     """
     Build base test cases for decode and prefill paths.
 
-    Used by:
+    This helper is only exercised by these tests.:
         - models/demos/deepseek_v3/tests/test_mla.py
         - models/demos/deepseek_v3/tests/test_decoder_block.py
         - models/demos/deepseek_v3/tests/test_model.py
 
-    This helper is only exercised by these tests.
-
     Args:
         users_per_row: Number of users per row (USERS_PER_ROW).
-        prefill_seq_lens: Iterable of prefill sequence lengths.
+        prefill_seq_len: Prefill sequence length to use when DEEPSEEK_MAX_SEQ_LEN_OVERRIDE is not set.
         include_decode_random_pos_ids: If True, include ("decode", 1, users_per_row, None).
 
-    environment variable DEEPSEEK_MAX_SEQ_LEN is primarily a CI override to expand prefill and decode coverage.
-    When set, we add specific prefill and decode position_ids (0 and max_seq_len - 1) in
-    addition to the default random prefill and decode position_ids cases.
+    The environment variable DEEPSEEK_MAX_SEQ_LEN_OVERRIDE is primarily a CI override to expand
+    prefill and decode coverage.
 
     Behavior:
-        - Adds decode cases:
-          - random position_ids (optional)
-          - position_id 0
-          - position_id max_seq_len - 1 when DEEPSEEK_MAX_SEQ_LEN is set
-        - Adds prefill cases:
-          - a direct prefill at DEEPSEEK_MAX_SEQ_LEN when set
-          - the standard prefill seq lens with skip marks when DEEPSEEK_MAX_SEQ_LEN is not set
+        - Decode cases:
+          - when DEEPSEEK_MAX_SEQ_LEN_OVERRIDE is set, additionally includes:
+            - position_id 0
+            - position_id max_seq_len - 1
+        - Prefill cases:
+          - when DEEPSEEK_MAX_SEQ_LEN_OVERRIDE is not set, includes one prefill case using
+            prefill_seq_len: ("prefill", prefill_seq_len, 1, None)
+          - when DEEPSEEK_MAX_SEQ_LEN_OVERRIDE is set, replaces the prefill list with a single case:
+            ("prefill", max_seq_len, 1, None)
 
     """
     base_cases = []
     if include_decode_random_pos_ids:
         base_cases += [("decode", 1, users_per_row, None)]
 
-    max_seq_len_env = os.getenv("DEEPSEEK_MAX_SEQ_LEN")
+    max_seq_len_env = os.getenv("DEEPSEEK_MAX_SEQ_LEN_OVERRIDE")
     if max_seq_len_env is None:
-        # If DEEPSEEK_MAX_SEQ_LEN is not set, use the default prefill sequence length.
-        base_cases += [("prefill", seq_len, 1, None) for seq_len in prefill_seq_lens]
+        # If DEEPSEEK_MAX_SEQ_LEN_OVERRIDE is not set, use the default prefill sequence length.
+        base_cases += [("prefill", prefill_seq_len, 1, None)]
     else:
-        # If DEEPSEEK_MAX_SEQ_LEN is set, use it to expand prefill and decode coverage.
+        # If DEEPSEEK_MAX_SEQ_LEN_OVERRIDE is set, use it to expand prefill and decode coverage.
         max_seq_len = int(max_seq_len_env)
         base_cases += [
             ("decode", 1, users_per_row, 0),  # decode position_id 0
@@ -170,7 +171,7 @@ def get_base_test_cases(users_per_row, prefill_seq_lens, include_decode_random_p
     return base_cases
 
 
-def build_test_cases_and_ids(users_per_row, prefill_seq_lens, include_decode_random_pos_ids=True):
+def build_test_cases_and_ids(users_per_row, prefill_seq_len, include_decode_random_pos_ids=True):
     """
     Build base test cases and return expanded cases with matching pytest IDs.
 
@@ -179,7 +180,7 @@ def build_test_cases_and_ids(users_per_row, prefill_seq_lens, include_decode_ran
       - expand_test_cases_with_position_ids_ranges
       - build_expanded_test_ids
     """
-    base_cases = get_base_test_cases(users_per_row, prefill_seq_lens, include_decode_random_pos_ids)
+    base_cases = get_base_test_cases(users_per_row, prefill_seq_len, include_decode_random_pos_ids)
     expanded_cases = expand_test_cases_with_position_ids_ranges(base_cases)
     expanded_ids = build_expanded_test_ids(expanded_cases)
     return expanded_cases, expanded_ids

models/demos/deepseek_v3/tests/test_decoder_block.py

@@ -10,9 +10,8 @@
 from transformers.configuration_utils import PretrainedConfig
 
 import ttnn
-from models.demos.deepseek_v3.conftest import PREFILL_SEQ_LENS
 from models.demos.deepseek_v3.reference.modeling_deepseek import DeepseekV3DecoderLayer
-from models.demos.deepseek_v3.tests.pytest_utils import build_test_cases_and_ids
+from models.demos.deepseek_v3.tests.pytest_utils import DEFAULT_PREFILL_SEQ_LEN, build_test_cases_and_ids
 from models.demos.deepseek_v3.tt.decoder_block.decoder_block_2d import DecoderBlock2D
 from models.demos.deepseek_v3.tt.decoder_block.decoder_block_2d_base import DecoderBlock2DBase
 from models.demos.deepseek_v3.tt.decoder_block.moe_decoder_block_2d import MoEDecoderBlock2D
@@ -196,7 +195,7 @@ def run_test_forward_pass_decoder2d(
 
 TEST_CASES, TEST_IDS = build_test_cases_and_ids(
     USERS_PER_ROW,
-    PREFILL_SEQ_LENS[0:1],  # list of default prefill sequence lengths to test
+    DEFAULT_PREFILL_SEQ_LEN,  # default prefill sequence length to test
     include_decode_random_pos_ids=True,  # include decode random position_ids case
 )
 

models/demos/deepseek_v3/tests/test_embedding.py

@@ -2,6 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
+import os
+
 import pytest
 import torch
 from loguru import logger
@@ -10,7 +12,7 @@
 from torch.nn import Embedding as EmbeddingReference
 
 import ttnn
-from models.demos.deepseek_v3.conftest import PREFILL_SEQ_LENS
+from models.demos.deepseek_v3.tests.pytest_utils import DEFAULT_PREFILL_SEQ_LEN
 from models.demos.deepseek_v3.tt.embedding.embedding1d import Embedding1D
 from models.demos.deepseek_v3.tt.embedding.embedding2d import Embedding2D
 from models.demos.deepseek_v3.utils.config_helpers import sub_state_dict
@@ -23,52 +25,27 @@
     run_module_forward,
 )
 
+_max_seq_len_env = os.getenv("DEEPSEEK_MAX_SEQ_LEN_OVERRIDE")
+_prefill_seq_len = int(_max_seq_len_env) if _max_seq_len_env is not None else DEFAULT_PREFILL_SEQ_LEN
+
 
 @pytest.mark.parametrize(
-    "device_params",
+    "EmbeddingClass,mode,batch_size_or_seq_len",
     [
-        {"fabric_config": ttnn.FabricConfig.FABRIC_1D},
+        pytest.param(Embedding1D, "decode", 32, marks=pytest.mark.requires_device(["TG"])),
+        pytest.param(Embedding2D, "decode", 128, marks=pytest.mark.requires_device(["TG", "DUAL", "QUAD"])),
+        pytest.param(Embedding1D, "prefill", _prefill_seq_len, marks=pytest.mark.requires_device(["TG"])),
+        pytest.param(
+            Embedding2D, "prefill", _prefill_seq_len, marks=pytest.mark.requires_device(["TG", "DUAL", "QUAD"])
+        ),
     ],
-    indirect=True,
 )
 @pytest.mark.parametrize(
-    "EmbeddingClass,mode,batch_size_or_seq_len",
+    "device_params",
     [
-        pytest.param(Embedding1D, "decode", 32, marks=pytest.mark.requires_device(["TG"])),
-        pytest.param(Embedding2D, "decode", 128, marks=pytest.mark.requires_device(["TG", "DUAL", "QUAD"])),
-    ]
-    + [
-        pytest.param(Embedding1D, "prefill", seq_len, marks=pytest.mark.requires_device(["TG"]))
-        if seq_len == 128
-        else pytest.param(
-            Embedding1D,
-            "prefill",
-            seq_len,
-            marks=[
-                pytest.mark.requires_device(["TG"]),
-                pytest.mark.skip(
-                    f"Skipping prefilling with seq_len={seq_len} since this would cause us to exceed our available CI workload time"
-                ),
-            ],
-        )
-        for seq_len in PREFILL_SEQ_LENS
-    ]
-    + [
-        pytest.param(Embedding2D, "prefill", seq_len, marks=pytest.mark.requires_device(["TG", "DUAL", "QUAD"]))
-        if seq_len == 128
-        else pytest.param(
-            Embedding2D,
-            "prefill",
-            seq_len,
-            marks=[
-                pytest.mark.requires_device(["TG", "DUAL", "QUAD"]),
-                pytest.mark.skip(
-                    f"Skipping prefilling with seq_len={seq_len} since this would cause us to exceed our available CI workload time"
-                ),
-            ],
-        )
-        for seq_len in PREFILL_SEQ_LENS
+        {"fabric_config": ttnn.FabricConfig.FABRIC_1D},
     ],
+    indirect=True,
 )
 @pytest.mark.parametrize(
     "generate_reference_io",

models/demos/deepseek_v3/tests/test_lm_head.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
+import os
 from pathlib import Path
 from typing import Any
 
@@ -11,7 +12,7 @@
 from loguru import logger
 
 import ttnn
-from models.demos.deepseek_v3.conftest import PREFILL_SEQ_LENS
+from models.demos.deepseek_v3.tests.pytest_utils import DEFAULT_PREFILL_SEQ_LEN
 from models.demos.deepseek_v3.tt.ccl import CCL
 from models.demos.deepseek_v3.tt.lm_head import LMHead
 from models.demos.deepseek_v3.utils.config_helpers import sub_state_dict
@@ -39,30 +40,23 @@ def forward(self, hidden_states):
         return self.lm_head(hidden_states)
 
 
+_max_seq_len_env = os.getenv("DEEPSEEK_MAX_SEQ_LEN_OVERRIDE")
+_prefill_seq_len = int(_max_seq_len_env) if _max_seq_len_env is not None else DEFAULT_PREFILL_SEQ_LEN
+
+
 @pytest.mark.parametrize(
-    "device_params",
+    "mode,seq_len",
     [
-        {"fabric_config": ttnn.FabricConfig.FABRIC_1D},
+        ("decode", 32),
+        ("prefill", _prefill_seq_len),
     ],
-    indirect=True,
 )
 @pytest.mark.parametrize(
-    "mode,seq_len",
+    "device_params",
     [
-        ("decode", 32),
-    ]
-    + [
-        ("prefill", seq_len)
-        if seq_len == 128
-        else pytest.param(
-            "prefill",
-            seq_len,
-            marks=pytest.mark.skip(
-                f"Skipping prefilling with seq_len={seq_len} since this would cause us to exceed our available CI workload time"
-            ),
-        )
-        for seq_len in PREFILL_SEQ_LENS
+        {"fabric_config": ttnn.FabricConfig.FABRIC_1D},
     ],
+    indirect=True,
 )
 @pytest.mark.requires_device(["TG"])
 def test_forward_pass(

models/demos/deepseek_v3/tests/test_mla.py

@@ -11,10 +11,8 @@
 
 import ttnn
 from models.common.utility_functions import comp_pcc
-from models.demos.deepseek_v3.conftest import PREFILL_SEQ_LENS
 from models.demos.deepseek_v3.reference.modeling_deepseek import DeepseekV3Attention
-from models.demos.deepseek_v3.tests.pytest_utils import build_test_cases_and_ids
-from models.demos.deepseek_v3.tt.mla.mla1d import MLA1D
+from models.demos.deepseek_v3.tests.pytest_utils import DEFAULT_PREFILL_SEQ_LEN, build_test_cases_and_ids
 from models.demos.deepseek_v3.tt.mla.mla2d import MLA2D
 from models.demos.deepseek_v3.utils.config_helpers import USERS_PER_ROW, sub_state_dict
 from models.demos.deepseek_v3.utils.run_config import create_run_config
@@ -312,7 +310,7 @@ def run_test_forward_pass_mla2d(
 
 TEST_CASES, TEST_IDS = build_test_cases_and_ids(
     USERS_PER_ROW,
-    PREFILL_SEQ_LENS[0:1],  # list of default prefill sequence lengths to test
+    DEFAULT_PREFILL_SEQ_LEN,  # default prefill sequence length to test
     include_decode_random_pos_ids=True,  # include decode random position_ids case
 )
 

models/demos/deepseek_v3/tests/test_mlp.py

@@ -2,14 +2,16 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
+import os
+
 import pytest
 import torch
 from loguru import logger
 
 import ttnn
 from models.common.utility_functions import comp_pcc
-from models.demos.deepseek_v3.conftest import PREFILL_SEQ_LENS
 from models.demos.deepseek_v3.reference.modeling_deepseek import DeepseekV3MLP
+from models.demos.deepseek_v3.tests.pytest_utils import DEFAULT_PREFILL_SEQ_LEN
 from models.demos.deepseek_v3.tt.mlp.mlp import MLP
 from models.demos.deepseek_v3.tt.mlp.mlp_dequant import MLPDequant
 from models.demos.deepseek_v3.tt.mlp.non_expert import NonExpert
@@ -128,6 +130,17 @@ def run_weight_conversion_test(MLPClass, hf_config, state_dict, tmp_path, refere
     ttnn.deallocate(w1_ttnn)
 
 
+_max_seq_len_env = os.getenv("DEEPSEEK_MAX_SEQ_LEN_OVERRIDE")
+_prefill_seq_len = int(_max_seq_len_env) if _max_seq_len_env is not None else DEFAULT_PREFILL_SEQ_LEN
+
+
+@pytest.mark.parametrize(
+    "mode,seq_len",
+    [
+        ("decode", 32),
+        ("prefill", _prefill_seq_len),
+    ],
+)
 @pytest.mark.parametrize("device_params", [{"fabric_config": ttnn.FabricConfig.FABRIC_1D}], indirect=True)
 @pytest.mark.parametrize(
     "MLPClass,module_path",
@@ -137,24 +150,6 @@ def run_weight_conversion_test(MLPClass, hf_config, state_dict, tmp_path, refere
         (SharedExpert, "model.layers.3.mlp.shared_experts"),
     ],
 )
-@pytest.mark.parametrize(
-    "mode,seq_len",
-    [
-        ("decode", 32),
-    ]
-    + [
-        ("prefill", seq_len)
-        if seq_len == 128
-        else pytest.param(
-            "prefill",
-            seq_len,
-            marks=pytest.mark.skip(
-                f"Skipping prefilling with seq_len={seq_len} since this would cause us to exceed our available CI workload time"
-            ),
-        )
-        for seq_len in PREFILL_SEQ_LENS
-    ],
-)
 def test_forward_pass(
     MLPClass,
     module_path,

models/demos/deepseek_v3/tests/test_model.py

@@ -8,9 +8,8 @@
 from transformers.configuration_utils import PretrainedConfig
 
 import ttnn
-from models.demos.deepseek_v3.conftest import PREFILL_SEQ_LENS
 from models.demos.deepseek_v3.reference.modeling_deepseek import DeepseekV3ForCausalLM
-from models.demos.deepseek_v3.tests.pytest_utils import build_test_cases_and_ids
+from models.demos.deepseek_v3.tests.pytest_utils import DEFAULT_PREFILL_SEQ_LEN, build_test_cases_and_ids
 from models.demos.deepseek_v3.tt.mla.mla2d import MLA2D
 from models.demos.deepseek_v3.tt.model.row_batched_model import RowBatchedModel
 from models.demos.deepseek_v3.utils.config_helpers import USERS_PER_ROW, sub_state_dict
@@ -200,7 +199,7 @@ def run_test_forward_pass_dpmodel(
 
 TEST_CASES, TEST_IDS = build_test_cases_and_ids(
     USERS_PER_ROW,
-    PREFILL_SEQ_LENS[0:1],  # list of default prefill sequence lengths to test
+    DEFAULT_PREFILL_SEQ_LEN,  # default prefill sequence length to test
     include_decode_random_pos_ids=False,  # TODO: Remove include_decode_random_pos_ids=False once non-zero position_ids case is working.
 )