From bc50e39353ae579177fa71cad39f63b88abac06e Mon Sep 17 00:00:00 2001
From: mtairum <mtairum@tenstorrent.com>
Date: Wed, 22 May 2024 16:21:31 +0000
Subject: [PATCH] #5337: Add eth dispatch config to select Mixtral tests

---
 models/demos/t3000/mixtral8x7b/demo/demo.py                  | 1 +
 models/demos/t3000/mixtral8x7b/tests/test_mixtral_decoder.py | 1 +
 models/demos/t3000/mixtral8x7b/tests/test_mixtral_mlp.py     | 1 +
 models/demos/t3000/mixtral8x7b/tests/test_mixtral_model.py   | 1 +
 models/demos/t3000/mixtral8x7b/tests/test_mixtral_moe.py     | 1 +
 models/demos/t3000/mixtral8x7b/tests/test_mixtral_perf.py    | 1 +
 6 files changed, 6 insertions(+)

diff --git a/models/demos/t3000/mixtral8x7b/demo/demo.py b/models/demos/t3000/mixtral8x7b/demo/demo.py
index f7808b20de4..8483d89f242 100644
--- a/models/demos/t3000/mixtral8x7b/demo/demo.py
+++ b/models/demos/t3000/mixtral8x7b/demo/demo.py
@@ -15,6 +15,7 @@
     os.environ["MIXTRAL_TOKENIZER_PATH"] = "/mnt/MLPerf/tt_dnn-models/Mistral/Mixtral-8x7B-v0.1/"
     os.environ["MIXTRAL_CACHE_PATH"] = "/mnt/MLPerf/tt_dnn-models/Mistral/Mixtral-8x7B-v0.1/"
     os.environ["TT_METAL_ASYNC_DEVICE_QUEUE"] = "1"
+    os.environ["WH_ARCH_YAML"] = "wormhole_b0_80_arch_eth_dispatch.yaml"
 
 import ttnn
 from ttnn import ReplicateTensorToMesh, ConcatMeshToTensor
diff --git a/models/demos/t3000/mixtral8x7b/tests/test_mixtral_decoder.py b/models/demos/t3000/mixtral8x7b/tests/test_mixtral_decoder.py
index d39608a7644..51d14165835 100644
--- a/models/demos/t3000/mixtral8x7b/tests/test_mixtral_decoder.py
+++ b/models/demos/t3000/mixtral8x7b/tests/test_mixtral_decoder.py
@@ -12,6 +12,7 @@
     os.environ["MIXTRAL_TOKENIZER_PATH"] = "/mnt/MLPerf/tt_dnn-models/Mistral/Mixtral-8x7B-v0.1/"
     os.environ["MIXTRAL_CACHE_PATH"] = "/mnt/MLPerf/tt_dnn-models/Mistral/Mixtral-8x7B-v0.1/"
     os.environ["TT_METAL_ASYNC_DEVICE_QUEUE"] = "1"
+    os.environ["WH_ARCH_YAML"] = "wormhole_b0_80_arch_eth_dispatch.yaml"
 
 import ttnn
 from models.demos.t3000.mixtral8x7b.tt.mixtral_common import prepare_inputs_ttnn, prepare_rotation_mat_ttnn
diff --git a/models/demos/t3000/mixtral8x7b/tests/test_mixtral_mlp.py b/models/demos/t3000/mixtral8x7b/tests/test_mixtral_mlp.py
index d7a057e2b53..b349b91ceee 100644
--- a/models/demos/t3000/mixtral8x7b/tests/test_mixtral_mlp.py
+++ b/models/demos/t3000/mixtral8x7b/tests/test_mixtral_mlp.py
@@ -11,6 +11,7 @@
     os.environ["MIXTRAL_TOKENIZER_PATH"] = "/mnt/MLPerf/tt_dnn-models/Mistral/Mixtral-8x7B-v0.1/"
     os.environ["MIXTRAL_CACHE_PATH"] = "/mnt/MLPerf/tt_dnn-models/Mistral/Mixtral-8x7B-v0.1/"
     os.environ["TT_METAL_ASYNC_DEVICE_QUEUE"] = "1"
+    os.environ["WH_ARCH_YAML"] = "wormhole_b0_80_arch_eth_dispatch.yaml"
 
 import ttnn
 from ttnn import ReplicateTensorToMesh, ConcatMeshToTensor
diff --git a/models/demos/t3000/mixtral8x7b/tests/test_mixtral_model.py b/models/demos/t3000/mixtral8x7b/tests/test_mixtral_model.py
index 7a9afecb583..51b6c49bee1 100644
--- a/models/demos/t3000/mixtral8x7b/tests/test_mixtral_model.py
+++ b/models/demos/t3000/mixtral8x7b/tests/test_mixtral_model.py
@@ -14,6 +14,7 @@
     os.environ["MIXTRAL_TOKENIZER_PATH"] = "/mnt/MLPerf/tt_dnn-models/Mistral/Mixtral-8x7B-v0.1/"
     os.environ["MIXTRAL_CACHE_PATH"] = "/mnt/MLPerf/tt_dnn-models/Mistral/Mixtral-8x7B-v0.1/"
     os.environ["TT_METAL_ASYNC_DEVICE_QUEUE"] = "1"
+    os.environ["WH_ARCH_YAML"] = "wormhole_b0_80_arch_eth_dispatch.yaml"
 
 import ttnn
 from ttnn import ReplicateTensorToMesh, ConcatMeshToTensor
diff --git a/models/demos/t3000/mixtral8x7b/tests/test_mixtral_moe.py b/models/demos/t3000/mixtral8x7b/tests/test_mixtral_moe.py
index 221d956de29..94d93875a6a 100644
--- a/models/demos/t3000/mixtral8x7b/tests/test_mixtral_moe.py
+++ b/models/demos/t3000/mixtral8x7b/tests/test_mixtral_moe.py
@@ -12,6 +12,7 @@
     os.environ["MIXTRAL_TOKENIZER_PATH"] = "/mnt/MLPerf/tt_dnn-models/Mistral/Mixtral-8x7B-v0.1/"
     os.environ["MIXTRAL_CACHE_PATH"] = "/mnt/MLPerf/tt_dnn-models/Mistral/Mixtral-8x7B-v0.1/"
     os.environ["TT_METAL_ASYNC_DEVICE_QUEUE"] = "1"
+    os.environ["WH_ARCH_YAML"] = "wormhole_b0_80_arch_eth_dispatch.yaml"
 
 import ttnn
 from ttnn import ReplicateTensorToMesh, ConcatMeshToTensor
diff --git a/models/demos/t3000/mixtral8x7b/tests/test_mixtral_perf.py b/models/demos/t3000/mixtral8x7b/tests/test_mixtral_perf.py
index fa208ab9786..50137139855 100644
--- a/models/demos/t3000/mixtral8x7b/tests/test_mixtral_perf.py
+++ b/models/demos/t3000/mixtral8x7b/tests/test_mixtral_perf.py
@@ -11,6 +11,7 @@
     os.environ["MIXTRAL_TOKENIZER_PATH"] = "/mnt/MLPerf/tt_dnn-models/Mistral/Mixtral-8x7B-v0.1/"
     os.environ["MIXTRAL_CACHE_PATH"] = "/mnt/MLPerf/tt_dnn-models/Mistral/Mixtral-8x7B-v0.1/"
     os.environ["TT_METAL_ASYNC_DEVICE_QUEUE"] = "1"
+    os.environ["WH_ARCH_YAML"] = "wormhole_b0_80_arch_eth_dispatch.yaml"
 
 import ttnn
 from ttnn import ConcatMeshToTensor