From d9b6648b0c9e3e2290b98e1d1e4eea031b5c43e6 Mon Sep 17 00:00:00 2001
From: Tommy Reilly <gnurizen@gmail.com>
Date: Tue, 30 Jun 2026 07:01:14 -0400
Subject: [PATCH] Add separate rate limit for graph launches (default 20/s)

Graph launches previously bypassed the callback rate limiter entirely, so a
high graph-replay rate could flood the cupti_events ringbuf: each graph launch
fans out to many kernel activity records, and graphs dominate GPU event volume
on graph-heavy workloads (e.g. TRT-LLM decode).

Give graph launches their own token bucket (default 20/sec,
PARCAGPU_GRAPH_RATE_LIMIT) instead of the bypass. This preserves the original
intent -- eager bursts can't starve graph sampling, since the budgets are now
separate -- while bounding graph bursts. A rate-limited graph launch is never
inserted into the graph correlation map, so its kernel activities are filtered
before they reach the ringbuf, throttling the dominant traffic at the source.

Apply both configured rates (eager and graph) from process-global atomics that
each thread copies into its thread_local bucket on first use. Previously
init_debug called setRate on the thread_local bucket directly, which only
affected the single thread that ran init_debug -- so PARCAGPU_RATE_LIMIT (and
the new graph knob) silently had no effect on other launching threads.
---
 src/cupti.cpp | 54 ++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 41 insertions(+), 13 deletions(-)

diff --git a/src/cupti.cpp b/src/cupti.cpp
index 1e1dd8d..a2c7f87 100644
--- a/src/cupti.cpp
+++ b/src/cupti.cpp
@@ -37,9 +37,27 @@ static std::atomic<uint32_t> g_bufferCycle{0};
 // runtime calls)
 thread_local uint32_t runtimeEnterCorrelationId = 0;
 
-// Thread-local rate limiter for callback probes (default 100/sec,
-// configurable via PARCAGPU_RATE_LIMIT).
+// Configured probe rates (tokens/sec), read from env once in init_debug.
+// Globals because the buckets below are thread_local: setRate from init_debug
+// would only reach the single thread that ran it.
+std::atomic<double> g_eagerRateLimit{100.0};
+std::atomic<double> g_graphRateLimit{20.0};
+
+// Thread-local rate limiters. Graph launches get a separate, lower bucket
+// (PARCAGPU_GRAPH_RATE_LIMIT) so graph replays don't flood cupti_events and
+// eager bursts can't starve graph sampling.
 thread_local TokenBucket callbackLimiter(100.0);
+thread_local TokenBucket graphCallbackLimiter(20.0);
+
+// Copy the configured rates into this thread's buckets, once per thread.
+static inline void applyRateLimitsOnce() {
+  thread_local bool applied = false;
+  if (!applied) {
+    callbackLimiter.setRate(g_eagerRateLimit.load(std::memory_order_relaxed));
+    graphCallbackLimiter.setRate(g_graphRateLimit.load(std::memory_order_relaxed));
+    applied = true;
+  }
+}
 
 // ---------------------------------------------------------------------------
 // PC sampling probabilistic control.
@@ -179,7 +197,15 @@ void init_debug() {
     if (rateEnv != nullptr) {
       double rate = atof(rateEnv);
       if (rate > 0) {
-        callbackLimiter.setRate(rate);
+        g_eagerRateLimit.store(rate, std::memory_order_relaxed);
+      }
+    }
+
+    const char *graphRateEnv = getenv("PARCAGPU_GRAPH_RATE_LIMIT");
+    if (graphRateEnv != nullptr) {
+      double rate = atof(graphRateEnv);
+      if (rate > 0) {
+        g_graphRateLimit.store(rate, std::memory_order_relaxed);
       }
     }
 
@@ -642,7 +668,7 @@ class CuptiProfiler : public proton::Singleton<CuptiProfiler> {
         return;
       }
 
-      // Check if this is a graph launch (never rate limit these)
+      // Check if this is a graph launch (rate-limited via a separate bucket).
       bool isGraphLaunch = false;
       if (signedCbid < 0) {
         // Driver API: cuGraphLaunch = 514, cuGraphLaunch_ptsz = 515
@@ -658,15 +684,17 @@ class CuptiProfiler : public proton::Singleton<CuptiProfiler> {
                  CUPTI_RUNTIME_TRACE_CBID_cudaGraphLaunch_ptsz_v10000);
       }
 
-      // Rate limit probes using token bucket.  Skip rate limiting for graph
-      // launches (they share one correlation ID across many kernels) and when
-      // PC sampling is active (every kernel needs its correlation callback so
-      // PC samples can be matched with CPU stacks on the agent side).
-      if (!isGraphLaunch && !g_pcSamplingState.active) {
-        if (!callbackLimiter.tryAcquire()) {
-          DEBUG_PRINTF(
-              "[PARCAGPU] Rate limited: skipping probe for correlationId=%u\n",
-              correlationId);
+      // Rate limit via per-class buckets (eager vs graph). Skip when PC
+      // sampling is active: every kernel needs its correlation callback to
+      // match PC samples to CPU stacks on the agent side.
+      if (!g_pcSamplingState.active) {
+        applyRateLimitsOnce();
+        TokenBucket &limiter =
+            isGraphLaunch ? graphCallbackLimiter : callbackLimiter;
+        if (!limiter.tryAcquire()) {
+          DEBUG_PRINTF("[PARCAGPU] Rate limited: skipping probe for "
+                       "correlationId=%u (graph=%d)\n",
+                       correlationId, isGraphLaunch);
           return;
         }
       }