[https://nvbugs/6482566][fix] Make MoE all-to-all completion-flag timeout phase-aware

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 <cerrno>
 #include <cooperative_groups.h>
 #include <cstdint>
+#include <cstdlib>
 #include <type_traits>
 
 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<long>(kMaxSec), static_cast<long>(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<long>(sSteadySec), static_cast<long>(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;

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;
 };

cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp

@@ -21,6 +21,7 @@
 #include "tensorrt_llm/thop/moeAlltoAllMeta.h"
 #include "tensorrt_llm/thop/thUtils.h"
 
+#include <atomic>
 #include <c10/cuda/CUDAStream.h>
 #include <torch/extension.h>
 #include <torch/types.h>
@@ -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<bool> 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<std::vector<torch::Tensor>, 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);
 }

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.