From 3ac1b2f0e9dc66643215733de090f7bc14f3491b Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 23 Apr 2026 20:16:44 -0700 Subject: [PATCH 1/4] [None][feat] WideEP FT: add active_rank_mask to NVLink AlltoAll kernels Eliminates the infinite-spin AlltoAll hang that turns a single GPU failure in a Wide-EP group into a 5-minute HangDetector fire + full restart. The dispatch and combine kernels now take a uint64[2] bitmask of currently-alive EP ranks; dead ranks are skipped on every completion-flag write/wait, peer recv_counter store, EPLB stats write, and per-token routing decision (dead-targeted slots collapse to the same -1 sentinel combine already uses for duplicates). The mask is optional on both torch ops; omitting it (or passing all-ones) produces bit-identical output to the pre-change kernel. kMaxRanks is bumped 64 -> 128 to cover NVL72 with headroom; kRankMaskWords = 2 names the kernel ABI explicitly. Tests cover (a) all-ones mask matches no-mask bit-for-bit, and (b) one rank masked dead -> surviving ranks complete dispatch+combine without hang, dead-targeted topk slots dropped, in tests/unittest/_torch/multi_gpu/test_moe_a2a_rank_mask.py. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../moeAlltoAllKernels.cu | 59 ++- .../communicationKernels/moeAlltoAllKernels.h | 30 +- cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp | 53 ++- .../_torch/custom_ops/cpp_custom_ops.py | 2 + .../communication/nvlink_one_sided.py | 2 +- .../multi_gpu/test_moe_a2a_rank_mask.py | 407 ++++++++++++++++++ 6 files changed, 543 insertions(+), 10 deletions(-) create mode 100644 tests/unittest/_torch/multi_gpu/test_moe_a2a_rank_mask.py diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu index 91cb5725fede..66e5fb30de99 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu @@ -210,6 +210,14 @@ __device__ __forceinline__ int compute_target_rank_id(int expert_id, int base, i return remainder + (expert_id - split) / base; } +// Test bit `rank` in a kRankMaskWords-wide little-endian uint64 bitmask. +// Word 0 covers ranks 0..63, word 1 covers ranks 64..127, etc. +// `rank >> 6` and `rank & 63` divide / modulo by 64. +__device__ __forceinline__ bool is_rank_active(uint64_t const* mask, int rank) +{ + return (mask[rank >> 6] >> (rank & 63)) & 1ULL; +} + // ============================================================================ // Helper Functions for Vectorized Memory Operations // ============================================================================ @@ -432,7 +440,12 @@ __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); - if (already_copied & (1ULL << target_rank)) + // 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. + bool const target_dead = !is_rank_active(ptrs.active_rank_mask, target_rank); + if ((already_copied & (1ULL << target_rank)) || target_dead) { if (thread_idx == 0) { @@ -511,10 +524,13 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ if (is_last_token) { -// Store send_counters to recv_counters +// Store send_counters to recv_counters. +// Skip masked target ranks: their symmetric memory may be inaccessible. #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; int send_count = ptrs.send_counters[target_rank]; ptrs.recv_counters[target_rank][rank_id] = send_count; } @@ -522,9 +538,12 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ if constexpr (ENABLE_EPLB) { // Write local stats into peer buffers before the release fence below. + // Skip masked target ranks for the same reason as above. #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; int* target_stats = ptrs.eplb_gathered_stats[target_rank]; for (int expert_id = lane_id; expert_id < eplb_stats_num_experts; expert_id += warpSize) { @@ -543,9 +562,13 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ #else asm volatile("fence.acq_rel.sys;"); #endif + // Signal completion to all active peers; skip dead ranks (their symmetric memory + // is unreachable). #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; 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)); @@ -555,9 +578,13 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ #endif } + // Wait for all active peers to signal; skip dead ranks (otherwise we would + // spin forever — this is the bug the rank-mask is here to prevent). #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; bool flag_set = false; auto s = clock64(); do @@ -605,6 +632,10 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) TLLM_CHECK(params.ep_size > 0 && params.ep_size <= kMaxRanks); 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); // Prepare kernel pointers struct DispatchKernelPointers kernel_ptrs = {}; @@ -642,6 +673,12 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) kernel_ptrs.topk_send_indices = params.topk_send_indices; kernel_ptrs.eplb_local_stats = params.eplb_local_stats; + // Copy active-rank bitmask into the kernel pointers struct + for (int w = 0; w < kRankMaskWords; ++w) + { + kernel_ptrs.active_rank_mask[w] = params.active_rank_mask[w]; + } + int const kBlockSize = tensorrt_llm::common::getEnvMoeA2ADispatchBlockSize(); // One block per token: grid_size == local_num_tokens. If 0, launch a single block to @@ -1153,9 +1190,13 @@ __global__ void moeA2ACombineKernel( if (blockIdx.x == 0) { + // Signal readiness to all active peers; skip dead ranks (their symmetric memory + // is unreachable). #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; 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 @@ -1165,9 +1206,13 @@ __global__ void moeA2ACombineKernel( } } + // Wait for all active peers to signal; skip dead ranks (otherwise we would spin + // forever — this is the bug the rank-mask is here to prevent). #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; bool flag_set = false; auto s = clock64(); do @@ -1273,6 +1318,10 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) TLLM_CHECK(params.ep_size > 0 && params.ep_size <= kMaxRanks); 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); // Configure kernel launch (one block per token). int const kBlockSize = tensorrt_llm::common::getEnvMoeA2ACombineBlockSize(); @@ -1306,6 +1355,12 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) kernel_ptrs.topk_target_ranks = params.topk_target_ranks; kernel_ptrs.topk_send_indices = params.topk_send_indices; + // Copy active-rank bitmask into the kernel pointers struct + for (int w = 0; w < kRankMaskWords; ++w) + { + kernel_ptrs.active_rank_mask[w] = params.active_rank_mask[w]; + } + // stride_per_token: byte distance between tokens in the recv buffer. // FP8 external payload: EPT × 1 (compact FP8 layout) // FP8 in-place / non-FP8: EPT × sizeof(PayloadT) (payload-dtype stride) diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h index 317ff4d2240c..0b872db4356f 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h @@ -26,9 +26,12 @@ namespace kernels::moe_comm { // Configuration constants -static constexpr int kMaxTopK = 22; // Maximum top-k experts per token -static constexpr int kMaxPayloads = 4; // Maximum number of different payload types -static constexpr int kMaxRanks = 64; // Maximum supported EP size +static constexpr int kMaxTopK = 22; // Maximum top-k experts per token +static constexpr int kMaxPayloads = 4; // Maximum number of different payload types +static constexpr int kMaxRanks = 128; // Maximum supported EP size (covers NVL72 with headroom) +static constexpr int kRankMaskWords = 2; // uint64 words to hold the active-rank bitmask + // (kRankMaskWords * 64 must be >= kMaxRanks) +static_assert(kRankMaskWords * 64 >= kMaxRanks, "active_rank_mask too small for kMaxRanks"); // Describes a single payload type to be communicated struct PayloadDescriptor @@ -65,6 +68,12 @@ struct DispatchKernelPointers // 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. + uint64_t active_rank_mask[kRankMaskWords]; }; // Combine kernel pointers - non-const output in src_data_ptrs[0], const recv buffers @@ -82,6 +91,11 @@ struct CombineKernelPointers // 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 + + // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Combine skips flag + // writes/waits to/from masked peers; per-token accumulation uses topk_send_indices[k] < 0 + // (set by dispatch) to skip dead-targeted slots, so no explicit mask check is needed there. + uint64_t active_rank_mask[kRankMaskWords]; }; // Dispatch phase parameters @@ -125,6 +139,11 @@ 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. + uint64_t active_rank_mask[kRankMaskWords]; + // CUDA stream cudaStream_t stream; }; @@ -170,6 +189,11 @@ 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. + uint64_t active_rank_mask[kRankMaskWords]; + // CUDA stream cudaStream_t stream; }; diff --git a/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp b/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp index 7a767976dabd..8f467519fab7 100644 --- a/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp +++ b/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp @@ -42,6 +42,40 @@ 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 void resolveActiveRankMask(torch::optional const& maskTensor, int64_t epRank, + uint64_t (&out)[tensorrt_llm::kernels::moe_comm::kRankMaskWords]) +{ + using tensorrt_llm::kernels::moe_comm::kRankMaskWords; + if (!maskTensor.has_value() || !maskTensor.value().defined()) + { + for (int w = 0; w < kRankMaskWords; ++w) + { + out[w] = ~uint64_t{0}; + } + return; + } + 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"); + TORCH_CHECK(t.dim() == 1, "active_rank_mask must be a 1D tensor"); + TORCH_CHECK(t.numel() == kRankMaskWords, "active_rank_mask must have exactly ", kRankMaskWords, " uint64 elements"); + TORCH_CHECK(t.is_contiguous(), "active_rank_mask must be contiguous"); + auto const* src = static_cast(t.const_data_ptr()); + for (int w = 0; w < kRankMaskWords; ++w) + { + out[w] = src[w]; + } + // Local rank's bit must be set; otherwise the kernel would be running on a "dead" rank. + TORCH_CHECK((out[epRank >> 6] >> (epRank & 63)) & 1ULL, "active_rank_mask must mark the local ep_rank (", epRank, + ") as active"); +} + // Calculate auxiliary data offsets MoeA2ADataOffsets calculateOffsets(int epSize, int maxNumTokens, int eplbStatsNumExperts) { @@ -181,7 +215,8 @@ torch::Tensor moeA2AInitializeOp(torch::Tensor const& workspace, int64_t epRank, 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) + int64_t epSize, int64_t topK, int64_t numExperts, torch::optional eplbLocalStats, + torch::optional activeRankMask) { using tensorrt_llm::kernels::moe_comm::PayloadDescriptor; using tensorrt_llm::kernels::moe_comm::MoeA2ADispatchParams; @@ -360,6 +395,10 @@ 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.stream = at::cuda::getCurrentCUDAStream(); // Prepare for dispatch (zero counters/indices and increment flag_val) @@ -413,7 +452,8 @@ 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 = false, + torch::optional activeRankMask = torch::nullopt) { using tensorrt_llm::kernels::moe_comm::MoeA2ACombineParams; using tensorrt_llm::kernels::moe_comm::moe_a2a_combine_launch; @@ -520,6 +560,9 @@ 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.stream = at::cuda::getCurrentCUDAStream(); moe_a2a_prepare_combine_launch(params); @@ -613,12 +656,14 @@ TORCH_LIBRARY_FRAGMENT(trtllm, module) "moe_a2a_dispatch(Tensor token_selected_experts, Tensor[] input_payloads, " "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(a!)[], int, Tensor(a!))"); + "Tensor? eplb_local_stats=None, " + "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"); + "bool payload_in_workspace, bool use_low_precision=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/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index a8e113278e75..35d6f69fd626 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -480,6 +480,7 @@ def _( top_k: int, num_experts: int, eplb_local_stats: Optional[torch.Tensor] = None, + active_rank_mask: Optional[torch.Tensor] = None, ) -> Tuple[List[torch.Tensor], int, torch.Tensor]: recv_tensors: List[torch.Tensor] = [] for payload in input_payloads: @@ -510,6 +511,7 @@ def _( combine_payload_offset: int, payload_in_workspace: bool, use_low_precision: 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/modules/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py index 3b634dd7072c..f9068f71d717 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 @@ -51,7 +51,7 @@ class NVLinkOneSided(Communication): """ # Constants from C++ (must match moeAlltoAllKernels.h) - MAX_RANKS = 64 + MAX_RANKS = 128 MAX_TOP_K = 8 MAX_PAYLOADS = 8 diff --git a/tests/unittest/_torch/multi_gpu/test_moe_a2a_rank_mask.py b/tests/unittest/_torch/multi_gpu/test_moe_a2a_rank_mask.py new file mode 100644 index 000000000000..9f843671dff9 --- /dev/null +++ b/tests/unittest/_torch/multi_gpu/test_moe_a2a_rank_mask.py @@ -0,0 +1,407 @@ +# 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. +"""Unit tests for the active_rank_mask parameter on the MoE AlltoAll kernels (PR 1a.2). + +Two scenarios are exercised: + +1. **All-active mask matches no-mask** — passing a mask with every bit set must produce + bit-identical output to omitting the mask. Regression guard for the kernel mod's + default behavior. + +2. **One rank masked completes without hanging** — with bit K cleared, the surviving + N-1 ranks must complete dispatch + combine without spinning on the dead rank's + completion flag, and any token routed to the dead rank's experts must be dropped + (topk_target_ranks[k] == -1) rather than silently corrupting peer memory. + +The dispatch/combine kernels require `MnnvlMemory` (multi-node NVLink, GB200), so these +tests skip on hardware that does not support MNNVL and on nodes with fewer GPUs than +`ep_size`. +""" + +import pickle +import sys +import traceback + +import cloudpickle +import pytest +import torch +from mpi4py import MPI + +import tensorrt_llm as tllm +from tensorrt_llm._mnnvl_utils import MnnvlMemory +from tensorrt_llm._torch.distributed import MoeAlltoAll +from tensorrt_llm.mapping import Mapping + +cloudpickle.register_pickle_by_value(sys.modules[__name__]) +MPI.pickle.__init__( + cloudpickle.dumps, + cloudpickle.loads, + pickle.HIGHEST_PROTOCOL, +) + + +# Must match cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h +# kRankMaskWords. The full mask is little-endian: word 0 covers ranks 0..63. +EP_MASK_NUM_WORDS = 2 + + +@pytest.fixture(autouse=True) +def setup_test(): + torch.manual_seed(0xA2A) + tllm.logger.set_level("error") + + +def _ep_mask_words(ep_size: int, dead_ranks: set[int]) -> torch.Tensor: + """Build the uint64[EP_MASK_NUM_WORDS] CPU tensor expected by the C++ op.""" + mask_int = ((1 << ep_size) - 1) & ~sum(1 << r for r in dead_ranks) + word_mask = (1 << 64) - 1 + words = [(mask_int >> (i * 64)) & word_mask for i in range(EP_MASK_NUM_WORDS)] + return torch.tensor(words, dtype=torch.uint64, device="cpu") + + +def _generate_token_selected_experts( + local_num_tokens: int, num_experts: int, top_k: int +) -> torch.Tensor: + return torch.randint( + 0, num_experts, (local_num_tokens, top_k), dtype=torch.int32, device="cuda" + ) + + +def _make_payload(local_num_tokens: int, hidden_size: int, rank: int) -> torch.Tensor: + """Deterministic per-rank payload so we can assert exact bit-for-bit equality.""" + base = torch.arange(local_num_tokens * hidden_size, dtype=torch.bfloat16, device="cuda").view( + local_num_tokens, hidden_size + ) + # Encode rank into the payload so cross-rank mismatches are immediately visible. + return base + (rank * 1000.0) + + +def _read_topk_target_ranks(moe_a2a: MoeAlltoAll, max_num_tokens: int, top_k: int) -> torch.Tensor: + """Read the kernel-written topk_target_ranks[max_num_tokens, top_k] from workspace.""" + offset = moe_a2a.metainfo[MoeAlltoAll._METAINFO_INDEX["TOPK_TARGET_RANKS_OFFSET_INDEX"]].item() + raw = moe_a2a.workspace[ + moe_a2a.ep_rank, + offset : offset + max_num_tokens * top_k * 4, + ] + return raw.view(torch.int32).view(max_num_tokens, top_k).cpu() + + +def _run_dispatch_combine( + moe_a2a: MoeAlltoAll, + token_selected_experts: torch.Tensor, + payload: torch.Tensor, + runtime_max_tokens_per_rank: int, + active_rank_mask: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Drive dispatch + combine via the raw C++ ops (so we can pass active_rank_mask). + + Returns ``(combined_output, topk_target_ranks_snapshot)``. ``payload`` doubles as + both the dispatched payload and the staged combine payload, which is the simplest + end-to-end exercise of the masking path on both kernels. + """ + recv_tensors, combine_payload_offset, _ = torch.ops.trtllm.moe_a2a_dispatch( + token_selected_experts, + [payload], + moe_a2a.workspace, + moe_a2a.metainfo, + runtime_max_tokens_per_rank, + moe_a2a.ep_rank, + moe_a2a.ep_size, + moe_a2a.top_k, + moe_a2a.num_experts, + None, # eplb_local_stats + active_rank_mask, + ) + + # Snapshot the kernel-written routing table BEFORE combine (combine reads it but + # may also reset workspace state on subsequent rounds). + topk_target_ranks = _read_topk_target_ranks(moe_a2a, runtime_max_tokens_per_rank, moe_a2a.top_k) + + combine_payload = recv_tensors[0] # [ep_size, max_tokens, hidden_size] + combined = torch.ops.trtllm.moe_a2a_combine( + combine_payload, + token_selected_experts.size(0), + moe_a2a.workspace, + moe_a2a.metainfo, + runtime_max_tokens_per_rank, + moe_a2a.ep_rank, + moe_a2a.ep_size, + moe_a2a.top_k, + int(combine_payload_offset), + False, # payload_in_workspace + False, # use_low_precision + active_rank_mask, + ) + return combined.cpu(), topk_target_ranks + + +# --------------------------------------------------------------------------- +# Worker: regression — all-active mask must match no-mask (bit-identical). +# --------------------------------------------------------------------------- + + +def _worker_all_active_matches_no_mask( + ep_size: int, + local_num_tokens: int, + top_k: int, + workspace_size_per_rank: int, + num_experts: int, + hidden_size: int, +): + rank = tllm.mpi_rank() + torch.cuda.set_device(rank) + try: + mapping = Mapping(rank=rank, tp_size=ep_size, moe_ep_size=ep_size, world_size=ep_size) + moe_a2a = MoeAlltoAll( + mapping=mapping, + max_num_tokens=local_num_tokens, + top_k=top_k, + num_slots=num_experts, + workspace_size_per_rank=workspace_size_per_rank, + ) + + # Same RNG seed across both runs => identical inputs. + torch.manual_seed(0xA2A + rank) + token_selected_experts = _generate_token_selected_experts( + local_num_tokens, num_experts, top_k + ) + payload = _make_payload(local_num_tokens, hidden_size, rank) + + out_no_mask, topk_no_mask = _run_dispatch_combine( + moe_a2a, token_selected_experts, payload, local_num_tokens, active_rank_mask=None + ) + out_all_active, topk_all_active = _run_dispatch_combine( + moe_a2a, + token_selected_experts, + payload, + local_num_tokens, + active_rank_mask=_ep_mask_words(ep_size, dead_ranks=set()), + ) + + return ( + torch.equal(out_no_mask, out_all_active), + torch.equal(topk_no_mask, topk_all_active), + ) + except Exception: + traceback.print_exc() + raise + + +# --------------------------------------------------------------------------- +# Worker: one rank masked — surviving ranks complete; dead-targeted slots dropped. +# --------------------------------------------------------------------------- + + +def _worker_one_rank_masked( + ep_size: int, + dead_rank: int, + local_num_tokens: int, + top_k: int, + workspace_size_per_rank: int, + num_experts: int, + hidden_size: int, +): + rank = tllm.mpi_rank() + torch.cuda.set_device(rank) + try: + mapping = Mapping(rank=rank, tp_size=ep_size, moe_ep_size=ep_size, world_size=ep_size) + # Every rank participates in workspace init (it has MPI barriers internally). + moe_a2a = MoeAlltoAll( + mapping=mapping, + max_num_tokens=local_num_tokens, + top_k=top_k, + num_slots=num_experts, + workspace_size_per_rank=workspace_size_per_rank, + ) + + if rank == dead_rank: + # Simulate a dead rank: do not call dispatch/combine. Wait at a final + # barrier so the surviving ranks have someone to synchronize with at + # the end of the test. (The kernel itself never observes us because + # the surviving ranks pass a mask with our bit cleared.) + MPI.COMM_WORLD.barrier() + return ("dead", None, None, None) + + torch.manual_seed(0xA2A + rank) + token_selected_experts = _generate_token_selected_experts( + local_num_tokens, num_experts, top_k + ) + payload = _make_payload(local_num_tokens, hidden_size, rank) + + # Build mask with dead_rank's bit cleared. + mask = _ep_mask_words(ep_size, dead_ranks={dead_rank}) + + # Compute the per-token target ranks the way the kernel does so we can + # cross-check the workspace afterwards. + num_experts_per_rank = num_experts // ep_size + expected_target_ranks = (token_selected_experts // num_experts_per_rank).cpu() + + combined, topk_target_ranks = _run_dispatch_combine( + moe_a2a, token_selected_experts, payload, local_num_tokens, active_rank_mask=mask + ) + + MPI.COMM_WORLD.barrier() + return ( + "alive", + combined, + topk_target_ranks, + expected_target_ranks, + ) + except Exception: + traceback.print_exc() + raise + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize( + "mpi_pool_executor,local_num_tokens,top_k", + [ + (4, 16, 2), + (4, 32, 4), + ], + indirect=["mpi_pool_executor"], +) +def test_all_active_mask_matches_no_mask(mpi_pool_executor, local_num_tokens, top_k): + """An all-ones active_rank_mask must produce identical output to omitting it.""" + try: + MnnvlMemory.initialize() + assert MnnvlMemory.supports_mnnvl() + except Exception: + pytest.skip("MNNVL not supported on this system") + + ep_size = mpi_pool_executor.num_workers + if ep_size > torch.cuda.device_count(): + pytest.skip( + f"Need at least {ep_size} GPUs but only {torch.cuda.device_count()} are available" + ) + + hidden_size = 1024 + num_experts = 32 + workspace_size_per_rank = 256 * 1024 * 1024 + + args = (ep_size, local_num_tokens, top_k, workspace_size_per_rank, num_experts, hidden_size) + results = list( + mpi_pool_executor.map( + _worker_all_active_matches_no_mask, + *zip(*[args] * ep_size), + ) + ) + + for rank, (output_eq, topk_eq) in enumerate(results): + assert output_eq, f"rank {rank}: combine output differs between no-mask and all-active mask" + assert topk_eq, f"rank {rank}: topk_target_ranks differ between no-mask and all-active mask" + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize( + "mpi_pool_executor,dead_rank,local_num_tokens,top_k", + [ + (4, 2, 16, 2), + (4, 0, 16, 4), # mask the lowest-numbered rank + (4, 3, 32, 4), # mask the highest-numbered rank + ], + indirect=["mpi_pool_executor"], +) +def test_one_rank_masked_completes(mpi_pool_executor, dead_rank, local_num_tokens, top_k): + """With one rank masked dead, surviving ranks complete dispatch+combine. + + Verifies: + * No hang (the test reaches the assertions). + * On every surviving rank, any topk slot whose expert mapped to the dead + rank is dropped (topk_target_ranks == -1). + * Slots whose expert mapped to a surviving rank are unchanged from what + the contiguous-partition routing rule predicts. + """ + try: + MnnvlMemory.initialize() + assert MnnvlMemory.supports_mnnvl() + except Exception: + pytest.skip("MNNVL not supported on this system") + + ep_size = mpi_pool_executor.num_workers + if ep_size > torch.cuda.device_count(): + pytest.skip( + f"Need at least {ep_size} GPUs but only {torch.cuda.device_count()} are available" + ) + assert 0 <= dead_rank < ep_size + + hidden_size = 1024 + num_experts = 32 + workspace_size_per_rank = 256 * 1024 * 1024 + + args = ( + ep_size, + dead_rank, + local_num_tokens, + top_k, + workspace_size_per_rank, + num_experts, + hidden_size, + ) + results = list( + mpi_pool_executor.map( + _worker_one_rank_masked, + *zip(*[args] * ep_size), + ) + ) + + saw_dead = False + for rank, (status, combined, topk_target, expected_target) in enumerate(results): + if status == "dead": + assert rank == dead_rank + saw_dead = True + continue + assert status == "alive" + # Combine produced an output of the expected shape on the surviving rank. + assert combined is not None + assert combined.shape == (local_num_tokens, hidden_size) + + # Per-token routing assertions, using only the live tokens (the workspace + # topk arrays are sized [max_num_tokens, top_k] and may have stale rows + # beyond local_num_tokens; we only care about the live region). + live_topk = topk_target[:local_num_tokens] + live_expected = expected_target[:local_num_tokens] + + # Slot-level rules: + # - If expected_target_rank == dead_rank, the kernel must have set -1. + # - Otherwise, kernel must record the same target rank (or -1 if the + # same target was already covered earlier in the same token's top-k + # list; the kernel uses -1 as a "duplicate" sentinel). + for token_idx in range(local_num_tokens): + seen_ranks: set[int] = set() + for k in range(top_k): + exp = int(live_expected[token_idx, k].item()) + got = int(live_topk[token_idx, k].item()) + if exp == 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 exp in seen_ranks: + # Duplicate target within this token — kernel sets -1. + assert got == -1 + else: + assert got == exp, ( + f"rank {rank} token {token_idx} k={k}: target rank mismatch " + f"(expected={exp}, got={got})" + ) + seen_ranks.add(exp) + + assert saw_dead, f"dead rank {dead_rank} did not appear in results" From 4cde71596c1b3781ad0442805c8238d971e26755 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 16 Jun 2026 14:54:13 -0700 Subject: [PATCH 2/4] Address active rank mask review comments Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../moeAlltoAllKernels.cu | 11 +- .../communicationKernels/moeAlltoAllKernels.h | 4 +- cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp | 10 ++ .../multi_gpu/test_moe_a2a_rank_mask.py | 164 ++++++++---------- 4 files changed, 97 insertions(+), 92 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu index 66e5fb30de99..0c4431b8eaef 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu @@ -424,7 +424,7 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ int* smem_topk_target_ranks = smem; int* smem_topk_send_indices = smem + TOP_K; - uint64_t already_copied = 0; + uint64_t already_copied[kRankMaskWords] = {}; // Precompute the ceil/floor partition parameters once per thread, outside the // per-token TOP_K loop. The fast path (remainder == 0) then collapses to a single // integer divide per call, matching the pre-PR uniform-partition cost exactly. @@ -444,8 +444,11 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ // 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 ((already_copied & (1ULL << target_rank)) || target_dead) + if (target_already_copied || target_dead) { if (thread_idx == 0) { @@ -470,7 +473,7 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [ smem_topk_target_ranks[k] = target_rank; smem_topk_send_indices[k] = dst_token_idx; } - already_copied |= 1ULL << target_rank; + already_copied[mask_word] |= mask_bit; } // Sync before dispatching data ThreadingPolicy::sync(); @@ -630,6 +633,7 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) // Validate parameters TLLM_CHECK(params.top_k > 0 && params.top_k <= kMaxTopK); TLLM_CHECK(params.ep_size > 0 && params.ep_size <= kMaxRanks); + 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; @@ -1316,6 +1320,7 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) // Validate parameters TLLM_CHECK(params.top_k > 0 && params.top_k <= kMaxTopK); TLLM_CHECK(params.ep_size > 0 && params.ep_size <= kMaxRanks); + 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; diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h index 0b872db4356f..9a6f3904c501 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h @@ -142,7 +142,7 @@ struct MoeA2ADispatchParams // 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. - uint64_t active_rank_mask[kRankMaskWords]; + uint64_t active_rank_mask[kRankMaskWords] = {~uint64_t{0}, ~uint64_t{0}}; // CUDA stream cudaStream_t stream; @@ -192,7 +192,7 @@ struct MoeA2ACombineParams // 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. - uint64_t active_rank_mask[kRankMaskWords]; + uint64_t active_rank_mask[kRankMaskWords] = {~uint64_t{0}, ~uint64_t{0}}; // CUDA stream cudaStream_t stream; diff --git a/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp b/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp index 8f467519fab7..fc45afd792bb 100644 --- a/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp +++ b/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp @@ -52,6 +52,9 @@ inline void resolveActiveRankMask(torch::optional const& maskTens uint64_t (&out)[tensorrt_llm::kernels::moe_comm::kRankMaskWords]) { using tensorrt_llm::kernels::moe_comm::kRankMaskWords; + 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) @@ -151,11 +154,14 @@ MoeA2ADataOffsets calculateOffsets(int epSize, int maxNumTokens, int eplbStatsNu torch::Tensor moeA2AInitializeOp(torch::Tensor const& workspace, int64_t epRank, int64_t epSize, int64_t maxNumTokens, torch::optional eplbStatsNumExperts) { + using tensorrt_llm::kernels::moe_comm::kMaxRanks; + // Validate inputs CHECK_TH_CUDA(workspace); CHECK_TYPE(workspace, torch::kUInt8); TORCH_CHECK(workspace.dim() == 2, "workspace must be a 2D tensor of shape [epSize, sizePerRank]"); TORCH_CHECK(workspace.size(0) == epSize, "workspace first dimension must equal epSize"); + TORCH_CHECK(epSize > 0 && epSize <= kMaxRanks, "epSize must be in the range (0, ", kMaxRanks, "]"); TORCH_CHECK(epRank >= 0 && epRank < epSize, "epRank must be in the range [0, epSize)"); // Initialize workspace to zero @@ -223,6 +229,7 @@ std::tuple, int64_t, torch::Tensor> moeA2ADispatchOp( using tensorrt_llm::kernels::moe_comm::moe_a2a_dispatch_launch; using tensorrt_llm::kernels::moe_comm::kMaxTopK; using tensorrt_llm::kernels::moe_comm::kMaxPayloads; + using tensorrt_llm::kernels::moe_comm::kMaxRanks; // Validate inputs CHECK_INPUT(tokenSelectedExperts, torch::kInt32); @@ -238,6 +245,7 @@ std::tuple, int64_t, torch::Tensor> moeA2ADispatchOp( int64_t localNumTokens = tokenSelectedExperts.size(0); TORCH_CHECK(runtimeMaxTokensPerRank > 0, "runtimeMaxTokensPerRank must be positive"); + TORCH_CHECK(epSize > 0 && epSize <= kMaxRanks, "epSize must be in the range (0, ", kMaxRanks, "]"); TORCH_CHECK(epRank >= 0 && epRank < epSize, "epRank must be in the range [0, epSize)"); TORCH_CHECK(topK > 0 && topK <= kMaxTopK, "topK must be in the range (0, kMaxTopK]"); TORCH_CHECK(!inputPayloads.empty(), "inputPayloads must not be empty"); @@ -458,6 +466,7 @@ torch::Tensor moeA2ACombineOp(torch::Tensor const& payload, int64_t localNumToke using tensorrt_llm::kernels::moe_comm::MoeA2ACombineParams; using tensorrt_llm::kernels::moe_comm::moe_a2a_combine_launch; using tensorrt_llm::kernels::moe_comm::kMaxTopK; + using tensorrt_llm::kernels::moe_comm::kMaxRanks; // Validate inputs CHECK_TH_CUDA(payload); @@ -471,6 +480,7 @@ torch::Tensor moeA2ACombineOp(torch::Tensor const& payload, int64_t localNumToke TORCH_CHECK(reinterpret_cast(payload.data_ptr()) % 16 == 0, "payload must be 16-byte aligned"); int64_t elementsPerToken = payload.size(2); TORCH_CHECK(elementsPerToken > 0, "elementsPerToken must be positive"); + TORCH_CHECK(epSize > 0 && epSize <= kMaxRanks, "epSize must be in the range (0, ", kMaxRanks, "]"); TORCH_CHECK(epRank >= 0 && epRank < epSize, "epRank must be in the range [0, epSize)"); TORCH_CHECK(topK > 0 && topK <= kMaxTopK, "topK must be in the range (0, kMaxTopK]"); diff --git a/tests/unittest/_torch/multi_gpu/test_moe_a2a_rank_mask.py b/tests/unittest/_torch/multi_gpu/test_moe_a2a_rank_mask.py index 9f843671dff9..7c1d89ca14bd 100644 --- a/tests/unittest/_torch/multi_gpu/test_moe_a2a_rank_mask.py +++ b/tests/unittest/_torch/multi_gpu/test_moe_a2a_rank_mask.py @@ -32,9 +32,9 @@ import pickle import sys -import traceback import cloudpickle +import pynvml import pytest import torch from mpi4py import MPI @@ -63,6 +63,16 @@ def setup_test(): tllm.logger.set_level("error") +def _skip_if_mnnvl_unsupported() -> None: + try: + MnnvlMemory.initialize() + supports_mnnvl = MnnvlMemory.supports_mnnvl() + except (RuntimeError, pynvml.NVMLError) as exc: + pytest.skip(f"MNNVL not supported on this system: {exc}") + if not supports_mnnvl: + pytest.skip("MNNVL not supported on this system") + + def _ep_mask_words(ep_size: int, dead_ranks: set[int]) -> torch.Tensor: """Build the uint64[EP_MASK_NUM_WORDS] CPU tensor expected by the C++ op.""" mask_int = ((1 << ep_size) - 1) & ~sum(1 << r for r in dead_ranks) @@ -162,41 +172,35 @@ def _worker_all_active_matches_no_mask( ): rank = tllm.mpi_rank() torch.cuda.set_device(rank) - try: - mapping = Mapping(rank=rank, tp_size=ep_size, moe_ep_size=ep_size, world_size=ep_size) - moe_a2a = MoeAlltoAll( - mapping=mapping, - max_num_tokens=local_num_tokens, - top_k=top_k, - num_slots=num_experts, - workspace_size_per_rank=workspace_size_per_rank, - ) + mapping = Mapping(rank=rank, tp_size=ep_size, moe_ep_size=ep_size, world_size=ep_size) + moe_a2a = MoeAlltoAll( + mapping=mapping, + max_num_tokens=local_num_tokens, + top_k=top_k, + num_slots=num_experts, + workspace_size_per_rank=workspace_size_per_rank, + ) - # Same RNG seed across both runs => identical inputs. - torch.manual_seed(0xA2A + rank) - token_selected_experts = _generate_token_selected_experts( - local_num_tokens, num_experts, top_k - ) - payload = _make_payload(local_num_tokens, hidden_size, rank) + # Same RNG seed across both runs => identical inputs. + torch.manual_seed(0xA2A + rank) + token_selected_experts = _generate_token_selected_experts(local_num_tokens, num_experts, top_k) + payload = _make_payload(local_num_tokens, hidden_size, rank) - out_no_mask, topk_no_mask = _run_dispatch_combine( - moe_a2a, token_selected_experts, payload, local_num_tokens, active_rank_mask=None - ) - out_all_active, topk_all_active = _run_dispatch_combine( - moe_a2a, - token_selected_experts, - payload, - local_num_tokens, - active_rank_mask=_ep_mask_words(ep_size, dead_ranks=set()), - ) + out_no_mask, topk_no_mask = _run_dispatch_combine( + moe_a2a, token_selected_experts, payload, local_num_tokens, active_rank_mask=None + ) + out_all_active, topk_all_active = _run_dispatch_combine( + moe_a2a, + token_selected_experts, + payload, + local_num_tokens, + active_rank_mask=_ep_mask_words(ep_size, dead_ranks=set()), + ) - return ( - torch.equal(out_no_mask, out_all_active), - torch.equal(topk_no_mask, topk_all_active), - ) - except Exception: - traceback.print_exc() - raise + return ( + torch.equal(out_no_mask, out_all_active), + torch.equal(topk_no_mask, topk_all_active), + ) # --------------------------------------------------------------------------- @@ -215,53 +219,47 @@ def _worker_one_rank_masked( ): rank = tllm.mpi_rank() torch.cuda.set_device(rank) - try: - mapping = Mapping(rank=rank, tp_size=ep_size, moe_ep_size=ep_size, world_size=ep_size) - # Every rank participates in workspace init (it has MPI barriers internally). - moe_a2a = MoeAlltoAll( - mapping=mapping, - max_num_tokens=local_num_tokens, - top_k=top_k, - num_slots=num_experts, - workspace_size_per_rank=workspace_size_per_rank, - ) + mapping = Mapping(rank=rank, tp_size=ep_size, moe_ep_size=ep_size, world_size=ep_size) + # Every rank participates in workspace init (it has MPI barriers internally). + moe_a2a = MoeAlltoAll( + mapping=mapping, + max_num_tokens=local_num_tokens, + top_k=top_k, + num_slots=num_experts, + workspace_size_per_rank=workspace_size_per_rank, + ) - if rank == dead_rank: - # Simulate a dead rank: do not call dispatch/combine. Wait at a final - # barrier so the surviving ranks have someone to synchronize with at - # the end of the test. (The kernel itself never observes us because - # the surviving ranks pass a mask with our bit cleared.) - MPI.COMM_WORLD.barrier() - return ("dead", None, None, None) - - torch.manual_seed(0xA2A + rank) - token_selected_experts = _generate_token_selected_experts( - local_num_tokens, num_experts, top_k - ) - payload = _make_payload(local_num_tokens, hidden_size, rank) + if rank == dead_rank: + # Simulate a dead rank: do not call dispatch/combine. Wait at a final + # barrier so the surviving ranks have someone to synchronize with at + # the end of the test. (The kernel itself never observes us because + # the surviving ranks pass a mask with our bit cleared.) + MPI.COMM_WORLD.barrier() + return ("dead", None, None, None) - # Build mask with dead_rank's bit cleared. - mask = _ep_mask_words(ep_size, dead_ranks={dead_rank}) + torch.manual_seed(0xA2A + rank) + token_selected_experts = _generate_token_selected_experts(local_num_tokens, num_experts, top_k) + payload = _make_payload(local_num_tokens, hidden_size, rank) - # Compute the per-token target ranks the way the kernel does so we can - # cross-check the workspace afterwards. - num_experts_per_rank = num_experts // ep_size - expected_target_ranks = (token_selected_experts // num_experts_per_rank).cpu() + # Build mask with dead_rank's bit cleared. + mask = _ep_mask_words(ep_size, dead_ranks={dead_rank}) - combined, topk_target_ranks = _run_dispatch_combine( - moe_a2a, token_selected_experts, payload, local_num_tokens, active_rank_mask=mask - ) + # Compute the per-token target ranks the way the kernel does so we can + # cross-check the workspace afterwards. + num_experts_per_rank = num_experts // ep_size + expected_target_ranks = (token_selected_experts // num_experts_per_rank).cpu() - MPI.COMM_WORLD.barrier() - return ( - "alive", - combined, - topk_target_ranks, - expected_target_ranks, - ) - except Exception: - traceback.print_exc() - raise + combined, topk_target_ranks = _run_dispatch_combine( + moe_a2a, token_selected_experts, payload, local_num_tokens, active_rank_mask=mask + ) + + MPI.COMM_WORLD.barrier() + return ( + "alive", + combined, + topk_target_ranks, + expected_target_ranks, + ) # --------------------------------------------------------------------------- @@ -280,11 +278,7 @@ def _worker_one_rank_masked( ) def test_all_active_mask_matches_no_mask(mpi_pool_executor, local_num_tokens, top_k): """An all-ones active_rank_mask must produce identical output to omitting it.""" - try: - MnnvlMemory.initialize() - assert MnnvlMemory.supports_mnnvl() - except Exception: - pytest.skip("MNNVL not supported on this system") + _skip_if_mnnvl_unsupported() ep_size = mpi_pool_executor.num_workers if ep_size > torch.cuda.device_count(): @@ -300,7 +294,7 @@ def test_all_active_mask_matches_no_mask(mpi_pool_executor, local_num_tokens, to results = list( mpi_pool_executor.map( _worker_all_active_matches_no_mask, - *zip(*[args] * ep_size), + *zip(*[args] * ep_size, strict=True), ) ) @@ -329,11 +323,7 @@ def test_one_rank_masked_completes(mpi_pool_executor, dead_rank, local_num_token * Slots whose expert mapped to a surviving rank are unchanged from what the contiguous-partition routing rule predicts. """ - try: - MnnvlMemory.initialize() - assert MnnvlMemory.supports_mnnvl() - except Exception: - pytest.skip("MNNVL not supported on this system") + _skip_if_mnnvl_unsupported() ep_size = mpi_pool_executor.num_workers if ep_size > torch.cuda.device_count(): @@ -358,7 +348,7 @@ def test_one_rank_masked_completes(mpi_pool_executor, dead_rank, local_num_token results = list( mpi_pool_executor.map( _worker_one_rank_masked, - *zip(*[args] * ep_size), + *zip(*[args] * ep_size, strict=True), ) ) From ec451e0ea7a0d9aa82cda2a019ef8f5cfd26a212 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 23 Jun 2026 10:35:31 -0700 Subject: [PATCH 3/4] Move rank mask tests into MoE comm suite Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../_torch/modules/moe/test_moe_comm.py | 369 ++++++++++++++++ .../multi_gpu/test_moe_a2a_rank_mask.py | 397 ------------------ 2 files changed, 369 insertions(+), 397 deletions(-) delete mode 100644 tests/unittest/_torch/multi_gpu/test_moe_a2a_rank_mask.py diff --git a/tests/unittest/_torch/modules/moe/test_moe_comm.py b/tests/unittest/_torch/modules/moe/test_moe_comm.py index c59751f42025..4d5ea9c3162f 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_comm.py +++ b/tests/unittest/_torch/modules/moe/test_moe_comm.py @@ -114,6 +114,10 @@ # to avoid _WORKSPACE singleton assertion failures. NVLINK_WORKSPACE_MB = "512" +# Must match kRankMaskWords in +# cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h. +EP_MASK_NUM_WORDS = 2 + # ============================================================================ # Test Configuration @@ -192,6 +196,84 @@ def _safe_cpu(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: return t.cpu() +def _ep_mask_words(ep_size: int, dead_ranks: Set[int]) -> torch.Tensor: + """Build the uint64[EP_MASK_NUM_WORDS] CPU tensor expected by moe_a2a ops.""" + mask_int = ((1 << ep_size) - 1) & ~sum(1 << rank for rank in dead_ranks) + word_mask = (1 << 64) - 1 + words = [(mask_int >> (i * 64)) & word_mask for i in range(EP_MASK_NUM_WORDS)] + return torch.tensor(words, dtype=torch.uint64, device="cpu") + + +def _make_rank_mask_payload(local_num_tokens: int, hidden_size: int, rank: int) -> torch.Tensor: + """Make deterministic per-rank payloads for exact equality assertions.""" + base = torch.arange(local_num_tokens * hidden_size, dtype=torch.bfloat16, device="cuda").view( + local_num_tokens, hidden_size + ) + return base + (rank * 1000.0) + + +def _read_nvlink_topk_target_ranks( + comm: NVLinkOneSided, + max_num_tokens: int, + top_k: int, +) -> torch.Tensor: + """Read topk_target_ranks[max_num_tokens, top_k] from NVLinkOneSided workspace.""" + from tensorrt_llm.bindings import internal as _tllm_internal + + offset_index = int(_tllm_internal.thop.MOE_A2A_TOPK_TARGET_RANKS_OFFSET_INDEX) + offset = comm.moe_a2a_metainfo[offset_index].item() + raw = comm.workspace[ + comm.ep_rank, + offset : offset + max_num_tokens * top_k * 4, + ] + return raw.view(torch.int32).view(max_num_tokens, top_k).cpu() + + +def _run_nvlink_rank_mask_dispatch_combine( + comm: NVLinkOneSided, + token_selected_experts: torch.Tensor, + payload: torch.Tensor, + runtime_max_tokens_per_rank: int, + active_rank_mask: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor]: + """Run raw NVLink one-sided dispatch/combine with an optional active rank mask.""" + recv_tensors, combine_payload_offset, _ = torch.ops.trtllm.moe_a2a_dispatch( + token_selected_experts, + [payload], + comm.workspace, + comm.moe_a2a_metainfo, + runtime_max_tokens_per_rank, + comm.ep_rank, + comm.ep_size, + comm.top_k, + comm.num_experts, + None, # eplb_local_stats + active_rank_mask, + ) + + topk_target_ranks = _read_nvlink_topk_target_ranks( + comm, + runtime_max_tokens_per_rank, + comm.top_k, + ) + + combined = torch.ops.trtllm.moe_a2a_combine( + recv_tensors[0], + token_selected_experts.size(0), + comm.workspace, + comm.moe_a2a_metainfo, + runtime_max_tokens_per_rank, + comm.ep_rank, + comm.ep_size, + comm.top_k, + int(combine_payload_offset), + False, # payload_in_workspace + False, # use_low_precision + active_rank_mask, + ) + return combined.cpu(), topk_target_ranks + + # ============================================================================ # Source Encoding Utilities # ============================================================================ @@ -856,6 +938,154 @@ def _worker_full_pipeline(config: CommTestConfig) -> dict: comm.destroy() +def _make_rank_mask_config( + ep_size: int, + local_num_tokens: int, + top_k: int, +) -> CommTestConfig: + """Build the small NVLinkOneSided config used by active-rank-mask tests.""" + return CommTestConfig( + comm_type=COMM_NVLINK_ONE_SIDED, + ep_size=ep_size, + num_experts=FIXED_NUM_EXPERTS, + top_k=top_k, + hidden_size=1024, + all_num_tokens=[local_num_tokens] * ep_size, + ) + + +def _worker_rank_mask_all_active_matches_no_mask(config: CommTestConfig) -> dict: + """Check that all-active active_rank_mask is bit-identical to no 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, + ) + 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( + 0, + config.num_experts, + (local_num_tokens, config.top_k), + dtype=torch.int32, + device="cuda", + ) + payload = _make_rank_mask_payload(local_num_tokens, config.hidden_size, rank) + + out_no_mask, topk_no_mask = _run_nvlink_rank_mask_dispatch_combine( + comm, + token_selected_experts, + payload, + local_num_tokens, + active_rank_mask=None, + ) + 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()), + ) + + return { + "output_eq": torch.equal(out_no_mask, out_all_active), + "topk_eq": torch.equal(topk_no_mask, topk_all_active), + } + except Exception: + traceback.print_exc() + raise + finally: + if comm is not None and hasattr(comm, "destroy"): + comm.destroy() + + +def _expected_target_ranks( + token_selected_experts: torch.Tensor, + num_experts: int, + ep_size: int, +) -> torch.Tensor: + """Map each selected expert to its target EP rank using the kernel partition rule.""" + token_selected_experts_cpu = token_selected_experts.cpu() + expected = torch.empty_like(token_selected_experts_cpu) + for token_idx in range(token_selected_experts_cpu.shape[0]): + for k in range(token_selected_experts_cpu.shape[1]): + expert_id = int(token_selected_experts_cpu[token_idx, k].item()) + expected[token_idx, k] = _expert_id_to_rank(expert_id, num_experts, ep_size) + return expected + + +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.""" + 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, + ) + # All ranks must initialize the symmetric workspace before the dead rank + # stops participating in dispatch/combine. + comm = create_comm_object(config.comm_type, mapping, config) + + if rank == dead_rank: + MPI.COMM_WORLD.barrier() + return {"status": "dead"} + + local_num_tokens = config.all_num_tokens[rank] + torch.manual_seed(0xA2A + rank) + token_selected_experts = torch.randint( + 0, + config.num_experts, + (local_num_tokens, config.top_k), + 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( + comm, + token_selected_experts, + payload, + local_num_tokens, + active_rank_mask=mask, + ) + expected_target_ranks = _expected_target_ranks( + token_selected_experts, + config.num_experts, + config.ep_size, + ) + + MPI.COMM_WORLD.barrier() + return { + "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() + + # ============================================================================ # Verification Functions # ============================================================================ @@ -1630,6 +1860,102 @@ def _run_full_test(mpi_pool_executor, config: CommTestConfig): verify_combine_results(all_results, config, rtol=0.02, atol=0.15) +def _skip_if_rank_mask_config_unsupported(config: CommTestConfig) -> None: + """Skip active-rank-mask tests when NVLinkOneSided cannot run locally.""" + skip_reason = check_platform_support(config.comm_type) + if skip_reason: + pytest.skip(skip_reason) + + skip_reason = check_feasibility(config.comm_type, config) + if skip_reason: + pytest.skip(skip_reason) + + if config.ep_size > torch.cuda.device_count(): + pytest.skip(f"Need {config.ep_size} GPUs but only {torch.cuda.device_count()} available") + + +def _run_rank_mask_all_active_test( + mpi_pool_executor, + 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) + + results = list( + mpi_pool_executor.map( + _worker_rank_mask_all_active_matches_no_mask, + *zip(*[(config,)] * config.ep_size), + ) + ) + + for rank, result in enumerate(results): + assert result["output_eq"], ( + f"rank {rank}: combine output differs between no-mask and all-active mask" + ) + assert result["topk_eq"], ( + f"rank {rank}: topk_target_ranks differ between no-mask and all-active mask" + ) + + +def _run_rank_mask_one_rank_masked_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_one_rank_masked, + *zip(*worker_args), + ) + ) + + saw_dead = False + for rank, result in enumerate(results): + if result["status"] == "dead": + assert rank == dead_rank + saw_dead = True + continue + + assert result["status"] == "alive" + combined = result["combined"] + topk_target_ranks = result["topk_target_ranks"] + expected_target_ranks = result["expected_target_ranks"] + + assert combined.shape == (local_num_tokens, config.hidden_size) + + live_topk = topk_target_ranks[:local_num_tokens] + live_expected = expected_target_ranks[:local_num_tokens] + for token_idx in range(local_num_tokens): + seen_ranks: Set[int] = set() + 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 got == -1 + else: + assert got == expected, ( + f"rank {rank} token {token_idx} k={k}: target rank mismatch " + f"(expected={expected}, got={got})" + ) + seen_ranks.add(expected) + + assert saw_dead, f"dead rank {dead_rank} did not appear in results" + + # ============================================================================ # Test Class # ============================================================================ @@ -1682,3 +2008,46 @@ def test_moe_comm_postquant(self, mpi_pool_executor, config: CommTestConfig): def test_moe_comm_non_divisible_ep(self, mpi_pool_executor, config: CommTestConfig): """Verify NVLinkOneSided with non-divisible EP (num_experts % ep_size != 0).""" _run_full_test(mpi_pool_executor, config) + + @pytest.mark.threadleak(enabled=False) + @pytest.mark.parametrize( + "mpi_pool_executor,local_num_tokens,top_k", + [ + (4, 16, 2), + (4, 32, 4), + ], + indirect=["mpi_pool_executor"], + ) + def test_moe_comm_rank_mask_all_active_matches_no_mask( + self, + mpi_pool_executor, + local_num_tokens: int, + top_k: int, + ): + """Verify all-active active_rank_mask matches omitted mask for NVLinkOneSided.""" + _run_rank_mask_all_active_test(mpi_pool_executor, 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), + (4, 0, 16, 4), + (4, 3, 32, 4), + ], + indirect=["mpi_pool_executor"], + ) + def test_moe_comm_rank_mask_one_rank_masked_completes( + self, + mpi_pool_executor, + dead_rank: int, + local_num_tokens: int, + top_k: int, + ): + """Verify masked-dead rank is skipped by raw NVLinkOneSided moe_a2a ops.""" + _run_rank_mask_one_rank_masked_test( + mpi_pool_executor, + dead_rank, + local_num_tokens, + top_k, + ) diff --git a/tests/unittest/_torch/multi_gpu/test_moe_a2a_rank_mask.py b/tests/unittest/_torch/multi_gpu/test_moe_a2a_rank_mask.py deleted file mode 100644 index 7c1d89ca14bd..000000000000 --- a/tests/unittest/_torch/multi_gpu/test_moe_a2a_rank_mask.py +++ /dev/null @@ -1,397 +0,0 @@ -# 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. -"""Unit tests for the active_rank_mask parameter on the MoE AlltoAll kernels (PR 1a.2). - -Two scenarios are exercised: - -1. **All-active mask matches no-mask** — passing a mask with every bit set must produce - bit-identical output to omitting the mask. Regression guard for the kernel mod's - default behavior. - -2. **One rank masked completes without hanging** — with bit K cleared, the surviving - N-1 ranks must complete dispatch + combine without spinning on the dead rank's - completion flag, and any token routed to the dead rank's experts must be dropped - (topk_target_ranks[k] == -1) rather than silently corrupting peer memory. - -The dispatch/combine kernels require `MnnvlMemory` (multi-node NVLink, GB200), so these -tests skip on hardware that does not support MNNVL and on nodes with fewer GPUs than -`ep_size`. -""" - -import pickle -import sys - -import cloudpickle -import pynvml -import pytest -import torch -from mpi4py import MPI - -import tensorrt_llm as tllm -from tensorrt_llm._mnnvl_utils import MnnvlMemory -from tensorrt_llm._torch.distributed import MoeAlltoAll -from tensorrt_llm.mapping import Mapping - -cloudpickle.register_pickle_by_value(sys.modules[__name__]) -MPI.pickle.__init__( - cloudpickle.dumps, - cloudpickle.loads, - pickle.HIGHEST_PROTOCOL, -) - - -# Must match cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h -# kRankMaskWords. The full mask is little-endian: word 0 covers ranks 0..63. -EP_MASK_NUM_WORDS = 2 - - -@pytest.fixture(autouse=True) -def setup_test(): - torch.manual_seed(0xA2A) - tllm.logger.set_level("error") - - -def _skip_if_mnnvl_unsupported() -> None: - try: - MnnvlMemory.initialize() - supports_mnnvl = MnnvlMemory.supports_mnnvl() - except (RuntimeError, pynvml.NVMLError) as exc: - pytest.skip(f"MNNVL not supported on this system: {exc}") - if not supports_mnnvl: - pytest.skip("MNNVL not supported on this system") - - -def _ep_mask_words(ep_size: int, dead_ranks: set[int]) -> torch.Tensor: - """Build the uint64[EP_MASK_NUM_WORDS] CPU tensor expected by the C++ op.""" - mask_int = ((1 << ep_size) - 1) & ~sum(1 << r for r in dead_ranks) - word_mask = (1 << 64) - 1 - words = [(mask_int >> (i * 64)) & word_mask for i in range(EP_MASK_NUM_WORDS)] - return torch.tensor(words, dtype=torch.uint64, device="cpu") - - -def _generate_token_selected_experts( - local_num_tokens: int, num_experts: int, top_k: int -) -> torch.Tensor: - return torch.randint( - 0, num_experts, (local_num_tokens, top_k), dtype=torch.int32, device="cuda" - ) - - -def _make_payload(local_num_tokens: int, hidden_size: int, rank: int) -> torch.Tensor: - """Deterministic per-rank payload so we can assert exact bit-for-bit equality.""" - base = torch.arange(local_num_tokens * hidden_size, dtype=torch.bfloat16, device="cuda").view( - local_num_tokens, hidden_size - ) - # Encode rank into the payload so cross-rank mismatches are immediately visible. - return base + (rank * 1000.0) - - -def _read_topk_target_ranks(moe_a2a: MoeAlltoAll, max_num_tokens: int, top_k: int) -> torch.Tensor: - """Read the kernel-written topk_target_ranks[max_num_tokens, top_k] from workspace.""" - offset = moe_a2a.metainfo[MoeAlltoAll._METAINFO_INDEX["TOPK_TARGET_RANKS_OFFSET_INDEX"]].item() - raw = moe_a2a.workspace[ - moe_a2a.ep_rank, - offset : offset + max_num_tokens * top_k * 4, - ] - return raw.view(torch.int32).view(max_num_tokens, top_k).cpu() - - -def _run_dispatch_combine( - moe_a2a: MoeAlltoAll, - token_selected_experts: torch.Tensor, - payload: torch.Tensor, - runtime_max_tokens_per_rank: int, - active_rank_mask: torch.Tensor | None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Drive dispatch + combine via the raw C++ ops (so we can pass active_rank_mask). - - Returns ``(combined_output, topk_target_ranks_snapshot)``. ``payload`` doubles as - both the dispatched payload and the staged combine payload, which is the simplest - end-to-end exercise of the masking path on both kernels. - """ - recv_tensors, combine_payload_offset, _ = torch.ops.trtllm.moe_a2a_dispatch( - token_selected_experts, - [payload], - moe_a2a.workspace, - moe_a2a.metainfo, - runtime_max_tokens_per_rank, - moe_a2a.ep_rank, - moe_a2a.ep_size, - moe_a2a.top_k, - moe_a2a.num_experts, - None, # eplb_local_stats - active_rank_mask, - ) - - # Snapshot the kernel-written routing table BEFORE combine (combine reads it but - # may also reset workspace state on subsequent rounds). - topk_target_ranks = _read_topk_target_ranks(moe_a2a, runtime_max_tokens_per_rank, moe_a2a.top_k) - - combine_payload = recv_tensors[0] # [ep_size, max_tokens, hidden_size] - combined = torch.ops.trtllm.moe_a2a_combine( - combine_payload, - token_selected_experts.size(0), - moe_a2a.workspace, - moe_a2a.metainfo, - runtime_max_tokens_per_rank, - moe_a2a.ep_rank, - moe_a2a.ep_size, - moe_a2a.top_k, - int(combine_payload_offset), - False, # payload_in_workspace - False, # use_low_precision - active_rank_mask, - ) - return combined.cpu(), topk_target_ranks - - -# --------------------------------------------------------------------------- -# Worker: regression — all-active mask must match no-mask (bit-identical). -# --------------------------------------------------------------------------- - - -def _worker_all_active_matches_no_mask( - ep_size: int, - local_num_tokens: int, - top_k: int, - workspace_size_per_rank: int, - num_experts: int, - hidden_size: int, -): - rank = tllm.mpi_rank() - torch.cuda.set_device(rank) - mapping = Mapping(rank=rank, tp_size=ep_size, moe_ep_size=ep_size, world_size=ep_size) - moe_a2a = MoeAlltoAll( - mapping=mapping, - max_num_tokens=local_num_tokens, - top_k=top_k, - num_slots=num_experts, - workspace_size_per_rank=workspace_size_per_rank, - ) - - # Same RNG seed across both runs => identical inputs. - torch.manual_seed(0xA2A + rank) - token_selected_experts = _generate_token_selected_experts(local_num_tokens, num_experts, top_k) - payload = _make_payload(local_num_tokens, hidden_size, rank) - - out_no_mask, topk_no_mask = _run_dispatch_combine( - moe_a2a, token_selected_experts, payload, local_num_tokens, active_rank_mask=None - ) - out_all_active, topk_all_active = _run_dispatch_combine( - moe_a2a, - token_selected_experts, - payload, - local_num_tokens, - active_rank_mask=_ep_mask_words(ep_size, dead_ranks=set()), - ) - - return ( - torch.equal(out_no_mask, out_all_active), - torch.equal(topk_no_mask, topk_all_active), - ) - - -# --------------------------------------------------------------------------- -# Worker: one rank masked — surviving ranks complete; dead-targeted slots dropped. -# --------------------------------------------------------------------------- - - -def _worker_one_rank_masked( - ep_size: int, - dead_rank: int, - local_num_tokens: int, - top_k: int, - workspace_size_per_rank: int, - num_experts: int, - hidden_size: int, -): - rank = tllm.mpi_rank() - torch.cuda.set_device(rank) - mapping = Mapping(rank=rank, tp_size=ep_size, moe_ep_size=ep_size, world_size=ep_size) - # Every rank participates in workspace init (it has MPI barriers internally). - moe_a2a = MoeAlltoAll( - mapping=mapping, - max_num_tokens=local_num_tokens, - top_k=top_k, - num_slots=num_experts, - workspace_size_per_rank=workspace_size_per_rank, - ) - - if rank == dead_rank: - # Simulate a dead rank: do not call dispatch/combine. Wait at a final - # barrier so the surviving ranks have someone to synchronize with at - # the end of the test. (The kernel itself never observes us because - # the surviving ranks pass a mask with our bit cleared.) - MPI.COMM_WORLD.barrier() - return ("dead", None, None, None) - - torch.manual_seed(0xA2A + rank) - token_selected_experts = _generate_token_selected_experts(local_num_tokens, num_experts, top_k) - payload = _make_payload(local_num_tokens, hidden_size, rank) - - # Build mask with dead_rank's bit cleared. - mask = _ep_mask_words(ep_size, dead_ranks={dead_rank}) - - # Compute the per-token target ranks the way the kernel does so we can - # cross-check the workspace afterwards. - num_experts_per_rank = num_experts // ep_size - expected_target_ranks = (token_selected_experts // num_experts_per_rank).cpu() - - combined, topk_target_ranks = _run_dispatch_combine( - moe_a2a, token_selected_experts, payload, local_num_tokens, active_rank_mask=mask - ) - - MPI.COMM_WORLD.barrier() - return ( - "alive", - combined, - topk_target_ranks, - expected_target_ranks, - ) - - -# --------------------------------------------------------------------------- -# Tests -# --------------------------------------------------------------------------- - - -@pytest.mark.threadleak(enabled=False) -@pytest.mark.parametrize( - "mpi_pool_executor,local_num_tokens,top_k", - [ - (4, 16, 2), - (4, 32, 4), - ], - indirect=["mpi_pool_executor"], -) -def test_all_active_mask_matches_no_mask(mpi_pool_executor, local_num_tokens, top_k): - """An all-ones active_rank_mask must produce identical output to omitting it.""" - _skip_if_mnnvl_unsupported() - - ep_size = mpi_pool_executor.num_workers - if ep_size > torch.cuda.device_count(): - pytest.skip( - f"Need at least {ep_size} GPUs but only {torch.cuda.device_count()} are available" - ) - - hidden_size = 1024 - num_experts = 32 - workspace_size_per_rank = 256 * 1024 * 1024 - - args = (ep_size, local_num_tokens, top_k, workspace_size_per_rank, num_experts, hidden_size) - results = list( - mpi_pool_executor.map( - _worker_all_active_matches_no_mask, - *zip(*[args] * ep_size, strict=True), - ) - ) - - for rank, (output_eq, topk_eq) in enumerate(results): - assert output_eq, f"rank {rank}: combine output differs between no-mask and all-active mask" - assert topk_eq, f"rank {rank}: topk_target_ranks differ between no-mask and all-active mask" - - -@pytest.mark.threadleak(enabled=False) -@pytest.mark.parametrize( - "mpi_pool_executor,dead_rank,local_num_tokens,top_k", - [ - (4, 2, 16, 2), - (4, 0, 16, 4), # mask the lowest-numbered rank - (4, 3, 32, 4), # mask the highest-numbered rank - ], - indirect=["mpi_pool_executor"], -) -def test_one_rank_masked_completes(mpi_pool_executor, dead_rank, local_num_tokens, top_k): - """With one rank masked dead, surviving ranks complete dispatch+combine. - - Verifies: - * No hang (the test reaches the assertions). - * On every surviving rank, any topk slot whose expert mapped to the dead - rank is dropped (topk_target_ranks == -1). - * Slots whose expert mapped to a surviving rank are unchanged from what - the contiguous-partition routing rule predicts. - """ - _skip_if_mnnvl_unsupported() - - ep_size = mpi_pool_executor.num_workers - if ep_size > torch.cuda.device_count(): - pytest.skip( - f"Need at least {ep_size} GPUs but only {torch.cuda.device_count()} are available" - ) - assert 0 <= dead_rank < ep_size - - hidden_size = 1024 - num_experts = 32 - workspace_size_per_rank = 256 * 1024 * 1024 - - args = ( - ep_size, - dead_rank, - local_num_tokens, - top_k, - workspace_size_per_rank, - num_experts, - hidden_size, - ) - results = list( - mpi_pool_executor.map( - _worker_one_rank_masked, - *zip(*[args] * ep_size, strict=True), - ) - ) - - saw_dead = False - for rank, (status, combined, topk_target, expected_target) in enumerate(results): - if status == "dead": - assert rank == dead_rank - saw_dead = True - continue - assert status == "alive" - # Combine produced an output of the expected shape on the surviving rank. - assert combined is not None - assert combined.shape == (local_num_tokens, hidden_size) - - # Per-token routing assertions, using only the live tokens (the workspace - # topk arrays are sized [max_num_tokens, top_k] and may have stale rows - # beyond local_num_tokens; we only care about the live region). - live_topk = topk_target[:local_num_tokens] - live_expected = expected_target[:local_num_tokens] - - # Slot-level rules: - # - If expected_target_rank == dead_rank, the kernel must have set -1. - # - Otherwise, kernel must record the same target rank (or -1 if the - # same target was already covered earlier in the same token's top-k - # list; the kernel uses -1 as a "duplicate" sentinel). - for token_idx in range(local_num_tokens): - seen_ranks: set[int] = set() - for k in range(top_k): - exp = int(live_expected[token_idx, k].item()) - got = int(live_topk[token_idx, k].item()) - if exp == 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 exp in seen_ranks: - # Duplicate target within this token — kernel sets -1. - assert got == -1 - else: - assert got == exp, ( - f"rank {rank} token {token_idx} k={k}: target rank mismatch " - f"(expected={exp}, got={got})" - ) - seen_ranks.add(exp) - - assert saw_dead, f"dead rank {dead_rank} did not appear in results" From b844c4f5dc27af7cc99430f447a46f32ab26a30b Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Wed, 24 Jun 2026 20:53:27 -0700 Subject: [PATCH 4/4] test: reuse EP health mask words in MoE comm test Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../unittest/_torch/modules/moe/test_moe_comm.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/tests/unittest/_torch/modules/moe/test_moe_comm.py b/tests/unittest/_torch/modules/moe/test_moe_comm.py index 4d5ea9c3162f..057f87a95f7a 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_comm.py +++ b/tests/unittest/_torch/modules/moe/test_moe_comm.py @@ -73,6 +73,7 @@ NVLinkTwoSidedFlashinfer, ) from tensorrt_llm._torch.modules.fused_moe.deep_ep_utils import deep_ep_installed +from tensorrt_llm._torch.modules.fused_moe.ep_group_health import EPGroupHealth from tensorrt_llm.deep_ep.buffer import Buffer from tensorrt_llm.mapping import Mapping @@ -114,11 +115,6 @@ # to avoid _WORKSPACE singleton assertion failures. NVLINK_WORKSPACE_MB = "512" -# Must match kRankMaskWords in -# cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h. -EP_MASK_NUM_WORDS = 2 - - # ============================================================================ # Test Configuration # ============================================================================ @@ -197,11 +193,11 @@ def _safe_cpu(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: def _ep_mask_words(ep_size: int, dead_ranks: Set[int]) -> torch.Tensor: - """Build the uint64[EP_MASK_NUM_WORDS] CPU tensor expected by moe_a2a ops.""" - mask_int = ((1 << ep_size) - 1) & ~sum(1 << rank for rank in dead_ranks) - word_mask = (1 << 64) - 1 - words = [(mask_int >> (i * 64)) & word_mask for i in range(EP_MASK_NUM_WORDS)] - return torch.tensor(words, dtype=torch.uint64, device="cpu") + """Build the CPU active-rank mask tensor expected by moe_a2a ops.""" + health = EPGroupHealth(ep_size) + for rank in dead_ranks: + health.mark_failed(rank) + return torch.tensor(health.get_mask_words(), dtype=torch.uint64, device="cpu") def _make_rank_mask_payload(local_num_tokens: int, hidden_size: int, rank: int) -> torch.Tensor: