diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu index 0f2d453c363a..472a5877a80d 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu @@ -392,7 +392,7 @@ __global__ void moeA2APrepareDispatchKernel( // Dispatch Kernels // ============================================================================ -template +template __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [local_num_tokens, TOP_K] const DispatchKernelPointers ptrs, // Struct containing all kernel pointers int num_payloads, // Number of payloads @@ -440,15 +440,17 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ // Supports the non-divisible case where num_experts % ep_size != 0. int target_rank = compute_target_rank_id(expert_id, ep_base, ep_remainder); - // Skip duplicates AND dead ranks: both produce the same -1 sentinel that combine - // checks via topk_send_indices[k] < 0. A token whose only target is dead is dropped - // from this collective; higher-layer logic (EPLB redistribution) is responsible - // for re-routing such tokens on subsequent iterations. int const mask_word = target_rank >> 6; uint64_t const mask_bit = 1ULL << (target_rank & 63); - bool const target_already_copied = already_copied[mask_word] & mask_bit; - bool const target_dead = !is_rank_active(ptrs.active_rank_mask, target_rank); - if (target_already_copied || target_dead) + bool const target_already_copied = (already_copied[mask_word] & mask_bit) != 0; + bool skip_target = target_already_copied; + if constexpr (ENABLE_RANK_MASK) + { + // This is a fail-closed safety guard until post-commit routing is enforced end to end. + // A masked route is not valid model output; the failed execution epoch must be discarded. + skip_target = skip_target || !is_rank_active(ptrs.active_rank_mask, target_rank); + } + if (skip_target) { if (thread_idx == 0) { @@ -532,8 +534,11 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ #pragma unroll 1 // No unroll as one iter is typically enough for (int target_rank = lane_id; target_rank < ep_size; target_rank += warpSize) { - if (!is_rank_active(ptrs.active_rank_mask, target_rank)) - continue; + if constexpr (ENABLE_RANK_MASK) + { + if (!is_rank_active(ptrs.active_rank_mask, target_rank)) + continue; + } int send_count = ptrs.send_counters[target_rank]; ptrs.recv_counters[target_rank][rank_id] = send_count; } @@ -545,8 +550,11 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ #pragma unroll 1 for (int target_rank = 0; target_rank < ep_size; ++target_rank) { - if (!is_rank_active(ptrs.active_rank_mask, target_rank)) - continue; + if constexpr (ENABLE_RANK_MASK) + { + if (!is_rank_active(ptrs.active_rank_mask, target_rank)) + continue; + } int* target_stats = ptrs.eplb_gathered_stats[target_rank]; for (int expert_id = lane_id; expert_id < eplb_stats_num_experts; expert_id += warpSize) { @@ -570,8 +578,11 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ #pragma unroll 1 // No unroll as one iter is typically enough for (int target_rank = lane_id; target_rank < ep_size; target_rank += warpSize) { - if (!is_rank_active(ptrs.active_rank_mask, target_rank)) - continue; + if constexpr (ENABLE_RANK_MASK) + { + if (!is_rank_active(ptrs.active_rank_mask, target_rank)) + continue; + } uint32_t* flag_addr = &ptrs.completion_flags[target_rank][rank_id]; asm volatile("st.relaxed.sys.u32 [%0], %1;" ::"l"(flag_addr), "r"(expected_value)); @@ -586,8 +597,11 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ #pragma unroll 1 // No unroll for (int peer_rank = lane_id; peer_rank < ep_size; peer_rank += warpSize) { - if (!is_rank_active(ptrs.active_rank_mask, peer_rank)) - continue; + if constexpr (ENABLE_RANK_MASK) + { + if (!is_rank_active(ptrs.active_rank_mask, peer_rank)) + continue; + } bool flag_set = false; auto s = clock64(); do @@ -636,10 +650,11 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) TLLM_CHECK(params.ep_rank >= 0 && params.ep_rank < params.ep_size); TLLM_CHECK(params.local_num_tokens >= 0); TLLM_CHECK(params.num_payloads > 0 && params.num_payloads <= kMaxPayloads); - // The local rank must always be marked active in its own view of the mask; - // otherwise the kernel itself would be running on a "dead" rank. - TLLM_CHECK_WITH_INFO((params.active_rank_mask[params.ep_rank >> 6] >> (params.ep_rank & 63)) & 1ULL, - "active_rank_mask must mark the local ep_rank (%d) as active", params.ep_rank); + if (params.enable_rank_mask) + { + TLLM_CHECK_WITH_INFO((params.active_rank_mask[params.ep_rank >> 6] >> (params.ep_rank & 63)) & 1ULL, + "active_rank_mask must mark the local ep_rank (%d) as active", params.ep_rank); + } // Prepare kernel pointers struct DispatchKernelPointers kernel_ptrs = {}; @@ -693,12 +708,15 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) grid_size = 1; } int shared_bytes = 2 * params.top_k * (int) sizeof(int); - SWITCH_BOOL(params.enable_eplb, EPLB_STATS, SWITCH_TOP_K(params.top_k, TOP_K, { - auto kernel_fn = moeA2ADispatchKernel; - launchWithPdlWhenEnabled("moeA2ADispatchKernel", kernel_fn, grid_size, kBlockSize, shared_bytes, params.stream, - params.token_selected_experts, kernel_ptrs, params.num_payloads, params.max_tokens_per_rank, - params.local_num_tokens, params.ep_rank, params.ep_size, params.num_experts, params.eplb_stats_num_experts); - })) + SWITCH_BOOL(params.enable_rank_mask, ENABLE_RANK_MASK, {SWITCH_BOOL(params.enable_eplb, EPLB_STATS, { + SWITCH_TOP_K(params.top_k, TOP_K, { + auto kernel_fn = moeA2ADispatchKernel; + launchWithPdlWhenEnabled("moeA2ADispatchKernel", kernel_fn, grid_size, kBlockSize, shared_bytes, + params.stream, params.token_selected_experts, kernel_ptrs, params.num_payloads, + params.max_tokens_per_rank, params.local_num_tokens, params.ep_rank, params.ep_size, params.num_experts, + params.eplb_stats_num_experts); + }); + })}) } // ============================================================================ @@ -741,7 +759,7 @@ __device__ void vectorized_combine_impl(T* dst_typed_base, int size_per_token, i { int target_rank = ptrs.topk_target_ranks[local_token_idx * TOP_K + k]; int dst_idx = ptrs.topk_send_indices[local_token_idx * TOP_K + k]; - if (dst_idx < 0 || !is_rank_active(ptrs.active_rank_mask, target_rank)) + if (dst_idx < 0) { acc[k].fill(0.0f); continue; @@ -766,12 +784,8 @@ __device__ void vectorized_combine_impl(T* dst_typed_base, int size_per_token, i #pragma unroll for (int k = 0; k < TOP_K; ++k) { - int target_rank = ptrs.topk_target_ranks[local_token_idx * TOP_K + k]; - int dst_idx = ptrs.topk_send_indices[local_token_idx * TOP_K + k]; - if (dst_idx < 0 || !is_rank_active(ptrs.active_rank_mask, target_rank)) - { + if (ptrs.topk_send_indices[local_token_idx * TOP_K + k] < 0) continue; // acc[k] already holds 0.0f from fill() above - } #pragma unroll for (int j = elems_per_vec - 1; j >= 0; --j) acc[k][j] = static_cast(reinterpret_cast(&acc[k])[j]); @@ -1158,7 +1172,7 @@ __global__ void moeA2APrepareCombineKernel(uint8_t* recv_buffer_bytes, void cons // Generic Combine Kernel Implementation (Templated by data type) // ============================================================================ -template +template __global__ void moeA2ACombineKernel( const CombineKernelPointers ptrs, // Combine-specific struct, src_data_ptrs[0] is output int max_tokens_per_rank, int elements_per_token, int local_num_tokens, int rank_id, int ep_size, @@ -1203,8 +1217,11 @@ __global__ void moeA2ACombineKernel( #pragma unroll 1 // No unroll for (int peer_rank = lane_id; peer_rank < ep_size; peer_rank += warpSize) { - if (!is_rank_active(ptrs.active_rank_mask, peer_rank)) - continue; + if constexpr (ENABLE_RANK_MASK) + { + if (!is_rank_active(ptrs.active_rank_mask, peer_rank)) + continue; + } uint32_t* flag_addr = &ptrs.completion_flags[peer_rank][rank_id]; asm volatile("st.relaxed.sys.u32 [%0], %1;" ::"l"(flag_addr), "r"(expected_value)); #if ENABLE_DEBUG_PRINT @@ -1219,8 +1236,11 @@ __global__ void moeA2ACombineKernel( #pragma unroll 1 // No unroll for (int peer_rank = lane_id; peer_rank < ep_size; peer_rank += warpSize) { - if (!is_rank_active(ptrs.active_rank_mask, peer_rank)) - continue; + if constexpr (ENABLE_RANK_MASK) + { + if (!is_rank_active(ptrs.active_rank_mask, peer_rank)) + continue; + } bool flag_set = false; auto s = clock64(); do @@ -1327,10 +1347,11 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) TLLM_CHECK(params.ep_rank >= 0 && params.ep_rank < params.ep_size); TLLM_CHECK(params.local_num_tokens >= 0); TLLM_CHECK(params.elements_per_token > 0); - // The local rank must always be marked active in its own view of the mask; - // otherwise the kernel itself would be running on a "dead" rank. - TLLM_CHECK_WITH_INFO((params.active_rank_mask[params.ep_rank >> 6] >> (params.ep_rank & 63)) & 1ULL, - "active_rank_mask must mark the local ep_rank (%d) as active", params.ep_rank); + if (params.enable_rank_mask) + { + TLLM_CHECK_WITH_INFO((params.active_rank_mask[params.ep_rank >> 6] >> (params.ep_rank & 63)) & 1ULL, + "active_rank_mask must mark the local ep_rank (%d) as active", params.ep_rank); + } // Configure kernel launch (one block per token). int const kBlockSize = tensorrt_llm::common::getEnvMoeA2ACombineBlockSize(); @@ -1385,14 +1406,16 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) auto const effective_dtype = params.use_low_precision ? nvinfer1::DataType::kFP8 : params.dtype; // Launch appropriate kernel with compact macros - SWITCH_DTYPE(effective_dtype, TKernelType, { - SWITCH_TOP_K(params.top_k, TOP_K, { - auto kernel_fn = moeA2ACombineKernel; - launchWithPdlWhenEnabled("moeA2ACombineKernel", kernel_fn, grid, kBlockSize, 0, params.stream, kernel_ptrs, - params.max_tokens_per_rank, params.elements_per_token, params.local_num_tokens, params.ep_rank, - params.ep_size, stride_per_token); + SWITCH_BOOL(params.enable_rank_mask, ENABLE_RANK_MASK, { + SWITCH_DTYPE(effective_dtype, TKernelType, { + SWITCH_TOP_K(params.top_k, TOP_K, { + auto kernel_fn = moeA2ACombineKernel; + launchWithPdlWhenEnabled("moeA2ACombineKernel", kernel_fn, grid, kBlockSize, 0, params.stream, + kernel_ptrs, params.max_tokens_per_rank, params.elements_per_token, params.local_num_tokens, + params.ep_rank, params.ep_size, stride_per_token); + }); }); - }); + }) } // Kernel to sanitize expert ids for invalid tokens diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h index 138ca92e71a8..177293684874 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h @@ -62,17 +62,16 @@ struct DispatchKernelPointers int* local_token_counter; // Atomic counter for completed tokens // Top-K compact routing info per local token (size: [local_num_tokens, top_k]) - int* topk_target_ranks; // target rank per k, -1 for duplicates - int* topk_send_indices; // dst index per k, -1 for duplicates + int* topk_target_ranks; // target rank per k, -1 for invalid or duplicate routes + int* topk_send_indices; // dst index per k, -1 for invalid or duplicate routes // Optional: Statistics for EPLB int const* eplb_local_stats; // [eplb_stats_num_experts] int* eplb_gathered_stats[kMaxRanks]; // [ep_size, eplb_stats_num_experts] per rank - // Active-rank bitmask: bit i set => rank i is alive and participates in this collective. - // Word 0 covers ranks 0..63; word 1 covers ranks 64..127. Tokens routed to a masked - // rank are dropped (topk_*[k] = -1); flag writes/waits to/from masked peers are skipped. - // The local rank's own bit must always be set; this is checked at launch time. + // Active-rank bitmask: bit i set => rank i participates in this collective. + // 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]; }; @@ -89,12 +88,11 @@ struct CombineKernelPointers uint32_t* flag_val; // The value of the flag for this round (stored on the local rank) // Top-K compact routing info per local token (size: [local_num_tokens, top_k]) - int const* topk_target_ranks; // target rank per k, -1 for duplicates - int const* topk_send_indices; // dst index per k, -1 for duplicates + int const* topk_target_ranks; // target rank per k, -1 for invalid or duplicate routes + int const* topk_send_indices; // dst index per k, -1 for invalid or duplicate routes - // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Combine skips flag - // writes/waits to/from masked peers and also skips per-token accumulation for ranks that - // become inactive between dispatch and combine. + // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Combine skips + // completion flag writes/waits to/from inactive peers. uint64_t active_rank_mask[kRankMaskWords]; }; @@ -139,9 +137,14 @@ struct MoeA2ADispatchParams int const* eplb_local_stats; // [eplb_stats_num_experts] int* eplb_gathered_stats[kMaxRanks]; // [ep_size, eplb_stats_num_experts] per rank - // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. The launch function - // copies these words into the kernel pointers struct. Defaults to all-ones for - // backwards-compatible "no masking" behavior. + // Whether to instantiate a kernel with active-rank checks. + // This is a launch-lifetime mode, independent of future execution-abort handling. + bool enable_rank_mask{false}; + + // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Used only when + // enable_rank_mask is true; defaults to all-ones for backwards-compatible behavior. + // The mask is copied by value into kernel arguments. Rank-mask mode must reject + // CUDA graph replay until generation-scoped invalidation and recapture are available. uint64_t active_rank_mask[kRankMaskWords] = {~uint64_t{0}, ~uint64_t{0}}; // CUDA stream @@ -189,9 +192,14 @@ struct MoeA2ACombineParams // rank has signaled the target rank void const* recv_buffers[kMaxRanks]; // Per-rank receive buffers (only for single payload) - // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. The launch function - // copies these words into the kernel pointers struct. Defaults to all-ones for - // backwards-compatible "no masking" behavior. + // Whether to instantiate a kernel with active-rank checks in peer synchronization. + // This is a launch-lifetime mode, independent of future execution-abort handling. + bool enable_rank_mask{false}; + + // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Used only when + // enable_rank_mask is true; defaults to all-ones for backwards-compatible behavior. + // The mask is copied by value into kernel arguments. Rank-mask mode must reject + // CUDA graph replay until generation-scoped invalidation and recapture are available. uint64_t active_rank_mask[kRankMaskWords] = {~uint64_t{0}, ~uint64_t{0}}; // CUDA stream diff --git a/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp b/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp index fc45afd792bb..373f936c4c5d 100644 --- a/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp +++ b/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp @@ -42,12 +42,13 @@ inline size_t alignOffset(size_t offset, size_t alignment) return (offset + alignment - 1) & ~(alignment - 1); } -// Resolve an optional rank-mask tensor into a fixed-width uint64 array. -// If the caller did not provide a mask, default to "all ranks active" (all bits set), which -// reproduces the pre-fault-tolerance behavior bit-for-bit. -// -// On failure (wrong dtype / device / shape), throws via TORCH_CHECK so the error surfaces -// at the Python op boundary rather than the kernel launch. +inline bool hasActiveRankMask(torch::optional const& maskTensor) +{ + return maskTensor.has_value() && maskTensor.value().defined(); +} + +// Resolve a provided rank-mask tensor into a fixed-width uint64 array. On failure +// (wrong dtype / device / shape), throw at the Python op boundary rather than launch. inline void resolveActiveRankMask(torch::optional const& maskTensor, int64_t epRank, uint64_t (&out)[tensorrt_llm::kernels::moe_comm::kRankMaskWords]) { @@ -55,14 +56,7 @@ inline void resolveActiveRankMask(torch::optional const& maskTens using tensorrt_llm::kernels::moe_comm::kMaxRanks; TORCH_CHECK( epRank >= 0 && epRank < kMaxRanks, "epRank must be in the range [0, ", kMaxRanks, ") for active_rank_mask"); - if (!maskTensor.has_value() || !maskTensor.value().defined()) - { - for (int w = 0; w < kRankMaskWords; ++w) - { - out[w] = ~uint64_t{0}; - } - return; - } + TORCH_CHECK(hasActiveRankMask(maskTensor), "active_rank_mask must be defined"); torch::Tensor const& t = maskTensor.value(); TORCH_CHECK(t.is_cpu(), "active_rank_mask must be a CPU tensor"); TORCH_CHECK(t.scalar_type() == torch::kUInt64, "active_rank_mask must have dtype uint64"); @@ -222,7 +216,7 @@ std::tuple, int64_t, torch::Tensor> moeA2ADispatchOp( torch::Tensor const& tokenSelectedExperts, std::vector const& inputPayloads, torch::Tensor const& workspace, torch::Tensor const& metainfo, int64_t runtimeMaxTokensPerRank, int64_t epRank, int64_t epSize, int64_t topK, int64_t numExperts, torch::optional eplbLocalStats, - torch::optional activeRankMask) + bool enableRankMask, torch::optional activeRankMask) { using tensorrt_llm::kernels::moe_comm::PayloadDescriptor; using tensorrt_llm::kernels::moe_comm::MoeA2ADispatchParams; @@ -403,9 +397,15 @@ std::tuple, int64_t, torch::Tensor> moeA2ADispatchOp( params.eplb_local_stats = nullptr; } - // Resolve the optional active-rank mask. Default (no mask) = all bits set, which - // exactly reproduces the pre-fault-tolerance kernel behavior. - resolveActiveRankMask(activeRankMask, epRank, params.active_rank_mask); + params.enable_rank_mask = enableRankMask; + if (params.enable_rank_mask) + { + resolveActiveRankMask(activeRankMask, epRank, params.active_rank_mask); + } + else + { + TORCH_CHECK(!hasActiveRankMask(activeRankMask), "active_rank_mask requires enable_rank_mask=True"); + } params.stream = at::cuda::getCurrentCUDAStream(); @@ -460,7 +460,7 @@ std::tuple, int64_t, torch::Tensor> moeA2ADispatchOp( // In both cases, the combine kernel reads from the workspace at 'combinePayloadOffset'. torch::Tensor moeA2ACombineOp(torch::Tensor const& payload, int64_t localNumTokens, torch::Tensor const& workspace, torch::Tensor const& metainfo, int64_t runtimeMaxTokensPerRank, int64_t epRank, int64_t epSize, int64_t topK, - int64_t combinePayloadOffset, bool payloadInWorkspace, bool useLowPrecision = false, + int64_t combinePayloadOffset, bool payloadInWorkspace, bool useLowPrecision, bool enableRankMask, torch::optional activeRankMask = torch::nullopt) { using tensorrt_llm::kernels::moe_comm::MoeA2ACombineParams; @@ -570,8 +570,15 @@ torch::Tensor moeA2ACombineOp(torch::Tensor const& payload, int64_t localNumToke params.recv_buffers[target_rank] = target_workspace_ptr + combinePayloadOffset; } - // Resolve the optional active-rank mask. Default (no mask) = all bits set. - resolveActiveRankMask(activeRankMask, epRank, params.active_rank_mask); + params.enable_rank_mask = enableRankMask; + if (params.enable_rank_mask) + { + resolveActiveRankMask(activeRankMask, epRank, params.active_rank_mask); + } + else + { + TORCH_CHECK(!hasActiveRankMask(activeRankMask), "active_rank_mask requires enable_rank_mask=True"); + } params.stream = at::cuda::getCurrentCUDAStream(); @@ -667,13 +674,13 @@ TORCH_LIBRARY_FRAGMENT(trtllm, module) "Tensor(a!->*) workspace, Tensor metainfo, int runtime_max_tokens_per_rank, " "int ep_rank, int ep_size, int top_k, int num_experts, " "Tensor? eplb_local_stats=None, " - "Tensor? active_rank_mask=None) -> (Tensor(a!)[], int, Tensor(a!))"); + "bool enable_rank_mask=False, Tensor? active_rank_mask=None) -> (Tensor(a!)[], int, Tensor(a!))"); module.def( "moe_a2a_combine(Tensor(a) payload, int local_num_tokens," "Tensor(a!) workspace, Tensor metainfo, int runtime_max_tokens_per_rank, " "int ep_rank, int ep_size, int top_k, int combine_payload_offset, " "bool payload_in_workspace, bool use_low_precision=False, " - "Tensor? active_rank_mask=None) -> Tensor"); + "bool enable_rank_mask=False, Tensor? active_rank_mask=None) -> Tensor"); module.def( "moe_a2a_initialize(Tensor(a!) workspace, int ep_rank, int ep_size, int max_num_tokens_per_rank, " "int? eplb_stats_num_experts=None) -> Tensor"); diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py index 2b71aaefa257..92434f26710b 100644 --- a/tensorrt_llm/_torch/alltoall_watchdog.py +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -53,6 +53,15 @@ _WORKSPACE_WATCHDOG_STATE_INIT_LOCK = threading.Lock() +def reject_rank_mask_cuda_graph_capture(rank_mask_enabled: bool) -> None: + """Reject graph capture until membership-scoped recapture is available.""" + if rank_mask_enabled and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "rank-mask mode does not support CUDA graphs until generation-scoped " + "invalidation and recapture are implemented" + ) + + def _normalize_completion_flag(value: int) -> int: return int(value) & _COMPLETION_FLAG_MASK diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 391779ac0c90..e82d81e219b2 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -496,6 +496,7 @@ def _( top_k: int, num_experts: int, eplb_local_stats: Optional[torch.Tensor] = None, + enable_rank_mask: bool = False, active_rank_mask: Optional[torch.Tensor] = None, ) -> Tuple[List[torch.Tensor], int, torch.Tensor]: recv_tensors: List[torch.Tensor] = [] @@ -527,6 +528,7 @@ def _( combine_payload_offset: int, payload_in_workspace: bool, use_low_precision: bool = False, + enable_rank_mask: bool = False, active_rank_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: return payload.new_empty((local_num_tokens, payload.shape[2])) diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index 2bf31ef680a4..24f98307c70b 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -19,7 +19,7 @@ DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, ActiveRankMaskSnapshot, AlltoAllWatchdog, AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, - EPGroupHealthLike) + EPGroupHealthLike, reject_rank_mask_cuda_graph_capture) from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger from tensorrt_llm.mapping import Mapping @@ -151,8 +151,9 @@ def __init__( Note: The terminology is mapped to `num_experts` in this class and the kernels. num_experts: (Optional) Number of experts for EPLB stats (must be <= num_slots). DO NOT provide this parameter if EPLB is not enabled. Note: The terminology is mapped to `eplb_stats_num_experts` in this class and the kernels. - ep_group_health: Optional read-only committed EP membership. When present, its mask is passed to the - CUDA kernels and defines the peers expected by the watchdog. Timeout detection never mutates it. + ep_group_health: Optional read-only committed EP membership. When present, rank-mask handling is + enabled in the CUDA kernels, and its mask defines the peers expected by the watchdog. Timeout + detection never mutates it. CUDA graphs are rejected until membership-scoped recapture lands. alltoall_watchdog_timeout_s: Optional timeout for the host-side AlltoAll watchdog. If None, the watchdog is disabled. alltoall_watchdog_poll_interval_s: Poll interval for the watchdog thread. @@ -234,6 +235,8 @@ def __init__( # Internal state self._state: _A2AState = _A2AState() self.ep_group_health = ep_group_health + # Keep the kernel specialization stable for this communicator's lifetime. + self._rank_mask_enabled = ep_group_health is not None workspace_state = self._WORKSPACE assert workspace_state is not None metainfo_index = self._METAINFO_INDEX @@ -291,14 +294,16 @@ def dispatch(self, invalid_token_expert_id: If not None, set the token_selected_experts of the invalid tokens to this expert id. This is used to notify the MoE to skip these tokens for GroupGEMM. expert_id_payload_index: The index of token_selected_experts in the input_payloads. Must be provided if invalid_token_expert_id is not None. eplb_local_stats: (Optional) [num_experts] tensor containing local statistics for EPLB - active_rank_mask: Optional uint64 CPU tensor overriding committed membership for this dispatch. When + active_rank_mask: Optional uint64 CPU tensor overriding committed membership in rank-mask mode. When omitted, the committed mask and generation are captured together. Combine reuses that mask and - fails closed if the committed generation changes first. + fails closed if the committed generation changes first. The masked kernel rejects inactive routes + before remote access; that sentinel is an internal abort artifact, not valid model output. Returns: recv_tensors: List of tensors received, each has shape [ep_size, max_tokens_per_rank, payload_num_elements_per_token] """ assert self._state.phase == "idle", "dispatch called twice without an intervening combine" + reject_rank_mask_cuda_graph_capture(self._rank_mask_enabled) assert runtime_max_tokens_per_rank <= self.max_num_tokens, "runtime_max_tokens_per_rank must not exceed max_num_tokens" if eplb_local_stats is not None: assert self.enable_eplb, "eplb_local_stats provided but enable_eplb is False" @@ -308,8 +313,13 @@ def dispatch(self, 0 ) == self.eplb_stats_num_experts, "eplb_local_stats size must match eplb_stats_num_experts" + requested_active_rank_mask = active_rank_mask + if (not self._rank_mask_enabled + and requested_active_rank_mask is not None): + raise ValueError( + "active_rank_mask requires committed EP group health") active_rank_mask_snapshot = self._watchdog_coordinator.capture_active_rank_mask( - active_rank_mask) + requested_active_rank_mask) active_rank_mask = active_rank_mask_snapshot.active_rank_mask recv_tensors, combine_payload_offset, eplb_gathered_stats = torch.ops.trtllm.moe_a2a_dispatch( token_selected_experts, @@ -322,6 +332,7 @@ def dispatch(self, self.top_k, self.num_experts, eplb_local_stats, + self._rank_mask_enabled, active_rank_mask, ) self._watchdog_coordinator.watch_collective(self._alltoall_watchdog, @@ -367,24 +378,32 @@ def combine( runtime_max_tokens_per_rank: Maximum of the number of tokens of each DP rank's local batch. payload_in_workspace: If True, 'payload' is a view into 'workspace' at 'combine_payload_offset' and no staging copy is needed. If False, the op stages 'payload' into the workspace region before combining. use_low_precision_combine: If True, quantize the combine payload to FP8 for NVLink transfer (halves NVLink bandwidth usage, output precision is preserved). - active_rank_mask: Optional uint64 CPU tensor. If supplied, it must match the mask captured by dispatch - for this collective. A committed-generation change since dispatch aborts the collective epoch. + active_rank_mask: Optional uint64 CPU tensor. In rank-mask mode, it must match the mask captured by + dispatch for this collective when supplied. A committed-generation change since dispatch aborts + the collective epoch. Returns: combined_output: [local_num_tokens, num_elements_per_token] tensor of combined results """ assert self._state.phase == "dispatched", "combine called before a successful dispatch" + reject_rank_mask_cuda_graph_capture(self._rank_mask_enabled) assert runtime_max_tokens_per_rank <= self.max_num_tokens, "runtime_max_tokens_per_rank must not exceed max_num_tokens" active_rank_mask_snapshot = self._state.active_rank_mask_snapshot assert active_rank_mask_snapshot is not None + requested_active_rank_mask = active_rank_mask + if (not self._rank_mask_enabled + and requested_active_rank_mask is not None): + raise ValueError( + "active_rank_mask requires committed EP group health") active_rank_mask = self._watchdog_coordinator.active_rank_mask_for_combine( - active_rank_mask_snapshot, active_rank_mask) + active_rank_mask_snapshot, requested_active_rank_mask) output = torch.ops.trtllm.moe_a2a_combine( payload, self._state.local_num_tokens, self.workspace, self.metainfo, runtime_max_tokens_per_rank, self.ep_rank, self.ep_size, self.top_k, self._state.combine_payload_offset, - payload_in_workspace, use_low_precision_combine, active_rank_mask) + payload_in_workspace, use_low_precision_combine, + self._rank_mask_enabled, active_rank_mask) self._watchdog_coordinator.watch_collective(self._alltoall_watchdog, "combine", active_rank_mask) diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py index 8f0b3c9486aa..55b7b2297a7d 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py @@ -38,6 +38,7 @@ AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, EPGroupHealthLike, + reject_rank_mask_cuda_graph_capture, ) from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger @@ -182,8 +183,9 @@ def __init__( use_low_precision_combine: If True, quantize the combine payload to FP8 for NVLink transfer (halves NVLink bandwidth usage, output precision is preserved). Corresponds to model_config.use_low_precision_moe_combine. - ep_group_health: Optional read-only committed EP membership. When present, its mask is passed to - the CUDA kernels and defines the peers expected by the watchdog. Timeout detection never mutates it. + ep_group_health: Optional read-only committed EP membership. When present, rank-mask handling is + enabled in the CUDA kernels, and its mask defines the peers expected by the watchdog. Timeout + detection never mutates it. CUDA graphs are rejected until membership-scoped recapture lands. alltoall_watchdog_timeout_s: Optional timeout for the host-side AlltoAll watchdog. If None, the watchdog is disabled. alltoall_watchdog_poll_interval_s: Poll interval for the watchdog thread. @@ -321,6 +323,8 @@ def __init__( self.moe_a2a_metainfo = workspace_state["metainfo"] self.max_num_tokens_per_rank = workspace_state["max_num_tokens_per_rank"] self.ep_group_health = ep_group_health + # Keep the kernel specialization stable for this communicator's lifetime. + self._rank_mask_enabled = ep_group_health is not None self._watchdog_coordinator = AlltoAllWatchdogCoordinator( workspace_state=workspace_state, workspace=self.workspace, @@ -425,9 +429,11 @@ def dispatch( token_final_scales: Router weights [local_num_tokens, top_k] all_rank_num_tokens: Token counts per rank [ep_size] use_dp_padding: Whether to use DP padding (optional) - **kwargs: Strategy-specific arguments. ``active_rank_mask`` may override the committed membership - for dispatch. Without an override, the committed mask and generation are captured together; - combine reuses that mask and fails closed if the generation changes first. + **kwargs: Strategy-specific arguments. In rank-mask mode, ``active_rank_mask`` may override the + committed membership for dispatch. Without an override, the committed mask and generation are + captured together; combine reuses that mask and fails closed if the generation changes first. + The masked kernel rejects inactive routes before remote access; that sentinel is an internal abort + artifact, not valid model output. Returns: Tuple of (hidden_states, hidden_states_sf, token_selected_slots, token_final_scales) @@ -435,6 +441,7 @@ def dispatch( """ if self._dispatch_state.get("phase") == "dispatched": raise RuntimeError("dispatch called twice without an intervening combine") + reject_rank_mask_cuda_graph_capture(self._rank_mask_enabled) # Calculate runtime_max_tokens_per_rank from all_rank_num_tokens runtime_max_tokens_per_rank = max(all_rank_num_tokens) @@ -458,8 +465,11 @@ def dispatch( assert eplb_local_stats.size(0) == self.eplb_stats_num_experts, ( "eplb_local_stats size must match eplb_stats_num_experts" ) + requested_active_rank_mask = kwargs.get("active_rank_mask") + if not self._rank_mask_enabled and requested_active_rank_mask is not None: + raise ValueError("active_rank_mask requires committed EP group health") active_rank_mask_snapshot = self._watchdog_coordinator.capture_active_rank_mask( - kwargs.get("active_rank_mask") + requested_active_rank_mask ) active_rank_mask = active_rank_mask_snapshot.active_rank_mask @@ -475,6 +485,7 @@ def dispatch( self.top_k, self.num_experts, eplb_local_stats, + self._rank_mask_enabled, active_rank_mask, ) ) @@ -550,9 +561,9 @@ def combine( final_hidden_states: Output from MoE computation Shape: [ep_size, max_tokens_per_rank, hidden_size] or [ep_size * max_tokens_per_rank, hidden_size] (will be reshaped) - **kwargs: Strategy-specific arguments. If ``active_rank_mask`` is supplied, it must match the mask - captured by dispatch for this collective. A committed-generation change since dispatch aborts - the collective epoch. + **kwargs: Strategy-specific arguments. In rank-mask mode, ``active_rank_mask`` must match the mask + captured by dispatch for this collective when supplied. A committed-generation change since + dispatch aborts the collective epoch. Returns: Combined output tensor [local_num_tokens, hidden_size] @@ -560,6 +571,7 @@ def combine( """ if self._dispatch_state.get("phase") != "dispatched": raise RuntimeError("combine called before a successful dispatch") + reject_rank_mask_cuda_graph_capture(self._rank_mask_enabled) local_num_tokens = self._dispatch_state.get("local_num_tokens") combine_payload_offset = self._dispatch_state.get("combine_payload_offset") @@ -590,9 +602,12 @@ def combine( active_rank_mask_snapshot = self._dispatch_state.get("active_rank_mask_snapshot") if not isinstance(active_rank_mask_snapshot, ActiveRankMaskSnapshot): raise RuntimeError("combine called but dispatch rank-mask snapshot is missing") + requested_active_rank_mask = kwargs.get("active_rank_mask") + if not self._rank_mask_enabled and requested_active_rank_mask is not None: + raise ValueError("active_rank_mask requires committed EP group health") active_rank_mask = self._watchdog_coordinator.active_rank_mask_for_combine( active_rank_mask_snapshot, - kwargs.get("active_rank_mask"), + requested_active_rank_mask, ) output = torch.ops.trtllm.moe_a2a_combine( final_hidden_states, @@ -606,6 +621,7 @@ def combine( int(combine_payload_offset), bool(self.payload_in_workspace), bool(self.use_low_precision_combine), + self._rank_mask_enabled, active_rank_mask, ) self._watchdog_coordinator.watch_collective( diff --git a/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py b/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py index fb1586ae1054..e0ee7e213a6d 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py +++ b/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py @@ -60,12 +60,19 @@ def get_wide_ep_ft_options( one process-local membership object shared by all MoE communication layers. The AlltoAll watchdog reads this object to determine expected peers and reports suspects through its ``on_timeout`` seam; it never mutates the - committed membership directly. + committed membership directly. Rank-mask mode rejects CUDA graphs until + generation-scoped invalidation and recapture are implemented. """ extra_attrs = getattr(model_config, "extra_attrs", {}) health = extra_attrs.get(_HEALTH_KEY) or extra_attrs.get("ep_group_health") - if health is None and _env_enabled(): + rank_mask_enabled = health is not None or _env_enabled() + if rank_mask_enabled and getattr(model_config, "use_cuda_graph", False): + raise ValueError( + "WideEP fault tolerance does not support CUDA graphs until generation-scoped " + "invalidation and recapture are implemented" + ) + if health is None and rank_mask_enabled: health = EPGroupHealth(model_config.mapping.moe_ep_size) extra_attrs[_HEALTH_KEY] = health diff --git a/tests/unittest/_torch/modules/moe/test_moe_comm.py b/tests/unittest/_torch/modules/moe/test_moe_comm.py index 0c90db578191..fe8f6cc2723b 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_comm.py +++ b/tests/unittest/_torch/modules/moe/test_moe_comm.py @@ -267,6 +267,7 @@ def _run_nvlink_rank_mask_dispatch( token_selected_experts: torch.Tensor, payload: torch.Tensor, runtime_max_tokens_per_rank: int, + enable_rank_mask: bool, active_rank_mask: Optional[torch.Tensor], ) -> Tuple[List[torch.Tensor], int, torch.Tensor, torch.Tensor]: """Run raw NVLink one-sided dispatch with an optional active rank mask.""" @@ -281,6 +282,7 @@ def _run_nvlink_rank_mask_dispatch( comm.top_k, comm.num_experts, None, # eplb_local_stats + enable_rank_mask, active_rank_mask, ) @@ -303,6 +305,7 @@ def _run_nvlink_rank_mask_combine( local_num_tokens: int, runtime_max_tokens_per_rank: int, combine_payload_offset: int, + enable_rank_mask: bool, active_rank_mask: Optional[torch.Tensor], ) -> torch.Tensor: """Run raw NVLink one-sided combine with an optional active rank mask.""" @@ -318,6 +321,7 @@ def _run_nvlink_rank_mask_combine( combine_payload_offset, False, # payload_in_workspace False, # use_low_precision + enable_rank_mask, active_rank_mask, ) @@ -327,15 +331,19 @@ def _run_nvlink_rank_mask_dispatch_combine( token_selected_experts: torch.Tensor, payload: torch.Tensor, runtime_max_tokens_per_rank: int, + enable_rank_mask: bool, active_rank_mask: Optional[torch.Tensor], -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Run raw NVLink one-sided dispatch/combine with an optional active rank mask.""" - recv_tensors, combine_payload_offset, topk_target_ranks, _ = _run_nvlink_rank_mask_dispatch( - comm, - token_selected_experts, - payload, - runtime_max_tokens_per_rank, - active_rank_mask, + recv_tensors, combine_payload_offset, topk_target_ranks, topk_send_indices = ( + _run_nvlink_rank_mask_dispatch( + comm, + token_selected_experts, + payload, + runtime_max_tokens_per_rank, + enable_rank_mask, + active_rank_mask, + ) ) combined = _run_nvlink_rank_mask_combine( comm, @@ -343,9 +351,10 @@ def _run_nvlink_rank_mask_dispatch_combine( token_selected_experts.size(0), runtime_max_tokens_per_rank, combine_payload_offset, + enable_rank_mask, active_rank_mask, ) - return combined.cpu(), topk_target_ranks + return combined.cpu(), topk_target_ranks, topk_send_indices def _expected_nvlink_rank_mask_combine_output( @@ -355,15 +364,14 @@ def _expected_nvlink_rank_mask_combine_output( topk_send_indices: torch.Tensor, local_num_tokens: int, runtime_max_tokens_per_rank: int, - dead_ranks: Set[int], ) -> torch.Tensor: - """Compute combine output from dispatched workspace while skipping dead ranks.""" + """Compute combine output from the routes recorded by dispatch.""" from tensorrt_llm.bindings import internal as _tllm_internal hidden_size = payload.shape[-1] expected = torch.zeros( (local_num_tokens, hidden_size), - dtype=payload.dtype, + dtype=torch.float32, device=payload.device, ) payload_offset_index = int(_tllm_internal.thop.MOE_A2A_PAYLOAD_DATA_OFFSET_INDEX) @@ -376,7 +384,7 @@ def _expected_nvlink_rank_mask_combine_output( for k in range(comm.top_k): target_rank = int(topk_target_ranks[token_idx, k].item()) dst_idx = int(topk_send_indices[token_idx, k].item()) - if dst_idx < 0 or target_rank in dead_ranks: + if dst_idx < 0: continue raw = comm.workspace[target_rank, payload_offset : payload_offset + bytes_per_rank] recv_payload = raw.view(payload.dtype).view( @@ -384,8 +392,8 @@ def _expected_nvlink_rank_mask_combine_output( runtime_max_tokens_per_rank, hidden_size, ) - expected[token_idx] += recv_payload[comm.ep_rank, dst_idx] - return expected.cpu() + expected[token_idx] += recv_payload[comm.ep_rank, dst_idx].float() + return expected.to(payload.dtype).cpu() # ============================================================================ @@ -1208,20 +1216,42 @@ def _worker_rank_mask_all_active_matches_no_mask(config: CommTestConfig) -> dict device="cuda", ) payload = _make_rank_mask_payload(local_num_tokens, config.hidden_size, rank) + all_active_mask = _ep_mask_words(config.ep_size, dead_ranks=set()) - out_no_mask, topk_no_mask = _run_nvlink_rank_mask_dispatch_combine( + with pytest.raises(RuntimeError, match="requires enable_rank_mask"): + _run_nvlink_rank_mask_dispatch( + comm, + token_selected_experts, + payload, + local_num_tokens, + enable_rank_mask=False, + active_rank_mask=all_active_mask, + ) + with pytest.raises(RuntimeError, match="active_rank_mask must be defined"): + _run_nvlink_rank_mask_dispatch( + comm, + token_selected_experts, + payload, + local_num_tokens, + enable_rank_mask=True, + active_rank_mask=None, + ) + + out_no_mask, topk_no_mask, _ = _run_nvlink_rank_mask_dispatch_combine( comm, token_selected_experts, payload, local_num_tokens, + enable_rank_mask=False, active_rank_mask=None, ) - out_all_active, topk_all_active = _run_nvlink_rank_mask_dispatch_combine( + out_all_active, topk_all_active, _ = _run_nvlink_rank_mask_dispatch_combine( comm, token_selected_experts, payload, local_num_tokens, - active_rank_mask=_ep_mask_words(config.ep_size, dead_ranks=set()), + enable_rank_mask=True, + active_rank_mask=all_active_mask, ) return { @@ -1256,7 +1286,7 @@ def _worker_rank_mask_one_rank_masked( config: CommTestConfig, dead_rank: int, ) -> dict: - """Run dispatch/combine with one EP rank omitted from active_rank_mask.""" + """Verify masked-route rejection, then run with survivor-only routing.""" rank = tllm.mpi_rank() torch.cuda.set_device(rank) @@ -1278,89 +1308,59 @@ def _worker_rank_mask_one_rank_masked( local_num_tokens = config.all_num_tokens[rank] torch.manual_seed(0xA2A + rank) - token_selected_experts = torch.randint( - 0, - config.num_experts, + mask = _ep_mask_words(config.ep_size, dead_ranks={dead_rank}) + dead_expert_id = next( + expert_id + for expert_id in range(config.num_experts) + if _expert_id_to_rank(expert_id, config.num_experts, config.ep_size) == dead_rank + ) + masked_routes = torch.full( (local_num_tokens, config.top_k), + dead_expert_id, dtype=torch.int32, device="cuda", ) - payload = _make_rank_mask_payload(local_num_tokens, config.hidden_size, rank) - mask = _ep_mask_words(config.ep_size, dead_ranks={dead_rank}) - - combined, topk_target_ranks = _run_nvlink_rank_mask_dispatch_combine( + _, _, masked_target_ranks, masked_send_indices = _run_nvlink_rank_mask_dispatch( comm, - token_selected_experts, - payload, + masked_routes, + _make_rank_mask_payload(local_num_tokens, config.hidden_size, rank), local_num_tokens, + enable_rank_mask=True, active_rank_mask=mask, ) - expected_target_ranks = _expected_target_ranks( - token_selected_experts, - config.num_experts, - config.ep_size, - ) - MPI.COMM_WORLD.barrier() - return { - "rank": rank, - "status": "alive", - "combined": combined, - "topk_target_ranks": topk_target_ranks, - "expected_target_ranks": expected_target_ranks, - } - except Exception: - traceback.print_exc() - raise - finally: - if comm is not None and hasattr(comm, "destroy"): - comm.destroy() - - -def _worker_rank_mask_inactive_before_combine( - config: CommTestConfig, - dead_rank: int, -) -> dict: - """Dispatch with all ranks active, then omit one rank from combine's active mask.""" - rank = tllm.mpi_rank() - torch.cuda.set_device(rank) - - comm = None - try: - mapping = Mapping( - rank=rank, - tp_size=config.ep_size, - moe_ep_size=config.ep_size, - world_size=config.ep_size, + live_expert_ids = torch.tensor( + [ + expert_id + for expert_id in range(config.num_experts) + if _expert_id_to_rank(expert_id, config.num_experts, config.ep_size) != dead_rank + ], + dtype=torch.int32, + device="cuda", ) - comm = create_comm_object(config.comm_type, mapping, config) - - local_num_tokens = config.all_num_tokens[rank] - torch.manual_seed(0xA2A + rank) - token_selected_experts = torch.randint( + live_expert_indices = torch.randint( 0, - config.num_experts, + live_expert_ids.numel(), (local_num_tokens, config.top_k), - dtype=torch.int32, + dtype=torch.int64, device="cuda", ) + token_selected_experts = live_expert_ids[live_expert_indices] payload = _make_rank_mask_payload(local_num_tokens, config.hidden_size, rank) - recv_tensors, combine_payload_offset, topk_target_ranks, topk_send_indices = ( - _run_nvlink_rank_mask_dispatch( - comm, - token_selected_experts, - payload, - local_num_tokens, - active_rank_mask=_ep_mask_words(config.ep_size, dead_ranks=set()), - ) + combined, topk_target_ranks, topk_send_indices = _run_nvlink_rank_mask_dispatch_combine( + comm, + token_selected_experts, + payload, + local_num_tokens, + enable_rank_mask=True, + active_rank_mask=mask, + ) + expected_target_ranks = _expected_target_ranks( + token_selected_experts, + config.num_experts, + config.ep_size, ) - - if rank == dead_rank: - MPI.COMM_WORLD.barrier() - return {"rank": rank, "status": "dead"} - - dead_ranks = {dead_rank} expected = _expected_nvlink_rank_mask_combine_output( comm, payload, @@ -1368,23 +1368,18 @@ def _worker_rank_mask_inactive_before_combine( topk_send_indices, local_num_tokens, local_num_tokens, - dead_ranks, ) - combined = _run_nvlink_rank_mask_combine( - comm, - recv_tensors[0], - local_num_tokens, - local_num_tokens, - combine_payload_offset, - active_rank_mask=_ep_mask_words(config.ep_size, dead_ranks=dead_ranks), - ).cpu() MPI.COMM_WORLD.barrier() return { "rank": rank, "status": "alive", + "masked_target_ranks": masked_target_ranks[:local_num_tokens], + "masked_send_indices": masked_send_indices[:local_num_tokens], "combined": combined, "expected": expected, + "topk_target_ranks": topk_target_ranks, + "expected_target_ranks": expected_target_ranks, } except Exception: traceback.print_exc() @@ -2369,11 +2364,21 @@ def _run_rank_mask_one_rank_masked_test( continue assert result["status"] == "alive" + assert torch.all(result["masked_target_ranks"] == -1), ( + f"rank {rank}: dispatch retained a route to masked rank {dead_rank}" + ) + assert torch.all(result["masked_send_indices"] == -1), ( + f"rank {rank}: dispatch allocated a send slot for masked rank {dead_rank}" + ) combined = result["combined"] + expected = result["expected"] topk_target_ranks = result["topk_target_ranks"] expected_target_ranks = result["expected_target_ranks"] assert combined.shape == (local_num_tokens, config.hidden_size) + assert torch.equal(combined, expected), ( + f"rank {rank}: combine output does not match the pre-routed live payloads" + ) live_topk = topk_target_ranks[:local_num_tokens] live_expected = expected_target_ranks[:local_num_tokens] @@ -2382,12 +2387,8 @@ def _run_rank_mask_one_rank_masked_test( for k in range(top_k): expected = int(live_expected[token_idx, k].item()) got = int(live_topk[token_idx, k].item()) - if expected == dead_rank: - assert got == -1, ( - f"rank {rank} token {token_idx} k={k}: token routed to dead " - f"rank {dead_rank} should have been dropped (got={got})" - ) - elif expected in seen_ranks: + assert expected != dead_rank + if expected in seen_ranks: assert got == -1 else: assert got == expected, ( @@ -2399,45 +2400,6 @@ def _run_rank_mask_one_rank_masked_test( assert saw_dead, f"dead rank {dead_rank} did not appear in results" -def _run_rank_mask_inactive_before_combine_test( - mpi_pool_executor, - dead_rank: int, - local_num_tokens: int, - top_k: int, -) -> None: - ep_size = mpi_pool_executor.num_workers - config = _make_rank_mask_config(ep_size, local_num_tokens, top_k) - _skip_if_rank_mask_config_unsupported(config) - assert 0 <= dead_rank < ep_size - - worker_args = [(config, dead_rank)] * config.ep_size - results = list( - mpi_pool_executor.map( - _worker_rank_mask_inactive_before_combine, - *zip(*worker_args), - ) - ) - - saw_dead = False - for result in results: - rank = result["rank"] - if result["status"] == "dead": - assert rank == dead_rank - saw_dead = True - continue - - assert result["status"] == "alive" - combined = result["combined"] - expected = result["expected"] - assert combined is not None - assert expected is not None - assert torch.equal(combined, expected), ( - f"rank {rank}: combine output included a rank masked inactive before combine" - ) - - assert saw_dead, f"dead rank {dead_rank} did not appear in results" - - # ============================================================================ # Test Class # ============================================================================ @@ -2526,33 +2488,10 @@ def test_moe_comm_rank_mask_one_rank_masked_completes( local_num_tokens: int, top_k: int, ) -> None: - """Verify masked-dead rank is skipped by raw NVLinkOneSided moe_a2a ops.""" + """Verify peer synchronization skips a masked rank after routing excludes it.""" _run_rank_mask_one_rank_masked_test( mpi_pool_executor, dead_rank, local_num_tokens, top_k, ) - - @pytest.mark.threadleak(enabled=False) - @pytest.mark.parametrize( - "mpi_pool_executor,dead_rank,local_num_tokens,top_k", - [ - (4, 2, 16, 2), - ], - indirect=["mpi_pool_executor"], - ) - def test_moe_comm_rank_mask_inactive_before_combine_skips_stale_dispatch_slots( - self, - mpi_pool_executor, - dead_rank: int, - local_num_tokens: int, - top_k: int, - ) -> None: - """Verify combine skips slots from a rank masked inactive after dispatch.""" - _run_rank_mask_inactive_before_combine_test( - mpi_pool_executor, - dead_rank, - local_num_tokens, - top_k, - ) diff --git a/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py index cdff07cbcea4..c5f4b260160f 100644 --- a/tests/unittest/_torch/modules/test_alltoall_watchdog.py +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -30,6 +30,7 @@ AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, CompletionFlagReadTimeout, + reject_rank_mask_cuda_graph_capture, ) from tensorrt_llm._torch.modules.fused_moe.ep_group_health import EPGroupHealth from tensorrt_llm._torch.modules.fused_moe.wide_ep_ft import get_wide_ep_ft_options @@ -197,6 +198,31 @@ def test_wide_ep_ft_options_create_shared_health_when_enabled( assert poll_again_s == poll_interval_s +def test_wide_ep_ft_options_reject_cuda_graphs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("TLLM_FAULT_TOLERANCE_MODE", "1") + model_config = SimpleNamespace( + extra_attrs={}, + mapping=SimpleNamespace(moe_ep_size=4), + use_cuda_graph=True, + ) + + with pytest.raises(ValueError, match="does not support CUDA graphs"): + get_wide_ep_ft_options(model_config) + assert model_config.extra_attrs == {} + + +def test_rank_mask_mode_rejects_direct_cuda_graph_capture( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + + reject_rank_mask_cuda_graph_capture(False) + with pytest.raises(RuntimeError, match="does not support CUDA graphs"): + reject_rank_mask_cuda_graph_capture(True) + + def test_wide_ep_ft_options_ignore_legacy_enable_flag(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("TLLM_FAULT_TOLERANCE_MODE", raising=False) monkeypatch.setenv("TRTLLM_ENABLE_WIDE_EP_FT", "1")