From 66b27b19bf830c781bc1ad12ca67f2021b87d56a Mon Sep 17 00:00:00 2001
From: chenjiahui9 <chenjiahui9@xiaomi.com>
Date: Tue, 12 May 2026 21:38:33 +0800
Subject: [PATCH] alltoall: use NCCL_WIN_DEFAULT for window registration

---
 src/common.cu | 15 +++++++++++----
 src/common.h  |  3 +++
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/src/common.cu b/src/common.cu
index b14c3d5..9f5b04b 100644
--- a/src/common.cu
+++ b/src/common.cu
@@ -24,6 +24,7 @@
 #pragma weak ncclDevCommCreate
 #pragma weak ncclDevCommDestroy
 #pragma weak ncclCommQueryProperties
+#pragma weak AlltoAllRunTest
 
 #define DIVUP(x, y) \
     (((x)+(y)-1)/(y))
@@ -825,8 +826,11 @@ testResult_t threadInit(struct threadArgs* args) {
   for (int i=0; i<args->nGpus; i++) {
 #if NCCL_VERSION_CODE >= NCCL_VERSION(2,27,0)
     if (test_ncclVersion >= NCCL_VERSION(2,27,0) && (local_register == SYMMETRIC_REGISTER)) {
-      NCCLCHECK(ncclCommWindowRegister(args->comms[i], args->sendbuffs[i], args->maxbytes, (ncclWindow_t*)&args->sendRegHandles[i], NCCL_WIN_COLL_SYMMETRIC));
-      NCCLCHECK(ncclCommWindowRegister(args->comms[i], args->recvbuffs[i], args->maxbytes, (ncclWindow_t*)&args->recvRegHandles[i], NCCL_WIN_COLL_SYMMETRIC));
+      // AlltoAll does not use the symk path (AllReduce/AllGather/ReduceScatter only),
+      // so NCCL_WIN_COLL_SYMMETRIC is unnecessary and NCCL_WIN_DEFAULT suffices.
+      int winFlags = (ncclTestEngine.runTest == AlltoAllRunTest) ? NCCL_WIN_DEFAULT : NCCL_WIN_COLL_SYMMETRIC;
+      NCCLCHECK(ncclCommWindowRegister(args->comms[i], args->sendbuffs[i], args->maxbytes, (ncclWindow_t*)&args->sendRegHandles[i], winFlags));
+      NCCLCHECK(ncclCommWindowRegister(args->comms[i], args->recvbuffs[i], args->maxbytes, (ncclWindow_t*)&args->recvRegHandles[i], winFlags));
     } else
 #endif
     {
@@ -1404,8 +1408,11 @@ testResult_t run() {
        TESTCHECK(AllocateBuffs(sendbuffs+i, sendBytes, recvbuffs+i, recvBytes, expected+i, (size_t)maxBytes));
 #if NCCL_VERSION_CODE >= NCCL_VERSION(2,27,0)
        if (test_ncclVersion >= NCCL_VERSION(2,27,0) && (local_register == SYMMETRIC_REGISTER)) {
-         NCCLCHECK(ncclCommWindowRegister(comms[i], sendbuffs[i], maxBytes, (ncclWindow_t*)&sendRegHandles[i], NCCL_WIN_COLL_SYMMETRIC));
-         NCCLCHECK(ncclCommWindowRegister(comms[i], recvbuffs[i], maxBytes, (ncclWindow_t*)&recvRegHandles[i], NCCL_WIN_COLL_SYMMETRIC));
+         // AlltoAll does not use the symk path (AllReduce/AllGather/ReduceScatter only),
+         // so NCCL_WIN_COLL_SYMMETRIC is unnecessary and NCCL_WIN_DEFAULT suffices.
+         int wf = (ncclTestEngine.runTest == AlltoAllRunTest) ? NCCL_WIN_DEFAULT : NCCL_WIN_COLL_SYMMETRIC;
+         NCCLCHECK(ncclCommWindowRegister(comms[i], sendbuffs[i], maxBytes, (ncclWindow_t*)&sendRegHandles[i], wf));
+         NCCLCHECK(ncclCommWindowRegister(comms[i], recvbuffs[i], maxBytes, (ncclWindow_t*)&recvRegHandles[i], wf));
        } else
 #endif
        {
diff --git a/src/common.h b/src/common.h
index 990fad6..d7ed191 100644
--- a/src/common.h
+++ b/src/common.h
@@ -124,6 +124,9 @@ struct testEngine {
 #endif
 };
 
+extern testResult_t AlltoAllRunTest(struct threadArgs* args, int root, ncclDataType_t type,
+    const char* typeName, ncclRedOp_t op, const char* opName);
+
 extern struct testEngine ncclTestEngine;
 
 struct threadArgs {