From 553f5ea60f8b87c92dba30411a6541bce9bbcab0 Mon Sep 17 00:00:00 2001 From: Darshan Sanghani Date: Mon, 8 Jun 2026 16:22:23 -0700 Subject: [PATCH] Add NCCL task protocol trace metadata Problem The executed protocol (LL/LL128/Simple) of a collective is not visible in profiles. The kernel symbol is named after the inlined LL body, not the tuned protocol, and the only way to see the real protocol today is NCCL_DEBUG_SUBSYS=TUNING, which is too noisy for normal profiling. Solution Push one NVTX range per scheduled collective task, named with its protocol, plus a compact structured payload (comm, rank, bytes, collective, algorithm, protocol, dtype, root, channels) stored as enum integers to keep the launch path cheap. nsys then reports the executed protocol per task, even when multiple protocols are fused into one launch. Limitations This surfaces the protocol in nsys. torch.profiler / Kineto does not decode external NVTX payloads; those consumers should read the protocol from the NCCL profiler plugin instead. This is part 1 of 2 addressing #2196. Part 2 removes the misleading protocol suffix from kernel symbols and should be merged after this change, so the executed protocol stays observable once the suffix is gone. Addresses #2196 Signed-off-by: Darshan Sanghani --- src/enqueue.cc | 59 ++++++++++++++++++++++++++++++ src/include/nvtx.h | 1 + src/include/nvtx_payload_schemas.h | 23 +++++++++++- 3 files changed, 81 insertions(+), 2 deletions(-) diff --git a/src/enqueue.cc b/src/enqueue.cc index 9f59b217c4..2167c6710a 100644 --- a/src/enqueue.cc +++ b/src/enqueue.cc @@ -17,6 +17,7 @@ #include "register_inline.h" #include "ce_coll.h" #include "nvtx.h" +#include "nvtx_payload_schemas.h" #include "scheduler.h" #include "compiler.h" #include "rma/rma.h" @@ -25,11 +26,67 @@ #include // PRIx64 #include #include // FLT_MAX +#include // std::extent NCCL_PARAM(L1SharedMemoryCarveout, "L1_SHARED_MEMORY_CARVEOUT", 0); NCCL_PARAM(AllgathervEnable, "ALLGATHERV_ENABLE", 1); NCCL_PARAM(SymCeThreshold, "SYM_CE_THRESHOLD", 8 * 1024 * 1024); +static const char* ncclNvtxProtoRangeName(int proto) { + switch (proto) { + case NCCL_PROTO_LL: return "NCCL Kernel Task Proto LL"; + case NCCL_PROTO_LL128: return "NCCL Kernel Task Proto LL128"; + case NCCL_PROTO_SIMPLE: return "NCCL Kernel Task Proto Simple"; + default: return "NCCL Kernel Task Proto Unknown"; + } +} + +static int ncclNvtxPushKernelTaskProtoRanges(struct ncclComm* comm, struct ncclKernelPlan* plan) { + if (ncclParamNvtxDisable()) return 0; + + constexpr uint64_t schemaId = + NVTX_PAYLOAD_ENTRY_TYPE_SCHEMA_ID_STATIC_START + NVTX_SID_KernelTask; + static const payload_schema schema{ + NcclNvtxParamsKernelTaskSchema, + std::extent::value - 1, + schemaId, + sizeof(NcclNvtxParamsKernelTask) + }; + + int ranges = 0; + struct ncclTaskColl* task = ncclIntruQueueHead(&plan->collTaskQueue); + while (task != nullptr) { + const NcclNvtxParamsKernelTask payload = { + comm->commHash, + comm->rank, + task->count * ncclTypeSize(task->datatype), + task->func, + task->algorithm, + task->protocol, + task->datatype, + task->root, + task->nChannels + }; + nvtxPayloadData_t payloadData = {}; + nvtxEventAttributes_t attr = {}; + attr.version = NVTX_VERSION; + attr.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE; + attr.messageType = NVTX_MESSAGE_TYPE_ASCII; + attr.message.ascii = ncclNvtxProtoRangeName(task->protocol); + NVTX_PAYLOAD_EVTATTR_SET_DATA(attr, &payloadData, schemaId, &payload, sizeof(payload)); + nvtxDomainRangePushEx(nvtx3::domain::get(), &attr); + ranges++; + task = task->next; + } + return ranges; +} + +static void ncclNvtxPopKernelTaskProtoRanges(int ranges) { + while (ranges-- > 0) { + nvtxDomainRangePop(nvtx3::domain::get()); + } +} + // Returns maximum kernel stack size of all CUDA kernels ncclResult_t ncclInitKernelsForDevice(int cudaArch, int maxSharedMem, size_t* maxStackSize) { ncclResult_t result = ncclSuccess; @@ -1697,6 +1754,7 @@ ncclResult_t ncclLaunchPrepare(struct ncclComm* comm) { } } + int nvtxTaskProtoRanges = ncclNvtxPushKernelTaskProtoRanges(comm, plan); if (!persistent && comm->sharedRes->persistentRefs) { status = CUDACLEARERROR(cudaEventQuery(comm->sharedRes->hostStream.serialEvent)); } @@ -1846,6 +1904,7 @@ ncclResult_t ncclLaunchKernel(struct ncclComm* comm, struct ncclKernelPlan* plan } do_return: + ncclNvtxPopKernelTaskProtoRanges(nvtxTaskProtoRanges); NCCLCHECK(ncclProfilerStopKernelLaunchEvent(plan)); return ret; } diff --git a/src/include/nvtx.h b/src/include/nvtx.h index ce86cf012d..656e780ac1 100644 --- a/src/include/nvtx.h +++ b/src/include/nvtx.h @@ -43,6 +43,7 @@ #define NVTX_SID_PutSignal 21 #define NVTX_SID_Signal 22 #define NVTX_SID_WaitSignal 23 +#define NVTX_SID_KernelTask 25 // When adding new schema IDs, DO NOT re-use/overlap with the enum schema ID below! // Define static schema ID for the reduction operation. diff --git a/src/include/nvtx_payload_schemas.h b/src/include/nvtx_payload_schemas.h index a4c54788d5..dfefd1e89d 100644 --- a/src/include/nvtx_payload_schemas.h +++ b/src/include/nvtx_payload_schemas.h @@ -5,8 +5,8 @@ * See LICENSE.txt for more license information *************************************************************************/ -/// Definitions of NVTX payload types and schemas used for the NVTX -/// instrumentation in init.cc and collectives.cc. +/// Definitions of NVTX payload types and schemas used for NCCL NVTX +/// instrumentation. #ifndef NVTX_PAYLOAD_SCHEMAS_H_ #define NVTX_PAYLOAD_SCHEMAS_H_ @@ -32,6 +32,11 @@ static constexpr char const* nccl_nvtxRankStr = "Rank"; static constexpr char const* nccl_nvtxNranksStr = "No. of ranks"; static constexpr char const* nccl_nvtxMsgSizeStr = "Message size [bytes]"; static constexpr char const* nccl_nvtxReductionOpStrpStr = "Reduction operation"; +static constexpr char const* nccl_nvtxCollectiveStr = "Collective"; +static constexpr char const* nccl_nvtxAlgorithmStr = "Algorithm"; +static constexpr char const* nccl_nvtxProtocolStr = "Protocol"; +static constexpr char const* nccl_nvtxDatatypeStr = "Datatype"; +static constexpr char const* nccl_nvtxChannelCountStr = "No. of channels"; NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES( NcclNvtxParamsCommInitAll, static constexpr, @@ -138,4 +143,18 @@ NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsWaitSignal, static con (int, npeers, TYPE_INT, "Number of peers"), (int, ctx, TYPE_INT, "Context ID"))) +NCCL_NVTX_DEFINE_STRUCT_WITH_SCHEMA_ENTRIES(NcclNvtxParamsKernelTask, static constexpr, + NCCL_NVTX_PAYLOAD_ENTRIES( + (uint64_t, comm, TYPE_UINT64, nccl_nvtxCommStr), + (int, rank, TYPE_INT, nccl_nvtxRankStr), + (size_t, bytes, TYPE_SIZE, nccl_nvtxMsgSizeStr), + (int, func, TYPE_INT, nccl_nvtxCollectiveStr), + (int, algorithm, TYPE_INT, nccl_nvtxAlgorithmStr), + (int, protocol, TYPE_INT, nccl_nvtxProtocolStr), + (int, datatype, TYPE_INT, nccl_nvtxDatatypeStr), + (int, root, TYPE_INT, "Root"), + (int, nChannels, TYPE_INT, nccl_nvtxChannelCountStr) + ) +) + #endif // end include guard