Add batch_size parameterization and lower to 16, restore default optimization level and prompt.

Author: ddilbazTT

benchmark/tt-xla/test_llms.py

@@ -166,7 +166,16 @@ def test_llm(
             json.dump(results, file, indent=2)
 
 
-def test_llm_tp(ModelLoaderModule, variant, output_file, num_layers=None, batch_size=None, optimization_level=None, request=None, **kwargs):
+def test_llm_tp(
+    ModelLoaderModule,
+    variant,
+    output_file,
+    num_layers=None,
+    batch_size=None,
+    optimization_level=None,
+    request=None,
+    **kwargs,
+):
     # Need to define arch since get_xla_device_arch() doesn't work when spmd is enabled
     arch = "wormhole_llmbox"
     mesh_config_fn = ModelLoaderModule.get_mesh_config
@@ -175,7 +184,7 @@ def test_llm_tp(ModelLoaderModule, variant, output_file, num_layers=None, batch_
         batch_size = DEFAULT_BATCH_SIZE
     if optimization_level is None:
         optimization_level = DEFAULT_OPTIMIZATION_LEVEL
-    
+
     test_llm(
         ModelLoaderModule=ModelLoaderModule,
         variant=variant,
@@ -619,12 +628,12 @@ def test_gpt_oss_20b_tp(output_file, num_layers, request):
 
     variant = ModelVariant.GPT_OSS_20B
     test_llm_tp(
-        ModelLoader, 
-        variant, 
-        output_file, 
+        ModelLoader,
+        variant,
+        output_file,
         num_layers=num_layers,
-        batch_size=16, # https://github.com/tenstorrent/tt-xla/issues/3251
+        batch_size=16,  # https://github.com/tenstorrent/tt-xla/issues/3251
         optimization_level=0,
         request=request,
-        required_pcc=0.86
+        required_pcc=0.86,
     )