From d7eaccb40f1a5c069eb51e6fd3a8aa81d901f91c Mon Sep 17 00:00:00 2001
From: Mamzi Bayatpour <mbayatpour@local>
Date: Tue, 6 Jan 2026 12:14:31 -0800
Subject: [PATCH] sharp: apply SHARP_COLLNET_OVERLAP_AG env profile

---
 src/sharp_plugin.c | 28 ++++++++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/src/sharp_plugin.c b/src/sharp_plugin.c
index 7ae2f29..c86b7df 100644
--- a/src/sharp_plugin.c
+++ b/src/sharp_plugin.c
@@ -8,12 +8,12 @@
 
 #include <pthread.h>
 #include <stdint.h>
+#include <stdlib.h>
 #include <string.h>
 #include <sys/types.h>
 #include <sys/time.h>
 #include <unistd.h>
 #include <assert.h>
-#include <string.h>
 
 #include "config.h"
 #include "core.h"
@@ -56,6 +56,28 @@ enum ncclSharpRequestType {
   NCCL_SHARP_REQ_IFLUSH,
 };
 
+static void ncclSharpSetEnv(const char* name, const char* value) {
+  /* We explicitly want the plugin profile to win over existing env vars. */
+  setenv(name, value, 1);
+}
+
+static void ncclSharpApplyLibEnvProfile(void) {
+  const char* overlapAgEnv = getenv("SHARP_COLLNET_OVERLAP_AG");
+  /* Only apply a profile if Megatron explicitly requests it. */
+  if (overlapAgEnv == NULL) return;
+
+  const int overlapAg = (strcmp(overlapAgEnv, "1") == 0);
+  if (overlapAg) {
+    /* For allgather overlap (multicast path) */
+    ncclSharpSetEnv("SHARP_COLL_ENABLE_MCAST", "1");
+    ncclSharpSetEnv("SHARP_COLL_ENABLE_SAT", "0");
+  } else {
+    /* For reduce-scatter (non-overlap default path) */
+    ncclSharpSetEnv("SHARP_COLL_ENABLE_MCAST", "0");
+    ncclSharpSetEnv("SHARP_COLL_ENABLE_SAT", "1");
+  }
+}
+
 struct ncclSharpRequest {
   int requestType;
   void *sharpRequest;
@@ -220,12 +242,14 @@ ncclResult_t ncclSharpInit(void **ctx, uint64_t commId, ncclDebugLogger_t logFun
   gettimeofday(&tval, NULL);
   srand((int) tval.tv_usec);
 
-  /* set SHARP COLL library default for plugin */
   setenv("SHARP_COLL_ENABLE_SAT", "1", 0);
   setenv("SHARP_COLL_NUM_COLL_GROUP_RESOURCE_ALLOC_THRESHOLD", "0", 0);
   setenv("SHARP_COLL_LOCK_ON_COMM_INIT", "1", 0);
   setenv("SHARP_COLL_LOG_LEVEL", "3", 0);
 
+  /* Apply Megatron-driven env profile (only if SHARP_COLLNET_OVERLAP_AG is set) */
+  ncclSharpApplyLibEnvProfile();
+
   if (ncclParamSharpDisableAG()) {
     INFO(NCCL_NET, "Disabled SHARP Allgather");
     ncclCollNetPlugin_v10.iallgather = NULL;