diff --git a/extension/llm/runner/targets.bzl b/extension/llm/runner/targets.bzl
index 8ae0c1ba3be..0c5ad4170e5 100644
--- a/extension/llm/runner/targets.bzl
+++ b/extension/llm/runner/targets.bzl
@@ -11,17 +11,6 @@ def define_common_targets():
         ],
     )
 
-    runtime.cxx_library(
-        name = "stats",
-        exported_headers = [
-            "stats.h",
-            "util.h",
-        ],
-        visibility = [
-            "@EXECUTORCH_CLIENTS",
-        ],
-    )
-
     runtime.cxx_library(
         name = "constants",
         exported_headers = [
@@ -35,6 +24,21 @@ def define_common_targets():
     for aten in get_aten_mode_options():
         aten_suffix = "_aten" if aten else ""
 
+        runtime.cxx_library(
+            name = "stats" + aten_suffix,
+            exported_headers = [
+                "stats.h",
+                "util.h",
+            ],
+            visibility = [
+                "@EXECUTORCH_CLIENTS",
+            ],
+            exported_deps = [
+                ":constants",
+                 "//executorch/extension/module:module" + aten_suffix,
+            ],
+        )
+
         runtime.cxx_library(
             name = "text_decoder_runner" + aten_suffix,
             exported_headers = ["text_decoder_runner.h"],
@@ -43,7 +47,7 @@ def define_common_targets():
                 "@EXECUTORCH_CLIENTS",
             ],
             exported_deps = [
-                ":stats",
+                ":stats" + aten_suffix,
                 "//executorch/kernels/portable/cpu/util:arange_util" + aten_suffix,
                 "//executorch/extension/llm/sampler:sampler" + aten_suffix,
                 "//executorch/extension/llm/runner/io_manager:io_manager" + aten_suffix,
diff --git a/extension/llm/runner/util.h b/extension/llm/runner/util.h
index 5ede889e034..513fd109255 100644
--- a/extension/llm/runner/util.h
+++ b/extension/llm/runner/util.h
@@ -7,11 +7,14 @@
  */
 
 #pragma once
+#include <executorch/extension/llm/runner/constants.h>
+#include <executorch/extension/module/module.h>
 #include <executorch/extension/tensor/tensor.h>
 #include <executorch/runtime/platform/compiler.h>
 #include <stdio.h>
 #include <time.h>
 #include <cctype>
+#include <vector>
 #if defined(__linux__) || defined(__ANDROID__) || defined(__unix__)
 #include <sys/resource.h>
 #endif