flashinfer/comm/mnnvl.py (1)

640-655: Minor polish: unused recvmsg outputs and predictable opId

Two small, non‑blocking cleanups:

• In IpcSocket.recv_fd(), the unpacked msg, flags, and addr from recvmsg are unused. Renaming them to _msg, _flags, _addr will make that explicit and silence linters:

_msg, ancdata, _flags, _addr = self.sock.recvmsg(...)

• opId for the socket name is generated with random.randint. Since it’s only used as a uniqueness hint (not security‑sensitive), this is fine; if you want to appease S311 you could switch to secrets.randbits(64) or document that it’s non‑cryptographic.

Both are optional, but would make static analysis quieter.

Also applies to: 885-889

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

23-25: Explicitly include <array> and <tuple>, and guard adjustGridConfig against smCount == 0

Within this header:

• LamportBufferLayout, LamportFlags, PackedVec, and several kernels use std::array.
• adjustGridConfig returns std::tuple<int, int, int> and callers use std::get.

But only <type_traits> is included; <array> and <tuple> are currently pulled in (if at all) via transitive includes, which is fragile.

Also, adjustGridConfig relies on GetCudaMultiProcessorCount():

int smCount = GetCudaMultiProcessorCount();
while (numTokens * clusterSize > smCount && clusterSize > 1 && blockSize <= 512) {
  ...
}

If GetCudaMultiProcessorCount() ever returns 0 (e.g., CUDA error or misconfiguration), this loop will keep shrinking clusterSize and inflating blockSize in a somewhat opaque way.

Suggestions:

• Add explicit includes at the top of the header:

#include <array>
#include <tuple>

• Make adjustGridConfig robust to a 0 or negative SM count by early‑clamping:

int smCount = GetCudaMultiProcessorCount();
if (smCount <= 0) {
  // Fall back to single-SM configuration
  clusterSize = 1;
  blockSize = std::min(threadsNeeded, 1024);
  return {blockSize, clusterSize, loadsPerThread};
}

This keeps the fused path predictable even if the helper cannot obtain a valid SM count.

Also applies to: 54-177, 143-163, 291-313, 348-359, 385-419, 449-497

---

509-651: Confirm lamport clear / wait protocol assumptions for oneshot kernel

The oneshot fused kernel uses LamportFlags as follows:

• Out‑of‑bounds threads call ctaArrive() then clearDirtyLamportBuf() and return.
• In‑bounds threads:
  • write their shard into the multicast lamport buffer,
  • call ctaArrive() again,
  • then call clearDirtyLamportBuf() and spin on the Lamport buffers until all entries are non‑negZero.

This protocol assumes:

• Every thread in the grid calls clearDirtyLamportBuf() exactly once per iteration.
• Buffer flags and bytesToClear are correctly initialized to match the configured numTokens * tokenDim * WorldSize.

Given that this is a direct Lamport port, the logic looks consistent, but the protocol is subtle. I’d recommend:

• Double‑checking the initialization of buffer_flags in MNNVLAllreduceFusionWorkspace matches the expectations here (current index, dirty index, bytes per buffer, and stage counts).
• Adding a brief comment near the kernel launch documenting that buffer_flags must follow the [cur, dirty, bytes_per_buffer, dirty_num_stages, bytes_to_clear[4], access_ptr] layout used by LamportFlags.

No code change strictly required, but the invariants are nontrivial and worth locking down in comments/tests.

---

754-885: Two‑shot path & RMSNorm fusion: validate world sizes and loads‑per‑thread bounds

The two‑shot kernels and dispatchers introduce several constraints:

• twoshotAllreduceFusionDispatch<T> only supports nRanks in {2, 4, 8, 16, 32, 64} and enforces tokenDim % (sizeof(float4) / sizeof(T)) == 0.
• rmsNormLamport is instantiated with LoadsPerThread in [1, 8] and uses float4 loads into shared memory; dynamic shared memory is sized as 3 * rnBlockSize * iters * sizeof(T) and indexed accordingly.

The implementation looks coherent, but a few invariants are implicit:

• MNNVLTwoShotStage::NUM_STAGES must stay in sync with the LamportFlags<float4> usage and the two bytes_to_clear entries in waitAndUpdate.
• rnLoadsPerThread retrieved from adjustGridConfig must remain in [1, 8]; the default: branch already errors if it’s out of range, which is good.
• rnClusterSize from adjustGridConfig is assumed to be <= 8 given __shared__ float sharedVal[8]; in the RMSNorm kernel.

Given these contracts, I’d suggest:

• Adding asserts (or comments) that rnClusterSize <= 8 when CGA is used, to guard future changes to adjustGridConfig.
• Extending tests to cover the corner cases where tokenDim is just at or above the supported boundary (e.g., maximum hidden size and multiple world sizes) so we don’t regress the FLASHINFER_CHECK conditions.

Functionally the code looks sound; this is mainly about making the implicit constraints explicit.

Also applies to: 898-959, 1062-1219

csrc/trtllm_mnnvl_allreduce.cu (1)

99-107: Error message still mentions “twoshot” even for oneshot path

Regardless of use_oneshot, the failure message says:

TVM_FFI_ICHECK(status == cudaSuccess)
    << "twoshot_allreduce_dispatch_world_size failed with error code "
    << cudaGetErrorString(status);

This is slightly misleading when the oneshot dispatch is used. Consider making the message neutral (e.g., “allreduce_fusion_dispatch failed…”) or switching on use_oneshot to provide a more accurate label. Behavior is otherwise fine.

tests/comm/test_trtllm_mnnvl_allreduce.py (2)

232-270: Use the same eps for reference RMSNorm as the fused kernel

In prepare_test_data, the fused reference path uses:

norm_out = rmsnorm(
    residual_out, norm_weight, torch.finfo(dtype).eps, enable_pdl=False
)

But the actual fused kernel is driven by the eps argument passed into row_linear_residual_norm_fusion_forward (eps = 1e-5 in run_mnnvl_ar_full).

To keep the reference as close as possible to the fused implementation (and not rely on loose tolerances), consider:

def prepare_test_data(..., fusion: bool, eps: float):
    ...
    if fusion:
        ...
        norm_out = rmsnorm(residual_out, norm_weight, eps, enable_pdl=False)

and threading eps through the call sites.

---

273-281: Annotate legacy_explicit_workspace_bytes as optional

Ruff’s RUF013 warning here is valid:

def run_mnnvl_ar_full(...,
    legacy_explicit_workspace_bytes: int = None,
    legacy_api: bool = False,
):

Changing the signature to make the optionality explicit improves readability and typing:

from typing import Optional

def run_mnnvl_ar_full(
    ...,
    legacy_explicit_workspace_bytes: Optional[int] = None,
    legacy_api: bool = False,
) -> None:
    ...

or, in Python 3.10+:

legacy_explicit_workspace_bytes: int | None = None

flashinfer/comm/trtllm_mnnvl_ar.py (1)

203-205: Drop debug print from hot path.
This unconditional print will spam stdout for every call to the fused kernel. Please remove it or guard it behind a proper debug logger.

-        print(
-            f"[Rank {rank}] Inside Kernel: multicast_buffer_ptr: {multicast_buffer_ptr:x}, buffer_ptrs_dev: {buffer_ptrs_dev:x}, buffer_ptr_local: {buffer_ptr_local:x}, buffer_flags_mnnvl: {buffer_flags_mnnvl}"
-        )

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

csrc/trtllm_mnnvl_allreduce.cu (1)

56-69: Guard RMSNorm fusion against missing residual_in and validate its shape.

The precondition at line 56 enforces presence of residual_out, gamma, and epsilon when rmsnorm_fusion=true, but omits residual_in. The kernel will dereference residualInPtr unconditionally when RMSNormFusion is true, causing undefined behavior if residual_in is absent.

Additionally, shape validation (lines 61-68) only covers residual_out and gamma; residual_in is not validated.

Extend the precondition to include residual_in:

-    TVM_FFI_ICHECK((residual_out.has_value() && gamma.has_value() && epsilon.has_value()) ||
+    TVM_FFI_ICHECK((residual_out.has_value() && residual_in.has_value() &&
+                    gamma.has_value() && epsilon.has_value()) ||
                    !rmsnorm_fusion)
-        << "residual_out, gamma, and epsilon must be provided if rmsnorm_fusion is true";
+        << "residual_out, residual_in, gamma, and epsilon must be provided if rmsnorm_fusion is true";

Add shape validation for residual_in within the if (rmsnorm_fusion) block:

     if (rmsnorm_fusion) {
       TVM_FFI_ICHECK(residual_out.value().size(0) == num_tokens &&
                      residual_out.value().size(1) == token_dim)
           << "residual_out shape mismatch: expected (" << input.size(0) << ", " << input.size(1)
           << ") but got (" << residual_out.value().size(0) << ", " << residual_out.value().size(1)
           << ")";
+      TVM_FFI_ICHECK(residual_in.value().size(0) == num_tokens &&
+                     residual_in.value().size(1) == token_dim)
+          << "residual_in shape mismatch: expected (" << num_tokens << ", " << token_dim
+          << ") but got (" << residual_in.value().size(0) << ", "
+          << residual_in.value().size(1) << ")";
       TVM_FFI_ICHECK(gamma.value().size(0) == token_dim)
           << "gamma must have the same shape as token dimension (" << token_dim << ") but got ("
           << gamma.value().size(0) << ")";
     }

flashinfer/comm/trtllm_mnnvl_ar.py (1)

331-332: Restore RMSNorm epsilon default to 1e-5.

Overriding epsilon with torch.finfo(input.dtype).eps replaces the kernel's built-in 1e-5 default (see line 91 in csrc/trtllm_mnnvl_allreduce.cu). For fp16 this becomes ~1e-3, materially changing RMSNorm results and breaking parity with TensorRT-LLM.

Apply this diff:

     if epsilon is None:
-        epsilon = torch.finfo(input.dtype).eps
+        epsilon = 1e-5

flashinfer/comm/trtllm_mnnvl_ar.py (1)

502-504: Clarify assertion for legacy API compatibility.

The assertion at lines 502-504 will fail with a cryptic message if wait_for_results=False is passed. Since this is deprecated legacy code, the assertion is reasonable, but consider improving the error message for clarity:

-    assert wait_for_results and (out is not None), (
-        "Calling the legacy trtllm_mnnvl_all_reduce with wait_for_results=False is not supported. Please use trtllm_mnnvl_allreduce instead."
-    )
+    if not wait_for_results or out is None:
+        raise ValueError(
+            "Legacy trtllm_mnnvl_all_reduce requires wait_for_results=True and a valid output tensor. "
+            "Please use the new trtllm_mnnvl_allreduce API instead."
+        )

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: raayandhar
Repo: flashinfer-ai/flashinfer PR: 2070
File: include/flashinfer/gemm/bf16_gemm_cutlass_template.h:145-160
Timestamp: 2025-11-12T03:35:17.583Z
Learning: In flashinfer GEMM implementations (e.g., include/flashinfer/gemm/bf16_gemm_cutlass_template.h, fp8_gemm_cutlass_template.h), it is acceptable to catch and silently ignore std::runtime_error exceptions in getWorkspaceSizeImpl when probing multiple GEMM configurations, as some configurations may legitimately fail due to SMEM constraints. This pattern should include a comment like "// Swallow errors when SMEM exceeds maximum allowed" to document the rationale.

flashinfer/comm/mnnvl.py (1)

566-664: Close remaining POSIX FDs in IPC path to avoid leaks

In the POSIX handle path of _alloc_mn_mcast_mem, a few FDs are still never closed:

• local_shareable_uc_handle returned by cuMemExportToShareableHandle (line 958) is used in the IPC ring allgather but never closed.
• During the ring, each rank sends its local_shareable_uc_handle to all peers, including itself. The self‑recv for p == group_rank populates all_shareable_uc_handles[self.group_rank], but that FD is never imported (due to if p != self.group_rank) and also never closed.

You already close imported POSIX FDs after cuMemImportFromShareableHandle and close the multicast FD after import; closing the remaining two FDs will complete the cleanup and prevent per‑allocation FD leaks in long‑running jobs.

One way to fix this:

        if (
            self._shareable_handle_type
            == cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC
        ):
            # All-gather fabric handles
            all_shareable_uc_handles = self.comm_backend.allgather(
                local_shareable_uc_handle.data
            )
        else:
            # Implement the allgather logic with ipc socket
            all_shareable_uc_handles = [None] * self.group_size
            for i in range(self.group_size):
                self.comm_backend.barrier()
                # Send to peer at offset i
                dest_rank = (self.group_rank + i) % self.group_size
                self._ipc_socket.send_fd(local_shareable_uc_handle, dest_rank)
                # Receive from peer at offset -i
                src_rank = (self.group_rank + self.group_size - i) % self.group_size
                all_shareable_uc_handles[src_rank] = self._ipc_socket.recv_fd()
+           if (
+               self._shareable_handle_type
+               == cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR
+           ):
+               # Close our exported UC handle FD and the self-received FD
+               os.close(local_shareable_uc_handle)
+               if all_shareable_uc_handles[self.group_rank] is not None:
+                   os.close(all_shareable_uc_handles[self.group_rank])

The existing per‑peer close after import:

if self._shareable_handle_type == CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR:
    os.close(all_shareable_uc_handles[p])

and the multicast close:

if self._shareable_handle_type == CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR:
    os.close(shareable_mc_handle)

can stay as‑is.

Together with the new __del__ logic calling self._ipc_socket.close(), this fully addresses the descriptor‑leak concern in the IPC path.

Also applies to: 957-1005, 1008-1055

tests/comm/test_trtllm_mnnvl_allreduce.py (1)

233-271: Align reference RMSNorm epsilon with kernel default (still using torch.finfo(dtype).eps)

prepare_test_data still uses torch.finfo(dtype).eps as the epsilon for the reference RMSNorm:

norm_out = rmsnorm(
    residual_out, norm_weight, torch.finfo(dtype).eps, enable_pdl=False
)

while the kernel and test harness default to eps = 1e-5 (see run_mnnvl_ar_full and the C++ FFI wrapper’s params.epsilon default). This inconsistency can mask subtle discrepancies behind loose tolerances or cause avoidable test drift.

To keep the reference path exactly aligned with the implementation, switch this to the same constant:

-        norm_out = rmsnorm(
-            residual_out, norm_weight, torch.finfo(dtype).eps, enable_pdl=False
-        )
+        norm_out = rmsnorm(
+            residual_out,
+            norm_weight,
+            1e-5,
+            enable_pdl=False,
+        )

(or better, reuse the same eps value passed into run_mnnvl_ar_full to avoid hard‑coding the constant twice).

csrc/trtllm_mnnvl_allreduce.cu (1)

100-114: Ensure epsilon defaults stay consistent with Python API and tests

Here params.epsilon falls back to 1e-5 when the Optional epsilon is not provided:

params.epsilon = epsilon.has_value() ? epsilon.value() : 1e-5;

The Python wrapper in flashinfer/comm/trtllm_mnnvl_ar.py and the tests in tests/comm/test_trtllm_mnnvl_allreduce.py should use the same default to avoid silent discrepancies between the kernel and reference paths. The core test harness already sets eps = 1e-5; the remaining mismatch is in the reference RMSNorm computation (see prepare_test_data), which still uses torch.finfo(dtype).eps.

flashinfer/comm/mnnvl.py (3)

132-150: Fix alloc_and_copy_to_cuda return type vs None behavior

alloc_and_copy_to_cuda is annotated as returning int but still returns None for an empty host_ptr_array:

def alloc_and_copy_to_cuda(host_ptr_array: List[int]) -> int:
    if not host_ptr_array:
        return None

Current call sites (signal_pads and uc_ptrs) always pass non‑empty lists, so behavior is correct, but the annotation is misleading and could hide bugs if the function gets reused.

Either make the return type explicit about None or enforce non‑emptiness by raising:

-def alloc_and_copy_to_cuda(host_ptr_array: List[int]) -> int:
-    if not host_ptr_array:
-        return None
+def alloc_and_copy_to_cuda(host_ptr_array: List[int]) -> int:
+    if not host_ptr_array:
+        raise ValueError("host_ptr_array must be non-empty")

(or change the annotation to Optional[int] if you prefer the sentinel behavior).

---

885-893: IPC opId bootstrap looks fine; consider documenting ordering guarantees

_init_ipc_socket uses an MPI‑like bcast to distribute a randomly chosen opId from rank 0, then uses it to construct IpcSocket endpoints on all ranks. This nicely avoids hard‑coding operation IDs and lines up with the C++ IPC model.

Given the reliance on collective barriers around send_fd/recv_fd, it would help future maintainers to mention in a comment here that all ranks are expected to participate in the same sequence of IPC operations for a given opId, and that mismatched usage will deadlock. The code is correct as written; this is just a documentation/clarity suggestion.

---

1143-1170: McastGPUBuffer workspace integration and pointer getters look consistent

The new comm_backend_for_handle_transfer parameter is threaded through to McastDeviceMemory, and the added get_unicast_ptr wrapper simply delegates to mcast_device_memory.get_unicast_ptr(rank). This lines up with how tests and get_allreduce_mnnvl_workspace use these pointers and keeps pointer access encapsulated.

The placeholder buffer‑view methods (get_multicast_buffer, get_unicast_buffer) are clearly marked NotImplementedError, so they won’t be hit accidentally. If you plan to expose tensor views later, you can implement them via create_tensor_from_cuda_memory.

Also applies to: 1209-1212

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: raayandhar
Repo: flashinfer-ai/flashinfer PR: 2070
File: include/flashinfer/gemm/bf16_gemm_cutlass_template.h:145-160
Timestamp: 2025-11-12T03:35:17.583Z
Learning: In flashinfer GEMM implementations (e.g., include/flashinfer/gemm/bf16_gemm_cutlass_template.h, fp8_gemm_cutlass_template.h), it is acceptable to catch and silently ignore std::runtime_error exceptions in getWorkspaceSizeImpl when probing multiple GEMM configurations, as some configurations may legitimately fail due to SMEM constraints. This pattern should include a comment like "// Swallow errors when SMEM exceeds maximum allowed" to document the rationale.

flashinfer/comm/trtllm_mnnvl_ar.py (1)

361-362: Critical: Restore epsilon default to 1e-5 to match kernel.

This epsilon fallback was flagged as critical in a previous review but remains unresolved. Using torch.finfo(input.dtype).eps sets epsilon to approximately 1e-3 for fp16, diverging from the kernel's built-in 1e-5 default (see csrc/trtllm_mnnvl_allreduce.cu line 96). This materially alters RMSNorm results and breaks compatibility with TensorRT-LLM.

Apply this fix:

-    if epsilon is None:
-        epsilon = torch.finfo(input.dtype).eps
+    if epsilon is None:
+        epsilon = 1e-5

flashinfer/comm/trtllm_mnnvl_ar.py (1)

118-136: Consider replacing @functools.cache on instance method.

Using @functools.cache on an instance method can prevent the instance from being garbage collected, leading to memory leaks. Since this method takes self as the first parameter, the cache will hold references to the instance.

Consider either:

1. Making this a standalone function that takes workspace parameters explicitly
2. Using @functools.lru_cache(maxsize=...) with a reasonable limit
3. Implementing manual caching in the instance if needed

Based on learnings

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: raayandhar
Repo: flashinfer-ai/flashinfer PR: 2070
File: include/flashinfer/gemm/bf16_gemm_cutlass_template.h:145-160
Timestamp: 2025-11-12T03:35:17.583Z
Learning: In flashinfer GEMM implementations (e.g., include/flashinfer/gemm/bf16_gemm_cutlass_template.h, fp8_gemm_cutlass_template.h), it is acceptable to catch and silently ignore std::runtime_error exceptions in getWorkspaceSizeImpl when probing multiple GEMM configurations, as some configurations may legitimately fail due to SMEM constraints. This pattern should include a comment like "// Swallow errors when SMEM exceeds maximum allowed" to document the rationale.

flashinfer/comm/trtllm_mnnvl_ar.py (1)

377-378: Restore RMSNorm epsilon default to 1e-5.

Overriding epsilon with torch.finfo(input.dtype).eps replaces the kernel's built-in 1e-5 default (see trtllm_mnnvl_allreduce_fusion in csrc/trtllm_mnnvl_allreduce.cu line ~35: params.epsilon = epsilon.has_value() ? epsilon.value() : 1e-5). For fp16 this becomes ~1e-3, materially changing RMSNorm results and breaking numerical parity.

Apply this diff to fix:

     if epsilon is None:
-        epsilon = torch.finfo(input.dtype).eps
+        epsilon = 1e-5

flashinfer/comm/trtllm_mnnvl_ar.py (1)

134-152: Consider alternatives to @functools.cache on instance methods.

Using @functools.cache (or @lru_cache) on instance methods can prevent garbage collection of instances because the cache holds references to bound methods, which in turn hold references to self. Since MNNVLAllreduceFusionWorkspace instances are likely long-lived in typical usage, this may be acceptable, but consider these alternatives:

• Use @functools.lru_cache(maxsize=128) to limit cache growth
• Move caching logic to a module-level cache keyed on relevant parameters
• Document the caching behavior and its memory implications

Based on learnings

Apply this diff if you want to limit cache size:

-    @functools.cache
+    @functools.lru_cache(maxsize=128)
     def is_buffer_size_sufficient(

flashinfer/comm/mnnvl.py (2)

640-654: Prefix unused unpacked variables with underscore.

The variables msg, flags, and addr from recvmsg are unpacked but never used. Prefix them with _ to indicate they're intentionally ignored.

Apply this diff:

-        msg, ancdata, flags, addr = self.sock.recvmsg(
+        _msg, ancdata, _flags, _addr = self.sock.recvmsg(

---

893-900: Consider using secrets module for opId generation.

While cryptographic randomness is not strictly required for socket naming, using secrets.randbelow(2**64) instead of random.randint provides better collision resistance if multiple jobs run concurrently on the same node.

Apply this diff:

+import secrets
+
     def _init_ipc_socket(self):
         if self.group_rank == 0:
-            # Gnerate the opId
-            opId = random.randint(0, 2**64 - 1)
+            # Generate the opId
+            opId = secrets.randbelow(2**64)

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: raayandhar
Repo: flashinfer-ai/flashinfer PR: 2070
File: include/flashinfer/gemm/bf16_gemm_cutlass_template.h:145-160
Timestamp: 2025-11-12T03:35:17.583Z
Learning: In flashinfer GEMM implementations (e.g., include/flashinfer/gemm/bf16_gemm_cutlass_template.h, fp8_gemm_cutlass_template.h), it is acceptable to catch and silently ignore std::runtime_error exceptions in getWorkspaceSizeImpl when probing multiple GEMM configurations, as some configurations may legitimately fail due to SMEM constraints. This pattern should include a comment like "// Swallow errors when SMEM exceeds maximum allowed" to document the rationale.