[None][fix] Merge back #19593 from release/1.3: Make MoE communication timeouts follow the warmup phase

3rdparty/CMakeLists.txt

@@ -143,15 +143,52 @@ foreach(DEP_IDX RANGE ${DEP_COUNT_MINUS_ONE})
     list(APPEND FETCH_ARGS SOURCE_SUBDIR "${DEP_SOURCE_SUBDIR}")
   endif()
 
-  if(DEP_PATCH_FILE AND NOT DEP_PATCH_FILE STREQUAL "")
+  # patch_file is one path or a list of paths, relative to this directory and
+  # applied in order.
+  set(_patch_files "")
+  string(
+    JSON
+    _patch_file_type
+    ERROR_VARIABLE
+    _err
+    TYPE
+    "${DEP_OBJECT}"
+    "patch_file")
+  if(_patch_file_type STREQUAL "ARRAY")
+    string(JSON _patch_count LENGTH "${DEP_OBJECT}" "patch_file")
+    if(_patch_count GREATER 0)
+      math(EXPR _patch_last "${_patch_count} - 1")
+      foreach(_patch_idx RANGE ${_patch_last})
+        string(JSON _patch GET "${DEP_OBJECT}" "patch_file" ${_patch_idx})
+        list(APPEND _patch_files "${CMAKE_CURRENT_SOURCE_DIR}/${_patch}")
+      endforeach()
+    endif()
+  elseif(DEP_PATCH_FILE AND NOT DEP_PATCH_FILE STREQUAL "")
+    list(APPEND _patch_files "${CMAKE_CURRENT_SOURCE_DIR}/${DEP_PATCH_FILE}")
+  endif()
+
+  if(_patch_files)
     set(_patch_command "${Python3_EXECUTABLE}"
                        "${CMAKE_CURRENT_SOURCE_DIR}/patch_dependency.py")
-    set(_patch_file "${CMAKE_CURRENT_SOURCE_DIR}/${DEP_PATCH_FILE}")
-    # Reconfigure, and rerun the patch step, when the patch changes.
+    # Reconfigure, and rerun the patch step, when a patch changes.
     set_property(
       DIRECTORY
       APPEND
-      PROPERTY CMAKE_CONFIGURE_DEPENDS "${_patch_file}")
+      PROPERTY CMAKE_CONFIGURE_DEPENDS ${_patch_files})
+    list(LENGTH _patch_files _patch_count)
+    if(_patch_count EQUAL 1)
+      list(GET _patch_files 0 _patch_file)
+    else()
+      # patch_dependency.py applies a single patch, so concatenate the patches
+      # in order. They must touch disjoint hunks.
+      set(_patch_content "")
+      foreach(_patch IN LISTS _patch_files)
+        file(READ "${_patch}" _patch_part)
+        string(APPEND _patch_content "${_patch_part}")
+      endforeach()
+      set(_patch_file "${CMAKE_CURRENT_BINARY_DIR}/patches/${DEP_NAME}.patch")
+      file(WRITE "${_patch_file}" "${_patch_content}")
+    endif()
     file(SHA256 "${_patch_file}" _patch_sha256)
     list(
       APPEND

3rdparty/fetch_content.json

@@ -27,21 +27,26 @@
       "git_repository": "${github_base_url}/deepseek-ai/DeepEP",
       "git_tag": "931243dab44d7b3976a67e253427b4db25b42988",
       "use_url": true,
-      "patch_file": "patches/deep_ep_intranode_combine_fix.patch"
+      "patch_file": [
+        "patches/deep_ep_intranode_combine_fix.patch",
+        "patches/deep_ep_runtime_timeout.patch"
+      ]
     },
     {
       "name": "nccl_extensions_download",
       "display_name": "nccl_extensions",
       "git_repository": "${github_base_url}/NVIDIA/nccl-extensions",
       "git_tag": "af42925a9cc6a818b04681eb8bdbedaed506aa30",
-      "use_url": true
+      "use_url": true,
+      "patch_file": "patches/nccl_ep_group_timeout.patch"
     },
     {
       "name": "deepgemm",
       "git_repository": "https://github.com/deepseek-ai/DeepGEMM",
       "git_tag": "b64107f2b9599ca76445b7f62eedf66bae1d095b",
       "git_submodules_recurse": true,
-      "source_subdir": "dont-add-this-project-with-add-subdirectory"
+      "source_subdir": "dont-add-this-project-with-add-subdirectory",
+      "patch_file": "patches/deepgemm_configurable_barrier_timeout.patch"
     },
     {
       "name": "eigen",

3rdparty/patches/deep_ep_runtime_timeout.patch

@@ -0,0 +1,95 @@
+--- a/csrc/deep_ep.cpp
++++ b/csrc/deep_ep.cpp
+@@ -1614,6 +1614,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+         .def("low_latency_combine_low_precision", &deep_ep::Buffer::low_latency_combine_low_precision)
+         .def("get_next_low_latency_combine_buffer", &deep_ep::Buffer::get_next_low_latency_combine_buffer);
+
++    m.def("set_timeout_seconds", &deep_ep::set_timeout_seconds);
++    m.def("get_timeout_seconds", &deep_ep::get_timeout_seconds);
+     m.def("is_sm90_compiled", deep_ep::is_sm90_compiled);
+     m.def("quantize_bf16_to_nvfp4", deep_ep::quantize_bf16_to_nvfp4);
+     m.def("dequantize_nvfp4_to_bf16", deep_ep::dequantize_nvfp4_to_bf16);
+--- a/csrc/kernels/api.cuh
++++ b/csrc/kernels/api.cuh
+@@ -4,6 +4,13 @@
+
+ namespace deep_ep {
+
++// Timeouts of the normal kernels and their CPU-side waits for subsequent launches on the current
++// device; 0 restores the compile-time defaults. The device timeout is read when a kernel runs, so
++// it also applies to replays of captured CUDA graphs. Must not be called while a stream captures.
++void set_timeout_seconds(int timeout_secs);
++
++int get_timeout_seconds();
++
+ // Intranode runtime
+ namespace intranode {
+
+--- a/csrc/kernels/configs.cuh
++++ b/csrc/kernels/configs.cuh
+@@ -10,12 +10,27 @@
+ #define NUM_WAIT_NANOSECONDS 500
+
+ #ifndef ENABLE_FAST_DEBUG
+-#define NUM_CPU_TIMEOUT_SECS 100
+-#define NUM_TIMEOUT_CYCLES 200000000000ull // 200G cycles ~= 100s
++#define DEFAULT_NUM_CPU_TIMEOUT_SECS 100
++#define DEFAULT_NUM_TIMEOUT_CYCLES 200000000000ull // 200G cycles ~= 100s
+ #else
+-#define NUM_CPU_TIMEOUT_SECS 10
+-#define NUM_TIMEOUT_CYCLES 20000000000ull // 20G cycles ~= 10s
++#define DEFAULT_NUM_CPU_TIMEOUT_SECS 10
++#define DEFAULT_NUM_TIMEOUT_CYCLES 20000000000ull // 20G cycles ~= 10s
+ #endif
++#define NUM_TIMEOUT_CYCLES_PER_SEC 2000000000ull
++
++// Both timeouts can change at runtime, see `set_timeout_seconds` in `runtime.cu`
++#define NUM_CPU_TIMEOUT_SECS (deep_ep::num_cpu_timeout_secs)
++#define NUM_TIMEOUT_CYCLES (deep_ep::num_timeout_cycles)
++
++namespace deep_ep {
++
++extern int num_cpu_timeout_secs;
++#ifdef __CUDACC__
++// NOTES: defined in `runtime.cu`, which requires relocatable device code
++extern __device__ unsigned long long num_timeout_cycles;
++#endif
++
++} // namespace deep_ep
+
+ #define LOW_LATENCY_SEND_PHASE 1
+ #define LOW_LATENCY_RECV_PHASE 2
+--- a/csrc/kernels/runtime.cu
++++ b/csrc/kernels/runtime.cu
+@@ -12,6 +12,23 @@
+
+ namespace deep_ep {
+
++int num_cpu_timeout_secs = DEFAULT_NUM_CPU_TIMEOUT_SECS;
++__device__ unsigned long long num_timeout_cycles = DEFAULT_NUM_TIMEOUT_CYCLES;
++
++void set_timeout_seconds(int timeout_secs) {
++    EP_HOST_ASSERT(0 <= timeout_secs and timeout_secs <= 86400);
++    const unsigned long long cycles = timeout_secs == 0 ? DEFAULT_NUM_TIMEOUT_CYCLES
++                                                        : static_cast<unsigned long long>(timeout_secs) * NUM_TIMEOUT_CYCLES_PER_SEC;
++    CUDA_CHECK(cudaMemcpyToSymbol(num_timeout_cycles, &cycles, sizeof(cycles)));
++    num_cpu_timeout_secs = timeout_secs == 0 ? DEFAULT_NUM_CPU_TIMEOUT_SECS : timeout_secs;
++}
++
++int get_timeout_seconds() {
++    unsigned long long cycles;
++    CUDA_CHECK(cudaMemcpyFromSymbol(&cycles, num_timeout_cycles, sizeof(cycles)));
++    return static_cast<int>(cycles / NUM_TIMEOUT_CYCLES_PER_SEC);
++}
++
+ namespace intranode {
+
+ template<int kNumRanks>
+--- a/deep_ep/__init__.py
++++ b/deep_ep/__init__.py
+@@ -4,4 +4,4 @@ from .utils import EventOverlap
+ from .buffer import Buffer
+
+ # noinspection PyUnresolvedReferences
+-from deep_ep_cpp import Config
++from deep_ep_cpp import Config, get_timeout_seconds, set_timeout_seconds