trace: event tracer for the RPC backend, rpc-server and llama-server

ggml/CMakeLists.txt

@@ -329,6 +329,7 @@ set(GGML_PUBLIC_HEADERS
     include/ggml-blas.h
     include/ggml-cann.h
     include/ggml-cpp.h
+    include/ggml-trace.h
     include/ggml-cuda.h
     include/ggml-opt.h
     include/ggml-metal.h

ggml/include/ggml-rpc.h

@@ -7,7 +7,7 @@ extern "C" {
 #endif
 
 #define RPC_PROTO_MAJOR_VERSION    5
-#define RPC_PROTO_MINOR_VERSION    1
+#define RPC_PROTO_MINOR_VERSION    2
 #define RPC_PROTO_PATCH_VERSION    0
 
 #ifdef  __cplusplus

ggml/include/ggml-trace.h

@@ -0,0 +1,53 @@
+// Event tracer for the RPC backend and llama.cpp. Off unless GGML_RPC_TRACE (or rpc-server
+// --trace) names a file; scripts/rpc_trace/merge.py aligns the JSON lines each process writes.
+
+#pragma once
+
+#include "ggml.h"
+#include "ggml-backend.h"
+
+#include <stdarg.h>
+#include <stdint.h>
+
+#ifdef  __cplusplus
+extern "C" {
+#endif
+
+    // 1 while a trace file is open; read it at the call sites so a disabled tracer is one branch
+    GGML_API int ggml_trace_flag;
+
+    // `path` NULL means GGML_RPC_TRACE; the first call with a usable path wins
+    GGML_API int  ggml_trace_open(const char * path, const char * role);
+    GGML_API void ggml_trace_close(void);
+
+    GGML_API int64_t ggml_trace_time_us(void);
+
+    GGML_API int ggml_trace_tid(void);
+
+    // tags every event this thread raises with a pipeline group, -1 means no group
+    GGML_API void ggml_trace_set_group(int group);
+    GGML_API int  ggml_trace_get_group(void);
+
+    // names the tensor or graph the next RPC commands belong to; `name` must outlive the call
+    GGML_API void ggml_trace_set_subject(const char * name, uint64_t uid);
+
+    // t1 == t0 is an instant; `fields` may be NULL and is inlined verbatim into the JSON object
+    GGML_API void ggml_trace_event(const char * phase, const char * name,
+                                   int64_t t0, int64_t t1, const char * fields);
+
+    GGML_API void ggml_trace_eventf(const char * phase, const char * name,
+                                    int64_t t0, int64_t t1, const char * fmt, ...);
+
+    // t1 client sends, t2 peer receives, t3 peer replies, t4 client receives (microseconds).
+    GGML_API void ggml_trace_clock_offset(const char * peer, int64_t t1, int64_t t2, int64_t t3, int64_t t4);
+
+    // bracket a submit with compute-stream events; begin returns 0 without hooks, none ever wait
+    GGML_API uint64_t ggml_trace_gpu_begin(ggml_backend_t backend, const char * name);
+    GGML_API void     ggml_trace_gpu_end  (ggml_backend_t backend, uint64_t tag);
+    GGML_API void     ggml_trace_gpu_flush(void);
+
+    GGML_API const char * ggml_trace_escape(char * dst, size_t dst_size, const char * src);
+
+#ifdef  __cplusplus
+}
+#endif

ggml/src/CMakeLists.txt

@@ -195,12 +195,14 @@ add_library(ggml-base
             ../include/ggml-backend.h
             ../include/ggml-cpp.h
             ../include/ggml-opt.h
+            ../include/ggml-trace.h
             ../include/gguf.h
             ggml.c
             ggml.cpp
             ggml-alloc.c
             ggml-backend.cpp
             ggml-backend-meta.cpp
+            ggml-trace.cpp
             ggml-opt.cpp
             ggml-threading.cpp
             ggml-threading.h

ggml/src/ggml-backend.cpp

@@ -10,6 +10,7 @@
 
 #include "ggml-backend.h"
 #include "ggml-backend-impl.h"
+#include "ggml-trace.h"
 #include "ggml-alloc.h"
 #include "ggml-impl.h"
 
@@ -492,10 +493,22 @@ void ggml_backend_tensor_copy(const struct ggml_tensor * src, struct ggml_tensor
         GGML_LOG_DEBUG("%s: warning: slow copy from %s to %s\n", __func__, ggml_backend_buffer_name(src->buffer), ggml_backend_buffer_name(dst->buffer));
 #endif // NDEBUG
         size_t nbytes = ggml_nbytes(src);
+        const int64_t t0 = ggml_trace_flag ? ggml_trace_time_us() : 0;
         void * data = malloc(nbytes);
+        const int64_t t1 = ggml_trace_flag ? ggml_trace_time_us() : 0;
         ggml_backend_tensor_get(src, data, 0, nbytes);
+        const int64_t t2 = ggml_trace_flag ? ggml_trace_time_us() : 0;
         ggml_backend_tensor_set(dst, data, 0, nbytes);
+        const int64_t t3 = ggml_trace_flag ? ggml_trace_time_us() : 0;
         free(data);
+        if (ggml_trace_flag) {
+            ggml_trace_eventf("sched", "copy_stage", t0, ggml_trace_time_us(),
+                              "\"tensor\":\"%s\",\"bytes\":%zu,\"src\":\"%s\",\"dst\":\"%s\","
+                              "\"malloc_us\":%lld,\"get_us\":%lld,\"set_us\":%lld",
+                              src->name, nbytes,
+                              ggml_backend_buffer_name(src->buffer), ggml_backend_buffer_name(dst->buffer),
+                              (long long) (t1 - t0), (long long) (t2 - t1), (long long) (t3 - t2));
+        }
     }
 }
 
@@ -1608,6 +1621,9 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
         int split_backend_id = split->backend_id;
         ggml_backend_t split_backend = sched->backends[split_backend_id];
 
+        const int64_t t_split0 = ggml_trace_flag ? ggml_trace_time_us() : 0;
+        int64_t       t_inputs = t_split0;
+
         // ensure the previous split's async work has completed before we start
         // this split, the allocator may have reused buffer regions across splits
         if (split->n_inputs == 0 && prev_backend_id >= 0 && prev_backend_id != split_backend_id) {
@@ -1741,6 +1757,10 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
             }
         }
 
+        if (ggml_trace_flag) { t_inputs = ggml_trace_time_us(); }
+
+        const uint64_t gpu_tag = ggml_trace_flag ? ggml_trace_gpu_begin(split_backend, ggml_backend_name(split_backend)) : 0;
+
         if (!sched->callback_eval) {
             enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph);
             if (ec != GGML_STATUS_SUCCESS) {
@@ -1780,11 +1800,25 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
             }
         }
 
+        if (ggml_trace_flag) {
+            ggml_trace_gpu_end(split_backend, gpu_tag);
+        }
+
         // record the event of this split
         if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
             ggml_backend_event_record(sched->events[split_backend_id][sched->cur_copy], split_backend);
         }
 
+        if (ggml_trace_flag) {
+            // note: the submit is asynchronous, so t1 is when the work was queued, not finished
+            ggml_trace_eventf("sched", "split", t_split0, ggml_trace_time_us(),
+                              "\"split\":%d,\"n_splits\":%d,\"backend\":\"%s\",\"n_inputs\":%d,"
+                              "\"n_nodes\":%d,\"inputs_us\":%lld,\"gpu_tag\":%llu",
+                              split_id, sched->n_splits, ggml_backend_name(split_backend),
+                              split->n_inputs, split->graph.n_nodes,
+                              (long long) (t_inputs - t_split0), (unsigned long long) gpu_tag);
+        }
+
         prev_backend_id = split_backend_id;
     }
 
@@ -1980,6 +2014,10 @@ void ggml_backend_sched_synchronize(ggml_backend_sched_t sched) {
     for (int i = 0; i < sched->n_backends; i++) {
         ggml_backend_synchronize(sched->backends[i]);
     }
+    if (ggml_trace_flag) {
+        // everything is idle here, the only safe point to collect completed GPU spans
+        ggml_trace_gpu_flush();
+    }
     if (!sched->is_alloc) {
         // if the graph is not already allocated, always use copy 0 after a synchronization
         // this ensures that during generation the same copy is used every time,

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -5472,6 +5472,173 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t
     GGML_UNUSED(reg);
 }
 
+// GPU timing marks for the event tracer: CUDA events on the compute stream, resolved against one
+// anchor event, reached through ggml_backend_reg_get_proc_address so no caller links CUDA itself.
+
+struct ggml_cuda_trace_mark {
+    uint64_t    tag;
+    int         kind;
+    cudaEvent_t event;
+};
+
+// elapsed time is only defined between events of the same device, so each device needs its own
+// anchor and its own queues. A single shared anchor bound to whichever device marked first made
+// every mark on the other devices undeliverable, and the tracer still handed out tags for them,
+// so a multi-GPU trace silently lost all work outside that one device.
+struct ggml_cuda_trace_device {
+    std::vector<ggml_cuda_trace_mark> pending;
+    std::vector<cudaEvent_t>          spare;
+    cudaEvent_t                       anchor    = nullptr;
+    int64_t                           anchor_us = 0;  // 0 until the anchor's wall clock is known
+    // the stream the anchor was recorded on, so poll can ask whether it is idle
+    cudaStream_t                      stream    = nullptr;
+    int                               device    = 0;
+    bool                              anchor_fixed = false;
+};
+
+struct ggml_cuda_trace_state {
+    std::mutex                             mutex;
+    std::map<int, ggml_cuda_trace_device>  devs;
+};
+
+static ggml_cuda_trace_state & ggml_cuda_trace() {
+    static ggml_cuda_trace_state state;
+    return state;
+}
+
+// kind 0 = start of a span, 1 = end
+extern "C" GGML_BACKEND_API void ggml_backend_cuda_trace_mark(ggml_backend_t backend, uint64_t tag, int kind);
+extern "C" void ggml_backend_cuda_trace_mark(ggml_backend_t backend, uint64_t tag, int kind) {
+    ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;
+    ggml_cuda_trace_state & st = ggml_cuda_trace();
+
+    std::lock_guard<std::mutex> lock(st.mutex);
+
+    ggml_cuda_trace_device & d = st.devs[cuda_ctx->device];
+
+    if (d.anchor == nullptr) {
+        ggml_cuda_set_device(cuda_ctx->device);
+        if (cudaEventCreate(&d.anchor) != cudaSuccess) {
+            d.anchor = nullptr;
+            return;
+        }
+        // every later mark is reported as anchor_us + elapsed(anchor, mark). The anchor is only
+        // recorded here, never waited on: synchronizing would drain whatever the scheduler has
+        // already queued on this stream, so switching the tracer on would change the execution it
+        // is supposed to observe. Its wall clock is established in poll, which ties it to real time
+        // through a probe event once the stream is idle, so waiting drains nothing. See there.
+        cudaEventRecord(d.anchor, cuda_ctx->stream());
+        d.stream = cuda_ctx->stream();
+        d.device = cuda_ctx->device;
+    }
+
+    cudaEvent_t event = nullptr;
+    if (!d.spare.empty()) {
+        event = d.spare.back();
+        d.spare.pop_back();
+    } else {
+        if (cudaEventCreate(&event) != cudaSuccess) {
+            return;
+        }
+    }
+
+    if (cudaEventRecord(event, cuda_ctx->stream()) != cudaSuccess) {
+        d.spare.push_back(event);
+        return;
+    }
+
+    d.pending.push_back({ tag, kind, event });
+}
+
+// never waits: returns the marks already completed, so the caller loops until it gets < `max`.
+extern "C" GGML_BACKEND_API int ggml_backend_cuda_trace_poll(uint64_t * tags, int * kinds, int64_t * t_us, int max);
+extern "C" int ggml_backend_cuda_trace_poll(uint64_t * tags, int * kinds, int64_t * t_us, int max) {
+    ggml_cuda_trace_state & st = ggml_cuda_trace();
+
+    std::lock_guard<std::mutex> lock(st.mutex);
+
+    int n = 0;
+    for (auto & entry : st.devs) {
+        ggml_cuda_trace_device & d = entry.second;
+        if (d.anchor == nullptr) {
+            continue;
+        }
+        // the anchor is recorded before any mark on this stream, so it always completes first.
+        // Until it has, its wall clock is unknown and the marks simply stay pending.
+        if (!d.anchor_fixed) {
+            if (cudaEventQuery(d.anchor) != cudaSuccess) {
+                continue;
+            }
+
+            // Taking ggml_time_us() here dates the anchor to this poll rather than to when it
+            // actually completed, and every mark is reported as anchor_us + elapsed(anchor, mark),
+            // so the whole device timeline shifts forward by however long the anchor had already
+            // been complete. In the RPC server that delay is a full graph, because
+            // ggml_backend_graph_compute() synchronizes before the serve loop polls again, which
+            // is exactly the case cross-device overlap and idle attribution are computed from.
+            //
+            // Tie GPU time to wall time properly instead: record a probe, wait for it, and measure
+            // back to the anchor. Synchronizing is only safe when the stream is already idle, since
+            // otherwise it would drain queued work and change the execution being observed, which
+            // is why the anchor itself is never waited on. When the stream is idle the probe
+            // completes immediately, so the wait returns at its completion and costs nothing. That
+            // is the normal state at poll time in the serve loop.
+            if (cudaStreamQuery(d.stream) == cudaSuccess) {
+                ggml_cuda_set_device(d.device);
+
+                cudaEvent_t probe = nullptr;
+                if (!d.spare.empty()) {
+                    probe = d.spare.back();
+                    d.spare.pop_back();
+                } else if (cudaEventCreate(&probe) != cudaSuccess) {
+                    probe = nullptr;
+                }
+
+                if (probe != nullptr) {
+                    float ms = 0.0f;
+                    if (cudaEventRecord(probe, d.stream) == cudaSuccess &&
+                        cudaEventSynchronize(probe) == cudaSuccess &&
+                        cudaEventElapsedTime(&ms, d.anchor, probe) == cudaSuccess) {
+                        d.anchor_us    = ggml_time_us() - (int64_t)(ms * 1000.0f);
+                        d.anchor_fixed = true;
+                    }
+                    d.spare.push_back(probe);
+                }
+            }
+
+            if (!d.anchor_fixed) {
+                // Stream busy, or the probe failed. Fall back to the previous approximation rather
+                // than stalling the marks. Freeze it either way: refining the anchor on a later
+                // poll would move marks reported after the change relative to marks reported
+                // before it, putting a step in the middle of one device's timeline, which is
+                // harder to reason about than a consistent offset.
+                d.anchor_us    = ggml_time_us();
+                d.anchor_fixed = true;
+            }
+        }
+
+        size_t keep = 0;
+        for (size_t i = 0; i < d.pending.size(); i++) {
+            ggml_cuda_trace_mark & mark = d.pending[i];
+            if (n < max && cudaEventQuery(mark.event) == cudaSuccess) {
+                float ms = 0.0f;
+                if (cudaEventElapsedTime(&ms, d.anchor, mark.event) == cudaSuccess) {
+                    tags [n] = mark.tag;
+                    kinds[n] = mark.kind;
+                    t_us [n] = d.anchor_us + (int64_t)(ms * 1000.0f);
+                    n++;
+                }
+                d.spare.push_back(mark.event);
+            } else {
+                d.pending[keep++] = mark;
+            }
+        }
+        d.pending.resize(keep);
+    }
+
+    return n;
+}
+
 static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) {
     GGML_UNUSED(reg);
     if (strcmp(name, "ggml_backend_comm_init") == 0) {
@@ -5492,6 +5659,12 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con
     if (strcmp(name, "ggml_backend_get_features") == 0) {
         return (void *)ggml_backend_cuda_get_features;
     }
+    if (strcmp(name, "ggml_backend_cuda_trace_mark") == 0) {
+        return (void *)ggml_backend_cuda_trace_mark;
+    }
+    if (strcmp(name, "ggml_backend_cuda_trace_poll") == 0) {
+        return (void *)ggml_backend_cuda_trace_poll;
+    }
     return nullptr;
 }
 

ggml/src/ggml-cuda/vendors/hip.h

@@ -60,7 +60,10 @@
 #define cudaErrorMemoryAllocation hipErrorOutOfMemory
 #define cudaErrorPeerAccessAlreadyEnabled hipErrorPeerAccessAlreadyEnabled
 #define cudaErrorPeerAccessNotEnabled hipErrorPeerAccessNotEnabled
+#define cudaEventCreate hipEventCreate
 #define cudaEventCreateWithFlags hipEventCreateWithFlags
+#define cudaEventElapsedTime hipEventElapsedTime
+#define cudaEventQuery hipEventQuery
 #define cudaEventDisableTiming hipEventDisableTiming
 #define cudaEventRecord hipEventRecord
 #define cudaEventSynchronize hipEventSynchronize

ggml/src/ggml-cuda/vendors/musa.h

@@ -48,7 +48,10 @@
 #define cudaErrorMemoryAllocation musaErrorMemoryAllocation
 #define cudaErrorPeerAccessAlreadyEnabled musaErrorPeerAccessAlreadyEnabled
 #define cudaErrorPeerAccessNotEnabled musaErrorPeerAccessNotEnabled
+#define cudaEventCreate musaEventCreate
 #define cudaEventCreateWithFlags musaEventCreateWithFlags
+#define cudaEventElapsedTime musaEventElapsedTime
+#define cudaEventQuery musaEventQuery
 #define cudaEventDisableTiming musaEventDisableTiming
 #define cudaEventRecord musaEventRecord
 #define cudaEventSynchronize musaEventSynchronize