tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py (2)

75-78:Guard against zero when capturing graphs.

If use_cuda_graph is True, assert all_rank_max_num_tokens > 0 to avoid capturing with a degenerate size.

Apply:

-        if use_cuda_graph:-            num_worst_tokens = all_rank_max_num_tokens * ep_size+        if use_cuda_graph:+            assert all_rank_max_num_tokens > 0, "all_rank_max_num_tokens must be > 0 when using CUDA graph"+            num_worst_tokens = all_rank_max_num_tokens * ep_size

---

217-223:typing.Map is undefined here; use Dict from typing.

Map isn’t imported; switch to Dict to avoid type-checker errors.

Apply:

-from typing import List, Optional, Tuple, Union+from typing import Dict, List, Optional, Tuple, Union@@-        self.buffers: Map[Mapping,-                          weakref.ReferenceType[VariableLengthBuffer]] = {}-        self.low_latency_buffers: Map[-            Mapping,-            weakref.ReferenceType[VariableLengthLowLatencyBuffer]] = {}+        self.buffers: Dict[Mapping,+                           weakref.ReferenceType[VariableLengthBuffer]] = {}+        self.low_latency_buffers: Dict[+            Mapping,+            weakref.ReferenceType[VariableLengthLowLatencyBuffer]] = {}

tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py (3)

217-219:Don’t overwrite NVSHMEM_QP_DEPTH if already set.

Respect user/env overrides and only set a default when unset. Keep the helpful issue link comment.

Apply:

-                os.environ['NVSHMEM_QP_DEPTH'] = str(-                    2 * (self.deep_ep_max_num_tokens + 1))+                if 'NVSHMEM_QP_DEPTH' not in os.environ:+                    os.environ['NVSHMEM_QP_DEPTH'] = str(+                        2 * (self.deep_ep_max_num_tokens + 1))

---

257-279:DeepEP feasibility: add divisibility check and clarify intranode detection.

Avoid false positives when moe_ep_size is not a multiple of local_mpi_size; also be explicit about internode divisibility.

Apply:

         def is_deepep_feasible(num_ranks: int) -> bool:             NUM_INTRANODE_SUPPORTED_RANKS = {2, 4, 8}             REQUIRED_LOCAL_MPI_SIZE = 8             NUM_INTERNODE_SUPPORTED_RDMA_RANKS = {2, 4, 8, 16}             mpi_size = local_mpi_size()             # Intranode cases-            if num_ranks == mpi_size and num_ranks in NUM_INTRANODE_SUPPORTED_RANKS:+            if num_ranks == mpi_size and num_ranks in NUM_INTRANODE_SUPPORTED_RANKS:                 return True             # Internode cases             if mpi_size != REQUIRED_LOCAL_MPI_SIZE:                 return False-            num_rdma_nodes = num_ranks // mpi_size+            if num_ranks % mpi_size != 0:+                return False+            num_rdma_nodes = num_ranks // mpi_size             return num_rdma_nodes in NUM_INTERNODE_SUPPORTED_RDMA_RANKS

---

301-305:Selection fallback to DeepEPLowLatency is sensible; log reason optionally.

Optional: log when DeepEP is installed but not feasible (ranks mismatch), to aid configuration debugging.

-                if is_deepep_feasible(mapping.moe_ep_size):-                    return AlltoallMethodType.DeepEP-                return AlltoallMethodType.DeepEPLowLatency+                if is_deepep_feasible(mapping.moe_ep_size):+                    return AlltoallMethodType.DeepEP+                logger.debug("DeepEP not feasible for moe_ep_size=%d (local_mpi_size=%d); using DeepEPLowLatency",+                             mapping.moe_ep_size, local_mpi_size())+                return AlltoallMethodType.DeepEPLowLatency

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Learnt from: timlee0212PR: NVIDIA/TensorRT-LLM#6886File: tensorrt_llm/_torch/models/modeling_deepseekv3.py:0-0Timestamp: 2025-08-14T06:36:40.701ZLearning: In DeepSeek V3 model (tensorrt_llm/_torch/models/modeling_deepseekv3.py), the disagreement between AllReduce.__init__ guard and _compute_mlp_tp_size logic for MNNVL usage is expected by design. The AllReduce component and MLP TP-size computation intentionally use different criteria for MNNVL availability decisions.