diff --git a/3rdparty/CMakeLists.txt b/3rdparty/CMakeLists.txt index 32e4b7a04931..b48fbf8ab55b 100644 --- a/3rdparty/CMakeLists.txt +++ b/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 diff --git a/3rdparty/fetch_content.json b/3rdparty/fetch_content.json index 4db7c8262044..24c079b57a5a 100644 --- a/3rdparty/fetch_content.json +++ b/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", diff --git a/3rdparty/patches/deep_ep_runtime_timeout.patch b/3rdparty/patches/deep_ep_runtime_timeout.patch new file mode 100644 index 000000000000..14b4847aab15 --- /dev/null +++ b/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(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(cycles / NUM_TIMEOUT_CYCLES_PER_SEC); ++} ++ + namespace intranode { + + template +--- 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 diff --git a/3rdparty/patches/deepgemm_configurable_barrier_timeout.patch b/3rdparty/patches/deepgemm_configurable_barrier_timeout.patch new file mode 100644 index 000000000000..53a637cf0ba9 --- /dev/null +++ b/3rdparty/patches/deepgemm_configurable_barrier_timeout.patch @@ -0,0 +1,479 @@ +--- a/csrc/apis/config.hpp ++++ b/csrc/apis/config.hpp +@@ -27,6 +27,12 @@ static void register_apis(pybind11::module_& m) { + m.def("get_pdl", []() { + return *jit->default_launch_options.enable_pdl; + }); ++ m.def("set_barrier_timeout_seconds", [&](const int& new_barrier_timeout_seconds) { ++ runtime->set_barrier_timeout_seconds(new_barrier_timeout_seconds); ++ }); ++ m.def("get_barrier_timeout_seconds", [&]() { ++ return runtime->get_barrier_timeout_seconds(); ++ }); + m.def("use_deterministic_algorithms", [&](const bool enabled) { + heuristics_runtime->use_deterministic_algorithms(enabled); + }); +--- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp ++++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp +@@ -153,6 +153,7 @@ static void __instantiate_kernel() {{ + y.data_ptr(), + cumulative_local_expert_recv_stats_ptr, + num_tokens, ++ runtime->get_barrier_timeout_seconds(), + layout::SymBuffer<>(sym_buffer_ptrs, rank_idx), + tensor_map_l1_acts, + tensor_map_l1_weights, +--- a/csrc/jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp ++++ b/csrc/jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp +@@ -202,6 +202,7 @@ static void __instantiate_kernel() {{ + y.data_ptr(), + cumulative_local_expert_recv_stats_ptr, + num_tokens, ++ runtime->get_barrier_timeout_seconds(), + routed_scaling_factor, + l1_alphas, l2_alphas, a2_scales, + layout::SymBuffer<>(sym_buffer_ptrs, rank_idx), +--- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp ++++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp +@@ -220,6 +220,7 @@ static void __instantiate_kernel() {{ + y.data_ptr(), + cumulative_local_expert_recv_stats_ptr, + num_tokens, ++ runtime->get_barrier_timeout_seconds(), + layout::SymBuffer<>(sym_buffer_ptrs, rank_idx), + tensor_map_l1_acts, + tensor_map_l1_acts_sf, +--- a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp ++++ b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp +@@ -37,6 +37,7 @@ public: + void* y; + int* cumulative_local_expert_recv_stats; + int num_tokens; ++ int barrier_timeout_seconds; + layout::SymBuffer<> sym_buffer_ptrs; + + CUtensorMap tensor_map_l1_acts; +@@ -107,6 +108,7 @@ static void __instantiate_kernel() {{ + args.y, + args.cumulative_local_expert_recv_stats, + args.num_tokens, ++ args.barrier_timeout_seconds, + args.sym_buffer_ptrs, + args.tensor_map_l1_acts, + args.tensor_map_l1_acts_sf, +@@ -255,6 +257,7 @@ static void sm90_fp8_mega_moe( + .y = y.data_ptr(), + .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr, + .num_tokens = num_tokens, ++ .barrier_timeout_seconds = runtime->get_barrier_timeout_seconds(), + .sym_buffer_ptrs = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx), + .tensor_map_l1_acts = tensor_map_l1_acts, + .tensor_map_l1_acts_sf = tensor_map_l1_acts_sf, +--- a/csrc/runtime/jit.hpp ++++ b/csrc/runtime/jit.hpp +@@ -6,6 +6,7 @@ + + #include + #include ++#include + + namespace deep_gemm { + +@@ -27,6 +28,12 @@ inline void init_jit(const std::string& library_root_path) { + emplace_back("--diag-suppress=39,161,174,177,186,940"); + runtime->default_compiler_options.nvcc_flags-> + emplace_back("--compiler-options=-Wno-deprecated-declarations,-Wno-abi"); ++ // Allow overriding the in-kernel barrier timeout default, see `comm/barrier.cuh` ++ // NOTES: compiler flags are part of the kernel cache key, so changing this value ++ // naturally invalidates the JIT cache ++ if (const auto barrier_timeout_seconds = deep_jit::get_env("DG_JIT_BARRIER_TIMEOUT_SECONDS", 0); barrier_timeout_seconds > 0) ++ runtime->default_compiler_options.nvcc_flags-> ++ emplace_back("-DDG_BARRIER_TIMEOUT_SECONDS=" + std::to_string(barrier_timeout_seconds)); + // Preserve the existing 12.x family target when DeepJIT derives a different minor target. + if (runtime->device.get_arch_major() == 12 and runtime->device.get_arch_minor() != 0) + runtime->default_compiler_options.arch = "120f"; +--- a/csrc/runtime/runtime.hpp ++++ b/csrc/runtime/runtime.hpp +@@ -18,6 +18,7 @@ namespace deep_gemm { + class Runtime { + public: + int num_sms = 0, tc_util = 0; ++ int barrier_timeout_seconds = 0; + + // cuBLASLt utils + // cuBLAS will select worse heuristics when workspace > 16 MiB with shape, e.g. m=128, n=7168, k=16384 +@@ -92,6 +93,18 @@ public: + return tc_util == 0 ? 100 : tc_util; + } + ++ // In-kernel barrier timeout of subsequent launches, see `comm/barrier.cuh` ++ // NOTES: 0 selects the JIT-time default (`DG_JIT_BARRIER_TIMEOUT_SECONDS`, or 60 s); the value ++ // is a kernel argument rather than part of the kernel signature, so changing it does not recompile ++ void set_barrier_timeout_seconds(const int& new_barrier_timeout_seconds) { ++ DG_HOST_ASSERT(0 <= new_barrier_timeout_seconds and new_barrier_timeout_seconds <= 86400); ++ barrier_timeout_seconds = new_barrier_timeout_seconds; ++ } ++ ++ int get_barrier_timeout_seconds() const { ++ return barrier_timeout_seconds; ++ } ++ + }; + + inline auto runtime = deep_jit::LazyInit([](){ return std::make_shared(); }); +--- a/deep_gemm/__init__.py ++++ b/deep_gemm/__init__.py +@@ -22,6 +22,8 @@ from ._C import ( + set_block_size_multiple_of, + set_pdl, + get_pdl, ++ set_barrier_timeout_seconds, ++ get_barrier_timeout_seconds, + use_deterministic_algorithms, + ) + +--- a/deep_gemm/include/deep_gemm/comm/barrier.cuh ++++ b/deep_gemm/include/deep_gemm/comm/barrier.cuh +@@ -9,15 +9,36 @@ + + namespace deep_gemm::comm { + +-// 60s timeout, at 2 GHz +-constexpr int64_t kNumTimeoutCycles = 60ll * 2000000000ll; ++// Barrier timeout in seconds, overridable at JIT compile time. The host-side JIT runtime ++// turns the `DG_JIT_BARRIER_TIMEOUT_SECONDS` environment variable into this definition, ++// so slow-but-healthy rendezvous (e.g. large expert-parallel groups whose first ++// collective follows a long warm-up) can be distinguished from a genuine hang. ++#ifndef DG_BARRIER_TIMEOUT_SECONDS ++#define DG_BARRIER_TIMEOUT_SECONDS 60 ++#endif ++ ++// Timeout in cycles, at 2 GHz ++constexpr int64_t kNumCyclesPerSecond = 2000000000ll; ++constexpr int64_t kNumTimeoutCycles = static_cast(DG_BARRIER_TIMEOUT_SECONDS) * kNumCyclesPerSecond; ++ ++// Kernels also take a runtime timeout in seconds (see `Runtime::set_barrier_timeout_seconds`), ++// which can change between launches without recompiling; 0 selects the default above ++CUTLASS_DEVICE int64_t get_num_timeout_cycles(const uint32_t& timeout_seconds) { ++ return timeout_seconds > 0 ? static_cast(timeout_seconds) * kNumCyclesPerSecond : kNumTimeoutCycles; ++} ++ ++// Timeout reported by the diagnostics, in seconds ++CUTLASS_DEVICE long long get_timeout_seconds_for_print(const uint32_t& timeout_seconds) { ++ return static_cast(get_num_timeout_cycles(timeout_seconds) / kNumCyclesPerSecond); ++} + + // Spin until `pred()` holds; on timeout, `print_timeout()` runs before the assertion + template +-CUTLASS_DEVICE void wait_until(const pred_t& pred, const print_timeout_t& print_timeout) { ++CUTLASS_DEVICE void wait_until(const uint32_t& timeout_seconds, const pred_t& pred, const print_timeout_t& print_timeout) { ++ const auto num_timeout_cycles = get_num_timeout_cycles(timeout_seconds); + const auto start_clock = clock64(); + while (not pred()) { +- if (clock64() - start_clock >= kNumTimeoutCycles) { ++ if (clock64() - start_clock >= num_timeout_cycles) { + print_timeout(); + DG_DEVICE_ASSERT(false and "Timeout"); + } +@@ -36,6 +57,7 @@ template + CUTLASS_DEVICE void grid_sync(const WorkspaceT& workspace, + const uint32_t& sm_idx, const uint32_t& thread_idx, ++ const uint32_t& timeout_seconds, + const sync_scope_t& sync_scope) { + // NOTES: the implementation idea is from `cooperative_groups::this_grid().sync()` + static constexpr uint32_t kFinishSumTag = 0x80000000u; +@@ -46,15 +68,16 @@ CUTLASS_DEVICE void grid_sync(const WorkspaceT& workspace, + count_ptr, sm_idx == 0 ? (kFinishSumTag - (kNumSMs - 1)) : 1); + uint32_t new_value; + if constexpr (kTimeoutPolicy == BarrierTimeoutPolicy::Diagnostic) { +- wait_until([&]() { return ((new_value = ptx::ld_acq(count_ptr)) ^ old_value) & kFinishSumTag; }, [&]() { +- printf("DeepGEMM grid sync timeout: sm=%u, thread=%u, grid_sync_idx=%u, old=%u, current=%u, expected_tag=%u\n", +- sm_idx, thread_idx, kGridSyncIndex, old_value, new_value, old_value ^ kFinishSumTag); ++ wait_until(timeout_seconds, [&]() { return ((new_value = ptx::ld_acq(count_ptr)) ^ old_value) & kFinishSumTag; }, [&]() { ++ printf("DeepGEMM grid sync timeout (%llds): sm=%u, thread=%u, grid_sync_idx=%u, old=%u, current=%u, expected_tag=%u\n", ++ get_timeout_seconds_for_print(timeout_seconds), sm_idx, thread_idx, kGridSyncIndex, old_value, new_value, old_value ^ kFinishSumTag); + }); + } else { ++ const auto num_timeout_cycles = get_num_timeout_cycles(timeout_seconds); + const auto start_clock = clock64(); + do { + new_value = ptx::ld_acq(count_ptr); +- if (clock64() - start_clock >= kNumTimeoutCycles) { ++ if (clock64() - start_clock >= num_timeout_cycles) { + handle_grid_sync_timeout( + sm_idx, thread_idx, kGridSyncIndex, old_value, new_value, + old_value ^ kFinishSumTag); +@@ -72,6 +95,7 @@ template & sym_buffer, + const uint32_t& sm_idx, const uint32_t& thread_idx, ++ const uint32_t& timeout_seconds, + const sync_scope_t& sync_scope, + const bool& sync_prologue = true, + const bool& sync_epilogue = true) { +@@ -79,7 +103,7 @@ CUTLASS_DEVICE void nvlink_barrier(const WorkspaceT& workspace, + + // Grid sync before NVLink signaling + if (sync_prologue) +- grid_sync(workspace, sm_idx, thread_idx, sync_scope); ++ grid_sync(workspace, sm_idx, thread_idx, timeout_seconds, sync_scope); + + // NVLink cross-rank barrier, only SM 0 participates + if (sm_idx == 0) { +@@ -98,14 +122,15 @@ CUTLASS_DEVICE void nvlink_barrier(const WorkspaceT& workspace, + ptx::red_add(counter_ptr, 1); + const int target = signal_sign ? 0 : static_cast(kNumRanks); + if constexpr (kTimeoutPolicy == BarrierTimeoutPolicy::Diagnostic) { +- wait_until([&]() { return ptx::ld_acq_sys(signal_ptr) == target; }, [&]() { +- printf("DeepGEMM NVLink barrier timeout: rank=%d, counter=%d, signal=%d, target=%d, phase=%d, sign=%d, tag=%d\n", +- sym_buffer.rank_idx, *counter_ptr, ptx::ld_acq_sys(signal_ptr), target, signal_phase, signal_sign, kTag); ++ wait_until(timeout_seconds, [&]() { return ptx::ld_acq_sys(signal_ptr) == target; }, [&]() { ++ printf("DeepGEMM NVLink barrier timeout (%llds): rank=%d, counter=%d, signal=%d, target=%d, phase=%d, sign=%d, tag=%d\n", ++ get_timeout_seconds_for_print(timeout_seconds), sym_buffer.rank_idx, *counter_ptr, ptx::ld_acq_sys(signal_ptr), target, signal_phase, signal_sign, kTag); + }); + } else { ++ const auto num_timeout_cycles = get_num_timeout_cycles(timeout_seconds); + const auto start_clock = clock64(); + while (ptx::ld_acq_sys(signal_ptr) != target) { +- if (clock64() - start_clock >= kNumTimeoutCycles) { ++ if (clock64() - start_clock >= num_timeout_cycles) { + handle_nvlink_barrier_timeout( + sym_buffer.rank_idx, *counter_ptr, ptx::ld_acq_sys(signal_ptr), target, + signal_phase, signal_sign, kTag); +@@ -117,7 +142,7 @@ CUTLASS_DEVICE void nvlink_barrier(const WorkspaceT& workspace, + + // Grid sync after NVLink completion + if (sync_epilogue) +- grid_sync(workspace, sm_idx, thread_idx, sync_scope); ++ grid_sync(workspace, sm_idx, thread_idx, timeout_seconds, sync_scope); + } + + } // namespace deep_gemm::comm +--- a/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe.cuh ++++ b/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe.cuh +@@ -53,6 +53,7 @@ CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 1) void + sm100_bf16_mega_moe_impl(void* y, + int* cumulative_local_expert_recv_stats, + const uint32_t num_tokens, ++ const uint32_t barrier_timeout_seconds, + const __grid_constant__ layout::SymBuffer sym_buffer, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights, +@@ -340,7 +341,7 @@ sm100_bf16_mega_moe_impl(void* y, + + // Grid sync + comm::grid_sync( +- workspace, sm_idx, thread_idx, ++ workspace, sm_idx, thread_idx, barrier_timeout_seconds, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); } + ); + +@@ -370,7 +371,7 @@ sm100_bf16_mega_moe_impl(void* y, + // Barrier before pulling + comm::nvlink_barrier( +- workspace, sym_buffer, sm_idx, thread_idx, ++ workspace, sym_buffer, sm_idx, thread_idx, barrier_timeout_seconds, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + /* After the grid sync above, there is no more writes by other SMs (except 0) */ false, + /* After the NVLink barrier, there is a grid sync */ true +@@ -611,7 +612,7 @@ sm100_bf16_mega_moe_impl(void* y, + // Wait for all ranks to finish cleaning + comm::nvlink_barrier( +- workspace, sym_buffer, sm_idx, thread_idx, ++ workspace, sym_buffer, sm_idx, thread_idx, barrier_timeout_seconds, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + /* Before the NVLink barrier, there is a grid sync */ true, + /* At the end of kernel does not need to sync */ false +@@ -1189,7 +1190,7 @@ sm100_bf16_mega_moe_impl(void* y, + + // All local L2 writes are done after this grid sync + comm::grid_sync( +- workspace, sm_idx, epilogue_thread_idx, ++ workspace, sm_idx, epilogue_thread_idx, barrier_timeout_seconds, + [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } + ); + +@@ -1216,8 +1217,10 @@ sm100_bf16_mega_moe_impl(void* y, + const bool is_routed = lane_idx < kNumTopk and stored_topk_slot_idx >= 0; + const auto peer_ready_ptr = workspace.get_combine_ready_grid_idx_ptr( + is_routed ? static_cast(stored_topk_slot_idx) / kNumExpertsPerRank : 0); +- comm::wait_until([&]() { return __all_sync(0xffffffff, not is_routed or ptx::ld_acq_sys(peer_ready_ptr) == grid_idx); }, +- [&]() { printf("DeepGEMM combine peers timeout: rank=%u, token=%u\n", sym_buffer.rank_idx, token_idx); }); ++ comm::wait_until(barrier_timeout_seconds, ++ [&]() { return __all_sync(0xffffffff, not is_routed or ptx::ld_acq_sys(peer_ready_ptr) == grid_idx); }, ++ [&]() { printf("DeepGEMM combine peers timeout (%llds): rank=%u, token=%u\n", ++ comm::get_timeout_seconds_for_print(barrier_timeout_seconds), sym_buffer.rank_idx, token_idx); }); + + const uint32_t chunk_byte_offset = chunk_idx * kNumChunkBytes; + +--- a/deep_gemm/include/deep_gemm/impls/sm100_fp4_fp4_mega_moe.cuh ++++ b/deep_gemm/include/deep_gemm/impls/sm100_fp4_fp4_mega_moe.cuh +@@ -107,6 +107,7 @@ CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 1) void + sm100_fp4_fp4_mega_moe_impl(void* y, + int* cumulative_local_expert_recv_stats, + const uint32_t num_tokens, ++ const uint32_t barrier_timeout_seconds, + const float routed_scaling_factor, + // Optional per-local-expert scales: L1 has separate gate/up factors + // (e.g. modelopt's per-projection `weight_scale_2`) +@@ -453,7 +454,7 @@ sm100_fp4_fp4_mega_moe_impl(void* y, + + // Grid sync + comm::grid_sync( +- workspace, sm_idx, thread_idx, ++ workspace, sm_idx, thread_idx, barrier_timeout_seconds, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); } + ); + +@@ -477,7 +478,7 @@ sm100_fp4_fp4_mega_moe_impl(void* y, + // Barrier before pulling + comm::nvlink_barrier( +- workspace, sym_buffer, sm_idx, thread_idx, ++ workspace, sym_buffer, sm_idx, thread_idx, barrier_timeout_seconds, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + /* After the grid sync above, there is no more writes by other SMs (except 0) */ false, + /* After the NVLink barrier, there is a grid sync */ true +@@ -737,7 +738,7 @@ sm100_fp4_fp4_mega_moe_impl(void* y, + // Wait for all ranks to finish cleaning + comm::nvlink_barrier( +- workspace, sym_buffer, sm_idx, thread_idx, ++ workspace, sym_buffer, sm_idx, thread_idx, barrier_timeout_seconds, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + /* Before the NVLink barrier, there is a grid sync */ true, + /* At the end of kernel does not need to sync */ false +@@ -1584,7 +1585,7 @@ sm100_fp4_fp4_mega_moe_impl(void* y, + // NVLink barrier (grid sync + cross-rank signal + grid sync): ~4 us + comm::nvlink_barrier( +- workspace, sym_buffer, sm_idx, epilogue_thread_idx, ++ workspace, sym_buffer, sm_idx, epilogue_thread_idx, barrier_timeout_seconds, + [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } + ); + +--- a/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh ++++ b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh +@@ -60,6 +60,7 @@ CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 1) void + sm100_fp8_fp4_mega_moe_impl(void* y, + int* cumulative_local_expert_recv_stats, + const uint32_t num_tokens, ++ const uint32_t barrier_timeout_seconds, + const __grid_constant__ layout::SymBuffer sym_buffer, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts_sf, +@@ -390,7 +391,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, + + // Grid sync + comm::grid_sync( +- workspace, sm_idx, thread_idx, ++ workspace, sm_idx, thread_idx, barrier_timeout_seconds, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); } + ); + +@@ -420,7 +421,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, + // Barrier before pulling + comm::nvlink_barrier( +- workspace, sym_buffer, sm_idx, thread_idx, ++ workspace, sym_buffer, sm_idx, thread_idx, barrier_timeout_seconds, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + /* After the grid sync above, there is no more writes by other SMs (except 0) */ false, + /* After the NVLink barrier, there is a grid sync */ true +@@ -680,7 +681,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, + // Wait for all ranks to finish cleaning + comm::nvlink_barrier( +- workspace, sym_buffer, sm_idx, thread_idx, ++ workspace, sym_buffer, sm_idx, thread_idx, barrier_timeout_seconds, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + /* Before the NVLink barrier, there is a grid sync */ true, + /* At the end of kernel does not need to sync */ false +@@ -1397,7 +1398,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, + + // All local L2 writes are done after this grid sync + comm::grid_sync( +- workspace, sm_idx, epilogue_thread_idx, ++ workspace, sm_idx, epilogue_thread_idx, barrier_timeout_seconds, + [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } + ); + +@@ -1424,8 +1425,10 @@ sm100_fp8_fp4_mega_moe_impl(void* y, + const bool is_routed = lane_idx < kNumTopk and stored_topk_slot_idx >= 0; + const auto peer_ready_ptr = workspace.get_combine_ready_grid_idx_ptr( + is_routed ? static_cast(stored_topk_slot_idx) / kNumExpertsPerRank : 0); +- comm::wait_until([&]() { return __all_sync(0xffffffff, not is_routed or ptx::ld_acq_sys(peer_ready_ptr) == grid_idx); }, +- [&]() { printf("DeepGEMM combine peers timeout: rank=%u, token=%u\n", sym_buffer.rank_idx, token_idx); }); ++ comm::wait_until(barrier_timeout_seconds, ++ [&]() { return __all_sync(0xffffffff, not is_routed or ptx::ld_acq_sys(peer_ready_ptr) == grid_idx); }, ++ [&]() { printf("DeepGEMM combine peers timeout (%llds): rank=%u, token=%u\n", ++ comm::get_timeout_seconds_for_print(barrier_timeout_seconds), sym_buffer.rank_idx, token_idx); }); + + const uint32_t chunk_byte_offset = chunk_idx * kNumChunkBytes; + +--- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh ++++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh +@@ -115,6 +115,7 @@ struct MegaMoEPhasePolicy { + void* y, \ + int* cumulative_local_expert_recv_stats, \ + const uint32_t num_tokens, \ ++ const uint32_t barrier_timeout_seconds, \ + const __grid_constant__ layout::SymBuffer sym_buffer, \ + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts, \ + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts_sf, \ +@@ -130,6 +131,7 @@ struct MegaMoEPhasePolicy { + void* y, \ + int* cumulative_local_expert_recv_stats, \ + const uint32_t num_tokens, \ ++ const uint32_t barrier_timeout_seconds, \ + const layout::SymBuffer& sym_buffer, \ + const cute::TmaDescriptor& tensor_map_l1_acts, \ + const cute::TmaDescriptor& tensor_map_l1_acts_sf, \ +@@ -142,7 +144,7 @@ struct MegaMoEPhasePolicy { + const float* __restrict__ l2_weights_sf + + #define DG_SM90_FP8_MOE_KERNEL_ARGS \ +- y, cumulative_local_expert_recv_stats, num_tokens, sym_buffer, \ ++ y, cumulative_local_expert_recv_stats, num_tokens, barrier_timeout_seconds, sym_buffer, \ + tensor_map_l1_acts, tensor_map_l1_acts_sf, tensor_map_l1_weights, \ + l1_weights_sf, tensor_map_l1_output, tensor_map_l2_acts, \ + tensor_map_l2_acts_sf, tensor_map_l2_weights, l2_weights_sf +@@ -543,7 +545,7 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { + kNumRanks, kNumSMs, kNumDispatchThreads, + kDispatchGridSyncIndex, kAfterWorkspaceCleanBarrierTag, + comm::BarrierTimeoutPolicy::TrapOnly>( +- workspace, sym_buffer, sm_idx, thread_idx, ++ workspace, sym_buffer, sm_idx, thread_idx, barrier_timeout_seconds, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + true, false); + return; +@@ -593,7 +595,7 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { + comm::grid_sync< + kNumSMs, kDispatchGridSyncIndex, + comm::BarrierTimeoutPolicy::TrapOnly>( +- workspace, sm_idx, thread_idx, ++ workspace, sm_idx, thread_idx, barrier_timeout_seconds, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); } + ); + +@@ -617,7 +619,7 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { + kNumRanks, kNumSMs, kNumDispatchThreads, + kDispatchGridSyncIndex, kBeforeDispatchPullBarrierTag, + comm::BarrierTimeoutPolicy::TrapOnly>( +- workspace, sym_buffer, sm_idx, thread_idx, ++ workspace, sym_buffer, sm_idx, thread_idx, barrier_timeout_seconds, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + false, true); + +@@ -2182,7 +2184,7 @@ sm90_fp8_mega_moe_core(DG_SM90_FP8_MOE_CORE_ARGS_DECL) { + kNumRanks, kNumSMs, kNumEpilogueThreads, + kEpilogueGridSyncIndex, kBeforeCombineReduceBarrierTag, + comm::BarrierTimeoutPolicy::TrapOnly>( +- workspace, sym_buffer, sm_idx, epilogue_thread_idx, ++ workspace, sym_buffer, sm_idx, epilogue_thread_idx, barrier_timeout_seconds, + [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } + ); + // Sync with dispatch (paired with dispatch's pre-cleanup sync) so that diff --git a/3rdparty/patches/nccl_ep_group_timeout.patch b/3rdparty/patches/nccl_ep_group_timeout.patch new file mode 100644 index 000000000000..5f4feaf55bab --- /dev/null +++ b/3rdparty/patches/nccl_ep_group_timeout.patch @@ -0,0 +1,280 @@ +--- a/nccl_ep/include/nccl_ep.h ++++ b/nccl_ep/include/nccl_ep.h +@@ -1068,6 +1068,32 @@ ncclResult_t ncclEpGetAsyncError(ncclEpGroup_t ep_group, int* error_out); + + ncclResult_t ncclEpErrorClear(ncclEpGroup_t ep_group); + ++// Set the timeout of the GPU-side wait loops for subsequent launches on this group. ++// Host-side only: each launch copies the timeout into its kernel arguments, so a ++// launch captured into a CUDA graph keeps the timeout that was set at capture. ++// Not thread-safe with concurrent launches on the same group. ++// ++// Arguments: ++// ep_group - [IN] EP group ++// timeout_ns - [IN] Timeout in nanoseconds, at millisecond resolution like ++// config.timeout_ns. 0 restores the timeout resolved by ++// ncclEpCreateGroup (NCCL_EP_TIMEOUT_MS > config.timeout_ns > default). ++// ++// Returns: ncclResult_t error code; ncclInvalidArgument for a nonzero timeout below ++// one millisecond or one that overflows the cycle count. ++ ++ncclResult_t ncclEpGroupSetTimeout(ncclEpGroup_t ep_group, uint64_t timeout_ns); ++ ++// Get the timeout of the GPU-side wait loops for subsequent launches on this group. ++// ++// Arguments: ++// ep_group - [IN] EP group ++// timeout_ns - [OUT] Timeout in nanoseconds, at millisecond resolution ++// ++// Returns: ncclResult_t error code ++ ++ncclResult_t ncclEpGroupGetTimeout(ncclEpGroup_t ep_group, uint64_t* timeout_ns); ++ + #ifdef __cplusplus + } + #endif +--- a/nccl_ep/nccl_ep.cc ++++ b/nccl_ep/nccl_ep.cc +@@ -788,6 +788,8 @@ struct ncclEpGroup { + int* mask_buffer = nullptr; // Device: int[nRanks], null if !enable_mask + int* async_error_flag = nullptr; // Host-pinned: 0 = ok, 1 = timeout occurred + uint64_t timeout_cycles = NUM_TIMEOUT_CYCLES; // GPU clock cycles for wait-loop timeout ++ uint64_t creation_timeout_cycles = NUM_TIMEOUT_CYCLES; // timeout_cycles resolved at creation ++ uint64_t clock_khz = 0; // SM clock rate that converts timeouts to cycles + + // Physical node properties (CUDA device assignment, IPC between co-located GPUs) + int gpus_per_node; // Physical GPUs per node (nRanks / nNodes) +@@ -1923,6 +1925,8 @@ ncclResult_t ncclEpCreateGroup(ncclEpGroup_t* out_ep_group, ncclComm_t comm, con + } + + ep_group->timeout_cycles = resolved; ++ ep_group->creation_timeout_cycles = resolved; ++ ep_group->clock_khz = clock_khz; + if (ep_group->rank == 0) { + uint64_t timeout_ms = resolved / (clock_khz * 1000ULL / 1000ULL); + char env_str[32]; +@@ -5488,6 +5492,30 @@ ncclResult_t ncclEpErrorClear(ncclEpGroup_t ep_group) { + __atomic_store_n(ep_group->async_error_flag, 0, __ATOMIC_RELEASE); + return ncclSuccess; + } ++ ++ncclResult_t ncclEpGroupSetTimeout(ncclEpGroup_t ep_group, uint64_t timeout_ns) { ++ EP_HOST_ASSERT(ep_group != nullptr); ++ EP_HOST_ASSERT(ep_group->clock_khz > 0); ++ if (timeout_ns == 0) { ++ ep_group->timeout_cycles = ep_group->creation_timeout_cycles; ++ return ncclSuccess; ++ } ++ // Same millisecond resolution and conversion as config.timeout_ns in ncclEpCreateGroup ++ const uint64_t timeout_ms = timeout_ns / 1000000ULL; ++ if (timeout_ms == 0 || timeout_ms > UINT64_MAX / ep_group->clock_khz) { ++ return ncclInvalidArgument; ++ } ++ ep_group->timeout_cycles = ep_group->clock_khz * timeout_ms; ++ return ncclSuccess; ++} ++ ++ncclResult_t ncclEpGroupGetTimeout(ncclEpGroup_t ep_group, uint64_t* timeout_ns) { ++ EP_HOST_ASSERT(ep_group != nullptr); ++ EP_HOST_ASSERT(timeout_ns != nullptr); ++ EP_HOST_ASSERT(ep_group->clock_khz > 0); ++ *timeout_ns = ep_group->timeout_cycles / ep_group->clock_khz * 1000000ULL; ++ return ncclSuccess; ++} + // ── Test-only internal helpers (declared in nccl_ep_test_internal.h) ───────── + // Exposed via a separate test header so library consumers cannot accidentally + // depend on these implementation details. +--- a/python/nccl/_extensions/bindings/_internal/nccl_ep.pxd ++++ b/python/nccl/_extensions/bindings/_internal/nccl_ep.pxd +@@ -5,6 +5,7 @@ + # This code was automatically generated with version 0.2.0. Do not modify it directly. + + from ..cynccl_ep cimport * ++from libc.stdint cimport uint64_t + + + ############################################################################### +@@ -29,3 +30,5 @@ cdef ncclResult_t _ncclEpMaskUpdate(ncclEpGroup_t ep_group, const int* mask, cud + cdef ncclResult_t _ncclEpMaskClean(ncclEpGroup_t ep_group, cudaStream_t stream) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil + cdef ncclResult_t _ncclEpGetAsyncError(ncclEpGroup_t ep_group, int* error_out) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil + cdef ncclResult_t _ncclEpErrorClear(ncclEpGroup_t ep_group) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil ++cdef ncclResult_t _ncclEpGroupSetTimeout(ncclEpGroup_t ep_group, uint64_t timeout_ns) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil ++cdef ncclResult_t _ncclEpGroupGetTimeout(ncclEpGroup_t ep_group, uint64_t* timeout_ns) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil +--- a/python/nccl/_extensions/bindings/_internal/nccl_ep_linux.pyx ++++ b/python/nccl/_extensions/bindings/_internal/nccl_ep_linux.pyx +@@ -43,7 +43,7 @@ cdef extern from "": + void* _cyb_dlsym "dlsym"(void*, const char*) nogil + const void * _cyb_RTLD_DEFAULT "RTLD_DEFAULT" + +-from libc.stdint cimport intptr_t ++from libc.stdint cimport intptr_t, uint64_t + + import threading as _cyb_threading + +@@ -143,6 +143,8 @@ cdef void* __ncclEpMaskUpdate = NULL + cdef void* __ncclEpMaskClean = NULL + cdef void* __ncclEpGetAsyncError = NULL + cdef void* __ncclEpErrorClear = NULL ++cdef void* __ncclEpGroupSetTimeout = NULL ++cdef void* __ncclEpGroupGetTimeout = NULL + + cdef int _init_nccl_ep() except -1 nogil: + global _cyb___py_nccl_ep_init +@@ -276,6 +278,20 @@ cdef int _init_nccl_ep() except -1 nogil: + handle = load_library() + __ncclEpErrorClear = _cyb_dlsym(handle, 'ncclEpErrorClear') + ++ global __ncclEpGroupSetTimeout ++ __ncclEpGroupSetTimeout = _cyb_dlsym(_cyb_RTLD_DEFAULT, 'ncclEpGroupSetTimeout') ++ if __ncclEpGroupSetTimeout == NULL: ++ if handle == NULL: ++ handle = load_library() ++ __ncclEpGroupSetTimeout = _cyb_dlsym(handle, 'ncclEpGroupSetTimeout') ++ ++ global __ncclEpGroupGetTimeout ++ __ncclEpGroupGetTimeout = _cyb_dlsym(_cyb_RTLD_DEFAULT, 'ncclEpGroupGetTimeout') ++ if __ncclEpGroupGetTimeout == NULL: ++ if handle == NULL: ++ handle = load_library() ++ __ncclEpGroupGetTimeout = _cyb_dlsym(handle, 'ncclEpGroupGetTimeout') ++ + _cyb_atomic_int_store(&_cyb___py_nccl_ep_init, 1) + return 0 + +@@ -346,6 +362,12 @@ cpdef dict _inspect_function_pointers(): + + global __ncclEpErrorClear + data["__ncclEpErrorClear"] = __ncclEpErrorClear ++ ++ global __ncclEpGroupSetTimeout ++ data["__ncclEpGroupSetTimeout"] = __ncclEpGroupSetTimeout ++ ++ global __ncclEpGroupGetTimeout ++ data["__ncclEpGroupGetTimeout"] = __ncclEpGroupGetTimeout + _cyb_func_ptrs = data + return data + +@@ -587,3 +609,23 @@ cdef ncclResult_t _ncclEpErrorClear(ncclEpGroup_t ep_group) except?_NCCLRESULT_T + raise FunctionNotFoundError("function ncclEpErrorClear is not found") + return (__ncclEpErrorClear)( + ep_group) ++ ++ ++cdef ncclResult_t _ncclEpGroupSetTimeout(ncclEpGroup_t ep_group, uint64_t timeout_ns) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil: ++ global __ncclEpGroupSetTimeout ++ _check_or_init_nccl_ep() ++ if __ncclEpGroupSetTimeout == NULL: ++ with gil: ++ raise FunctionNotFoundError("function ncclEpGroupSetTimeout is not found") ++ return (__ncclEpGroupSetTimeout)( ++ ep_group, timeout_ns) ++ ++ ++cdef ncclResult_t _ncclEpGroupGetTimeout(ncclEpGroup_t ep_group, uint64_t* timeout_ns) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil: ++ global __ncclEpGroupGetTimeout ++ _check_or_init_nccl_ep() ++ if __ncclEpGroupGetTimeout == NULL: ++ with gil: ++ raise FunctionNotFoundError("function ncclEpGroupGetTimeout is not found") ++ return (__ncclEpGroupGetTimeout)( ++ ep_group, timeout_ns) +--- a/python/nccl/_extensions/bindings/cynccl_ep.pxd ++++ b/python/nccl/_extensions/bindings/cynccl_ep.pxd +@@ -215,3 +215,5 @@ cdef ncclResult_t ncclEpMaskUpdate(ncclEpGroup_t ep_group, const int* mask, cuda + cdef ncclResult_t ncclEpMaskClean(ncclEpGroup_t ep_group, cudaStream_t stream) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil + cdef ncclResult_t ncclEpGetAsyncError(ncclEpGroup_t ep_group, int* error_out) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil + cdef ncclResult_t ncclEpErrorClear(ncclEpGroup_t ep_group) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil ++cdef ncclResult_t ncclEpGroupSetTimeout(ncclEpGroup_t ep_group, uint64_t timeout_ns) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil ++cdef ncclResult_t ncclEpGroupGetTimeout(ncclEpGroup_t ep_group, uint64_t* timeout_ns) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil +--- a/python/nccl/_extensions/bindings/cynccl_ep.pyx ++++ b/python/nccl/_extensions/bindings/cynccl_ep.pyx +@@ -81,3 +81,11 @@ cdef ncclResult_t ncclEpGetAsyncError(ncclEpGroup_t ep_group, int* error_out) ex + + cdef ncclResult_t ncclEpErrorClear(ncclEpGroup_t ep_group) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil: + return _nccl_ep._ncclEpErrorClear(ep_group) ++ ++ ++cdef ncclResult_t ncclEpGroupSetTimeout(ncclEpGroup_t ep_group, uint64_t timeout_ns) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil: ++ return _nccl_ep._ncclEpGroupSetTimeout(ep_group, timeout_ns) ++ ++ ++cdef ncclResult_t ncclEpGroupGetTimeout(ncclEpGroup_t ep_group, uint64_t* timeout_ns) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil: ++ return _nccl_ep._ncclEpGroupGetTimeout(ep_group, timeout_ns) +--- a/python/nccl/_extensions/bindings/nccl_ep.pxd ++++ b/python/nccl/_extensions/bindings/nccl_ep.pxd +@@ -8,7 +8,7 @@ + + # <<<< PREAMBLE CONTENT >>>> + +-from libc.stdint cimport intptr_t ++from libc.stdint cimport intptr_t, uint64_t + + + # <<<< END OF PREAMBLE CONTENT >>>> +@@ -66,4 +66,6 @@ cpdef mask_update(intptr_t ep_group, intptr_t mask, intptr_t stream) + cpdef mask_clean(intptr_t ep_group, intptr_t stream) + cpdef int get_async_error(intptr_t ep_group) except? -1 + cpdef error_clear(intptr_t ep_group) ++cpdef group_set_timeout(intptr_t ep_group, uint64_t timeout_ns) ++cpdef uint64_t group_get_timeout(intptr_t ep_group) except? 0 + cpdef object get_library_path() +--- a/python/nccl/_extensions/bindings/nccl_ep.pyx ++++ b/python/nccl/_extensions/bindings/nccl_ep.pyx +@@ -10,7 +10,7 @@ + cimport cpython as _cyb_cpython + cimport cpython.buffer as _cyb_cpython_buffer + from cython cimport view as _cyb_view +-from libc.stdint cimport intptr_t ++from libc.stdint cimport intptr_t, uint64_t + from libc.stdlib cimport ( + calloc as _cyb_calloc, + free as _cyb_free, +@@ -2810,6 +2810,20 @@ cpdef error_clear(intptr_t ep_group): + check_status(__status__) + + ++cpdef group_set_timeout(intptr_t ep_group, uint64_t timeout_ns): ++ with nogil: ++ __status__ = ncclEpGroupSetTimeout(ep_group, timeout_ns) ++ check_status(__status__) ++ ++ ++cpdef uint64_t group_get_timeout(intptr_t ep_group) except? 0: ++ cdef uint64_t timeout_ns ++ with nogil: ++ __status__ = ncclEpGroupGetTimeout(ep_group, &timeout_ns) ++ check_status(__status__) ++ return timeout_ns ++ ++ + # Hand-written: reports the .so the symbols resolved to; not a C entry point. + cpdef object get_library_path(): + from ._internal.nccl_ep import _inspect_loaded_library_path +--- a/python/nccl/ep/group.py ++++ b/python/nccl/ep/group.py +@@ -193,6 +193,25 @@ class Group: + ) + return Handle(ptr) + ++ def set_timeout_ns(self, timeout_ns: int) -> None: ++ """Set the timeout of the GPU-side wait loops for subsequent launches. ++ ++ Each launch copies the timeout into its kernel arguments, so a launch ++ captured into a CUDA graph keeps the timeout that was set at capture. ++ ++ Args: ++ timeout_ns: Timeout in nanoseconds, at millisecond resolution. ++ ``0`` restores the timeout resolved by :meth:`create` ++ (``NCCL_EP_TIMEOUT_MS`` > ``GroupConfig.timeout_ns`` > default). ++ """ ++ self._check_valid("set_timeout_ns") ++ _ep_bindings.group_set_timeout(self._ptr, timeout_ns) ++ ++ def get_timeout_ns(self) -> int: ++ """Timeout of the GPU-side wait loops for subsequent launches, in nanoseconds.""" ++ self._check_valid("get_timeout_ns") ++ return _ep_bindings.group_get_timeout(self._ptr) ++ + def destroy(self) -> None: + """Release the group. Subsequent operations on this object are invalid.""" + if self._ptr: diff --git a/cpp/tensorrt_llm/thop/moe/communication/moeAlltoAllOp.cpp b/cpp/tensorrt_llm/thop/moe/communication/moeAlltoAllOp.cpp index fd15b798e71d..edb2a5ad08c6 100644 --- a/cpp/tensorrt_llm/thop/moe/communication/moeAlltoAllOp.cpp +++ b/cpp/tensorrt_llm/thop/moe/communication/moeAlltoAllOp.cpp @@ -40,15 +40,29 @@ namespace torch_ext namespace moe_comm { -// Process-wide budget copied into subsequent dispatch/combine launch arguments. +// Process-wide completion-flag budget copied into subsequent dispatch/combine launch arguments, +// so a captured CUDA graph keeps the budget that was current at capture. +// +// No collective separates a rank's first-touch JIT/autotune work from its dispatch launch, so +// this device-side budget is in effect a deadline on the slowest peer's host-side progress. +// Python raises it during warmup (tensorrt_llm/_torch/moe/fused_moe/moe_comm_timeout_guard.py). static std::atomic gTimeoutCycles{tensorrt_llm::kernels::moe_comm::kDefaultTimeoutCycles}; +// Set the budget of subsequent launches in seconds; 0 restores the default. void moeA2ASetTimeoutOp(int64_t timeoutSec) { constexpr int64_t kMaxTimeoutSec = 24 * 60 * 60; - TORCH_CHECK(timeoutSec > 0 && timeoutSec <= kMaxTimeoutSec, "MoE all-to-all timeout must be in 1..", kMaxTimeoutSec, - " seconds"); - gTimeoutCycles.store(timeoutSec * tensorrt_llm::kernels::moe_comm::kAssumedClockHz, std::memory_order_relaxed); + TORCH_CHECK(timeoutSec >= 0 && timeoutSec <= kMaxTimeoutSec, "MoE all-to-all timeout must be in 0..", + kMaxTimeoutSec, " seconds (0 restores the default), got ", timeoutSec); + int64_t const timeoutCycles = timeoutSec == 0 ? tensorrt_llm::kernels::moe_comm::kDefaultTimeoutCycles + : timeoutSec * tensorrt_llm::kernels::moe_comm::kAssumedClockHz; + gTimeoutCycles.store(timeoutCycles, std::memory_order_relaxed); +} + +// Return the budget of subsequent launches in seconds. +int64_t moeA2AGetTimeoutOp() +{ + return gTimeoutCycles.load(std::memory_order_relaxed) / tensorrt_llm::kernels::moe_comm::kAssumedClockHz; } static constexpr size_t CACHELINE_ALIGNMENT = 128; @@ -1083,6 +1097,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, module) "moe_a2a_get_combine_payload_tensor(Tensor(a) workspace, Tensor metainfo, int ep_rank, int ep_size, " "int runtime_max_tokens_per_rank, ScalarType out_dtype, int hidden_size) -> Tensor(a)"); module.def("moe_a2a_set_timeout(int timeout_sec) -> ()", &tensorrt_llm::torch_ext::moe_comm::moeA2ASetTimeoutOp); + module.def("moe_a2a_get_timeout() -> int", &tensorrt_llm::torch_ext::moe_comm::moeA2AGetTimeoutOp); module.def( "moe_a2a_get_workspace_layout(int ep_size, int max_num_tokens_per_rank, int top_k, " "int dispatch_payload_bytes, int combine_input_bytes, int combine_recv_bytes, " diff --git a/docs/source/developer-guide/overview.md b/docs/source/developer-guide/overview.md index 2631d6130d61..8cdd2e122bc3 100644 --- a/docs/source/developer-guide/overview.md +++ b/docs/source/developer-guide/overview.md @@ -186,6 +186,15 @@ The value is baked into the kernels at JIT compile time and participates in the JIT cache key, so changing it recompiles the affected kernels rather than silently reusing a cached binary. +TensorRT-LLM overrides this default at runtime, without recompiling. While the +engine warms up, the barriers use `TRTLLM_MOE_COMM_WARMUP_TIMEOUT_SEC` (default +1800 seconds), because one-time per-rank work such as kernel JIT compilation can +delay a healthy rank by minutes. While serving, they use +`TRTLLM_MOE_COMM_TIMEOUT_SEC` when it is set, so `DG_JIT_BARRIER_TIMEOUT_SECONDS` +only sets the serving budget when `TRTLLM_MOE_COMM_TIMEOUT_SEC` is unset. Both +variables apply to every MoE communication kernel; see "Communication timeouts" +in `tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md`. + This is primarily a diagnostic knob: raising it distinguishes a slow-but-healthy rendezvous (all ranks eventually arrive and the run proceeds) from a genuine hang or a rank-count mismatch (the same `signal`/`target` split reappears at the diff --git a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index b5f3ecb95e1f..99891bd3d3a4 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -142,6 +142,25 @@ EPLB hooks fire only at the first/last chunk of the first/last `repeat_idx`. Mul No external `Communication.dispatch` / `.combine`. Zero-token chunks still launch the kernel so peer EP ranks can cross the in-kernel NVLink barrier. +The DeepGEMM kernel's in-kernel barriers trap when a peer is late; their timeout is relaxed during warmup, see [Communication timeouts](#communication-timeouts). + +### Communication timeouts + +Several MoE communication kernels wait for their peers inside the kernel and trap once a peer is late by more than a device-side timeout. The trap is a sticky CUDA error, so every waiting rank dies. During warmup, one-time per-rank work (kernel JIT, autotuning, module loading) can delay a healthy rank by minutes with no collective in between, so these timeouts are relaxed while the engine warms up: + +- `moe_comm_timeout_guard.py` owns the process-wide `MoECommTimeoutGuard`. It resolves two budgets, `TRTLLM_MOE_COMM_WARMUP_TIMEOUT_SEC` (default 1800 s) and `TRTLLM_MOE_COMM_TIMEOUT_SEC` (default: each library's native timeout), and on each transition pushes the current one to every registered `MoECommTimeoutProxy`. `CommunicationFactory.create_strategy` resolves the budgets before it tries any strategy, so an invalid value fails there instead of reading as an unavailable backend. The NVLinkOneSided-only names `TRTLLM_MOE_A2A_[WARMUP_]TIMEOUT_SEC` and `TRTLLM_NVLINK_ONE_SIDED_A2A_[WARMUP_]TIMEOUT_SEC` are deprecated aliases that now apply to every backend. +- The model engine's `is_warmup` setter calls `set_moe_comm_warmup`. The CUDA-graph pass that captures runs inside `moe_comm_serving_timeouts()`, because captured graphs keep the timeouts that were current at capture. Switching while a CUDA graph is being captured raises. +- Each backend with such a timeout registers a proxy when it is constructed. The guard holds proxies weakly, so the backend must keep its proxy alive: + +| Backend | Proxy applies the timeout through | Native serving timeout | +|---------|-----------------------------------|------------------------| +| NVLinkOneSided | `torch.ops.trtllm.moe_a2a_set_timeout`; copied into launch arguments | 300 s | +| DeepGEMM MegaMoE | `deep_gemm.set_barrier_timeout_seconds` (`3rdparty/patches/deepgemm_configurable_barrier_timeout.patch`); a kernel argument | 60 s, or `DG_JIT_BARRIER_TIMEOUT_SECONDS` | +| DeepEP normal kernels | `deep_ep.set_timeout_seconds` (`3rdparty/patches/deep_ep_runtime_timeout.patch`); a device variable read at execution | ~100 s | +| NcclEP | `nccl.ep.Group.set_timeout_ns`, one proxy per group (`3rdparty/patches/nccl_ep_group_timeout.patch`); copied into launch arguments | `NCCL_EP_TIMEOUT_MS`, else `GroupConfig.timeout_ns`, else ~100 s | + +Timeouts are nominal: NVLinkOneSided, DeepGEMM and DeepEP convert seconds at an assumed 2 GHz SM clock, while NcclEP uses the device's clock rate. + ### Core Design Principles 1. **Composition over inheritance** — Backend, Communication, EPLB, and Scheduler are independent, composable components @@ -194,6 +213,7 @@ Still on old path (standalone, with embedded communication): | `impl_blocks.py` | The blocks `MoE` and `MoEImplBase` share — `MoEExecutionContractMixin` (scheduler-facing declarations, `forward_fake`) and `MoEWeightOwnerMixin` (`create_weights` / `load_weights` / `_check_configs`) | | `quantization.py` | Quantization method implementations (`FusedMoEMethod` subclasses: weight creation, loading, quant/dequant ops per quant mode) | | `routing.py` | Routing methods (`TopKRouting`, etc.) | +| `moe_comm_timeout_guard.py` | Device-side timeouts of communication kernels, relaxed during warmup — budget resolution and `MoECommTimeoutProxy` registration (see [Communication timeouts](#communication-timeouts)) | | `moe_load_balancer.py` | EPLB implementation | | `moe_op_backend.py` | Op backend registry for TRTLLMGen (flashinfer/trtllm ops) | @@ -764,6 +784,7 @@ The `XXFusedMoE` names remain supported as module-level aliases — `DeepGemmFus - **Do NOT substitute a backend without recording it** — A degradation must be visible in the `MoEResolutionReport`, not only in a log line - **Do NOT pick `scheduler_kind` opportunistically** — Use `EXTERNAL_COMM` (default) unless your backend's fused kernel genuinely owns cross-rank exchange via SymmBuffer / equivalent in-kernel collective; `FUSED_COMM` brings hard invariants (no host comm, lockstep launches, no multi-stream overlap) - **Schedulers MUST NOT write `moe.repeat_idx`** — `repeat_idx` is wrapper state advanced once per `forward_impl` regardless of chunk count +- **Do NOT set a communication kernel's timeout directly or through environment variables at runtime** — A backend whose kernels wait for peers under a timeout registers a `MoECommTimeoutProxy` with the timeout guard when it is constructed and unregisters it before releasing what the proxy controls; only the guard decides the value (see [Communication timeouts](#communication-timeouts)). Writing environment variables at runtime races with NCCL threads that read them - **Do NOT allocate symmetric memory from `run_moe` in `FUSED_COMM` backends** — Symmetric-memory rendezvous is a build-time collective and is unsafe under PP / layer-skip or CUDA graph capture; allocate from `create_weights()` after `ConfigurableMoE` has synchronized EPLB-derived attributes. See `mega_moe/mega_moe_deepgemm.py` for the DG pattern and `mega_moe/mega_moe_cute_dsl.py:_alloc_symm_provider` for the NVSHMEM-equivalent provider. - **Do NOT add a new `FUSED_COMM` backend without a zero-token `quantize_input` regression test** — `FusedCommMoEScheduler` calls `quantize_input` for every chunk (including zero-token chunks) so each backend must return its own empty-tensor layout. - **Do NOT use a dataclass for an autotuner tactic without a tested `__repr__` round-trip** — `AutoTuner` serializes tactic values through `json.dumps`/`json.loads` and `eval(repr(tactic))`; a plain dataclass fails the `eval(repr(...))` check. Prefer a JSON-friendly **tuple of primitives or lists of primitives** (lists are JSON-friendly; tuples round-trip via `eval(repr(...))`). See the tactic-representation comment block in `tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py` for the 10-field tactic pattern (with legacy 8-field compatibility) (mma_tiler/cluster_shape as `list[int]`, `epi_flag_batch` as a nested `(int, int)` tuple, the rest as `bool`/`int`/`str`; `_unpack_tactic` is the single source of truth for the field order). The fallback tactic is the token-aware `default_megamoe_tactic(num_tokens)` helper, selected by `Sm100MegaMoENvfp4Runner.forward(tactic=-1)`, not a separate `fallback_tactic()` method. diff --git a/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py b/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py index caea659aff04..bc757c5fb622 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py +++ b/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py @@ -28,7 +28,12 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm.logger import logger -from ..nccl_ep_utils import nccl_ep_supports_version +from ..moe_comm_timeout_guard import get_moe_comm_timeout_budgets +from ..nccl_ep_utils import ( + is_nccl_ep_installed, + nccl_ep_supports_group_timeout, + nccl_ep_supports_version, +) from ..wide_ep_ft import get_wide_ep_ft_options from .allgather_reducescatter import AllGatherReduceScatter from .base import Communication @@ -125,7 +130,14 @@ def create_strategy( Note: Most parameters are extracted from model_config. Only MoE-specific parameters (num_experts, num_slots, top_k, expert_size_per_partition) need to be provided separately. + + Raises: + ValueError: If the MoE communication timeout environment variables are invalid. """ + # Backend constructors apply these budgets, and auto-selection below treats a constructor + # error as "backend unavailable", so an invalid variable must fail here instead. + get_moe_comm_timeout_budgets() + # Extract parameters from model_config mapping = model_config.mapping if mapping.has_cp_helix(): @@ -499,4 +511,9 @@ def _get_nccl_ep_unavailable_reason( f"{required_smem} bytes of dynamic shared memory, but the current device " f"supports only {max_dynamic_smem} bytes." ) + if is_nccl_ep_installed() and not nccl_ep_supports_group_timeout(): + return ( + "NcclEP requires nccl.ep.Group.set_timeout_ns; rebuild so nccl-extensions " + "carries 3rdparty/patches/nccl_ep_group_timeout.patch." + ) return None diff --git a/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py index c60052e1f9d9..c49b4fb75c55 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py +++ b/tensorrt_llm/_torch/moe/fused_moe/communication/nvlink_one_sided.py @@ -55,6 +55,7 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.math_utils import pad_up +from ..moe_comm_timeout_guard import register_moe_comm_timeout_proxy from .base import Communication _CFT_DEFAULT_MAX_BATCH_FOR_DISPATCH = 128 @@ -64,38 +65,6 @@ FORCE_CFT_ENV = "TRTLLM_NVLINK_ONE_SIDED_A2A_FORCE_CFT" _CFT_ALIGNMENT_BYTES = 16 _CFT_MIN_DRIVER_BRANCH = 615 -_TIMEOUT_ENV = "TRTLLM_NVLINK_ONE_SIDED_A2A_TIMEOUT_SEC" -_WARMUP_TIMEOUT_ENV = "TRTLLM_NVLINK_ONE_SIDED_A2A_WARMUP_TIMEOUT_SEC" -_DEFAULT_TIMEOUT_SEC = 300 -_DEFAULT_WARMUP_TIMEOUT_SEC = 1800 -_MAX_TIMEOUT_SEC = 24 * 60 * 60 - - -def get_timeout_seconds(in_warmup: bool = False) -> int: - """Resolve the nominal collective timeout for the engine's execution phase. - - First-touch JIT compilation, autotuning and module loading can delay peer - ranks by minutes during warmup, so that phase gets a larger default budget. - """ - name = _WARMUP_TIMEOUT_ENV if in_warmup else _TIMEOUT_ENV - default = _DEFAULT_WARMUP_TIMEOUT_SEC if in_warmup else _DEFAULT_TIMEOUT_SEC - value = os.environ.get(name) - if not value: - return default - try: - if re.fullmatch(r"\s*[+-]?[0-9]+", value) is not None: - seconds = int(value) - if 0 < seconds <= _MAX_TIMEOUT_SEC: - return seconds - except ValueError: - # Very long integer strings may exceed Python's conversion limit. - pass - tllm_logger.warning_once( - f'Ignoring invalid {name}="{value}" (expected 1..{_MAX_TIMEOUT_SEC} seconds); ' - f"using {default} s", - key=f"{name}_invalid_{value}", - ) - return default def get_force_cft() -> bool | None: @@ -267,6 +236,19 @@ def _get_cft_max_batch_for_combine() -> int: return _get_cft_max_batch(_CFT_MAX_BATCH_FOR_COMBINE_ENV, _CFT_DEFAULT_MAX_BATCH_FOR_COMBINE) +class _NVLinkOneSidedTimeoutProxy: + """Applies MoE communication timeouts to the process-wide completion-flag budget.""" + + name = "NVLinkOneSided" + + def set_timeout_seconds(self, seconds: int | None) -> None: + # The op restores its native default for 0. + torch.ops.trtllm.moe_a2a_set_timeout(0 if seconds is None else seconds) + + +_TIMEOUT_PROXY = _NVLinkOneSidedTimeoutProxy() + + class NVLinkOneSided(Communication): """ NVLINK one-sided comm AllToAll strategy for throughput scenarios. @@ -295,24 +277,6 @@ class NVLinkOneSided(Communication): # such model: its top_k varies across MoE layers. Remove once the endpoints can # be bound per workspace. _CFT_OWNER_SHAPE: Tuple[object, ...] | None = None - _timeout_initialized = False - - @staticmethod - def set_timeout(timeout_sec: int) -> None: - """Set the process-wide timeout for subsequent one-sided A2A launches. - - Args: - timeout_sec: Integer in [1, 86400], nominal seconds at an assumed - 2 GHz SM clock. Applies to both dispatch/combine and CFT/fence. - - Existing CUDA graphs retain the budget recorded at capture time. - """ - if isinstance(timeout_sec, bool) or not isinstance(timeout_sec, int): - raise TypeError("timeout_sec must be an integer number of seconds") - if not 0 < timeout_sec <= _MAX_TIMEOUT_SEC: - raise ValueError(f"timeout_sec must be in 1..{_MAX_TIMEOUT_SEC} seconds") - torch.ops.trtllm.moe_a2a_set_timeout(timeout_sec) - NVLinkOneSided._timeout_initialized = True # MetaInfo indices - initialized from C++ constants FLAG_VAL_OFFSET_INDEX = None @@ -490,11 +454,8 @@ def __init__( raise RuntimeError( f"NVLinkOneSided supports at most {self.MAX_RANKS} EP ranks, got ep_size={self.ep_size}." ) - - # Standalone callers have no engine phase transitions. Initialize their - # steady-state budget without overwriting an explicit or warmup timeout. - if not NVLinkOneSided._timeout_initialized: - NVLinkOneSided.set_timeout(get_timeout_seconds()) + # The completion-flag budget is process-wide, so every instance shares one proxy. + register_moe_comm_timeout_proxy(_TIMEOUT_PROXY) # Store needed parameters self.num_experts = num_slots diff --git a/tensorrt_llm/_torch/moe/fused_moe/deep_ep_utils.py b/tensorrt_llm/_torch/moe/fused_moe/deep_ep_utils.py index 0dcd115eee64..702b1c3de7a9 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/deep_ep_utils.py +++ b/tensorrt_llm/_torch/moe/fused_moe/deep_ep_utils.py @@ -9,18 +9,40 @@ from tensorrt_llm._utils import mpi_comm from tensorrt_llm.mapping import Mapping +from .moe_comm_timeout_guard import register_moe_comm_timeout_proxy + try: + from tensorrt_llm import deep_ep from tensorrt_llm.deep_ep import Buffer deep_ep_installed = True except ImportError: deep_ep_installed = False +class _DeepEPTimeoutProxy: + """Applies MoE communication timeouts to DeepEP's normal kernels. + + DeepEP keeps the kernel timeout and its CPU-side wait timeout + process-wide, so one proxy serves every buffer. Low-latency kernels do not + use them. + """ + + name = "DeepEP" + + def set_timeout_seconds(self, seconds: Optional[int]) -> None: + # DeepEP restores its compile-time defaults for 0. + deep_ep.set_timeout_seconds(0 if seconds is None else seconds) + + +_TIMEOUT_PROXY = _DeepEPTimeoutProxy() + + class VariableLengthBuffer: """ A wrapper of deep_ep.Buffer that accepts future size change """ def __init__(self, mapping: Mapping): + register_moe_comm_timeout_proxy(_TIMEOUT_PROXY) self.comm = mpi_comm().Split(mapping.pp_rank, mapping.moe_ep_rank) self.buffer = None diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py index f752d8f82cf9..37f50886ee56 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py @@ -24,6 +24,7 @@ import os import socket +from types import ModuleType from typing import Dict, List, Optional, Tuple import torch @@ -53,6 +54,7 @@ from ..impl_environment import MoEDep from ..impl_identity import MoEImplDescriptor, MoEImplId, register_moe_impl from ..interface import MoESchedulerKind, MoEWeightLoadingMode, _reject +from ..moe_comm_timeout_guard import register_moe_comm_timeout_proxy from ..quantization import W4A8MXFP4MXFP8MegaMoEDeepGemmMethod, _import_deep_gemm from ..routing import BaseMoeRoutingMethod @@ -81,6 +83,34 @@ _MEGA_MOE_SYMM_BUFFER_CACHE: Dict[tuple, Tuple[object, object]] = {} +class _DeepGemmTimeoutProxy: + """Applies MoE communication timeouts to DeepGEMM's in-kernel barriers. + + The timeout lives in DeepGEMM's process-wide runtime and each ``fp8_fp4_mega_moe`` + launch passes it to the kernel as an argument, so one proxy serves every layer. + """ + + name = "DeepGEMM MegaMoE" + + def __init__(self, dg: ModuleType) -> None: + self._dg = dg + + def set_timeout_seconds(self, seconds: int | None) -> None: + # DeepGEMM restores its JIT-time default for 0. + self._dg.set_barrier_timeout_seconds(0 if seconds is None else seconds) + + +# Kept alive here because the timeout guard holds its proxies weakly. +_TIMEOUT_PROXY: Optional[_DeepGemmTimeoutProxy] = None + + +def _register_timeout_proxy(dg: ModuleType) -> None: + global _TIMEOUT_PROXY + if _TIMEOUT_PROXY is None: + _TIMEOUT_PROXY = _DeepGemmTimeoutProxy(dg) + register_moe_comm_timeout_proxy(_TIMEOUT_PROXY) + + def _free_symm_buffer(buffered: object) -> int: """Release one DG SymmBuffer's symmetric memory. Returns its size in bytes. @@ -530,6 +560,7 @@ def __init__( # paying that on every forward (``run_moe`` path) shows up in host-side # CPU overhead even though the underlying ``import`` is cached by Python. self._dg = _import_deep_gemm() + _register_timeout_proxy(self._dg) # NVLink SymmBuffer activation workspace. Allocation is a # model-build-period collective (``symm_mem.rendezvous`` over the diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_comm_timeout_guard.py b/tensorrt_llm/_torch/moe/fused_moe/moe_comm_timeout_guard.py new file mode 100644 index 000000000000..b8ffbbf29baa --- /dev/null +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_comm_timeout_guard.py @@ -0,0 +1,344 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Warmup-aware device-side timeouts for MoE communication kernels. + +Several MoE communication kernels wait for their peers inside the kernel and trap once a peer is +late by more than a device-side timeout. The trap is a sticky CUDA error, so every waiting rank +dies. During warmup, one-time per-rank work (JIT compilation, autotuning, module loading) can +delay a healthy rank by minutes, and no collective separates that work from the next launch. + +The model engine reports each warmup transition here. The guard applies the warmup timeout to +every registered backend while the engine warms up and the serving timeout otherwise. Each +backend registers a ``MoECommTimeoutProxy`` when it is constructed. + +Environment variables (integer seconds in ``1..86400``): + TRTLLM_MOE_COMM_WARMUP_TIMEOUT_SEC: Warmup timeout for every backend; defaults to 1800. + TRTLLM_MOE_COMM_TIMEOUT_SEC: Serving timeout for every backend; when unset, each backend + keeps its native timeout. + +The NVLinkOneSided-only names ``TRTLLM_MOE_A2A_[WARMUP_]TIMEOUT_SEC`` and +``TRTLLM_NVLINK_ONE_SIDED_A2A_[WARMUP_]TIMEOUT_SEC`` are deprecated aliases of these variables. +""" + +import contextlib +import dataclasses +import os +import threading +import weakref +from collections.abc import Callable, Iterator, Mapping +from typing import Protocol + +import torch + +from tensorrt_llm.logger import logger + +__all__ = [ + "DEFAULT_WARMUP_TIMEOUT_SEC", + "MAX_TIMEOUT_SEC", + "SERVING_TIMEOUT_ENV", + "WARMUP_TIMEOUT_ENV", + "MoECommTimeoutBudgets", + "MoECommTimeoutGuard", + "MoECommTimeoutProxy", + "get_moe_comm_timeout_budgets", + "moe_comm_serving_timeouts", + "register_moe_comm_timeout_proxy", + "resolve_moe_comm_timeout_budgets", + "set_moe_comm_warmup", + "unregister_moe_comm_timeout_proxy", +] + +WARMUP_TIMEOUT_ENV = "TRTLLM_MOE_COMM_WARMUP_TIMEOUT_SEC" +SERVING_TIMEOUT_ENV = "TRTLLM_MOE_COMM_TIMEOUT_SEC" +DEFAULT_WARMUP_TIMEOUT_SEC = 1800 +MAX_TIMEOUT_SEC = 24 * 60 * 60 + +# The replacements broaden the scope: each deprecated name configured NVLinkOneSided only. +_DEPRECATED_ENV_ALIASES: dict[str, tuple[str, ...]] = { + WARMUP_TIMEOUT_ENV: ( + "TRTLLM_MOE_A2A_WARMUP_TIMEOUT_SEC", + "TRTLLM_NVLINK_ONE_SIDED_A2A_WARMUP_TIMEOUT_SEC", + ), + SERVING_TIMEOUT_ENV: ( + "TRTLLM_MOE_A2A_TIMEOUT_SEC", + "TRTLLM_NVLINK_ONE_SIDED_A2A_TIMEOUT_SEC", + ), +} + + +class MoECommTimeoutProxy(Protocol): + """Stands in for the device-side timeout of one MoE communication backend.""" + + name: str + """Backend name used in logs.""" + + def set_timeout_seconds(self, seconds: int | None) -> None: + """Apply a timeout to the launches issued after this call. + + Args: + seconds: Timeout in seconds, or ``None`` to restore the backend's native default. + """ + ... + + +@dataclasses.dataclass(frozen=True) +class MoECommTimeoutBudgets: + """Timeouts for warmup and for serving. + + Attributes: + warmup_seconds: Timeout applied to every backend during warmup. + serving_seconds: Timeout applied to every backend while serving; ``None`` keeps each + backend's native default. + """ + + warmup_seconds: int + serving_seconds: int | None + + def select(self, in_warmup: bool) -> int | None: + """Return the warmup or serving timeout; ``None`` selects native defaults.""" + return self.warmup_seconds if in_warmup else self.serving_seconds + + +def resolve_moe_comm_timeout_budgets(environ: Mapping[str, str]) -> MoECommTimeoutBudgets: + """Resolve the timeout budgets from environment variables. + + Args: + environ: The environment to read, usually ``os.environ``. + + Returns: + The budgets for warmup and serving. + + Raises: + ValueError: If a value is not an integer in ``1..86400``, a deprecated alias conflicts + with its replacement, or the warmup budget is shorter than the serving budget. + """ + warmup_seconds = _read_timeout_seconds(environ, WARMUP_TIMEOUT_ENV) + serving_seconds = _read_timeout_seconds(environ, SERVING_TIMEOUT_ENV) + if warmup_seconds is None: + warmup_seconds = DEFAULT_WARMUP_TIMEOUT_SEC + if serving_seconds is not None and warmup_seconds < serving_seconds: + raise ValueError( + f"{WARMUP_TIMEOUT_ENV} resolves to {warmup_seconds} s, which is shorter than " + f"{SERVING_TIMEOUT_ENV}={serving_seconds} s; warmup needs at least the serving budget" + ) + return MoECommTimeoutBudgets(warmup_seconds=warmup_seconds, serving_seconds=serving_seconds) + + +def _is_capturing_cuda_graph() -> bool: + return torch.cuda.is_available() and torch.cuda.is_current_stream_capturing() + + +class MoECommTimeoutGuard: + """Applies the warmup or serving timeout to every registered proxy. + + The guard starts in serving. Budgets are resolved from the environment the first time a proxy + needs one and then stay fixed. Proxies are held weakly; their owners keep them alive and + unregister them before releasing the resources the proxies control. + + Args: + environ: The environment holding the budget variables. + is_capturing: Returns whether the current CUDA stream is capturing a graph. + """ + + def __init__( + self, + environ: Mapping[str, str], + is_capturing: Callable[[], bool] = _is_capturing_cuda_graph, + ) -> None: + self._environ = environ + self._is_capturing = is_capturing + self._lock = threading.Lock() + self._in_warmup = False + self._budgets: MoECommTimeoutBudgets | None = None + self._proxies: weakref.WeakSet[MoECommTimeoutProxy] = weakref.WeakSet() + + @property + def in_warmup(self) -> bool: + """Whether the warmup timeout is in effect.""" + return self._in_warmup + + def budgets(self) -> MoECommTimeoutBudgets: + """Return the budgets, resolving them from the environment on first use. + + Raises: + ValueError: If the budget environment variables are invalid. + """ + with self._lock: + return self._resolve_budgets() + + def register(self, proxy: MoECommTimeoutProxy) -> None: + """Apply the current timeout to ``proxy`` and keep it updated on warmup transitions. + + A proxy whose setter raises is not registered, and the error propagates. + + Args: + proxy: The backend to control. + + Raises: + ValueError: If the budget environment variables are invalid. + """ + with self._lock: + proxy.set_timeout_seconds(self._resolve_budgets().select(self._in_warmup)) + self._proxies.add(proxy) + + def unregister(self, proxy: MoECommTimeoutProxy) -> None: + """Stop updating ``proxy``; unknown proxies are ignored. + + Args: + proxy: The backend to release. + """ + with self._lock: + self._proxies.discard(proxy) + + def set_warmup(self, in_warmup: bool) -> None: + """Apply the warmup or serving timeout to every registered proxy. + + Args: + in_warmup: Whether the engine enters warmup. + + Raises: + RuntimeError: If the current CUDA stream is capturing a graph. Proxies may issue + host-synchronizing CUDA calls, which are illegal during capture. + ValueError: If the budget environment variables are invalid. + """ + if self._is_capturing(): + raise RuntimeError( + "Cannot switch MoE communication timeouts while a CUDA graph is being captured" + ) + with self._lock: + if in_warmup == self._in_warmup: + return + self._in_warmup = in_warmup + proxies = list(self._proxies) + if not proxies: + return + seconds = self._resolve_budgets().select(in_warmup) + for proxy in proxies: + proxy.set_timeout_seconds(seconds) + applied = "native defaults" if seconds is None else f"{seconds} s" + names = ", ".join(proxy.name for proxy in proxies) + stage = "warmup" if in_warmup else "serving" + logger.info(f"MoE communication timeouts for {stage}: {applied} ({names})") + + @contextlib.contextmanager + def serving_timeouts(self) -> Iterator[None]: + """Apply the serving timeout in the enclosed block, then restore the previous timeout. + + A captured CUDA graph keeps the timeouts that were current at capture, and serving + replays it, so the pass that captures uses the serving timeout even during warmup. + """ + previous = self._in_warmup + self.set_warmup(False) + try: + yield + finally: + self.set_warmup(previous) + + def _resolve_budgets(self) -> MoECommTimeoutBudgets: + if self._budgets is None: + self._budgets = resolve_moe_comm_timeout_budgets(self._environ) + serving = self._budgets.serving_seconds + logger.info( + f"MoE communication timeouts: warmup={self._budgets.warmup_seconds} s, " + f"serving={'native defaults' if serving is None else f'{serving} s'}" + ) + return self._budgets + + +def _read_timeout_seconds(environ: Mapping[str, str], name: str) -> int | None: + resolved_name, resolved = name, _parse_timeout_seconds(environ, name) + for alias in _DEPRECATED_ENV_ALIASES[name]: + alias_value = _parse_timeout_seconds(environ, alias) + if alias_value is None: + continue + logger.warning_once( + f"{alias} is deprecated; use {name}, which applies to every MoE communication kernel.", + key=f"deprecated_env_{alias}", + ) + if resolved is None: + resolved_name, resolved = alias, alias_value + elif alias_value != resolved: + raise ValueError( + f"{resolved_name}={resolved} conflicts with the deprecated {alias}={alias_value}" + ) + return resolved + + +def _parse_timeout_seconds(environ: Mapping[str, str], name: str) -> int | None: + raw = environ.get(name) + if raw is None or not raw.strip(): + return None + error = ValueError( + f"{name}={raw!r} must be an integer number of seconds in 1..{MAX_TIMEOUT_SEC}" + ) + try: + seconds = int(raw) + except ValueError: + raise error from None + if not 1 <= seconds <= MAX_TIMEOUT_SEC: + raise error + return seconds + + +_DEFAULT_GUARD = MoECommTimeoutGuard(environ=os.environ) + + +def get_moe_comm_timeout_budgets() -> MoECommTimeoutBudgets: + """Return the process-wide budgets, resolving them from ``os.environ`` on first use. + + Backend selection calls this before trying any backend, because selection treats a + constructor error as "backend unavailable" and would otherwise hide an invalid variable. + + Raises: + ValueError: If the budget environment variables are invalid. + """ + return _DEFAULT_GUARD.budgets() + + +def register_moe_comm_timeout_proxy(proxy: MoECommTimeoutProxy) -> None: + """Register ``proxy`` with the process-wide guard; see ``MoECommTimeoutGuard.register``. + + Args: + proxy: The backend to control. + """ + _DEFAULT_GUARD.register(proxy) + + +def unregister_moe_comm_timeout_proxy(proxy: MoECommTimeoutProxy) -> None: + """Unregister ``proxy`` from the process-wide guard. + + Args: + proxy: The backend to release. + """ + _DEFAULT_GUARD.unregister(proxy) + + +def set_moe_comm_warmup(in_warmup: bool) -> None: + """Switch the process-wide guard; see ``MoECommTimeoutGuard.set_warmup``. + + Args: + in_warmup: Whether the engine enters warmup. + """ + _DEFAULT_GUARD.set_warmup(in_warmup) + + +@contextlib.contextmanager +def moe_comm_serving_timeouts() -> Iterator[None]: + """Apply the serving timeouts in the enclosed block. + + See ``MoECommTimeoutGuard.serving_timeouts``. + """ + with _DEFAULT_GUARD.serving_timeouts(): + yield diff --git a/tensorrt_llm/_torch/moe/fused_moe/nccl_ep_utils.py b/tensorrt_llm/_torch/moe/fused_moe/nccl_ep_utils.py index f729b0124d40..036484c889ad 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/nccl_ep_utils.py +++ b/tensorrt_llm/_torch/moe/fused_moe/nccl_ep_utils.py @@ -21,7 +21,7 @@ the persistent FP8 scales receive buffer. """ -from typing import Optional +from typing import TYPE_CHECKING, Optional import torch from packaging.version import InvalidVersion, Version @@ -29,6 +29,14 @@ from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping +from .moe_comm_timeout_guard import ( + register_moe_comm_timeout_proxy, + unregister_moe_comm_timeout_proxy, +) + +if TYPE_CHECKING: + from nccl.ep import Group + _MIN_NCCL_EP_INT32_TOPK_VERSION = "0.2" _NCCL_RUNTIME_ERRORS = (RuntimeError, OSError) _NCCL_AVAILABILITY_ERRORS = (ImportError,) + _NCCL_RUNTIME_ERRORS @@ -85,6 +93,28 @@ def _nccl_ep_supports_int32_topk_idx() -> bool: return nccl_ep_supports_version(_MIN_NCCL_EP_INT32_TOPK_VERSION) +def nccl_ep_supports_group_timeout() -> bool: + """Return True when ``nccl.ep.Group`` can change its wait-loop timeout after creation.""" + try: + from nccl.ep import Group + except _NCCL_AVAILABILITY_ERRORS: + return False + return hasattr(Group, "set_timeout_ns") + + +class _NcclEpGroupTimeoutProxy: + """Applies MoE communication timeouts to one NCCL EP group.""" + + name = "NcclEP" + + def __init__(self, group: "Group") -> None: + self._group = group + + def set_timeout_seconds(self, seconds: int | None) -> None: + # The group restores the timeout it resolved at creation for 0. + self._group.set_timeout_ns(0 if seconds is None else seconds * 1_000_000_000) + + # Singleton EP context keyed by (ep_size, ep_rank, max_tokens, num_experts, # hidden, max_top_k, use_internal_fp8_dispatch, layout). _ep_group_cache: dict = {} @@ -226,6 +256,10 @@ def __init__( max_token_bytes=max_token_bytes, ) self.ep_group = Group.create(self.comm, cfg) + # Held here because the timeout guard holds its proxies weakly; ``destroy`` unregisters + # the proxy before the group goes away. + self._timeout_proxy = _NcclEpGroupTimeoutProxy(self.ep_group) + register_moe_comm_timeout_proxy(self._timeout_proxy) logger.info( f"NCCL EP group created: ep_size={self.ep_size}, " @@ -361,6 +395,7 @@ def destroy(self): (the recommended nccl4py pattern), then ``Free`` on the MPI comm. """ if self.ep_group is not None: + unregister_moe_comm_timeout_proxy(self._timeout_proxy) try: self.ep_group.destroy() except _NCCL_RUNTIME_ERRORS as e: diff --git a/tensorrt_llm/_torch/moe/fused_moe/quantization.py b/tensorrt_llm/_torch/moe/fused_moe/quantization.py index 79e0946fee4d..021abd760120 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/moe/fused_moe/quantization.py @@ -7264,6 +7264,11 @@ def _import_deep_gemm(): f"tensorrt_llm.deep_gemm missing mega_moe symbols {missing}; " f"upgrade the TRT-LLM bundled DeepGEMM to a release that " f"includes fp8_fp4_mega_moe.") + if not hasattr(_dg, "set_barrier_timeout_seconds"): + raise _MegaMoEUnavailable( + "tensorrt_llm.deep_gemm lacks set_barrier_timeout_seconds; rebuild " + "so the bundled DeepGEMM carries " + "3rdparty/patches/deepgemm_configurable_barrier_timeout.patch.") mega_moe_params = inspect.signature(_dg.fp8_fp4_mega_moe).parameters missing_situ_params = [ diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py index 5432b7905438..e446642f19b9 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py @@ -4,7 +4,6 @@ """Shared helpers used by multiple model-runner families.""" import bisect -import contextlib import math from typing import TYPE_CHECKING, Any @@ -221,42 +220,3 @@ def resolve_mrope_position_deltas_cache(model: torch.nn.Module | None) -> torch. if cache is None: cache = getattr(getattr(model, "draft_model", None), "mrope_position_deltas_cache", None) return cache - - -@contextlib.contextmanager -def moe_a2a_steady_state_budget_for_capture(): - """Force the steady-state MoE all-to-all budget across CUDA-graph capture. - - The budget is a kernel launch argument, so it is frozen into each captured - graph. Capture happens inside the warmup window, so without this a replay - would keep warmup's relaxed deadline for the life of the process. - """ - set_moe_a2a_warmup(False) - try: - yield - finally: - set_moe_a2a_warmup(True) - - -def set_moe_a2a_warmup(in_warmup: bool) -> None: - """Select the MoE all-to-all completion-flag budget for the current phase. - - No-op when the op is unavailable (older bindings). - """ - from tensorrt_llm._torch.moe.fused_moe.communication.nvlink_one_sided import ( - NVLinkOneSided, - get_timeout_seconds, - ) - - timeout_sec = get_timeout_seconds(in_warmup) - try: - NVLinkOneSided.set_timeout(timeout_sec) - logger.info( - f"moe_a2a completion-flag budget: in_warmup={in_warmup}, " - f"timeout={timeout_sec} s (nominal, at an assumed 2 GHz clock64 rate)" - ) - except (AttributeError, RuntimeError) as e: - logger.warning( - f"moe_a2a_set_timeout unavailable, the all-to-all timeout " - f"budget was not switched: {type(e).__name__}: {e}" - ) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index f09883c04bcf..60f10b8c55ef 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -27,6 +27,7 @@ from tensorrt_llm._torch.models.modeling_multimodal_mixin import _build_request_multimodal_input from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM, timing_metric from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata +from tensorrt_llm._torch.moe.fused_moe.moe_comm_timeout_guard import moe_comm_serving_timeouts from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import ( MoeLoadBalancer, MoeLoadBalancerIterContext, @@ -108,7 +109,6 @@ get_position_id_offset, get_top_level_model, make_scheduled_inputs, - moe_a2a_steady_state_budget_for_capture, prepare_multimodal_indices, resolve_mrope_position_deltas_cache, ship_multimodal_indices, @@ -773,25 +773,8 @@ def _warmup_scheduled(self, resource_manager: ResourceManager, kv_cache_manager) # NVSHMEM). gc.collect() torch.cuda.empty_cache() - # Warm up every graph shape before capturing any graph. Attention - # kernels can switch implementations at smaller batch sizes and require - # a larger workspace, so the first pass grows the workspace to its - # maximum size. The second pass runs the final per-shape warmup and - # captures without resizing the workspace. - # Capture with the steady-state MoE all-to-all budget: the timeout is a - # launch argument and is baked into every later replay. - with ( - self._warmup_timer.phase("cuda_graph_capture"), - moe_a2a_steady_state_budget_for_capture(), - ): - with self.cuda_graph_runner.allow_capture(): - self.cuda_graph_runner.is_warmup_only = True - try: - self._run_cuda_graph_warmup(resource_manager) - finally: - self.cuda_graph_runner.is_warmup_only = False - self.cuda_graph_runner.padding_dummy_requests = {} - self._run_cuda_graph_warmup(resource_manager) + with self._warmup_timer.phase("cuda_graph_capture"): + self._warmup_and_capture_cuda_graphs(resource_manager) log_mem_snapshot("warmup/after_cuda_graph_capture") # Pre-compile DeepGEMM paged_mqa_logits_metadata for every 32-aligned # batch bucket the runtime can produce (max_batch_size scaled by the @@ -979,7 +962,8 @@ def _prewarm_cute_dsl_indexer_q(self) -> None: completion-flag deadline. It is a partial mitigation only: other first-touch compiles remain inside collective-bearing forwards, and some sit on the all-to-all path itself and cannot be pre-compiled this way. - The runtime budget (``moeA2AGetTimeoutCycles``) covers the general case. + The warmup MoE communication timeout (see ``moe_comm_timeout_guard.py``) + covers the general case. Only the fallback tactics are compiled -- what an eager, cache-miss forward selects. The runner's kernel cache key excludes m/n/k, so one @@ -1908,6 +1892,27 @@ def _get_graphs_to_capture(self, cuda_graph_batch_sizes: list[int]) -> list[tupl draft_lengths.append(0) return [(bs, draft_len) for bs in cuda_graph_batch_sizes for draft_len in draft_lengths] + def _warmup_and_capture_cuda_graphs(self, resource_manager: ResourceManager) -> None: + """Warm up every CUDA-graph shape, then capture the graphs. + + Attention kernels can switch implementations at smaller batch sizes and + require a larger workspace, so the first pass grows the workspace to its + maximum size without capturing. The second pass runs the final + per-shape warmup and captures without resizing the workspace. + """ + with self.cuda_graph_runner.allow_capture(): + self.cuda_graph_runner.is_warmup_only = True + try: + self._run_cuda_graph_warmup(resource_manager) + finally: + self.cuda_graph_runner.is_warmup_only = False + self.cuda_graph_runner.padding_dummy_requests = {} + # Captured graphs keep the MoE communication timeouts that were + # current at capture, and serving replays them, so capture with the + # serving timeouts. The warmup-only pass above keeps the warmup ones. + with moe_comm_serving_timeouts(): + self._run_cuda_graph_warmup(resource_manager) + def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): """Warm up or capture CUDA graphs for the configured graph shapes.""" is_warmup_only = self.cuda_graph_runner.is_warmup_only diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b9e44d5ad892..7ea413f05d95 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -39,6 +39,7 @@ from ..models.modeling_multimodal_mixin import MultimodalModelMixin from ..models.modeling_utils import DecoderModelForCausalLM, timing_metric from ..moe.expert_statistic import ExpertStatistic +from ..moe.fused_moe.moe_comm_timeout_guard import set_moe_comm_warmup from ..moe.fused_moe.moe_load_balancer import MoeLoadBalancer from ..route_capture import ROUTE_CAPTURE_ATTR, RouteCapture from ..speculative import SpecMetadata, update_spec_config_from_loaded_model @@ -55,8 +56,7 @@ setup_mm_encoder_attn_metadata) from .engine.runners import resolve_runner_type from .engine.runners.common import (make_scheduled_inputs, - resolve_mrope_position_deltas_cache, - set_moe_a2a_warmup) + resolve_mrope_position_deltas_cache) from .engine.runners.decoder import DecoderRunner, DecoderRunnerConfig from .engine.runners.encoder import EncoderRunner, EncoderRunnerConfig from .engine.runners.encoder_decoder import (EncoderDecoderRunner, @@ -914,9 +914,9 @@ def is_warmup(self, value: bool): self._is_warmup = value # This setter is the one choke point every warmup transition passes - # through, including PyExecutor's, so select the MoE all-to-all budget - # here rather than in set_warmup_flag(). - set_moe_a2a_warmup(value) + # through, including PyExecutor's, so switch the MoE communication + # timeouts here rather than in set_warmup_flag(). + set_moe_comm_warmup(value) self.moe_load_balancer_iter_info = (not value, not value) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 180b6562f90a..e4393ac57b1c 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -187,6 +187,7 @@ l0_b200: - unittest/_torch/moe/test_cutedsl_fc1_weight_layout.py - unittest/_torch/moe/test_moe_load_balancer.py - unittest/_torch/moe/test_moe_routing.py + - unittest/_torch/moe/test_moe_comm_timeout_guard.py - unittest/_torch/moe/fused_moe/test_deepgemm_fused_gather_finalize.py - unittest/_torch/moe/fused_moe/test_deepgemm_fused_expand_quant.py # ------------- MoE: test_moe_backend (by backend) --------------- @@ -203,6 +204,7 @@ l0_b200: - unittest/_torch/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fused_shared_experts - unittest/_torch/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fuse_shared_expert_layout - unittest/_torch/moe/test_moe_backend.py::test_import_deep_gemm_rejects_pre_situ_mega_moe_api + - unittest/_torch/moe/test_moe_backend.py::test_import_deep_gemm_rejects_mega_moe_without_barrier_timeout - unittest/_torch/moe/test_moe_backend.py::test_megamoe_streaming_reload_resets_slot_claims - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_selects_padded_quant_method - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_fc31_scale_c_drops_dequant_scale diff --git a/tests/unittest/_torch/executor/test_moe_comm_timeout_warmup.py b/tests/unittest/_torch/executor/test_moe_comm_timeout_warmup.py new file mode 100644 index 000000000000..9b5a1741210e --- /dev/null +++ b/tests/unittest/_torch/executor/test_moe_comm_timeout_warmup.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The model engine must switch the MoE communication timeouts at its warmup transitions. + +A relaxed warmup timeout that latches on into serving would leave hang detection permanently +slow, and CUDA graphs record their launch arguments, so capture must use the serving timeouts. +""" + +import contextlib +from collections.abc import Iterator + +import pytest + +from tensorrt_llm._torch.moe.fused_moe import moe_comm_timeout_guard +from tensorrt_llm._torch.moe.fused_moe.moe_comm_timeout_guard import ( + DEFAULT_WARMUP_TIMEOUT_SEC, + MoECommTimeoutGuard, +) +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner +from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine + +pytestmark = pytest.mark.cpu_only + + +@pytest.fixture +def guard(monkeypatch: pytest.MonkeyPatch) -> MoECommTimeoutGuard: + guard = MoECommTimeoutGuard(environ={}, is_capturing=lambda: False) + monkeypatch.setattr(moe_comm_timeout_guard, "_DEFAULT_GUARD", guard) + return guard + + +class _RecordingProxy: + name = "recording" + + def __init__(self) -> None: + self.calls: list[int | None] = [] + + def set_timeout_seconds(self, seconds: int | None) -> None: + self.calls.append(seconds) + + +class _WarmupFlagStub: + """Borrows the engine's ``is_warmup`` property without building a model. + + The setter only touches ``_is_warmup``, the MoE communication timeouts, and + ``moe_load_balancer_iter_info``, which is a no-op without a load balancer. + """ + + is_warmup = PyTorchModelEngine.is_warmup + moe_load_balancer_iter_info = PyTorchModelEngine.moe_load_balancer_iter_info + moe_load_balancer = None + + +class _GraphRunnerStub: + def __init__(self) -> None: + self.is_warmup_only = False + self.padding_dummy_requests: dict[str, object] = {"stale": object()} + + @contextlib.contextmanager + def allow_capture(self) -> Iterator[None]: + yield + + +class _CaptureRunnerStub: + """Records the runner mode and the timeout guard state of each CUDA-graph warmup pass.""" + + def __init__(self, guard: MoECommTimeoutGuard) -> None: + self.cuda_graph_runner = _GraphRunnerStub() + self.passes: list[tuple[bool, bool]] = [] + self._guard = guard + + def _run_cuda_graph_warmup(self, resource_manager: object) -> None: + self.passes.append((self.cuda_graph_runner.is_warmup_only, self._guard.in_warmup)) + + +def test_is_warmup_setter_switches_the_timeouts_both_ways(guard: MoECommTimeoutGuard) -> None: + proxy = _RecordingProxy() + guard.register(proxy) + stub = _WarmupFlagStub() + + stub.is_warmup = True + assert guard.in_warmup + + stub.is_warmup = False + assert not guard.in_warmup + assert not stub.is_warmup + assert proxy.calls == [None, DEFAULT_WARMUP_TIMEOUT_SEC, None] + + +def test_only_the_capturing_pass_uses_the_serving_timeouts(guard: MoECommTimeoutGuard) -> None: + runner = _CaptureRunnerStub(guard) + guard.set_warmup(True) + + DecoderRunner._warmup_and_capture_cuda_graphs(runner, resource_manager=None) + + assert runner.passes == [(True, True), (False, False)] + assert guard.in_warmup + assert runner.cuda_graph_runner.padding_dummy_requests == {} diff --git a/tests/unittest/_torch/moe/multi_gpu/test_nvlink_one_sided.py b/tests/unittest/_torch/moe/multi_gpu/test_nvlink_one_sided.py index 45983d643448..afe7ba316bd1 100644 --- a/tests/unittest/_torch/moe/multi_gpu/test_nvlink_one_sided.py +++ b/tests/unittest/_torch/moe/multi_gpu/test_nvlink_one_sided.py @@ -532,8 +532,8 @@ def _run(case: Case, pools: dict[tuple[int, bool], MPIPoolExecutor]) -> None: "TRTLLM_NVLINK_ONE_SIDED_A2A_CFT_MAX_BATCH_FOR_COMBINE": "128", "TRTLLM_NVLINK_ONE_SIDED_A2A_WORKSPACE_MB": "", "TRTLLM_ENABLE_PDL": "1" if case.pdl else "0", - "TRTLLM_NVLINK_ONE_SIDED_A2A_TIMEOUT_SEC": "30", - "TRTLLM_NVLINK_ONE_SIDED_A2A_WARMUP_TIMEOUT_SEC": "60", + "TRTLLM_MOE_COMM_TIMEOUT_SEC": "30", + "TRTLLM_MOE_COMM_WARMUP_TIMEOUT_SEC": "60", } key = (case.ep_size, case.pdl) if key not in pools: diff --git a/tests/unittest/_torch/moe/test_communication_factory.py b/tests/unittest/_torch/moe/test_communication_factory.py index 23c7f3a82f78..2c0381325b11 100644 --- a/tests/unittest/_torch/moe/test_communication_factory.py +++ b/tests/unittest/_torch/moe/test_communication_factory.py @@ -23,7 +23,7 @@ import torch from packaging.version import Version -from tensorrt_llm._torch.moe.fused_moe import nccl_ep_utils +from tensorrt_llm._torch.moe.fused_moe import moe_comm_timeout_guard, nccl_ep_utils from tensorrt_llm._torch.moe.fused_moe.communication import communication_factory from tensorrt_llm._torch.moe.fused_moe.communication import nvlink_one_sided as one_sided_module from tensorrt_llm._torch.moe.fused_moe.communication import nvlink_two_sided as two_sided_module @@ -68,6 +68,31 @@ def _strategy_unavailable(*args, **kwargs): def _mock_nccl_ep_v02(monkeypatch: pytest.MonkeyPatch) -> None: """Model the v0.2 runtime contract required by these factory tests.""" monkeypatch.setattr(communication_factory, "nccl_ep_supports_version", lambda _: True) + monkeypatch.setattr(communication_factory, "nccl_ep_supports_group_timeout", lambda: True) + + +def test_create_strategy_rejects_invalid_comm_timeout_before_selection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Selection treats constructor errors as "unavailable", so validation must come first.""" + guard = moe_comm_timeout_guard.MoECommTimeoutGuard( + environ={moe_comm_timeout_guard.SERVING_TIMEOUT_ENV: "soon"} + ) + monkeypatch.setattr(moe_comm_timeout_guard, "_DEFAULT_GUARD", guard) + one_sided = Mock(side_effect=AssertionError("backend construction reached")) + monkeypatch.setattr(communication_factory, "NVLinkOneSided", one_sided) + + with pytest.raises(ValueError, match=moe_comm_timeout_guard.SERVING_TIMEOUT_ENV): + communication_factory.CommunicationFactory.create_strategy( + _make_model_config(), + num_experts=32, + num_slots=32, + top_k=8, + expert_size_per_partition=16, + hidden_size=4096, + ) + + one_sided.assert_not_called() @pytest.mark.parametrize("use_flashinfer", [False, True]) @@ -232,6 +257,33 @@ def test_nccl_ep_invalid_version_is_unavailable(monkeypatch: pytest.MonkeyPatch) ) +@pytest.mark.parametrize(("has_setter", "supported"), [(False, False), (True, True)]) +def test_nccl_ep_group_timeout_probe( + monkeypatch: pytest.MonkeyPatch, + has_setter: bool, + supported: bool, +): + methods = {"set_timeout_ns": lambda self, timeout_ns: None} if has_setter else {} + fake_ep = SimpleNamespace(Group=type("Group", (), methods)) + monkeypatch.setitem(sys.modules, "nccl", SimpleNamespace(ep=fake_ep)) + monkeypatch.setitem(sys.modules, "nccl.ep", fake_ep) + + assert nccl_ep_utils.nccl_ep_supports_group_timeout() is supported + + +def test_nccl_ep_without_runtime_group_timeout_is_unavailable(monkeypatch: pytest.MonkeyPatch): + """A wheel without the carried nccl-extensions patch cannot follow warmup timeouts.""" + _mock_nccl_ep_v02(monkeypatch) + monkeypatch.setattr(communication_factory, "is_nccl_ep_installed", lambda: True) + monkeypatch.setattr(communication_factory, "nccl_ep_supports_group_timeout", lambda: False) + + reason = communication_factory.CommunicationFactory._get_nccl_ep_unavailable_reason( + torch.bfloat16, None, 32, 4096, 1024, 1024, 8 + ) + + assert reason is not None and "set_timeout_ns" in reason + + class _FakeNcclEP: def __init__( self, diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index efe174a325b5..ba98994088f3 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -154,6 +154,7 @@ def per_token_cast_to_fp8(*, use_packed_ue8m0=False): transform_sf_into_required_layout=lambda: None, transform_weights_for_mega_moe=lambda: None, per_token_cast_to_fp8=per_token_cast_to_fp8, + set_barrier_timeout_seconds=lambda seconds: None, ) monkeypatch.setattr(tensorrt_llm, "deep_gemm", deep_gemm) @@ -164,6 +165,35 @@ def per_token_cast_to_fp8(*, use_packed_ue8m0=False): quantization_module._import_deep_gemm() +def test_import_deep_gemm_rejects_mega_moe_without_barrier_timeout(monkeypatch): + import tensorrt_llm + import tensorrt_llm._torch.moe.fused_moe.quantization as quantization_module + + def fp8_fp4_mega_moe(*, situ_beta=None, situ_linear_beta=None): + pass + + def per_token_cast_to_fp8(*, use_packed_ue8m0=False): + pass + + deep_gemm = SimpleNamespace( + fp8_fp4_mega_moe=fp8_fp4_mega_moe, + get_symm_buffer_for_mega_moe=lambda: None, + transform_sf_into_required_layout=lambda: None, + transform_weights_for_mega_moe=lambda: None, + per_token_cast_to_fp8=per_token_cast_to_fp8, + ) + monkeypatch.setattr(tensorrt_llm, "deep_gemm", deep_gemm) + + with pytest.raises( + quantization_module._MegaMoEUnavailable, + match="lacks set_barrier_timeout_seconds", + ): + quantization_module._import_deep_gemm() + + deep_gemm.set_barrier_timeout_seconds = lambda seconds: None + assert quantization_module._import_deep_gemm() is deep_gemm + + def test_fp8_block_scale_moe_fallback_tactic_is_explicit_and_deterministic(): valid_tactics = [ [8, 0], diff --git a/tests/unittest/_torch/moe/test_moe_comm_timeout_guard.py b/tests/unittest/_torch/moe/test_moe_comm_timeout_guard.py new file mode 100644 index 000000000000..59aaab3d06a7 --- /dev/null +++ b/tests/unittest/_torch/moe/test_moe_comm_timeout_guard.py @@ -0,0 +1,422 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the MoE communication timeout guard and its backend proxies. + +The budget and guard tests are CPU-only. The proxy tests call each backend's native setter, so +they need the built extensions, and the DeepGEMM and DeepEP round trips also need a GPU. They +import the backends inside the tests, because the CPU-only stage collects this whole file. +""" + +import gc +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.moe.fused_moe import moe_comm_timeout_guard +from tensorrt_llm._torch.moe.fused_moe.moe_comm_timeout_guard import ( + DEFAULT_WARMUP_TIMEOUT_SEC, + MoECommTimeoutBudgets, + MoECommTimeoutGuard, + moe_comm_serving_timeouts, + register_moe_comm_timeout_proxy, + resolve_moe_comm_timeout_budgets, + set_moe_comm_warmup, +) + +SERVING_ENV = "TRTLLM_MOE_COMM_TIMEOUT_SEC" +WARMUP_ENV = "TRTLLM_MOE_COMM_WARMUP_TIMEOUT_SEC" +SERVING_ALIASES = ["TRTLLM_MOE_A2A_TIMEOUT_SEC", "TRTLLM_NVLINK_ONE_SIDED_A2A_TIMEOUT_SEC"] +WARMUP_ALIASES = [ + "TRTLLM_MOE_A2A_WARMUP_TIMEOUT_SEC", + "TRTLLM_NVLINK_ONE_SIDED_A2A_WARMUP_TIMEOUT_SEC", +] + +_NVLINK_ONE_SIDED_DEFAULT_SEC = 300 +_DEEP_EP_DEFAULT_SEC = 100 + + +class _RecordingProxy: + def __init__(self, name: str = "recording") -> None: + self.name = name + self.calls: list[int | None] = [] + + def set_timeout_seconds(self, seconds: int | None) -> None: + self.calls.append(seconds) + + +class _FailingProxy: + name = "failing" + + def set_timeout_seconds(self, seconds: int | None) -> None: + raise RuntimeError("setter unavailable") + + +class _CaptureState: + def __init__(self) -> None: + self.active = False + + def __call__(self) -> bool: + return self.active + + +def _make_guard( + environ: dict[str, str] | None = None, capture: _CaptureState | None = None +) -> MoECommTimeoutGuard: + return MoECommTimeoutGuard( + environ={} if environ is None else environ, + is_capturing=_CaptureState() if capture is None else capture, + ) + + +@pytest.mark.cpu_only +class TestBudgetResolution: + def test_defaults_relax_warmup_and_keep_native_serving(self) -> None: + assert resolve_moe_comm_timeout_budgets({}) == MoECommTimeoutBudgets( + warmup_seconds=DEFAULT_WARMUP_TIMEOUT_SEC, serving_seconds=None + ) + + def test_new_names_set_both_budgets(self) -> None: + budgets = resolve_moe_comm_timeout_budgets({WARMUP_ENV: "900", SERVING_ENV: "120"}) + + assert budgets == MoECommTimeoutBudgets(warmup_seconds=900, serving_seconds=120) + + def test_empty_value_counts_as_unset(self) -> None: + assert resolve_moe_comm_timeout_budgets({SERVING_ENV: " "}).serving_seconds is None + + @pytest.mark.parametrize("raw", ["abc", "1.5", "0", "-3", "86401", "1e3"]) + def test_invalid_value_raises_and_names_the_variable(self, raw: str) -> None: + with pytest.raises(ValueError, match=SERVING_ENV): + resolve_moe_comm_timeout_budgets({SERVING_ENV: raw}) + + @pytest.mark.parametrize( + ("warmup_alias", "serving_alias"), list(zip(WARMUP_ALIASES, SERVING_ALIASES)) + ) + def test_deprecated_aliases_apply_and_warn( + self, monkeypatch: pytest.MonkeyPatch, warmup_alias: str, serving_alias: str + ) -> None: + warned: list[str] = [] + monkeypatch.setattr( + moe_comm_timeout_guard.logger, "warning_once", lambda *msg, key: warned.append(key) + ) + + budgets = resolve_moe_comm_timeout_budgets({warmup_alias: "1000", serving_alias: "200"}) + + assert budgets == MoECommTimeoutBudgets(warmup_seconds=1000, serving_seconds=200) + assert sorted(warned) == sorted( + [f"deprecated_env_{warmup_alias}", f"deprecated_env_{serving_alias}"] + ) + + @pytest.mark.parametrize("serving_alias", SERVING_ALIASES) + def test_alias_equal_to_new_name_is_accepted(self, serving_alias: str) -> None: + budgets = resolve_moe_comm_timeout_budgets({SERVING_ENV: "200", serving_alias: "200"}) + + assert budgets.serving_seconds == 200 + + @pytest.mark.parametrize("serving_alias", SERVING_ALIASES) + def test_alias_conflicting_with_new_name_raises(self, serving_alias: str) -> None: + with pytest.raises(ValueError, match=serving_alias): + resolve_moe_comm_timeout_budgets({SERVING_ENV: "100", serving_alias: "200"}) + + def test_conflicting_aliases_raise(self) -> None: + with pytest.raises(ValueError, match=WARMUP_ALIASES[1]): + resolve_moe_comm_timeout_budgets({WARMUP_ALIASES[0]: "900", WARMUP_ALIASES[1]: "1200"}) + + @pytest.mark.parametrize("warmup_alias", WARMUP_ALIASES) + def test_invalid_alias_value_raises_and_names_the_alias(self, warmup_alias: str) -> None: + with pytest.raises(ValueError, match=warmup_alias): + resolve_moe_comm_timeout_budgets({warmup_alias: "abc"}) + + def test_warmup_shorter_than_serving_raises(self) -> None: + with pytest.raises(ValueError, match=WARMUP_ENV): + resolve_moe_comm_timeout_budgets({WARMUP_ENV: "100", SERVING_ENV: "200"}) + + def test_budgets_select_warmup_or_serving(self) -> None: + budgets = MoECommTimeoutBudgets(warmup_seconds=1800, serving_seconds=None) + + assert budgets.select(in_warmup=True) == 1800 + assert budgets.select(in_warmup=False) is None + + +@pytest.mark.cpu_only +class TestGuard: + def test_guard_starts_in_serving(self) -> None: + guard = _make_guard(environ={SERVING_ENV: "90"}) + proxy = _RecordingProxy() + + guard.register(proxy) + + assert not guard.in_warmup + assert proxy.calls == [90] + + def test_registration_applies_the_current_timeout_at_once(self) -> None: + guard = _make_guard() + guard.set_warmup(True) + proxy = _RecordingProxy() + + guard.register(proxy) + + assert proxy.calls == [DEFAULT_WARMUP_TIMEOUT_SEC] + + def test_transitions_reach_every_proxy(self) -> None: + guard = _make_guard(environ={SERVING_ENV: "90"}) + first, second = _RecordingProxy("first"), _RecordingProxy("second") + guard.register(first) + guard.register(second) + + guard.set_warmup(True) + guard.set_warmup(False) + + assert first.calls == [90, DEFAULT_WARMUP_TIMEOUT_SEC, 90] + assert second.calls == [90, DEFAULT_WARMUP_TIMEOUT_SEC, 90] + + def test_serving_without_override_restores_native_defaults(self) -> None: + guard = _make_guard() + guard.set_warmup(True) + proxy = _RecordingProxy() + guard.register(proxy) + + guard.set_warmup(False) + + assert proxy.calls == [DEFAULT_WARMUP_TIMEOUT_SEC, None] + + def test_repeating_the_current_state_pushes_nothing(self) -> None: + guard = _make_guard() + proxy = _RecordingProxy() + guard.register(proxy) + + guard.set_warmup(False) + guard.set_warmup(True) + guard.set_warmup(True) + + assert proxy.calls == [None, DEFAULT_WARMUP_TIMEOUT_SEC] + + def test_registering_a_proxy_twice_keeps_one_entry(self) -> None: + guard = _make_guard() + proxy = _RecordingProxy() + guard.register(proxy) + guard.register(proxy) + + guard.set_warmup(True) + + assert proxy.calls == [None, None, DEFAULT_WARMUP_TIMEOUT_SEC] + + def test_unregistered_proxy_stops_receiving_transitions(self) -> None: + guard = _make_guard() + proxy = _RecordingProxy() + guard.register(proxy) + + guard.unregister(proxy) + guard.set_warmup(True) + + assert proxy.calls == [None] + + def test_collected_proxy_is_dropped(self) -> None: + guard = _make_guard() + calls: list[int | None] = [] + + class _TemporaryProxy: + name = "temporary" + + def set_timeout_seconds(self, seconds: int | None) -> None: + calls.append(seconds) + + proxy = _TemporaryProxy() + guard.register(proxy) + del proxy + gc.collect() + + guard.set_warmup(True) + + assert calls == [None] + + def test_proxy_whose_setter_fails_is_not_registered(self) -> None: + guard = _make_guard() + with pytest.raises(RuntimeError, match="setter unavailable"): + guard.register(_FailingProxy()) + healthy = _RecordingProxy() + guard.register(healthy) + + guard.set_warmup(True) + + assert healthy.calls == [None, DEFAULT_WARMUP_TIMEOUT_SEC] + + def test_budgets_resolve_once(self) -> None: + environ = {SERVING_ENV: "90"} + guard = _make_guard(environ=environ) + proxy = _RecordingProxy() + guard.register(proxy) + + environ[SERVING_ENV] = "30" + guard.set_warmup(True) + guard.set_warmup(False) + + assert proxy.calls == [90, DEFAULT_WARMUP_TIMEOUT_SEC, 90] + + def test_invalid_environment_fails_the_first_registration(self) -> None: + guard = _make_guard(environ={SERVING_ENV: "never"}) + + with pytest.raises(ValueError, match=SERVING_ENV): + guard.register(_RecordingProxy()) + + def test_transitions_without_proxies_leave_the_environment_unread(self) -> None: + guard = _make_guard(environ={SERVING_ENV: "never"}) + + guard.set_warmup(True) + guard.set_warmup(False) + + with pytest.raises(ValueError, match=SERVING_ENV): + guard.budgets() + + def test_budgets_are_available_before_any_registration(self) -> None: + environ = {WARMUP_ENV: "900"} + guard = _make_guard(environ=environ) + + first = guard.budgets() + environ[WARMUP_ENV] = "1200" + + assert first == MoECommTimeoutBudgets(warmup_seconds=900, serving_seconds=None) + assert guard.budgets() == first + + def test_switching_during_cuda_graph_capture_raises(self) -> None: + capture = _CaptureState() + guard = _make_guard(capture=capture) + proxy = _RecordingProxy() + guard.register(proxy) + capture.active = True + + with pytest.raises(RuntimeError, match="captured"): + guard.set_warmup(True) + + assert not guard.in_warmup + assert proxy.calls == [None] + + def test_serving_timeouts_apply_serving_and_restore_warmup(self) -> None: + guard = _make_guard(environ={SERVING_ENV: "90"}) + guard.set_warmup(True) + proxy = _RecordingProxy() + guard.register(proxy) + + with guard.serving_timeouts(): + assert not guard.in_warmup + assert proxy.calls == [DEFAULT_WARMUP_TIMEOUT_SEC, 90] + + assert guard.in_warmup + assert proxy.calls == [DEFAULT_WARMUP_TIMEOUT_SEC, 90, DEFAULT_WARMUP_TIMEOUT_SEC] + + def test_serving_timeouts_restore_when_the_block_raises(self) -> None: + guard = _make_guard() + guard.set_warmup(True) + + with pytest.raises(KeyError), guard.serving_timeouts(): + raise KeyError("boom") + + assert guard.in_warmup + + def test_module_functions_use_the_process_wide_guard( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + guard = _make_guard(environ={SERVING_ENV: "90"}) + monkeypatch.setattr(moe_comm_timeout_guard, "_DEFAULT_GUARD", guard) + proxy = _RecordingProxy() + + register_moe_comm_timeout_proxy(proxy) + set_moe_comm_warmup(True) + with moe_comm_serving_timeouts(): + pass + + assert proxy.calls == [90, DEFAULT_WARMUP_TIMEOUT_SEC, 90, DEFAULT_WARMUP_TIMEOUT_SEC] + + +class TestBackendProxies: + def test_nvlink_one_sided_proxy_sets_the_completion_flag_budget(self) -> None: + from tensorrt_llm._torch.moe.fused_moe.communication import nvlink_one_sided + + proxy = nvlink_one_sided._TIMEOUT_PROXY + try: + proxy.set_timeout_seconds(1800) + assert torch.ops.trtllm.moe_a2a_get_timeout() == 1800 + proxy.set_timeout_seconds(None) + assert torch.ops.trtllm.moe_a2a_get_timeout() == _NVLINK_ONE_SIDED_DEFAULT_SEC + finally: + proxy.set_timeout_seconds(None) + + @pytest.mark.parametrize("timeout_sec", [-1, 86401]) + def test_moe_a2a_set_timeout_rejects_out_of_range(self, timeout_sec: int) -> None: + before = torch.ops.trtllm.moe_a2a_get_timeout() + + with pytest.raises(RuntimeError, match="MoE all-to-all timeout"): + torch.ops.trtllm.moe_a2a_set_timeout(timeout_sec) + + assert torch.ops.trtllm.moe_a2a_get_timeout() == before + + def test_deepgemm_proxy_maps_the_native_default_to_zero(self) -> None: + from tensorrt_llm._torch.moe.fused_moe.mega_moe.mega_moe_deepgemm import ( + _DeepGemmTimeoutProxy, + ) + + calls: list[int] = [] + proxy = _DeepGemmTimeoutProxy(SimpleNamespace(set_barrier_timeout_seconds=calls.append)) + + proxy.set_timeout_seconds(1800) + proxy.set_timeout_seconds(None) + + assert calls == [1800, 0] + + def test_deepgemm_proxy_sets_the_barrier_timeout(self) -> None: + if not torch.cuda.is_available(): + pytest.skip("DeepGEMM's device runtime needs a GPU") + dg = pytest.importorskip("tensorrt_llm.deep_gemm") + from tensorrt_llm._torch.moe.fused_moe.mega_moe.mega_moe_deepgemm import ( + _DeepGemmTimeoutProxy, + ) + + proxy = _DeepGemmTimeoutProxy(dg) + try: + proxy.set_timeout_seconds(1800) + assert dg.get_barrier_timeout_seconds() == 1800 + proxy.set_timeout_seconds(None) + assert dg.get_barrier_timeout_seconds() == 0 + with pytest.raises(RuntimeError): + dg.set_barrier_timeout_seconds(86401) + finally: + proxy.set_timeout_seconds(None) + + def test_deep_ep_proxy_sets_the_kernel_timeout(self) -> None: + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9: + pytest.skip("DeepEP is built for SM90 and newer") + from tensorrt_llm._torch.moe.fused_moe import deep_ep_utils + + if not deep_ep_utils.deep_ep_installed: + pytest.skip("DeepEP is not installed") + proxy = deep_ep_utils._TIMEOUT_PROXY + try: + proxy.set_timeout_seconds(1800) + assert deep_ep_utils.deep_ep.get_timeout_seconds() == 1800 + proxy.set_timeout_seconds(None) + assert deep_ep_utils.deep_ep.get_timeout_seconds() == _DEEP_EP_DEFAULT_SEC + finally: + proxy.set_timeout_seconds(None) + + def test_nccl_ep_proxy_converts_seconds_to_group_nanoseconds(self) -> None: + from tensorrt_llm._torch.moe.fused_moe.nccl_ep_utils import _NcclEpGroupTimeoutProxy + + calls: list[int] = [] + proxy = _NcclEpGroupTimeoutProxy(SimpleNamespace(set_timeout_ns=calls.append)) + + proxy.set_timeout_seconds(1800) + proxy.set_timeout_seconds(None) + + assert calls == [1_800_000_000_000, 0]