From 17b7abca9b27d12d801a8b12bdc4fe70391e1d74 Mon Sep 17 00:00:00 2001 From: Keita Watanabe Date: Mon, 24 Aug 2026 01:17:35 +0000 Subject: [PATCH] fix: resolve the GIN context count on the single-node path On a single node `scaleout_active` is false, and the else arm copied `num_allocated_qps` verbatim instead of calling `resolve_gin_context_cnt()`. When the caller leaves that value at its default of 0 and the unordered hybrid kernels are in use, the Python auto-fill is deliberately skipped so that C++ resolves the count, so a single-node run kept 0 all the way into the kernels: nccl.cu else arm -> gin_context_cnt = 0 elastic.py get_theoretical_num_qps -> min(num_sms * 16 + 1, 0) = 0 comm.cuh get_qp_mode -> kNumQPs == 1 fast path skipped, kNumSMs <= kNumAvailableQPs false qp_mapping.cuh -> balanced_partition(idx, 512, 0) -> n / 0 Nothing rejects q == 0. Compiled for the host that specialization exits on SIGFPE; in a GPU run it produced corrupted output and no reported fault: the combine reported 623 MB in 1.148 us, and its output differed from the reference in 99.98% of elements, surfacing as 'AssertionError: Diff: nan' in tests/elastic/test_ep.py. Resolve the count on both branches and write it back unconditionally, since Python reads it via get_num_allocated_qps() and caps the per-launch QP count with it. Add a static_assert so an illegal specialization fails to compile rather than miscomputing, and assert the delegation invariant in Python at the point the contract is broken. --- csrc/kernels/backend/nccl.cu | 5 ++--- deep_ep/buffers/elastic.py | 5 +++++ deep_ep/include/deep_ep/common/comm.cuh | 6 ++++++ deep_ep/include/deep_ep/common/qp_mapping.cuh | 10 ++++++++++ 4 files changed, 23 insertions(+), 3 deletions(-) diff --git a/csrc/kernels/backend/nccl.cu b/csrc/kernels/backend/nccl.cu index d86b0a07..4495b9ac 100644 --- a/csrc/kernels/backend/nccl.cu +++ b/csrc/kernels/backend/nccl.cu @@ -126,7 +126,7 @@ NCCLSymmetricMemoryContext::NCCLSymmetricMemoryContext(const int64_t& nccl_comm, if (scaleout_active) { gin_config = elastic::gin_alloc::make_gin_resources(resolve_gin_context_cnt()); } else { - gin_config.gin_context_cnt = this->num_allocated_qps; + gin_config.gin_context_cnt = resolve_gin_context_cnt(); gin_config.gin_indexed_signals_cnt = 0; } @@ -134,8 +134,7 @@ NCCLSymmetricMemoryContext::NCCLSymmetricMemoryContext(const int64_t& nccl_comm, "GIN indexed-signal budget cannot give each peer rail team a dedicated " "signal; reduce num_allocated_qps to raise the per-context signal count"); - if (scaleout_active) - this->num_allocated_qps = gin_config.gin_context_cnt; + this->num_allocated_qps = gin_config.gin_context_cnt; if (get_env("EP_BUFFER_DEBUG")) printf("GIN layout: gin_context_cnt=%d, gin_indexed_signals_cnt=%d, num_qp=%d\n", diff --git a/deep_ep/buffers/elastic.py b/deep_ep/buffers/elastic.py index e7b478ef..b21458df 100644 --- a/deep_ep/buffers/elastic.py +++ b/deep_ep/buffers/elastic.py @@ -377,6 +377,11 @@ def __init__(self, self.explicitly_destroy) self.num_allocated_qps = self.runtime.get_num_allocated_qps() + # The unordered hybrid path intentionally passes 0 down and relies on C++ to + # resolve it (see the automatic QP count above). A 0 read back here means that + # delegation did not happen; it would surface far downstream as `kNumQPs == 0`. + assert self.num_allocated_qps >= 1, \ + f'runtime returned num_allocated_qps={self.num_allocated_qps}: the QP count was never resolved' # Logical rank indices self.num_scaleout_ranks, self.num_scaleup_ranks = self.get_logical_domain_size() diff --git a/deep_ep/include/deep_ep/common/comm.cuh b/deep_ep/include/deep_ep/common/comm.cuh index 44302744..5cbb9e08 100644 --- a/deep_ep/include/deep_ep/common/comm.cuh +++ b/deep_ep/include/deep_ep/common/comm.cuh @@ -77,6 +77,12 @@ __device__ __forceinline__ void timeout_while(const func_t& func, const int64_t& template __device__ __forceinline__ std::pair get_qp_mode( const int& sm_idx, const int& channel_in_sm_idx, const bool& is_notify_warp = false) { + static_assert(kNumQPs >= 1, + "kNumQPs must be >= 1. With 0 the kNumQPs == 1 fast path is skipped, " + "kNumSMs <= kNumAvailableQPs is false, and balanced_partition() is called " + "with q == 0, i.e. n / 0. Compiled for the host this specialization exits on " + "SIGFPE; in a GPU run it produced corrupted output and no reported fault."); + constexpr auto kSharingCTA = NCCL_GIN_RESOURCE_SHARING_CTA; constexpr auto kSharingGrid = kNumSMs == 1 ? NCCL_GIN_RESOURCE_SHARING_CTA : NCCL_GIN_RESOURCE_SHARING_GPU; diff --git a/deep_ep/include/deep_ep/common/qp_mapping.cuh b/deep_ep/include/deep_ep/common/qp_mapping.cuh index 4e9652ee..12dfedd4 100644 --- a/deep_ep/include/deep_ep/common/qp_mapping.cuh +++ b/deep_ep/include/deep_ep/common/qp_mapping.cuh @@ -62,6 +62,11 @@ QP_MAPPING_HD constexpr QPSlot balanced_partition(int idx, int n, int q) { template QP_MAPPING_HD constexpr int channel_to_qp(int sm_idx, int channel_in_sm_idx, bool is_notify_warp = false) { + static_assert(kNumQPs >= 1, + "kNumQPs must be >= 1. With 0 the kNumQPs == 1 fast path is skipped, " + "kNumSMs <= kNumAvailableQPs is false, and balanced_partition() is called " + "with q == 0, i.e. n / 0. Compiled for the host this specialization exits on " + "SIGFPE; in a GPU run it produced corrupted output and no reported fault."); // Only one QP if constexpr (kNumQPs == 1) return 0; @@ -95,6 +100,11 @@ QP_MAPPING_HD constexpr int channel_to_qp(int sm_idx, int channel_in_sm_idx, // `signal_id < ceil(channels / qps)` -- within the tuner's signal budget. template QP_MAPPING_HD constexpr int channel_to_signal_id(int sm_idx, int channel_in_sm_idx) { + static_assert(kNumQPs >= 1, + "kNumQPs must be >= 1. With 0 the kNumQPs == 1 fast path is skipped, " + "kNumSMs <= kNumAvailableQPs is false, and balanced_partition() is called " + "with q == 0, i.e. n / 0. Compiled for the host this specialization exits on " + "SIGFPE; in a GPU run it produced corrupted output and no reported fault."); if constexpr (kNumQPs == 1) return sm_idx * kNumChannelsPerSM + channel_in_sm_idx;