fix(diag): print the NIC queue state next to every timeout

README.md

@@ -356,6 +356,8 @@ The library provides some environment variables, which may be useful:
     - `EP_JIT_DUMP_SASS`: `0` or `1`, dump SASS output, `0` by default
 - Debug and profiling
     - `EP_GIN_GDAKI_DEBUG`: `0` or `1`, enable NCCL Gin GDAKI debugging output, `0` by default
+    - `EP_DEBUG_CHECKS`: debug checks compiled into the kernels, off when unset or `0`. `1` or `all` turns on every check, a comma-separated list of the names below turns on single checks
+        - `queue`: after a network wait of the unordered kernels times out, a short timed flush of the waiting warp's GIN context prints whether this rank's NIC had completed every write posted on it (`NIC drained: 1`) or still held some (`0`). GIN flush timeouts then also print the context id and wait one second before trapping
     - `EP_USE_NVIDIA_TOOLS`: `0` or `1`, skip internal profiling when running under external NVIDIA tools, `0` by default
     - `EP_DISABLE_BARRIER_PROFILING`: `0` or `1`, disable barrier-based communication profiling in benchmarks, `0` by default
 - Build

csrc/jit/compiler.hpp

@@ -64,6 +64,8 @@ class Compiler {
             flags += " -Xcompiler -rdynamic -lineinfo";
         if (get_env("EP_GIN_GDAKI_DEBUG", 0))
             flags += " -DNCCL_DEVICE_GIN_GDAKI_ENABLE_DEBUG=1";
+        if (const auto debug_checks = get_debug_checks(); debug_checks != 0)
+            flags += fmt::format(" -DEP_DEBUG_CHECKS={}", debug_checks);
         flags += fmt::format(" -I {}/include", nccl_root.c_str());
 
         // Some special flags for EP

csrc/utils/system.hpp

@@ -5,10 +5,13 @@
 #include <filesystem>
 #include <functional>
 #include <random>
+#include <sstream>
 #include <string>
+#include <vector>
 #include <memory>
 #include <unistd.h>
 
+#include <deep_ep/common/compiled.cuh>
 #include <deep_ep/common/exception.cuh>
 
 #include "format.hpp"
@@ -34,6 +37,45 @@ static dtype_t get_env(const std::string& name, const dtype_t& default_value = d
     }
 }
 
+// `EP_DEBUG_CHECKS` becomes the bit mask the kernels are compiled with, in three steps.
+//   environment string  ->  check names   ->  one bit per name  ->  mask
+//   ""  or "0"               no names                              0
+//   "1" or "all"             every name in `kDebugCheckNames`      all bits
+//   "queue,nanny"            {"queue", "nanny"}                    kDebugCheckGinQueue | kDebugCheckNanny
+// An unknown name stops here, before any kernel is compiled.
+static std::vector<std::string> debug_check_names(const std::string& value) {
+    std::vector<std::string> names;
+    if (value.empty() or value == "0")
+        return names;
+    if (value == "1" or value == "all") {
+        for (const auto& check: kDebugCheckNames)
+            names.emplace_back(check.name);
+        return names;
+    }
+    std::stringstream list(value);
+    for (std::string name; std::getline(list, name, ',');)
+        names.push_back(name);
+    return names;
+}
+
+static int debug_check_bit(const std::string& name) {
+    for (const auto& check: kDebugCheckNames)
+        if (name == check.name)
+            return check.bit;
+    EP_HOST_UNREACHABLE(fmt::format("Unknown check in EP_DEBUG_CHECKS: {}", name));
+}
+
+static int debug_check_mask(const std::vector<std::string>& names) {
+    int mask = 0;
+    for (const auto& name: names)
+        mask |= debug_check_bit(name);
+    return mask;
+}
+
+static int get_debug_checks() {
+    return debug_check_mask(debug_check_names(get_env<std::string>("EP_DEBUG_CHECKS")));
+}
+
 static std::tuple<int, std::string> call_external_command(std::string command) {
     command = command + " 2>&1";
     const auto deleter = [](FILE* f) { if (f) pclose(f); };

deep_ep/include/deep_ep/common/comm.cuh

@@ -83,12 +83,37 @@ template <int64_t kNumTimeoutCycles, typename gin_t, typename coop_t>
 __device__ __forceinline__ void flush_with_timeout(const gin_t& gin, const coop_t& coop) {
     if (gin.flush(coop, cuda::memory_order_acquire, ncclGin_None{},
                   static_cast<uint64_t>(kNumTimeoutCycles)) != ncclSuccess) {
-        printf("DeepEP Gin flush timeout, sm: %d, thread: %d\n",
-               static_cast<int>(blockIdx.x), static_cast<int>(threadIdx.x));
+        if constexpr (kReportGinQueue) {
+            // Name the context and give the other warps' timeout reports a second to print before the trap
+            printf("DeepEP Gin flush timeout, sm: %d, thread: %d, ctx: %d\n",
+                   static_cast<int>(blockIdx.x), static_cast<int>(threadIdx.x), static_cast<int>(gin.contextId));
+            const auto stop = clock64() + kNumOneSecCycles;
+            while (clock64() < stop) {}
+        } else {
+            printf("DeepEP Gin flush timeout, sm: %d, thread: %d\n",
+                   static_cast<int>(blockIdx.x), static_cast<int>(threadIdx.x));
+        }
         ptx::trap();
     }
 }
 
+// With the `queue` check on, a timed flush after a network-wait timeout tells whether this rank's NIC still
+// holds writes. Compiled out otherwise, since the inlined flush changes the register allocation of the kernels.
+__device__ __forceinline__ void report_gin_queue(const ncclGin& gin,
+                                                 const int& scaleout_rank_idx, const int& scaleup_rank_idx) {
+    if constexpr (kReportGinQueue) {
+        const auto mask = __activemask();
+        if (threadIdx.x % 32 == __ffs(mask) - 1) {
+            const bool drained = gin.flush(ncclCoopThread(), cuda::memory_order_acquire, ncclGin_None{},
+                                           static_cast<uint64_t>(kNumOneSecCycles / 10)) == ncclSuccess;
+            printf("DeepEP GIN queue state, scale-out: %d, scale-up: %d, sm: %d, thread: %d, ctx: %d, NIC drained: %d\n",
+                   scaleout_rank_idx, scaleup_rank_idx, static_cast<int>(blockIdx.x), static_cast<int>(threadIdx.x),
+                   static_cast<int>(gin.contextId), static_cast<int>(drained));
+        }
+        __syncwarp(mask);
+    }
+}
+
 // Channels map onto QPs with a BALANCED CONTIGUOUS (block) partition: fill a QP with
 // consecutive channels before spilling to the next (so an SM's channels stay grouped
 // on the same GIN context), and when the channels do not divide the QPs evenly,
@@ -259,6 +284,7 @@ __forceinline__ __device__ void gin_barrier_wo_local_sync(
                 printf("DeepEP Gin barrier timeout, tag: %d, scaleout: %d, scaleup: %d, thread: %d, "
                        "slot: %d, target: %lu\n", kTag, scaleout_rank_idx, scaleup_rank_idx, thread_idx,
                        static_cast<int>(signal_idx), target);
+                report_gin_queue(gin, scaleout_rank_idx, scaleup_rank_idx);
                 // Wait another 1 second to let all threads print information and
                 // trap similar to timeout_while
                 const auto stop = clock64() + kNumOneSecCycles;

deep_ep/include/deep_ep/common/compiled.cuh

@@ -66,6 +66,10 @@ __device__ __forceinline__ longlong4_t make_longlong4_t(
 #define EP_NUM_TOPK_IDX_BITS 64
 #endif
 
+#ifndef EP_DEBUG_CHECKS
+#define EP_DEBUG_CHECKS 0
+#endif
+
 namespace deep_ep {
 
 #ifndef DISABLE_SM90_FEATURES
@@ -74,6 +78,18 @@ constexpr bool kEnableSM90Features = true;
 constexpr bool kEnableSM90Features = false;
 #endif
 
+// Debug checks compiled into the kernels. The host maps the names set in `EP_DEBUG_CHECKS` to these bits
+// (`get_debug_checks` in `csrc/utils/system.hpp`) and passes the mask to the JIT. The default build has none.
+constexpr int kDebugChecks = EP_DEBUG_CHECKS;
+constexpr int kDebugCheckGinQueue = 1 << 0;
+constexpr bool kReportGinQueue = (kDebugChecks & kDebugCheckGinQueue) != 0;
+
+struct DebugCheckName {
+    const char* name;
+    int bit;
+};
+constexpr DebugCheckName kDebugCheckNames[] = {{"queue", kDebugCheckGinQueue}};
+
 template <int kNumBits> struct int_with_bits;
 template <> struct int_with_bits<8>  { using type = int8_t;  };
 template <> struct int_with_bits<16> { using type = int16_t; };

deep_ep/include/deep_ep/impls/hybrid_combine_unordered.cuh

@@ -444,6 +444,7 @@ hybrid_unordered_combine_impl(nv_bfloat16* x,
                            "channel: %d, forward warp: %d, dst: %d, head: %u, tail: %u, depth: %d\n",
                            scaleout_rank_idx, scaleup_rank_idx, channel_idx, forward_warp_idx,
                            dst, my_head, ring_tail, kProxyRingDepth);
+                    comm::report_gin_queue(gin.gin, scaleout_rank_idx, scaleup_rank_idx);
                 }
                 return false;
             });
@@ -778,6 +779,7 @@ hybrid_unordered_combine_impl(nv_bfloat16* x,
                            "channel: %d, signal: %lu, target: %lu\n",
                            scaleout_rank_idx, kNumScaleoutRanks, scaleup_rank_idx, kNumScaleupRanks,
                            channel_idx, signal, target);
+                    comm::report_gin_queue(gin.gin, scaleout_rank_idx, scaleup_rank_idx);
                 }
                 return false;
             });

deep_ep/include/deep_ep/impls/hybrid_dispatch_unordered.cuh

@@ -406,6 +406,7 @@ hybrid_unordered_dispatch_impl(
                                    is_expert_reduction ? "expert" : "rank",
                                    scaleout_rank_idx, scaleup_rank_idx, thread_idx, j,
                                    decoded);
+                            comm::report_gin_queue(gin.gin, scaleout_rank_idx, scaleup_rank_idx);
                         }
                         return false;
                     });
@@ -930,16 +931,19 @@ hybrid_unordered_dispatch_impl(
                 __syncwarp();
 
                 // Timeout
-                if (is_last_check and lane_idx < kNumScaleoutRanks)
-                    printf("DeepEP hybrid dispatch (forwarding) timeout, scale-out: %d, scale-up: %d, "
-                           "channel: %d, lane: %d, old scale-out tail: %d, scale-out tail: (%d, %d), "
-                           "arrival anomaly: %d, part: %d, num valid: %d, "
-                           "completion missing: %d, missing part: %d, valid headers: %d, completions: %d\n",
-                           scaleout_rank_idx, scaleup_rank_idx,
-                           channel_idx, lane_idx, stored_scaleout_old_tail_idx,
-                           stored_finish_flag, stored_scaleout_tail_idx,
-                           static_cast<int>(anomaly_detected), anomaly_part, anomaly_num_valid,
-                           static_cast<int>(missing_part >= 0), missing_part, missing_num_valid, missing_landed);
+                if (is_last_check) {
+                    if (lane_idx < kNumScaleoutRanks)
+                        printf("DeepEP hybrid dispatch (forwarding) timeout, scale-out: %d, scale-up: %d, "
+                               "channel: %d, lane: %d, old scale-out tail: %d, scale-out tail: (%d, %d), "
+                               "arrival anomaly: %d, part: %d, num valid: %d, "
+                               "completion missing: %d, missing part: %d, valid headers: %d, completions: %d\n",
+                               scaleout_rank_idx, scaleup_rank_idx,
+                               channel_idx, lane_idx, stored_scaleout_old_tail_idx,
+                               stored_finish_flag, stored_scaleout_tail_idx,
+                               static_cast<int>(anomaly_detected), anomaly_part, anomaly_num_valid,
+                               static_cast<int>(missing_part >= 0), missing_part, missing_num_valid, missing_landed);
+                    comm::report_gin_queue(gin.gin, scaleout_rank_idx, scaleup_rank_idx);
+                }
                 return false;
             });
 
@@ -1114,10 +1118,12 @@ hybrid_unordered_dispatch_impl(
                         comm::signal_delta(gin.gin.readSignal(part_signal_id(k)), rx_part_base[k]);
                     if (landed >= static_cast<int64_t>(expected_k))
                         return true;
-                    if (is_last_check)
+                    if (is_last_check) {
                         printf("DeepEP hybrid dispatch (epilogue shadow drain) timeout, "
                                "channel: %d, part: %d, landed: %lld, expected: %d\n",
                                channel_idx, k, (long long)landed, expected_k);
+                        comm::report_gin_queue(gin.gin, scaleout_rank_idx, scaleup_rank_idx);
+                    }
                     return false;
                 });
                 gin.gin.increaseSignalShadow(part_signal_id(k), static_cast<uint64_t>(expected_k));