diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu index 74e40dbb2b81..8559653b6a5b 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu @@ -21,8 +21,10 @@ #include "tensorrt_llm/common/vec_dtypes.cuh" #include "tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h" #include "tensorrt_llm/kernels/quantization.cuh" +#include #include #include +#include #include TRTLLM_NAMESPACE_BEGIN @@ -32,6 +34,54 @@ namespace kernels::moe_comm using tensorrt_llm::common::launchWithPdlWhenEnabled; +// Resolve the completion-flag wait budget; see the header. Seconds are converted at +// an assumed 2 GHz SM clock, so they are nominal rather than wall-clock. +int64_t moeA2AGetTimeoutCycles(bool is_warmup) +{ + static constexpr int64_t kAssumedClockHz = 2000ll * 1000ll * 1000ll; + static constexpr int64_t kDefaultTimeoutSec = 300; + // Warmup contains one-time per-rank costs (JIT compilation, autotuning, module + // loading) that can run for minutes and are not synchronized against this + // collective, so it needs a larger budget than steady state. + static constexpr int64_t kDefaultWarmupTimeoutSec = 1800; + + // Reject trailing garbage, out-of-range values and anything that would overflow + // the cycle multiplication. + auto const readEnv = [](char const* name, int64_t fallback) -> int64_t + { + static constexpr int64_t kMaxSec = 24 * 60 * 60; // 1 day; * 2e9 stays well inside int64 + char const* v = std::getenv(name); + if (v == nullptr || *v == '\0') + { + return fallback; + } + errno = 0; + char* end = nullptr; + int64_t parsed = std::strtoll(v, &end, 10); + bool const trailingGarbage = (end == v) || (*end != '\0'); + if (trailingGarbage || errno == ERANGE || parsed <= 0 || parsed > kMaxSec) + { + TLLM_LOG_WARNING("Ignoring invalid %s=\"%s\" (expected 1..%ld seconds); using %ld s", name, v, + static_cast(kMaxSec), static_cast(fallback)); + return fallback; + } + return parsed; + }; + + static int64_t const sSteadySec = readEnv("TRTLLM_MOE_A2A_TIMEOUT_SEC", kDefaultTimeoutSec); + static int64_t const sWarmupSec = readEnv("TRTLLM_MOE_A2A_WARMUP_TIMEOUT_SEC", kDefaultWarmupTimeoutSec); + static bool const sLogged = []() + { + TLLM_LOG_INFO( + "MoE all-to-all completion-flag budget: steady=%ld s, warmup=%ld s (nominal, at an " + "assumed 2 GHz clock64 rate)", + static_cast(sSteadySec), static_cast(sWarmupSec)); + return true; + }(); + (void) sLogged; + return (is_warmup ? sWarmupSec : sSteadySec) * kAssumedClockHz; +} + #define ENABLE_DEBUG_PRINT 0 #define DISABLE_SYNC_FOR_PROFILING 0 @@ -161,10 +211,10 @@ using tensorrt_llm::common::launchWithPdlWhenEnabled; } #if DISABLE_TIMEOUT -#define check_timeout(s) false +#define check_timeout(s, budget) false #else -// 300 * 2000 MHz - should be high enough on any GPU but will prevent a hang -#define check_timeout(s) ((clock64() - (s)) > (300ll * 2000ll * 1000ll * 1000ll)) +// `budget` is in clock64() cycles, resolved on the host by moeA2AGetTimeoutCycles(). +#define check_timeout(s, budget) ((clock64() - (s)) > (budget)) #endif // ============================================================================ @@ -618,7 +668,7 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ rank_id, peer_rank, flag_value, expected_value, flag_ptr); #endif flag_set = flag_value == expected_value; - } while (!flag_set && !check_timeout(s)); + } while (!flag_set && !check_timeout(s, ptrs.timeout_cycles)); if (__builtin_expect(!flag_set, 0)) { @@ -659,6 +709,7 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) // Prepare kernel pointers struct DispatchKernelPointers kernel_ptrs = {}; + kernel_ptrs.timeout_cycles = params.timeout_cycles; // Fill source data pointers and payload sizes for (int i = 0; i < params.num_payloads; i++) @@ -1258,7 +1309,7 @@ __global__ void moeA2ACombineKernel( rank_id, peer_rank, flag_value, expected_value, flag_ptr); #endif flag_set = flag_value == expected_value; - } while (!flag_set && !check_timeout(s)); + } while (!flag_set && !check_timeout(s, ptrs.timeout_cycles)); if (__builtin_expect(!flag_set, 0)) { @@ -1365,6 +1416,7 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) // Prepare kernel pointers struct for combine CombineKernelPointers kernel_ptrs = {}; // Zero-initialize + kernel_ptrs.timeout_cycles = params.timeout_cycles; // Set output data pointer in src_data_ptrs[0] kernel_ptrs.src_data_ptrs[0] = params.output_data; diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h index 5184878ffc51..0b25e94a9e26 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h @@ -33,6 +33,9 @@ static constexpr int kRankMaskWords = 2; // uint64 words to hold the active-rank // (kRankMaskWords * 64 must be >= kMaxRanks) static_assert(kRankMaskWords * 64 >= kMaxRanks, "active_rank_mask too small for kMaxRanks"); +// Default completion-flag wait budget: 300 s at an assumed 2 GHz clock64 rate. +static constexpr int64_t kDefaultTimeoutCycles = 300ll * 2000ll * 1000ll * 1000ll; + // Describes a single payload type to be communicated struct PayloadDescriptor { @@ -73,6 +76,9 @@ struct DispatchKernelPointers // Word 0 covers ranks 0..63; word 1 covers ranks 64..127. The masked kernel // rejects inactive route targets and skips their peer counters, stats, and flags. uint64_t active_rank_mask[kRankMaskWords]; + + // Completion-flag wait budget in clock64() cycles; see moeA2AGetTimeoutCycles(). + int64_t timeout_cycles{kDefaultTimeoutCycles}; }; // Combine kernel pointers - non-const output in src_data_ptrs[0], const recv buffers @@ -94,6 +100,9 @@ struct CombineKernelPointers // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Combine skips // completion flag writes/waits to/from inactive peers. uint64_t active_rank_mask[kRankMaskWords]; + + // Completion-flag wait budget in clock64() cycles; see moeA2AGetTimeoutCycles(). + int64_t timeout_cycles{kDefaultTimeoutCycles}; }; // Dispatch phase parameters @@ -147,10 +156,22 @@ struct MoeA2ADispatchParams // CUDA graph replay until generation-scoped invalidation and recapture are available. uint64_t active_rank_mask[kRankMaskWords] = {~uint64_t{0}, ~uint64_t{0}}; + // Completion-flag wait budget in clock64() cycles; see moeA2AGetTimeoutCycles(). + int64_t timeout_cycles{kDefaultTimeoutCycles}; + // CUDA stream cudaStream_t stream; }; +// Resolve the completion-flag wait budget, in clock64() cycles. +// +// 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. Warmup therefore uses a larger budget than steady state. +// Overridable via TRTLLM_MOE_A2A_TIMEOUT_SEC / TRTLLM_MOE_A2A_WARMUP_TIMEOUT_SEC. +// See nvbugs/6482566. +int64_t moeA2AGetTimeoutCycles(bool is_warmup); + // Dispatch kernels void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params); // Prepare for dispatch: zero send_counters, local_token_counter and increment flag_val @@ -202,6 +223,9 @@ struct MoeA2ACombineParams // CUDA graph replay until generation-scoped invalidation and recapture are available. uint64_t active_rank_mask[kRankMaskWords] = {~uint64_t{0}, ~uint64_t{0}}; + // Completion-flag wait budget in clock64() cycles; see moeA2AGetTimeoutCycles(). + int64_t timeout_cycles{kDefaultTimeoutCycles}; + // CUDA stream cudaStream_t stream; }; diff --git a/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp b/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp index e985eb943ee1..81b3dc33fe28 100644 --- a/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp +++ b/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp @@ -21,6 +21,7 @@ #include "tensorrt_llm/thop/moeAlltoAllMeta.h" #include "tensorrt_llm/thop/thUtils.h" +#include #include #include #include @@ -34,6 +35,15 @@ namespace torch_ext namespace moe_comm { +// Whether the engine is in its startup warmup phase, which uses a larger +// completion-flag budget. See moeA2AGetTimeoutCycles(). +static std::atomic gInWarmup{false}; + +void moeA2ASetWarmupOp(bool in_warmup) +{ + gInWarmup.store(in_warmup, std::memory_order_relaxed); +} + static constexpr size_t CACHELINE_ALIGNMENT = 128; // TODO: Is Alignment necessary? @@ -409,6 +419,8 @@ std::tuple, int64_t, torch::Tensor> moeA2ADispatchOp( } params.stream = at::cuda::getCurrentCUDAStream(); + params.timeout_cycles + = tensorrt_llm::kernels::moe_comm::moeA2AGetTimeoutCycles(gInWarmup.load(std::memory_order_relaxed)); // Prepare for dispatch (zero counters/indices and increment flag_val) moe_a2a_prepare_dispatch_launch(params); @@ -582,6 +594,8 @@ torch::Tensor moeA2ACombineOp(torch::Tensor const& payload, int64_t localNumToke } params.stream = at::cuda::getCurrentCUDAStream(); + params.timeout_cycles + = tensorrt_llm::kernels::moe_comm::moeA2AGetTimeoutCycles(gInWarmup.load(std::memory_order_relaxed)); moe_a2a_prepare_combine_launch(params); @@ -692,6 +706,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, module) "moe_a2a_get_combine_payload_tensor(Tensor(a) workspace, int ep_rank, int ep_size, int " "runtime_max_tokens_per_rank, " "int combine_payload_offset, ScalarType out_dtype, int hidden_size) -> Tensor(a)"); + module.def("moe_a2a_set_warmup(bool in_warmup) -> ()", &tensorrt_llm::torch_ext::moe_comm::moeA2ASetWarmupOp); module.def("moe_a2a_get_aux_data_size(int ep_size, int max_num_tokens, int? eplb_stats_num_experts=None) -> int", &tensorrt_llm::torch_ext::moe_comm::moeA2AGetAuxDataSizeOp); } diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 390ea7e6f094..92f79e2505ca 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -301,6 +301,35 @@ def _configure_deep_gemm_pdl() -> None: _DEEP_GEMM_PDL_CONFIGURED = True +@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). + """ + try: + torch.ops.trtllm.moe_a2a_set_warmup(in_warmup) + logger.info(f"moe_a2a completion-flag budget: in_warmup={in_warmup}") + except (AttributeError, RuntimeError) as e: + logger.warning( + f"moe_a2a_set_warmup unavailable, the all-to-all timeout " + f"budget was not switched: {type(e).__name__}: {e}") + + class PyTorchModelEngine(ModelEngine): def __init__( @@ -953,6 +982,11 @@ def is_warmup(self): 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) + self.moe_load_balancer_iter_info = (not value, not value) @property @@ -1192,6 +1226,11 @@ def warmup(self, resource_manager: ResourceManager) -> None: and not isinstance(kv_cache_manager, MambaHybridCacheManager)) log_mem_snapshot("warmup/before_warmup") + # Compile the DSv4 indexer-Q CuTe DSL kernels before the first + # collective-bearing forward, so their JIT cost is not charged against the + # MoE all-to-all completion-flag deadline. + self._prewarm_cute_dsl_indexer_q() + log_mem_snapshot("warmup/after_cute_dsl_indexer_q") if not is_enc_dec: self._run_attention_warmup(resource_manager, can_run_general_warmup) @@ -1234,14 +1273,17 @@ def warmup(self, resource_manager: ResourceManager) -> None: # 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. - with self.cuda_graph_runner.allow_capture(): - self.cuda_graph_runner.is_warmup_only = True - try: + # Capture with the steady-state MoE all-to-all budget: the timeout is a + # launch argument and is baked into every later replay. + with _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) - finally: - self.cuda_graph_runner.is_warmup_only = False - self.cuda_graph_runner.padding_dummy_requests = {} - self._run_cuda_graph_warmup(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 @@ -1367,6 +1409,63 @@ def _warmup_dg_paged_mqa_logits_metadata(self) -> None: f"skipping bucket. {type(e).__name__}: {e}") torch.cuda.synchronize() + def _prewarm_cute_dsl_indexer_q(self) -> None: + """Pre-compile the DSv4 indexer-Q CuTe DSL kernels, then barrier. + + Runs before any collective-bearing forward so this op's first-touch + ``cute.compile`` is not charged against the MoE all-to-all + 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. + + 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 + compile per tactic covers every shape. Uses the real module and weights, + so it cannot drift from what the model runs. + + No-op on non-DSA models. See nvbugs/6482566. + """ + try: + from ..attention_backend.sparse.deepseek_v4.deepseek_v4 import \ + DeepseekV4Indexer + except ImportError: + return + + indexer = next( + (m + for m in self.model.modules() if isinstance(m, DeepseekV4Indexer) + and getattr(m, "wq_b", None) is not None), None) + if indexer is None: + return + + weight = indexer.wq_b.weight + # _fallback_tactic() branches on m at 4 and 8, so these three token + # counts cover every fallback tactic it can return. + with torch.inference_mode(): + for num_tokens in (4, 8, 16): + try: + qr = torch.zeros((num_tokens, weight.shape[1]), + dtype=torch.bfloat16, + device=weight.device) + position_ids = torch.zeros((num_tokens, ), + dtype=torch.int32, + device=weight.device) + indexer._project_and_quantize_q(qr, position_ids) + except Exception as e: + # Never fail startup for a prewarm miss; the kernel would + # simply be compiled later, as it is today. + logger.warning( + f"indexer-Q CuTe DSL prewarm skipped for {num_tokens} " + f"tokens. {type(e).__name__}: {e}") + torch.cuda.synchronize() + + # Hold every rank here until the slowest has finished compiling, so the + # first MoE all-to-all dispatch is entered without JIT skew. + if self.mapping.tp_size > 1: + self.dist.tp_allgather(1) + logger.info("indexer-Q CuTe DSL prewarm complete") + def _warmup_cute_dsl_radix_topk(self) -> None: """Pre-compile the DSA radix-filter CuTe DSL decode top-k for every cluster_size band during warmup, before serving. diff --git a/tests/unittest/_torch/misc/test_moe_a2a_warmup_timeout.py b/tests/unittest/_torch/misc/test_moe_a2a_warmup_timeout.py new file mode 100644 index 000000000000..0067bf2726b6 --- /dev/null +++ b/tests/unittest/_torch/misc/test_moe_a2a_warmup_timeout.py @@ -0,0 +1,78 @@ +import unittest +from unittest import mock + +from tensorrt_llm._torch.pyexecutor import model_engine +from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine + + +class _WarmupFlagStub: + """Minimal object that reuses the engine's is_warmup property. + + Building a real PyTorchModelEngine needs a model and a device; the property + itself only touches _is_warmup, the MoE all-to-all budget selector, and + moe_load_balancer_iter_info (a no-op without a balancer), so a stub exercises + the real code path without either. + """ + + is_warmup = PyTorchModelEngine.is_warmup + moe_load_balancer_iter_info = PyTorchModelEngine.moe_load_balancer_iter_info + + +class TestMoeA2AWarmupBudget(unittest.TestCase): + """The MoE all-to-all completion-flag budget must track the warmup phase. + + The kernel-side deadline is only safe if it is raised for warmup *and* + lowered again afterwards; a budget that latches on would leave the hang + watchdog permanently relaxed in steady state. See nvbugs/6482566. + """ + + def test_set_warmup_forwards_value_to_op(self): + with mock.patch.object( + model_engine.torch.ops.trtllm, "moe_a2a_set_warmup", create=True + ) as op: + model_engine._set_moe_a2a_warmup(True) + model_engine._set_moe_a2a_warmup(False) + self.assertEqual([c.args[0] for c in op.call_args_list], [True, False]) + + def test_missing_op_is_tolerated(self): + """An older C++ build without the op must not break startup.""" + with mock.patch.object( + model_engine.torch.ops.trtllm, + "moe_a2a_set_warmup", + create=True, + side_effect=AttributeError("no such op"), + ): + model_engine._set_moe_a2a_warmup(True) # must not raise + + def test_capture_context_selects_steady_state_then_restores(self): + """CUDA graphs bake the budget in at capture time. + + Capture runs inside the warmup window, so the context manager must hand + the kernel the steady-state budget and restore warmup afterwards. + """ + seen = [] + with mock.patch.object(model_engine, "_set_moe_a2a_warmup", side_effect=seen.append): + with model_engine._moe_a2a_steady_state_budget_for_capture(): + self.assertEqual(seen, [False]) + self.assertEqual(seen, [False, True]) + + def test_is_warmup_setter_switches_budget_both_ways(self): + """Regression: the budget must not latch on after warmup. + + PyExecutor sets is_warmup=True before calling warmup() and False after, + both through this setter. Selecting the budget anywhere else (e.g. only + in set_warmup_flag) leaves the relaxed warmup budget in force for the + whole serving lifetime. + """ + stub = _WarmupFlagStub() + seen = [] + with mock.patch.object(model_engine, "_set_moe_a2a_warmup", side_effect=seen.append): + stub.is_warmup = True + stub.is_warmup = False + + self.assertEqual(seen, [True, False]) + self.assertFalse(stub.is_warmup) + + +if __name__ == "__main__": + unittest.main()