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-trace.h

@@ -0,0 +1,83 @@
+// Event tracer for the RPC backend and llama.cpp.
+//
+// The tracer is off unless the environment variable GGML_RPC_TRACE is set to a file path (the
+// rpc-server also accepts --trace <path>). When it is off the only cost at a call site is one
+// load and one branch on ggml_trace_flag; no clock is read and nothing is formatted.
+//
+// Every process writes JSON lines to its own file: one header object, then one object per event.
+// scripts/rpc_trace/merge.py aligns the files of the two nodes with the clock offset measured by
+// the client at connect time and emits a Chrome trace plus a per decode step summary.
+
+#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 directly at the call sites, do not call a function:
+    //
+    //     const int64_t t0 = ggml_trace_flag ? ggml_time_us() : 0;
+    //
+    GGML_API int ggml_trace_flag;
+
+    // Opens the trace file. `path` may be NULL, then GGML_RPC_TRACE is used. `role` names the
+    // process in the merged trace ("client", "rpc-server"). Safe to call more than once, the
+    // first call with a usable path wins. Returns 1 if tracing is on afterwards.
+    GGML_API int  ggml_trace_open(const char * path, const char * role);
+    GGML_API void ggml_trace_close(void);
+
+    // Monotonic microseconds in the same clock as every timestamp written to the trace.
+    GGML_API int64_t ggml_trace_time_us(void);
+
+    // Stable small integer for the calling thread, used as the Chrome trace tid.
+    GGML_API int ggml_trace_tid(void);
+
+    // llama-server tags the threads that drive one pipeline group so that every event raised
+    // underneath, down to the individual RPC command, carries the group id. -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 of this thread belong to. The pointer must
+    // stay valid until it is replaced; the RPC backend passes tensor->name.
+    GGML_API void ggml_trace_set_subject(const char * name, uint64_t uid);
+
+    // Writes one event. `phase` is a free form category ("rpc.client", "sched", ...), `name` the
+    // event name. t0/t1 are microseconds from ggml_trace_time_us(); t1 == t0 makes it an instant.
+    // `fields` is appended verbatim inside the JSON object and may be NULL, for example
+    //     "\"bytes\":128,\"cmd\":\"GRAPH_COMPUTE\""
+    GGML_API void ggml_trace_event(const char * phase, const char * name,
+                                   int64_t t0, int64_t t1, const char * fields);
+
+    // Same, with printf formatting for the extra fields.
+    GGML_API void ggml_trace_eventf(const char * phase, const char * name,
+                                    int64_t t0, int64_t t1, const char * fmt, ...);
+
+    // Records the result of the four timestamp exchange with the peer into the header of this
+    // trace file, so the merge tool can put both nodes on one time line.
+    //   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);
+
+    // GPU spans.
+    //
+    // A backend submit only queues the work, so host timestamps around it say nothing about when
+    // the kernels ran. If the backend exposes the timing hooks through its registry (CUDA does),
+    // these bracket the submit with two events on the compute stream and report them later on the
+    // host monotonic scale. ggml_trace_gpu_begin returns a tag, 0 if the backend cannot do it.
+    // Nothing here ever waits on the GPU: the spans are emitted by ggml_trace_gpu_flush once both
+    // of their events have completed, which the caller does from a point where it is idle anyway.
+    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);
+
+    // Escapes a string for the JSON output. Returns `dst`.
+    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"
 
@@ -491,11 +492,26 @@ void ggml_backend_tensor_copy(const struct ggml_tensor * src, struct ggml_tensor
 #ifndef NDEBUG
         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
+        // The staging path taken between two backends that cannot copy to each other directly.
+        // For a layer split over RPC this is the hidden state going GPU -> host -> peer, so the
+        // trace breaks it into the host allocation, the read back and the send.
         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 +1624,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 +1760,12 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
             }
         }
 
+        if (ggml_trace_flag) { t_inputs = ggml_trace_time_us(); }
+
+        // GPU span around the submit of this split, so the trace shows when the kernels really
+        // ran and not only when the launch returned
+        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 +1805,26 @@ 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 on most backends, so t1 is when the work was
+            //       queued, not when it finished; the GPU rows in the merged trace say that
+            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 +2020,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, so this is where the completed GPU spans are collected
+        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,117 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t
     GGML_UNUSED(reg);
 }
 
+// -----------------------------------------------------------------------------
+// GPU timing marks for the event tracer (see ggml/include/ggml-trace.h)
+//
+// The RPC server needs to know when the kernels of a graph really ran on the GPU, not when the
+// launch returned, so it brackets ggml_backend_graph_compute with two CUDA events recorded on
+// the backend's compute stream. The events are resolved later, without blocking, against one
+// anchor event whose completion time on the host was measured once, which puts the GPU marks on
+// the same monotonic microsecond scale as every other event in the trace.
+//
+// Reached through ggml_backend_reg_get_proc_address, so the RPC backend does not have to link
+// against CUDA.
+// -----------------------------------------------------------------------------
+
+struct ggml_cuda_trace_mark {
+    uint64_t    tag;
+    int         kind;
+    cudaEvent_t event;
+};
+
+struct ggml_cuda_trace_state {
+    std::mutex                        mutex;
+    std::vector<ggml_cuda_trace_mark> pending;
+    std::vector<cudaEvent_t>          spare;
+    cudaEvent_t                       anchor    = nullptr;
+    int64_t                           anchor_us = 0;
+    int                               device    = -1;
+};
+
+static ggml_cuda_trace_state & ggml_cuda_trace() {
+    static ggml_cuda_trace_state state;
+    return state;
+}
+
+// records a mark on the compute stream of `backend`; 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);
+
+    if (st.anchor == nullptr) {
+        ggml_cuda_set_device(cuda_ctx->device);
+        if (cudaEventCreate(&st.anchor) != cudaSuccess) {
+            st.anchor = nullptr;
+            return;
+        }
+        st.device = cuda_ctx->device;
+        // the host time of the instant the anchor completed on the GPU; every later mark is
+        // reported as anchor_us + elapsed(anchor, mark)
+        cudaEventRecord(st.anchor, cuda_ctx->stream());
+        cudaEventSynchronize(st.anchor);
+        st.anchor_us = ggml_time_us();
+    }
+
+    // elapsed time is only defined between events of the same device
+    if (cuda_ctx->device != st.device) {
+        return;
+    }
+
+    cudaEvent_t event = nullptr;
+    if (!st.spare.empty()) {
+        event = st.spare.back();
+        st.spare.pop_back();
+    } else {
+        if (cudaEventCreate(&event) != cudaSuccess) {
+            return;
+        }
+    }
+
+    if (cudaEventRecord(event, cuda_ctx->stream()) != cudaSuccess) {
+        st.spare.push_back(event);
+        return;
+    }
+
+    st.pending.push_back({ tag, kind, event });
+}
+
+// collects the marks whose events have completed, without waiting for any of them. Returns the
+// number written; the caller loops until it gets less than `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);
+    if (st.anchor == nullptr) {
+        return 0;
+    }
+
+    int n = 0;
+    size_t keep = 0;
+    for (size_t i = 0; i < st.pending.size(); i++) {
+        ggml_cuda_trace_mark & mark = st.pending[i];
+        if (n < max && cudaEventQuery(mark.event) == cudaSuccess) {
+            float ms = 0.0f;
+            if (cudaEventElapsedTime(&ms, st.anchor, mark.event) == cudaSuccess) {
+                tags [n] = mark.tag;
+                kinds[n] = mark.kind;
+                t_us [n] = st.anchor_us + (int64_t)(ms * 1000.0f);
+                n++;
+            }
+            st.spare.push_back(mark.event);
+        } else {
+            st.pending[keep++] = mark;
+        }
+    }
+    st.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 +5603,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