flashinfer/comm/trtllm_ar.py (2)

503-506: Return typing is fine; consider a stronger metadata type for clarity.

Union return is OK. To improve API clarity, define a TypedDict for metadata and use it in the signature and docstrings. This avoids untyped dict usage downstream.

-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union, TypedDict
+
+class AllReduceFusionWorkspaceMeta(TypedDict):
+    tp_rank: int
+    tp_size: int
+    max_token_num: int
+    hidden_dim: int
+    use_fp32_lamport: bool
+    buffer_size: int
+    flag_size: int
+    lamport_comm_size: int
+    lamport_buffer_size: int
@@
-) -> Union[
-    Tuple[List[List[int]], torch.Tensor], Tuple[List[List[int]], torch.Tensor, dict]
-]:
+) -> Union[
+    Tuple[List[List[int]], torch.Tensor],
+    Tuple[List[List[int]], torch.Tensor, AllReduceFusionWorkspaceMeta],
+]:

---

621-636: Good: metadata emission. Consider adding device and validating key presence.

Including creation params is useful. Add a 'device' (CUDA index) for extra safety and validate required keys on consumption to avoid KeyError if a caller passes a partial dict.

         metadata = {
             "tp_rank": tp_rank,
             "tp_size": tp_size,
             "max_token_num": max_token_num,
             "hidden_dim": hidden_dim,
             "use_fp32_lamport": use_fp32_lamport,
+            "device": torch.cuda.current_device(),
             "buffer_size": buffer_size,
             "flag_size": flag_size,
             "lamport_comm_size": lamport_comm_size,
             "lamport_buffer_size": lamport_buffer_size,
         }

tests/comm/test_trtllm_allreduce_fusion.py (1)

60-71: Good: create workspace with metadata; add quick negative-path checks here.

Leverage the new validation with a small trio of negative assertions right after creating metadata to guard against regressions (no heavy kernels executed).

         ipc_handles, workspace_tensor, workspace_metadata = (
             comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(
                 rank,
                 world_size,
                 MAX_TOKEN_NUM,
                 hidden_dim,
                 group=group,
                 use_fp32_lamport=lamport_use_fp32,
                 create_metadata=True,  # Get metadata for validation
             )
         )
+
+        # Negative-path validation: token_num overflow
+        bad_token_num = MAX_TOKEN_NUM + 1
+        with pytest.raises(ValueError, match="token_num .* exceeds"):
+            comm.trtllm_allreduce_fusion(
+                allreduce_in=torch.empty(bad_token_num * hidden_dim, dtype=dtype, device=device),
+                world_size=world_size,
+                world_rank=rank,
+                token_num=bad_token_num,
+                hidden_dim=hidden_dim,
+                workspace_ptrs=workspace_tensor,
+                launch_with_pdl=True,
+                trigger_completion_at_end=True,
+                fp32_acc=False,
+                pattern_code=comm.AllReduceFusionPattern.kAllReduce,
+                use_oneshot=None,
+                allreduce_out=None,
+                residual_in=None,
+                residual_out=None,
+                norm_out=None,
+                quant_out=None,
+                scale_out=None,
+                rms_gamma=None,
+                rms_eps=None,
+                scale_factor=None,
+                layout_code=None,
+                metadata=workspace_metadata,
+            )
+
+        # Negative-path validation: world_size mismatch
+        meta_ws_mismatch = dict(workspace_metadata)
+        meta_ws_mismatch["tp_size"] = world_size + 1
+        with pytest.raises(ValueError, match="world_size .* does not match"):
+            comm.trtllm_allreduce_fusion(
+                allreduce_in=torch.empty(1 * hidden_dim, dtype=dtype, device=device),
+                world_size=world_size,
+                world_rank=rank,
+                token_num=1,
+                hidden_dim=hidden_dim,
+                workspace_ptrs=workspace_tensor,
+                launch_with_pdl=True,
+                trigger_completion_at_end=True,
+                fp32_acc=False,
+                pattern_code=comm.AllReduceFusionPattern.kAllReduce,
+                use_oneshot=None,
+                allreduce_out=None,
+                residual_in=None,
+                residual_out=None,
+                norm_out=None,
+                quant_out=None,
+                scale_out=None,
+                rms_gamma=None,
+                rms_eps=None,
+                scale_factor=None,
+                layout_code=None,
+                metadata=meta_ws_mismatch,
+            )
+
+        # Negative-path validation: hidden_dim mismatch
+        meta_hd_mismatch = dict(workspace_metadata)
+        meta_hd_mismatch["hidden_dim"] = hidden_dim + 64
+        with pytest.raises(ValueError, match="hidden_dim .* does not match"):
+            comm.trtllm_allreduce_fusion(
+                allreduce_in=torch.empty(1 * hidden_dim, dtype=dtype, device=device),
+                world_size=world_size,
+                world_rank=rank,
+                token_num=1,
+                hidden_dim=hidden_dim,
+                workspace_ptrs=workspace_tensor,
+                launch_with_pdl=True,
+                trigger_completion_at_end=True,
+                fp32_acc=False,
+                pattern_code=comm.AllReduceFusionPattern.kAllReduce,
+                use_oneshot=None,
+                allreduce_out=None,
+                residual_in=None,
+                residual_out=None,
+                norm_out=None,
+                quant_out=None,
+                scale_out=None,
+                rms_gamma=None,
+                rms_eps=None,
+                scale_factor=None,
+                layout_code=None,
+                metadata=meta_hd_mismatch,
+            )

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

flashinfer/comm/trtllm_ar.py (3)

510-513: Consider @overload for clearer type signatures.

While the Union return type maintains compatibility, static analysis tools and IDEs struggle with variable-length tuple unions. Using typing.overload would provide explicit signatures for each case, improving type safety and developer experience.

Example:

@overload
def trtllm_create_ipc_workspace_for_all_reduce_fusion(
    tp_rank: int,
    tp_size: int,
    max_token_num: int,
    hidden_dim: int,
    use_fp32_lamport: bool = False,
    group: Optional[ProcessGroup] = None,
    create_metadata: Literal[False] = False,
) -> Tuple[List[List[int]], torch.Tensor]: ...

@overload
def trtllm_create_ipc_workspace_for_all_reduce_fusion(
    tp_rank: int,
    tp_size: int,
    max_token_num: int,
    hidden_dim: int,
    use_fp32_lamport: bool = False,
    group: Optional[ProcessGroup] = None,
    create_metadata: Literal[True] = ...,
) -> Tuple[List[List[int]], torch.Tensor, dict]: ...

---

628-642: Consider TypedDict or dataclass for type-safe metadata.

Raw dicts are vulnerable to typos in both creation and consumption sites. Defining a structured type would enable static type checking, IDE autocompletion, and early error detection.

Example with TypedDict:

from typing import TypedDict

class WorkspaceMetadata(TypedDict):
    tp_rank: int
    tp_size: int
    max_token_num: int
    hidden_dim: int
    use_fp32_lamport: bool
    buffer_size: int
    flag_size: int
    lamport_comm_size: int
    lamport_buffer_size: int

---

913-917: Major: Size calculation should use lamport dtype from metadata.

When metadata is provided, the required lamport size should be calculated using metadata["use_fp32_lamport"] rather than allreduce_in.dtype. The workspace's lamport buffer dtype determines safety, not the input tensor dtype. Using the input dtype can underestimate required bytes and incorrectly enable oneshot mode, leading to buffer overruns.

Apply this diff:

+    # Determine element size: use metadata's lamport dtype when available
+    if metadata is not None:
+        elem_bytes = 4 if metadata["use_fp32_lamport"] else 2
+    else:
+        elem_bytes = 4 if allreduce_in.dtype == torch.float32 else 2
+    
     required_lamport_comm_size = (
-        token_num * hidden_dim * 2 * world_size
-        if allreduce_in.dtype != torch.float32
-        else token_num * hidden_dim * 4 * world_size
+        token_num * hidden_dim * elem_bytes * world_size
     )

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

flashinfer/comm/trtllm_ar.py (1)

893-898: Critical: Inverted logic in use_fp32_lamport validation.

Line 894 uses == which raises an error when metadata and dtype correctly match, allowing mismatches to pass. This is inverted and defeats the validation purpose.

Apply this diff to fix the inverted condition:

-        if metadata["use_fp32_lamport"] == (allreduce_in.dtype == torch.float32):
+        if metadata["use_fp32_lamport"] != (allreduce_in.dtype == torch.float32):

flashinfer/comm/trtllm_ar.py (2)

512-514: Consider typing.overload for better type safety (optional).

As noted in a previous review, using @overload decorators would provide clearer type signatures for different create_metadata values and improve IDE autocompletion. While the current Union approach works, overloads make it explicit what callers receive based on the boolean flag.

Based on learnings

---

630-640: Consider TypedDict or dataclass for metadata (recommended).

As noted in a previous review, defining a structured type for the metadata dict would prevent typos in keys and enable static type checking. This would make the API more robust when metadata is created here and consumed in trtllm_allreduce_fusion.

Based on learnings

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

flashinfer/comm/trtllm_ar.py (1)

859-905: Comprehensive workspace metadata validation addresses issue #1986.

The validation logic correctly prevents illegal memory access by checking token_num against max_token_num and ensuring all workspace parameters match the runtime configuration. The error accumulation pattern provides excellent user feedback by reporting all mismatches in a single error message.

The critical issue from past reviews regarding inverted logic at line 894 has been fixed—the condition now correctly uses != to detect mismatches between use_fp32_lamport and the input dtype.

Optional: Minor Pythonic improvement for key checking.

Lines 863-865 could be slightly more concise using set operations:

-        for key in required_keys:
-            if key not in metadata:
-                errors.append(f"Workspace metadata is missing required key: {key}")
+        missing_keys = set(required_keys) - set(metadata.keys())
+        if missing_keys:
+            errors.append(f"Workspace metadata is missing required keys: {sorted(missing_keys)}")

This is purely stylistic and doesn't affect correctness.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro