[RFC] Device Profiler Hook

[dolpm]

RFC: Device-Side Profiler Hook

Summary

Add an optional device-side profiler hook: a __device__ callback that
ncclKernelMain's profiler() invokes at each kernel work-item boundary
(START / STOP). A profiler plugin .so provides the hook + opaque context
per device, NCCL discovers it through the same plugin-.so path as the host
profiler, and from then on calls it directly from the kernel with a small POD
event. The plugin owns its telemetry end-to-end — buffer, format, atomicity,
depth, and draining.

The hook is additive and independent of the legacy path: it fires per work
item whenever registered, needs no proxy op or host callback, and does not
disable the ring. A plugin gets graph-launch-overhead-free kernel timing simply
by using the hook and not enabling ncclProfileKernelCh — so no proxy op /
host callback is armed at all — which the legacy ring alone cannot offer.

This is an alternative to the fixed workStarted / workCompleted ring +
proxy-delivery path for consumers that want it, while leaving that path fully
intact for everyone else.

Motivation

Today NCCL's device profiler writes two fixed-size per-channel rings
(workStarted, workCompleted), each entry a {timestamp, counter} pair, and
delivers them to the host profiler plugin via a proxy op. That design has two
limits:

1. Fixed schema. The ring records only a timestamp and work counter. A
   consumer wanting richer per-work telemetry (func identity, custom fields,
   its own ring depth/format) can't get it without patching core.
2. Graph overhead. Proxy delivery adds a CUDA host callback to captured
   CUDA graphs, which perturbs exactly the kernel timing a profiler is trying to
   measure.

We want a path where a plugin can capture whatever it needs, directly on the
device, with zero added host-side work in the graph.

Design

Event ABI (src/include/profiler_dev.h, new)

A POD event with a fixed layout for ABI stability:

typedef struct {
  uint64_t timestamp;   // globaltimer() ns (device clock)
  uint64_t workCounter; // monotonic per-channel work counter
  uint32_t funcId;      // collective identity (index into ncclDevFuncTable)
  uint8_t  channelId;
  uint8_t  phase;       // NCCL_PROFILER_DEV_START / _STOP
  uint16_t rsvd;
} ncclProfilerDevEvent_t;

typedef void (*ncclProfilerDevHook_t)(const ncclProfilerDevEvent_t* ev, void* devCtx);

The hook is called on lane 0, once per profiler-enabled work item. It must
be cheap and must not touch the collective's data or synchronization.

A helper resolves a funcId to a human-readable name (e.g.
"AllReduce_Sum_f32_RING_LL"):

const char* ncclProfilerDevFuncName(int funcId);

Discovery (src/plugin/profiler/profiler_v7.{h,cc}, new)

The device hook is provided by the profiler plugin .so as a new method on the
versioned profiler interface. ncclProfiler_v7 reuses the v6 event descriptor /
state args verbatim (no new host-side events) and adds one method:

typedef struct {
  const char* name;
  ncclResult_t (*init)(...);              // same as v6
  ncclResult_t (*startEvent)(...);        // same as v6
  ncclResult_t (*stopEvent)(...);         // same as v6
  ncclResult_t (*recordEventState)(...);  // same as v6
  ncclResult_t (*finalize)(...);          // same as v6
  // Return the plugin's device hook (a __device__ fn ptr valid on `cudaDev`)
  // and opaque context. *devHook == nullptr => no hook on this device.
  ncclResult_t (*getDeviceHook)(int cudaDev, void** devHook, void** devCtx);
} ncclProfiler_v7_t;

• ncclProfiler_v7_t becomes the default ncclProfiler_t. profiler.cc tries
  getNcclProfiler_v7() first in the existing version-fallback chain
  (v7 → v6 → … → v1).
• The v6 shim now adapts into v7: host callbacks pass straight through (the
  descriptor is identical) and getDeviceHook is set to nullptr. v5-and-older
  shims already adapt into ncclProfiler_t; their zero-initialized struct leaves
  getDeviceHook null. So a plugin built against any version ≤ v6 simply has no
  device hook.
• ncclProfilerDevGetHook(device, …) (internal) forwards to
  ncclProfiler->getDeviceHook when non-null. The plugin resolves its own
  __device__ symbol on device (a __device__ function has a distinct address
  per device — the plugin owns the device code, so resolution lives where it
  belongs).
• No new exported libnccl symbol. Control stays inverted (NCCL pulls from a
  passive plugin), exactly like the rest of the profiler. ncclProfilerDevFuncName
  remains exported for plugins to resolve funcId → name.

Propagation (src/init.cc)

At comm init (devCommSetup, after ncclProfilerPluginInit has loaded the
plugin), NCCL pulls the hook for the comm's device into the device comm:

ncclProfilerDevGetHook(comm->cudaDev, &comm->profiler.devHook, &comm->profiler.devCtx);
tmpCommAndChans.comm.profilerDevHook = comm->profiler.devHook;
tmpCommAndChans.comm.profilerDevCtx  = comm->profiler.devCtx;

nullptr (no plugin, or plugin returns none) => legacy ring behavior, unchanged.

ncclKernelComm (src/include/device.h) gains two void* fields
(profilerDevHook, profilerDevCtx), typed as void* to keep device.h free
of the hook typedef; common.h casts them. ncclProfilerState
(src/include/profiler.h) gains the matching host-side devHook / devCtx.

Kernel path (src/device/common.h)

profiler() is restructured into a single loop over work items. For each
profiler-enabled item:

• If a hook is registered, build a stack ncclProfilerDevEvent_t and call the
  hook.
• Independently, if the item is profilerEnabled, write the legacy
  workStarted / workCompleted ring as before.

The hook fires whenever registered and is independent of the host profiler
plugin; the legacy ring stays gated by profilerEnabled (the plugin's per-work
flag). Both can run simultaneously.

Func-name table (src/device/generate.py)

Codegen now emits ncclDevFuncName[] — one human-readable name per primary
func id, indexed the same as ncclDevKernelForFunc[] — plus
ncclDevFuncIdCount. This backs ncclProfilerDevFuncName() so a consumer can
resolve funcId without reproducing the codegen's variant enumeration.

Example / testing (plugins/profiler/example)

The example profiler plugin is bumped to v7 and gains a real device hook
(plugin_dev.cu): a __device__ callback that counts events per device into its
own device context and, at finalize, prints the counts and resolves the last
funcId via ncclProfilerDevFuncName(). It is opt-in via
NCCL_PROFILER_EXAMPLE_DEV_HOOK, so the example's host-side behavior is
unchanged by default. It is built with nvcc + relocatable device code (the
__device__ hook needs a real, callable global address for the kernel to invoke
it indirectly). This ships as a separate commit on top of the core change.

Validated end to end on 2× GB200 via all_reduce_perf: with the hook enabled it
fires per work item (e.g. ~6.7k events/device for a size sweep), resolves func
names (e.g. AllReduce_Sum_f32_RING_LL), and reports #wrong = 0; with the hook
disabled (default) behavior and results are unchanged.

Compatibility

• Fully backward compatible. No hook registered (nullptr) => existing
  workStarted / workCompleted ring + proxy path is unchanged.
• New ABI (ncclProfilerDevEvent_t) is POD with fixed layout and a reserved
  field for forward evolution.
• Only additive changes to ncclKernelComm / ncclProfilerState; no existing
  fields removed or reordered.

Alternatives considered

• Exported push API (ncclProfilerDevSetHook). The consumer resolves its own
  __device__ fn ptr and pushes it into a process-wide per-device store in
  libnccl before creating comms. Rejected as the OSS ABI: it adds a new exported
  libnccl symbol and a second registration mechanism unlike the established
  plugin-discovery model. Note the tradeoff — push is a more direct fit for a
  consumer that generates its hook at runtime (e.g. NVRTC), since pull requires
  such a consumer to wrap the runtime-resolved pointer behind a plugin .so +
  its own setter. We chose pull for ABI cleanliness; a runtime consumer builds a
  thin plugin whose getDeviceHook returns its runtime-set pointer.
• Dedicated ncclProfilerDev_v1 symbol on the same .so, independent of the
  host profiler version. Rejected: its only real advantage is letting a plugin
  provide a device hook without being the host profiler — a marginal case, since
  only one profiler plugin loads and a device profiler is a profiler. Folding
  the method onto ncclProfiler_v7 keeps a single versioned interface and reuses
  the existing fallback machinery; v7 reuses the v6 descriptor by typedef, so there
  is no struct duplication.
• Extend the existing ring schema. Rejected: still incurs proxy/host-callback
  graph overhead, and bakes one more fixed format into core.
• Deliver richer events through the proxy. Rejected for the same graph-overhead
  reason; the whole point is to avoid host callbacks in captured graphs.

Open questions

• The hook is pulled once at comm init; hot re-registration after comms exist is
  not supported (existing device comms keep their copy).

[kwen2501]

Thank you for the contribution. We are reviewing the PR.