diff --git a/extension/android/benchmark/android-llm-device-farm-test-spec.yml b/extension/android/benchmark/android-llm-device-farm-test-spec.yml
index 4e3274ce66f..ffb528767a5 100644
--- a/extension/android/benchmark/android-llm-device-farm-test-spec.yml
+++ b/extension/android/benchmark/android-llm-device-farm-test-spec.yml
@@ -10,18 +10,21 @@ phases:
     commands:
       # Prepare the model and the tokenizer
       - adb -s $DEVICEFARM_DEVICE_UDID shell "ls -la /sdcard/"
-      - adb -s $DEVICEFARM_DEVICE_UDID shell "mkdir -p /data/local/tmp/llama/"
-      - adb -s $DEVICEFARM_DEVICE_UDID shell "mv /sdcard/*.bin /data/local/tmp/llama/"
-      - adb -s $DEVICEFARM_DEVICE_UDID shell "mv /sdcard/*.pte /data/local/tmp/llama/"
-      - adb -s $DEVICEFARM_DEVICE_UDID shell "chmod 664 /data/local/tmp/llama/*.bin"
-      - adb -s $DEVICEFARM_DEVICE_UDID shell "chmod 664 /data/local/tmp/llama/*.pte"
-      - adb -s $DEVICEFARM_DEVICE_UDID shell "ls -la /data/local/tmp/llama/"
+      - adb -s $DEVICEFARM_DEVICE_UDID shell "mkdir -p /data/local/tmp/minibench/"
+      - adb -s $DEVICEFARM_DEVICE_UDID shell "mv /sdcard/*.bin /data/local/tmp/minibench/"
+      - adb -s $DEVICEFARM_DEVICE_UDID shell "mv /sdcard/*.pte /data/local/tmp/minibench/"
+      - adb -s $DEVICEFARM_DEVICE_UDID shell "chmod 664 /data/local/tmp/minibench/*.bin"
+      - adb -s $DEVICEFARM_DEVICE_UDID shell "chmod 664 /data/local/tmp/minibench/*.pte"
+      - adb -s $DEVICEFARM_DEVICE_UDID shell "ls -la /data/local/tmp/minibench/"
+      - adb -s $DEVICEFARM_DEVICE_UDID shell "run-as org.pytorch.minibench rm -rf files"
 
   test:
     commands:
       # By default, the following ADB command is used by Device Farm to run your Instrumentation test.
       # Please refer to Android's documentation for more options on running instrumentation tests with adb:
       # https://developer.android.com/studio/test/command-line#run-tests-with-adb
+
+      # Run the Instrumentation test for sanity check
       - echo "Starting the Instrumentation test"
       - |
         adb -s $DEVICEFARM_DEVICE_UDID shell "am instrument -r -w --no-window-animation \
@@ -67,17 +70,33 @@ phases:
         fi;
 
       # Run the new generic benchmark activity https://developer.android.com/tools/adb#am
-      - echo "Run LLM benchmark"
+      - echo "Determine model type"
+      - |
+        BIN_FOUND="$(adb -s $DEVICEFARM_DEVICE_UDID shell find /data/local/tmp/minibench/ -name '*.bin')"
+        if [ -z "$BIN_FOUND" ]; then
+          echo "No tokenizer files found in /data/local/tmp/minibench/"
+        else
+          echo "tokenizer files found in /data/local/tmp/minibench/"
+        fi
+
+      - echo "Run benchmark"
       - |
-        adb -s $DEVICEFARM_DEVICE_UDID shell am start -W -n org.pytorch.minibench/.LlmBenchmarkActivity \
-        --es "model_dir" "/data/local/tmp/llama" \
-        --es "tokenizer_path" "/data/local/tmp/llama/tokenizer.bin"
+        adb -s $DEVICEFARM_DEVICE_UDID shell am force-stop org.pytorch.minibench
+        if [ -z "$BIN_FOUND" ]; then
+          adb -s $DEVICEFARM_DEVICE_UDID shell am start -W -n org.pytorch.minibench/.BenchmarkActivity \
+            --es "model_dir" "/data/local/tmp/minibench"
+        else
+          adb -s $DEVICEFARM_DEVICE_UDID shell am start -W -n org.pytorch.minibench/.LlmBenchmarkActivity \
+            --es "model_dir" "/data/local/tmp/minibench" \
+            --es "tokenizer_path" "/data/local/tmp/minibench/tokenizer.bin"
+        fi
+
 
   post_test:
     commands:
-      - echo "Gather LLM benchmark results"
+      - echo "Gather benchmark results"
       - |
-        BENCHMARK_RESULTS=""
+        BENCHMARK_RESULTS=$(adb -s $DEVICEFARM_DEVICE_UDID shell run-as org.pytorch.minibench cat files/benchmark_results.json)
         ATTEMPT=0
         MAX_ATTEMPT=10
         while [ -z "${BENCHMARK_RESULTS}" ] && [ $ATTEMPT -lt $MAX_ATTEMPT ]; do
diff --git a/extension/android/benchmark/app/src/main/java/org/pytorch/minibench/BenchmarkActivity.java b/extension/android/benchmark/app/src/main/java/org/pytorch/minibench/BenchmarkActivity.java
index a79f668f80b..9ede7d69184 100644
--- a/extension/android/benchmark/app/src/main/java/org/pytorch/minibench/BenchmarkActivity.java
+++ b/extension/android/benchmark/app/src/main/java/org/pytorch/minibench/BenchmarkActivity.java
@@ -11,10 +11,14 @@
 import android.app.Activity;
 import android.content.Intent;
 import android.os.Bundle;
+import com.google.gson.Gson;
 import java.io.File;
 import java.io.FileWriter;
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
 import org.pytorch.executorch.Module;
 
 public class BenchmarkActivity extends Activity {
@@ -32,20 +36,39 @@ protected void onCreate(Bundle savedInstanceState) {
     int numIter = intent.getIntExtra("num_iter", 10);
 
     // TODO: Format the string with a parsable format
-    StringBuilder resultText = new StringBuilder();
+    Stats stats = new Stats();
 
     Module module = Module.load(model.getPath());
     for (int i = 0; i < numIter; i++) {
       long start = System.currentTimeMillis();
       module.forward();
       long forwardMs = System.currentTimeMillis() - start;
-      resultText.append(forwardMs).append(";");
+      stats.latency.add(forwardMs);
     }
 
+    // TODO (huydhn): Remove txt files here once the JSON format is ready
     try (FileWriter writer = new FileWriter(getFilesDir() + "/benchmark_results.txt")) {
-      writer.write(resultText.toString());
+      writer.write(stats.toString());
     } catch (IOException e) {
       e.printStackTrace();
     }
+
+    // TODO (huydhn): Figure out on what the final JSON results looks like, we need something
+    // with the same number of fields as https://github.com/pytorch/pytorch/pull/135042
+    try (FileWriter writer = new FileWriter(getFilesDir() + "/benchmark_results.json")) {
+      Gson gson = new Gson();
+      writer.write(gson.toJson(stats));
+    } catch (IOException e) {
+      e.printStackTrace();
+    }
+  }
+}
+
+class Stats {
+  List<Long> latency = new ArrayList<>();
+
+  @Override
+  public String toString() {
+    return "latency: " + latency.stream().map(Object::toString).collect(Collectors.joining(""));
   }
 }