include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh (2)

218-242: Lamport sync helpers are fragile w.r.t. early-exit / partial participation
LamportFlags::ctaArrive() uses __syncthreads() / cluster.sync(), which will deadlock if any threads in the block exit early before reaching it. Today you mostly avoid OOB by shaping blockSize, but this is a sharp edge for future tuning (e.g., cluster sizing, loadsPerThread, partial tiles). Consider reworking ctaArrive() to avoid requiring full-block participation (e.g., a single designated thread does the atomic without a block-wide barrier), or enforce the invariant with stronger checks/docs at the call sites.

Also applies to: 385-418, 509-651, 760-888

---

924-950: Minor: offsets array is oversized (stack/register pressure)
In rmsNormLamport, offsets is declared as LoadsPerThread * kELTS_PER_LOAD but indexed as offsets[i] for i < LoadsPerThread. Tighten this to LoadsPerThread.

-  uint32_t offsets[LoadsPerThread * kELTS_PER_LOAD];
+  uint32_t offsets[LoadsPerThread];

csrc/trtllm_mnnvl_allreduce.cu (1)

29-35: Add basic range/null guards for pointer-like int64 args
A null/invalid multicast_buffer_ptr, buffer_ptrs_dev, buffer_ptr_local, or empty buffer_flags_mnnvl will turn into UB quickly. Even lightweight > 0 / non-empty checks would make failures actionable.

Also applies to: 79-90

flashinfer/comm/mnnvl.py (1)

640-656: Minor: unused recvmsg tuple parts (ruff RUF059)
Prefix unused msg/flags/addr with _ to avoid lint noise.

tests/comm/test_allreduce_negative.py (1)

36-70: Consider extracting duplicated fixture setup to reduce code repetition.

Both TestMNNVLUnsupportedPatterns and TestMNNVLMissingRequiredParameters have nearly identical setup fixtures. Consider extracting the common workspace creation and teardown logic into a shared fixture or base class to improve maintainability.

@pytest.fixture
def mnnvl_workspace():
    """Shared fixture for MNNVL workspace setup and teardown."""
    rank, world_size, gpus_per_node = setup_mpi_and_cuda()
    
    workspace = create_allreduce_fusion_workspace(
        backend="mnnvl",
        world_size=world_size,
        rank=rank,
        max_token_num=128,
        hidden_dim=2880,
        dtype=torch.float16,
        topology="single_node",
        gpus_per_node=gpus_per_node,
    )
    
    yield workspace, rank, world_size, gpus_per_node
    
    if workspace is not None:
        workspace.destroy()
    trtllm_mnnvl_ar.mpi_barrier()

flashinfer/comm/workspace_base.py (1)

52-61: Consider using a more specific type for use_oneshot parameter.

The use_oneshot: Optional[Any] parameter is very permissive. Based on the relevant code snippets, the TRTLLM backend uses bool while MNNVL uses MNNVLAllreduceFusionStrategy. Consider using a Union type or documenting the expected types per backend.

+from typing import Optional, Any, Union
+
 @abstractmethod
 def is_buffer_size_sufficient(
     self,
     tp_size: int,
     num_tokens: int,
     hidden_dim: int,
     dtype: torch.dtype,
-    use_oneshot: Optional[Any] = None,
+    strategy: Optional[Any] = None,  # Backend-specific: bool for TRTLLM, MNNVLAllreduceFusionStrategy for MNNVL
 ) -> bool:
     pass

tests/comm/test_trtllm_mnnvl_allreduce.py (2)

17-103: Consider inlining the nested func function.

The nested func function doesn't capture any variables from the outer scope that aren't already passed as parameters, and it's only called once. Consider inlining it for clarity.

---

233-271: Fix type annotation for reference_output.

The variable is initialized to None but typed as Tuple[torch.Tensor, ...]. Use Optional[Tuple[torch.Tensor, ...]] or remove the annotation since Python will infer it correctly from the subsequent assignments.

     x_local = x_full[rank, :, :]
-    reference_output: Tuple[torch.Tensor, ...] = None
+    reference_output: Optional[Tuple[torch.Tensor, ...]] = None
     if fusion:

Note: Optional is already imported from typing.

tests/comm/test_allreduce_unified_api.py (4)

26-46: Hardcoded port may cause conflicts in parallel test runs.

The hardcoded MASTER_PORT = "29500" could cause port conflicts if multiple test instances run concurrently on the same machine.

Consider using a dynamic port or environment variable fallback:

-    os.environ["MASTER_PORT"] = "29500"
+    os.environ["MASTER_PORT"] = os.environ.get("MASTER_PORT", "29500")

---

133-146: Consider documenting why tolerances are set so high.

The tolerance values (rtol=0.05, atol=0.15) are quite loose for numerical validation. While this may be necessary for distributed bf16/fp16 operations, a brief comment explaining this would help future maintainers understand this isn't masking bugs.

---

190-197: Remove unused monkeypatch parameter.

The monkeypatch parameter is declared but never used in this function, as flagged by static analysis.

 def run_allreduce_test(
-    monkeypatch,
     seq_lens: list[int],
     fusion: bool,
     dtype: torch.dtype,
     hidden_size: int,
     backend: str,
 ):

Also update the call site in test_allreduce_unified accordingly.

---

314-334: Consider adding pytest markers for CI optimization.

The parametrized test creates 144 combinations (6 seq_lens × 2 fusion × 2 dtype × 2 hidden_size × 3 backend). Consider adding @pytest.mark.slow or similar markers to allow selective test execution in CI.

flashinfer/comm/trtllm_mnnvl_ar.py (2)

79-81: Use logging instead of print for production code.

Line 81 uses print() while other parts of this class use logging.debug(). For consistency and proper log level control, use logging:

-        print("Allocating MNNVL Allreduce Fusion Workspace...")
+        logging.info("Allocating MNNVL Allreduce Fusion Workspace...")

---

179-197: functools.cache on instance method can cause memory leaks.

Using @functools.cache on an instance method binds self to the cache key, preventing the instance from being garbage collected even after all other references are dropped. Since workspaces may be created and destroyed frequently, this could accumulate memory.

Consider using lru_cache with a bounded size, or move the caching to the static method get_required_buffer_size_bytes which doesn't have this issue:

-    @functools.cache
     def is_buffer_size_sufficient(
         self,
         tp_size: int,
         num_tokens: int,
         hidden_dim: int,
         dtype: torch.dtype,
         strategy: MNNVLAllreduceFusionStrategy = MNNVLAllreduceFusionStrategy.AUTO,
     ) -> bool:
-        """
-        Calculate the required buffer size for a given problem size.
-        """
+        """Check if the buffer is sufficient for a given problem size."""
         required_buffer_size = self.get_required_buffer_size_bytes(
             tp_size, num_tokens, hidden_dim, dtype, strategy
         )
-        if required_buffer_size > self.buffer_size_bytes:
-            return False
-        else:
-            return True
+        return required_buffer_size <= self.buffer_size_bytes

flashinfer/comm/allreduce.py (4)

146-161: Use logging instead of print, and address unused parameter.

1. Line 160 uses print() instead of logging for consistency with other modules.
2. The use_oneshot parameter is unused (static analysis ARG002) but may be kept for API consistency with MNNVLAllReduceFusionWorkspace.is_buffer_size_sufficient.

+import logging
+
 ...
     def is_buffer_size_sufficient(
         self,
         tp_size: int,
         num_tokens: int,
         hidden_dim: int,
         dtype: torch.dtype,
-        use_oneshot: Optional[Any] = None,
+        use_oneshot: Optional[Any] = None,  # Unused, kept for API consistency with MNNVL
     ) -> bool:
         try:
             check_trtllm_allreduce_fusion_workspace_metadata(
                 num_tokens, hidden_dim, tp_size, dtype, self.metadata
             )
             return True
         except ValueError as e:
-            print(f"Workspace is insufficient for problem size. {e}")
+            logging.debug(f"Workspace is insufficient for problem size: {e}")
             return False

---

200-217: Simplify _mnnvl_workspace_check - currently always returns True.

The function has a redundant conditional structure since both branches return True. If this is intentional scaffolding for future checks, consider adding a comment; otherwise simplify:

 def _mnnvl_workspace_check(
     backend: str,
     world_size: int,
     rank: int,
     max_token_num: int,
     hidden_dim: int,
     dtype: torch.dtype,
     topology: Literal["single_node", "multi_node"],
 ) -> bool:
     """
     Check if mnnvl backend CAN be used for workspace creation.
-
     """
-
-    if topology == "multi_node":
-        return True
-
-    return True
+    # MNNVL supports both single-node and multi-node topologies
+    return True

---

286-297: Use explicit Optional type hints for parameters defaulting to None.

Parameters with None defaults should explicitly declare Optional types for PEP 484 compliance and better IDE support.

 def create_allreduce_fusion_workspace(
     backend: Literal["trtllm", "mnnvl", "auto"] = "auto",
-    world_size: int = None,
-    rank: int = None,
-    max_token_num: int = None,
-    hidden_dim: int = None,
-    dtype: torch.dtype = None,
+    world_size: Optional[int] = None,
+    rank: Optional[int] = None,
+    max_token_num: Optional[int] = None,
+    hidden_dim: Optional[int] = None,
+    dtype: Optional[torch.dtype] = None,
     topology: Literal["single_node", "multi_node"] = "single_node",
     process_group: Optional["torch.distributed.ProcessGroup"] = None,
-    gpus_per_node: int = None,
+    gpus_per_node: Optional[int] = None,
     comm_backend: Optional[CommBackend] = None,
 ) -> AllReduceFusionWorkspace:

---

682-693: Prefix unused variable with underscore.

The residual_result variable is unpacked but never used, as flagged by static analysis.

             # Call the MNNVL fusion function
-            norm_result, residual_result = trtllm_mnnvl_fused_allreduce_add_rmsnorm(
+            norm_result, _residual_result = trtllm_mnnvl_fused_allreduce_add_rmsnorm(
                 input=input,
                 residual_in=residual_in,
                 gamma=rms_gamma,
                 workspace=workspace,
                 epsilon=rms_eps,
                 output=norm_out,
                 residual_out=residual_out,
                 launch_with_pdl=launch_with_pdl,
             )
             return norm_result

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro