Fix NPKit for SendRecv (#1651)

This commit is contained in:
Bertan Dogancay
2025-04-21 12:34:47 -04:00
committed by GitHub
orang tua 9a55ff60a9
melakukan ac8ec4c08c
2 mengubah file dengan 3 tambahan dan 2 penghapusan
+2 -2
Melihat File
@@ -23,7 +23,7 @@ struct RunWorkBatch<ncclFuncSendRecv, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPL
#if defined(ENABLE_NPKIT)
bool isNpKitThread = (tid == 0);
int npKitCtxIdx = blockIdx.x + group;
int npKitCtxIdx = blockIdx.x * NCCL_MAX_DEV_WORK_P2P_ELEMENTS + group;
#endif
#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_CPU)
@@ -81,7 +81,7 @@ struct RunWorkBatch<ncclFuncSendRecv, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPL
#if defined(ENABLE_NPKIT)
bool isNpKitThread = (tid == 0);
int npKitCtxIdx = blockIdx.x + group;
int npKitCtxIdx = blockIdx.x * NCCL_MAX_DEV_WORK_P2P_ELEMENTS + group;
#endif
#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_CPU)
+1
Melihat File
@@ -383,6 +383,7 @@ constexpr size_t ncclDevWorkSize(enum ncclDevWorkType type) {
#define NCCL_MAX_DEV_WORK_BATCH_BYTES 128
#define NCCL_MAX_DEV_WORK_BATCH_COLLS (NCCL_MAX_DEV_WORK_BATCH_BYTES/sizeof(ncclDevWorkColl))
#define NCCL_MAX_DEV_WORK_P2P_PER_BATCH 1
#define NCCL_MAX_DEV_WORK_P2P_ELEMENTS 2
struct alignas(16) ncclDevWorkBatch {
union {
struct {