From dd1c1d1ed2bb2bff34927d7d3bd5a72069eba94a Mon Sep 17 00:00:00 2001
From: wuyanan7 <wuyanan7@xiaomi.com>
Date: Mon, 1 Jun 2026 11:52:28 +0800
Subject: [PATCH] fix: change release scope from sys to gpu for dispatch update
 tail

---
 deep_ep/include/deep_ep/common/handle.cuh     | 24 +++++++++++++++++++
 deep_ep/include/deep_ep/common/ptx.cuh        |  8 +++++++
 .../include/deep_ep/impls/hybrid_dispatch.cuh |  2 +-
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/deep_ep/include/deep_ep/common/handle.cuh b/deep_ep/include/deep_ep/common/handle.cuh
index d1dd7ffc9..68a2dbacc 100644
--- a/deep_ep/include/deep_ep/common/handle.cuh
+++ b/deep_ep/include/deep_ep/common/handle.cuh
@@ -116,6 +116,30 @@ struct NCCLGin {
         }
     }
 
+    template <typename team_t, typename dtype_t>
+    __device__ __forceinline__
+    void red_add_rel_gpu(dtype_t* sym_ptr, const dtype_t& value, const int& dst_rank_idx,
+                     const int& extra_options = 0) const {
+        const auto dst_ptr = get_sym_ptr<team_t>(sym_ptr, dst_rank_idx);
+        // Use symmetric pointers as much as possible, RDMA otherwise
+        if (dst_ptr != nullptr) {
+            // NOTES: local rank (or even NVLink-connected) for tag rail can also bypass
+            ptx::red_add_rel_gpu(dst_ptr, value);
+        } else {
+            EP_DEVICE_ASSERT((not std::is_same_v<team_t, ncclTeamTagLsa>));
+            EP_DEVICE_ASSERT((std::is_same_v<dtype_t, int64_t>) or (std::is_same_v<dtype_t, uint64_t>));
+            // TODO(NCCL): support all dtypes
+            gin.signal(TEAM_WORLD_RAIL(), dst_rank_idx,
+                       ncclGin_VASignalAdd(nccl_window, reinterpret_cast<int64_t>(sym_ptr) - lsa_base_ptr, static_cast<uint64_t>(value)),
+                       ncclCoopThread(),
+                       ncclGin_None(),
+                       cuda::thread_scope_thread,
+                       cuda::thread_scope_device,
+                       ncclGinOptFlagsDefault | extra_options);
+        }
+    }
+
+
     __device__ __forceinline__
     void wait(ncclGinRequest_t& request) const {
         gin.wait(request);
diff --git a/deep_ep/include/deep_ep/common/ptx.cuh b/deep_ep/include/deep_ep/common/ptx.cuh
index 7d2a4b4c7..492761463 100644
--- a/deep_ep/include/deep_ep/common/ptx.cuh
+++ b/deep_ep/include/deep_ep/common/ptx.cuh
@@ -269,6 +269,14 @@ __forceinline__ __device__ void red_add_rel_sys(const int64_t* ptr, const int64_
     asm volatile("red.release.sys.global.add.u64 [%0], %1;" :: "l"(ptr), "l"(value));
 }
 
+__forceinline__ __device__ void red_add_rel_gpu(const int* ptr, const int& value) {
+    asm volatile("red.release.gpu.global.add.s32 [%0], %1;" :: "l"(ptr), "r"(value));
+}
+
+__forceinline__ __device__ void red_add_rel_gpu(const int64_t* ptr, const int64_t& value) {
+    asm volatile("red.release.gpu.global.add.u64 [%0], %1;" :: "l"(ptr), "l"(value));
+}
+
 template <typename dtype_t>
 __forceinline__ __device__ dtype_t ld_acquire_sys(const dtype_t* ptr) {
     if constexpr (sizeof(dtype_t) == 4) {
diff --git a/deep_ep/include/deep_ep/impls/hybrid_dispatch.cuh b/deep_ep/include/deep_ep/impls/hybrid_dispatch.cuh
index 7d75b0238..31a42db34 100644
--- a/deep_ep/include/deep_ep/impls/hybrid_dispatch.cuh
+++ b/deep_ep/include/deep_ep/impls/hybrid_dispatch.cuh
@@ -339,7 +339,7 @@ hybrid_dispatch_impl(
 
                 // NOTES: the "release" scope will be `sys` for the local rank (we may involve NVLink so not `gpu`)
                 // For RDMA requests, "release" is ensured by "atomic"
-                gin.red_add_rel<ncclTeamTagRail>(ptr, signaled_tail - old_signaled_tail, lane_idx);
+                gin.red_add_rel_gpu<ncclTeamTagRail>(ptr, signaled_tail - old_signaled_tail, lane_idx);
                 stored_old_scaleout_tail = stored_scaleout_tail;
             }
             __syncwarp();