From ac8ec4c08cf0a8c3fc23e0eafcfd6b92eb1d4e73 Mon Sep 17 00:00:00 2001
From: Bertan Dogancay <111835151+BertanDogancay@users.noreply.github.com>
Date: Mon, 21 Apr 2025 12:34:47 -0400
Subject: [PATCH] Fix NPKit for SendRecv (#1651)

---
 src/device/sendrecv.h | 4 ++--
 src/include/device.h  | 1 +
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/device/sendrecv.h b/src/device/sendrecv.h
index 0967c2bc43..c016023cf6 100644
--- a/src/device/sendrecv.h
+++ b/src/device/sendrecv.h
@@ -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)
diff --git a/src/include/device.h b/src/include/device.h
index 34fd1e5fba..4c202be14a 100644
--- a/src/include/device.h
+++ b/src/include/device.h
@@ -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 {