diff --git a/services/gpuservice/gpumem/GpuMem.cpp b/services/gpuservice/gpumem/GpuMem.cpp
index 245edb8..c78322e 100644
--- a/services/gpuservice/gpumem/GpuMem.cpp
+++ b/services/gpuservice/gpumem/GpuMem.cpp
@@ -25,6 +25,7 @@
 #include <libbpf_android.h>
 #include <log/log.h>
 #include <unistd.h>
+#include <utils/Timers.h>
 #include <utils/Trace.h>
 
 #include <unordered_map>
@@ -128,4 +129,24 @@
     }
 }
 
+void GpuMem::traverseGpuMemTotals(const std::function<void(int64_t ts, uint32_t gpuId, uint32_t pid,
+                                                           uint64_t size)>& callback) {
+    auto res = mGpuMemTotalMap.getFirstKey();
+    if (!res.ok()) return;
+    uint64_t key = res.value();
+    while (true) {
+        uint32_t gpu_id = key >> 32;
+        uint32_t pid = key;
+
+        res = mGpuMemTotalMap.readValue(key);
+        if (!res.ok()) break;
+        uint64_t size = res.value();
+
+        callback(systemTime(), gpu_id, pid, size);
+        res = mGpuMemTotalMap.getNextKey(key);
+        if (!res.ok()) break;
+        key = res.value();
+    }
+}
+
 } // namespace android
diff --git a/services/gpuservice/gpumem/include/gpumem/GpuMem.h b/services/gpuservice/gpumem/include/gpumem/GpuMem.h
index 49a9f95..de691e2 100644
--- a/services/gpuservice/gpumem/include/gpumem/GpuMem.h
+++ b/services/gpuservice/gpumem/include/gpumem/GpuMem.h
@@ -20,6 +20,8 @@
 #include <utils/String16.h>
 #include <utils/Vector.h>
 
+#include <functional>
+
 namespace android {
 
 class GpuMem {
@@ -33,27 +35,9 @@
     void dump(const Vector<String16>& args, std::string* result);
     bool isInitialized() { return mInitialized.load(); }
 
-    // Traverse the map and send each value read back to the callback function.
-    // Used for tracing.
-    template <typename lambda>
-    void traceGpuMemTotals(lambda tracerCallback) {
-        auto res = mGpuMemTotalMap.getFirstKey();
-        if (!res.ok()) return;
-        uint64_t key = res.value();
-        while (true) {
-            uint32_t gpu_id = key >> 32;
-            uint32_t pid = key;
-
-            res = mGpuMemTotalMap.readValue(key);
-            if (!res.ok()) break;
-            uint64_t size = res.value();
-
-            tracerCallback(gpu_id, pid, size);
-            res = mGpuMemTotalMap.getNextKey(key);
-            if (!res.ok()) break;
-            key = res.value();
-        }
-    }
+    // Traverse the gpu memory total map to feed the callback function.
+    void traverseGpuMemTotals(const std::function<void(int64_t ts, uint32_t gpuId, uint32_t pid,
+                                                       uint64_t size)>& callback);
 
 private:
     // Friend class for testing.
diff --git a/services/gpuservice/tests/unittests/GpuMemTest.cpp b/services/gpuservice/tests/unittests/GpuMemTest.cpp
index abaf30a..45e8367 100644
--- a/services/gpuservice/tests/unittests/GpuMemTest.cpp
+++ b/services/gpuservice/tests/unittests/GpuMemTest.cpp
@@ -41,6 +41,8 @@
 constexpr uint64_t TEST_PROC_VAL_1 = 234;
 constexpr uint64_t TEST_PROC_KEY_2 = 4294967298; // (1 << 32) + 2
 constexpr uint64_t TEST_PROC_VAL_2 = 345;
+constexpr uint32_t TEST_KEY_MASK = 0x1 | 0x2 | 0x4;
+constexpr uint32_t TEST_KEY_COUNT = 3;
 
 class GpuMemTest : public testing::Test {
 public:
@@ -143,5 +145,37 @@
                                        TEST_PROC_VAL_2)));
 }
 
+TEST_F(GpuMemTest, traverseGpuMemTotals) {
+    SKIP_IF_BPF_NOT_SUPPORTED;
+    ASSERT_RESULT_OK(mTestMap.writeValue(TEST_GLOBAL_KEY, TEST_GLOBAL_VAL, BPF_ANY));
+    ASSERT_RESULT_OK(mTestMap.writeValue(TEST_PROC_KEY_1, TEST_PROC_VAL_1, BPF_ANY));
+    ASSERT_RESULT_OK(mTestMap.writeValue(TEST_PROC_KEY_2, TEST_PROC_VAL_2, BPF_ANY));
+    mTestableGpuMem.setGpuMemTotalMap(mTestMap);
+
+    static uint32_t sMask = 0;
+    static uint32_t sCount = 0;
+    mGpuMem->traverseGpuMemTotals([](int64_t, uint32_t gpuId, uint32_t pid, uint64_t size) {
+        const uint64_t key = ((uint64_t)gpuId << 32) | pid;
+        switch (key) {
+            case TEST_GLOBAL_KEY:
+                EXPECT_EQ(size, TEST_GLOBAL_VAL);
+                sMask |= 0x1;
+                break;
+            case TEST_PROC_KEY_1:
+                EXPECT_EQ(size, TEST_PROC_VAL_1);
+                sMask |= 0x2;
+                break;
+            case TEST_PROC_KEY_2:
+                EXPECT_EQ(size, TEST_PROC_VAL_2);
+                sMask |= 0x4;
+                break;
+        }
+        sCount++;
+    });
+
+    EXPECT_EQ(sMask, TEST_KEY_MASK);
+    EXPECT_EQ(sCount, TEST_KEY_COUNT);
+}
+
 } // namespace
 } // namespace android
diff --git a/services/gpuservice/tracing/GpuMemTracer.cpp b/services/gpuservice/tracing/GpuMemTracer.cpp
index 9e01a64..6366e1d 100644
--- a/services/gpuservice/tracing/GpuMemTracer.cpp
+++ b/services/gpuservice/tracing/GpuMemTracer.cpp
@@ -23,9 +23,7 @@
 #include <gpumem/GpuMem.h>
 #include <perfetto/trace/android/gpu_mem_event.pbzero.h>
 #include <unistd.h>
-#include <utils/Timers.h>
 
-#include <algorithm>
 #include <thread>
 
 PERFETTO_DEFINE_DATA_SOURCE_STATIC_MEMBERS(android::GpuMemTracer::GpuMemDataSource);
@@ -79,10 +77,10 @@
         ALOGE("Cannot trace without GpuMem initialization");
         return;
     }
-    mGpuMem->traceGpuMemTotals([](uint32_t gpuId, uint32_t pid, uint64_t size) {
+    mGpuMem->traverseGpuMemTotals([](int64_t ts, uint32_t gpuId, uint32_t pid, uint64_t size) {
         GpuMemDataSource::Trace([&](GpuMemDataSource::TraceContext ctx) {
             auto packet = ctx.NewTracePacket();
-            packet->set_timestamp(systemTime());
+            packet->set_timestamp(ts);
             auto* event = packet->set_gpu_mem_total_event();
             event->set_gpu_id(gpuId);
             event->set_pid(pid);