fix: address EPLB mask review feedback

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

Author: chienchunhung

cpp/tensorrt_llm/runtime/moeLoadBalancer/moeLoadBalancer.cpp

@@ -59,13 +59,16 @@ namespace
 
 bool isRankMasked(std::vector<uint8_t> const* deadRankMask, int rank)
 {
-    return deadRankMask != nullptr && rank >= 0 && rank < static_cast<int>(deadRankMask->size())
-        && ((*deadRankMask)[rank] != 0);
+    return deadRankMask != nullptr && ((*deadRankMask)[rank] != 0);
 }
 
 int getActiveSlotCount(
     tensorrt_llm::kernels::MoeLoadBalanceMetaInfo const& metaInfo, std::vector<uint8_t> const* deadRankMask)
 {
+    TLLM_CHECK_WITH_INFO(deadRankMask == nullptr || static_cast<int>(deadRankMask->size()) == metaInfo.epSize,
+        "deadRankMask size (%ld) must match epSize (%d)",
+        deadRankMask == nullptr ? 0L : static_cast<long>(deadRankMask->size()), metaInfo.epSize);
+
     int activeRankCount = 0;
     for (int rank = 0; rank < metaInfo.epSize; ++rank)
     {
@@ -163,7 +166,8 @@ void doPlacement(tensorrt_llm::kernels::MoeLoadBalanceMetaInfo metaInfo, float*
 
     for (int expertId = 0; expertId < metaInfo.expertCount; ++expertId)
     {
-        assert(replicaCount[expertId] > 0); // Ensure replica count is positive
+        TLLM_CHECK_WITH_INFO(replicaCount[expertId] > 0, "Replica count (%d) for expert %d must be positive",
+            replicaCount[expertId], expertId);
         double slotSize = expertLoadFactor[expertId] / static_cast<double>(replicaCount[expertId]);
         for (int replicaId = 0; replicaId < replicaCount[expertId]; ++replicaId)
         {
@@ -172,7 +176,9 @@ void doPlacement(tensorrt_llm::kernels::MoeLoadBalanceMetaInfo metaInfo, float*
         }
     }
 
-    assert(static_cast<int>(allReplicas.size()) == totalSlotCount);
+    TLLM_CHECK_WITH_INFO(static_cast<int>(allReplicas.size()) == totalSlotCount,
+        "Replica count sum (%ld) must match active slot count (%d)", static_cast<long>(allReplicas.size()),
+        totalSlotCount);
 
     // 2. Sort replicas by slotSize descending
     std::sort(allReplicas.begin(), allReplicas.end());
@@ -1124,14 +1130,14 @@ void MoeLoadBalancer::reconfigureMaskOnly(std::vector<int> const& deadRanks)
     {
         layer->validateMaskOnly(candidateDeadRankMask);
     }
-    {
-        std::lock_guard<std::mutex> maskLock(mDeadRankMaskMutex);
-        mDeadRankMask = candidateDeadRankMask;
-    }
     for (auto& layer : mLayers)
     {
         layer->reconfigureMaskOnly(candidateDeadRankMask);
     }
+    {
+        std::lock_guard<std::mutex> maskLock(mDeadRankMaskMutex);
+        mDeadRankMask = candidateDeadRankMask;
+    }
 }
 
 void MoeLoadBalancer::workerThread()

cpp/tests/unit_tests/runtime/moeLoadBalancerTest.cpp

@@ -290,6 +290,109 @@ INSTANTIATE_TEST_SUITE_P(PlacementTests, MoePlacementTest,
         return name;
     });
 
+TEST(MoeLoadBalancerMaskOnlyTest, DynamicPlacementHonorsDeadRankMask)
+{
+    constexpr int kExpertCount = 4;
+    constexpr int kTopK = 2;
+    constexpr int kEpRank = 0;
+    constexpr int kEpSize = 4;
+    constexpr int kSlotCountPerRank = 2;
+    constexpr int kDeadRank = 2;
+    constexpr int kActiveSlotCount = (kEpSize - 1) * kSlotCountPerRank;
+
+    tensorrt_llm::kernels::MoeLoadBalanceMetaInfo metaInfo{kExpertCount, kTopK, kEpRank, kEpSize, kSlotCountPerRank};
+    std::vector<float> expertLoadFactor{8.0F, 4.0F, 2.0F, 1.0F};
+    std::vector<uint8_t> deadRankMask{0, 0, 1, 0};
+
+    MoePlacementCpuInfo cpuPlacement;
+    doReplication(metaInfo, expertLoadFactor.data(), &cpuPlacement, &deadRankMask);
+
+    int replicaSum = 0;
+    for (int replicaCount : cpuPlacement.expertReplicaCount)
+    {
+        replicaSum += replicaCount;
+    }
+    EXPECT_EQ(replicaSum, kActiveSlotCount);
+
+    cpuPlacement.rankExpertIds.resize(kEpSize);
+    for (int rank = 0; rank < kEpSize; ++rank)
+    {
+        cpuPlacement.rankExpertIds[rank].resize(kSlotCountPerRank, 99);
+    }
+
+    doPlacement(metaInfo, expertLoadFactor.data(), &cpuPlacement, &deadRankMask);
+
+    std::vector<int> placedReplicas(kExpertCount, 0);
+    int assignedSlotCount = 0;
+    for (int rank = 0; rank < kEpSize; ++rank)
+    {
+        for (int slot = 0; slot < kSlotCountPerRank; ++slot)
+        {
+            int const expertId = cpuPlacement.rankExpertIds[rank][slot];
+            if (rank == kDeadRank)
+            {
+                EXPECT_EQ(expertId, -1);
+                continue;
+            }
+
+            EXPECT_GE(expertId, 0);
+            EXPECT_LT(expertId, kExpertCount);
+            if (expertId >= 0 && expertId < kExpertCount)
+            {
+                ++placedReplicas[expertId];
+                ++assignedSlotCount;
+            }
+        }
+    }
+
+    EXPECT_EQ(assignedSlotCount, kActiveSlotCount);
+    for (int expertId = 0; expertId < kExpertCount; ++expertId)
+    {
+        EXPECT_EQ(placedReplicas[expertId], cpuPlacement.expertReplicaCount[expertId]);
+    }
+}
+
+TEST(MoeLoadBalancerMaskOnlyTest, DynamicPlacementRejectsMismatchedDeadRankMask)
+{
+    constexpr int kExpertCount = 4;
+    constexpr int kTopK = 2;
+    constexpr int kEpRank = 0;
+    constexpr int kEpSize = 4;
+    constexpr int kSlotCountPerRank = 2;
+
+    tensorrt_llm::kernels::MoeLoadBalanceMetaInfo metaInfo{kExpertCount, kTopK, kEpRank, kEpSize, kSlotCountPerRank};
+    std::vector<float> expertLoadFactor{1.0F, 1.0F, 1.0F, 1.0F};
+    std::vector<uint8_t> deadRankMask{0, 1};
+
+    MoePlacementCpuInfo cpuPlacement;
+    EXPECT_THROW(doReplication(metaInfo, expertLoadFactor.data(), &cpuPlacement, &deadRankMask),
+        tensorrt_llm::common::TllmException);
+}
+
+TEST(MoeLoadBalancerMaskOnlyTest, DynamicPlacementRejectsReplicaCountMismatch)
+{
+    constexpr int kExpertCount = 4;
+    constexpr int kTopK = 2;
+    constexpr int kEpRank = 0;
+    constexpr int kEpSize = 4;
+    constexpr int kSlotCountPerRank = 2;
+
+    tensorrt_llm::kernels::MoeLoadBalanceMetaInfo metaInfo{kExpertCount, kTopK, kEpRank, kEpSize, kSlotCountPerRank};
+    std::vector<float> expertLoadFactor{1.0F, 1.0F, 1.0F, 1.0F};
+    std::vector<uint8_t> deadRankMask{0, 0, 1, 0};
+
+    MoePlacementCpuInfo cpuPlacement;
+    cpuPlacement.expertReplicaCount = {1, 1, 1, 1};
+    cpuPlacement.rankExpertIds.resize(kEpSize);
+    for (int rank = 0; rank < kEpSize; ++rank)
+    {
+        cpuPlacement.rankExpertIds[rank].resize(kSlotCountPerRank, -1);
+    }
+
+    EXPECT_THROW(doPlacement(metaInfo, expertLoadFactor.data(), &cpuPlacement, &deadRankMask),
+        tensorrt_llm::common::TllmException);
+}
+
 TEST(MoeLoadBalancerMaskOnlyTest, ReconfigureMaskOnlyRemovesDeadRankSlots)
 {
     setenv("TLLM_HOST_ACCESSIBLE_ALLOW_MANAGED_FALLBACK", "1", 1);

tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py

@@ -989,7 +989,7 @@ def set_iter_info(self, enable_statistic: Optional[bool],
         if enable_update_weights is not None:
             self.enable_update_weights = enable_update_weights
 
-    def reconfigure_mask_only(self, dead_ranks: List[int]):
+    def reconfigure_mask_only(self, dead_ranks: list[int]) -> None:
         """
         Reconfigure EPLB routing so slots on dead EP ranks are unreachable.