Add new tests: 4b dq, kv sdpa (#2378)

Summary:
Pull Request resolved: #2378

Add new tests: 4b dq, kv sdpa

Reviewed By: mergennachin, manuelcandales

Differential Revision: D54806721

fbshipit-source-id: a18d8f91f3cf181f2a1e69297e742634686d2e4e

Author: Michael Gschwind

examples/models/llama2/export_llama_lib.py

@@ -499,13 +499,20 @@ def _export_llama(modelname, args) -> str:  # noqa: C901
     if builder.dtype == DType.fp16:
         modelname = f"{modelname}_h"
 
-    builder.save_to_pte(modelname)
-
     if args.output_name:
         modelname = args.output_name
-        if modelname[-4:] == ".pte":
+        if modelname.endswith(".pte"):
+            output_file = modelname
             modelname = modelname[:-4]
+            print(f"modelname: {modelname}")
+            print(f"output_file: {output_file}")
+        else:
+            output_file = f"{builder.output_dir}/{modelname}.pte"
+            print(f"modelname: {modelname}")
+            print(f"output_file: {output_file}")
+    else:
+        output_file = f"{builder.output_dir}/{modelname}.pte"
 
-    output_file = f"{builder.output_dir}/{modelname}.pte"
+    builder.save_to_pte(output_file)
 
     return output_file

examples/portable/utils.py

@@ -101,7 +101,11 @@ def export_to_exec_prog(
 def save_pte_program(
     prog: ExecutorchProgramManager, model_name: str, output_dir: str = ""
 ) -> None:
-    filename = os.path.join(output_dir, f"{model_name}.pte")
+    if model_name.endswith(".pte"):
+        filename = model_name
+    else:
+        filename = os.path.join(output_dir, f"{model_name}.pte")
+
     try:
         with open(filename, "wb") as file:
             prog.write_to_file(file)