add rdma_comm_bytes log

Author: WT1W

magi_attention/functional/dist_attn.py

@@ -1413,8 +1413,10 @@ def _fetch_remote_kv(
             * num_tensors,
             input=local_kv if self.concat_kv else local_kv[0],
         )
-        internode_output_seqlen: int = group_cast_args.get(
-            "internode_output_seqlen", -1
+        internode_output_seqlen: int = group_cast_args.get("internode_output_seqlen", 0)
+        group_cast_kv_rdma_bytes = self.compute_group_comm_bytes(
+            comm_tokens=internode_output_seqlen * num_tensors,
+            input=local_kv if self.concat_kv else local_kv[0],
         )
 
         with nvtx.add_nvtx_event(
@@ -1426,7 +1428,7 @@ def _fetch_remote_kv(
                 f"{output_kv_shape=} | "
                 f"{output_kv_dtype=} |"
                 f"{num_tensors=} | "
-                f"{internode_output_seqlen=}"
+                f"{group_cast_kv_rdma_bytes=}"
             )
         ):
             # launch group cast kernel
@@ -1498,8 +1500,10 @@ def _fetch_remote_q(
             ].group_cast_comm_tokens,
             input=local_q,
         )
-        internode_output_seqlen: int = group_cast_args.get(
-            "internode_output_seqlen", -1
+        internode_output_seqlen: int = group_cast_args.get("internode_output_seqlen", 0)
+        group_cast_q_rdma_bytes = self.compute_group_comm_bytes(
+            comm_tokens=internode_output_seqlen,
+            input=local_q,
         )
 
         with nvtx.add_nvtx_event(
@@ -1511,7 +1515,7 @@ def _fetch_remote_q(
                 f"output_q.shape={remote_q_buffer.shape} | "
                 f"output_q.dtype={remote_q_buffer.dtype} | "
                 f"num_tensors=1 | "
-                f"{internode_output_seqlen=}"
+                f"{group_cast_q_rdma_bytes=}"
             )
         ):
             # launch group cast kernel
@@ -1627,8 +1631,20 @@ def _fetch_remote_qo_do_lse(
             )
 
             group_cast_qo_do_lse_bytes = group_cast_qo_do_bytes + group_cast_lse_bytes
+
             internode_output_seqlen: int = group_cast_args.get(
-                "internode_output_seqlen", -1
+                "internode_output_seqlen", 0
+            )
+            group_cast_qo_do_rdma_bytes = self.compute_group_comm_bytes(
+                comm_tokens=internode_output_seqlen * 3,
+                input=local_qo_do[0],
+            )
+            group_cast_lse_rdma_bytes = self.compute_group_comm_bytes(
+                comm_tokens=internode_output_seqlen,
+                input=local_lse,
+            )
+            group_cast_qo_do_lse_rdma_bytes = (
+                group_cast_qo_do_rdma_bytes + group_cast_lse_rdma_bytes
             )
 
             with nvtx.add_nvtx_event(
@@ -1645,7 +1661,7 @@ def _fetch_remote_qo_do_lse(
                     f"output_lse_shape={remote_lse_buffer.shape} | "
                     f"output_lse_dtype={remote_lse_buffer.dtype} | "
                     f"num_tensors_lse=1 | "
-                    f"{internode_output_seqlen=}"
+                    f"{group_cast_qo_do_lse_rdma_bytes=}"
                 )
             ):
                 # launch group cast kernel
@@ -1696,9 +1712,6 @@ def _fetch_remote_qo_do_lse(
                 ].group_cast_comm_tokens,
                 lse=local_lse,
             )
-            internode_output_seqlen_lse: int = group_cast_args_lse.get(
-                "internode_output_seqlen", -1
-            )
 
             with nvtx.add_nvtx_event(
                 (
@@ -1708,8 +1721,7 @@ def _fetch_remote_qo_do_lse(
                     f"input_lse.dtype={local_lse.dtype} | "
                     f"output_lse.shape={remote_lse_buffer.shape} | "
                     f"output_lse.dtype={remote_lse_buffer.dtype} | "
-                    f"num_tensors=1 | "
-                    f"{internode_output_seqlen_lse=}"
+                    f"num_tensors=1"
                 )
             ):
                 # launch group cast kernel for lse
@@ -1746,9 +1758,6 @@ def _fetch_remote_qo_do_lse(
                 ].group_cast_comm_tokens,
                 input=local_qo_do,
             )
-            internode_output_seqlen_qo_do: int = group_cast_args_qo_do.get(
-                "internode_output_seqlen", -1
-            )
 
             with nvtx.add_nvtx_event(
                 (
@@ -1758,8 +1767,7 @@ def _fetch_remote_qo_do_lse(
                     f"input_qo_do.dtype={local_qo_do.dtype} | "
                     f"output_qo_do.shape={remote_qo_do_buffer.shape} | "  # type: ignore
                     f"output_qo_do.dtype={remote_qo_do_buffer.dtype} | "  # type: ignore
-                    f"num_tensors=1 | "
-                    f"{internode_output_seqlen_qo_do=}"
+                    f"num_tensors=1"
                 )
             ):
                 # launch group cast kernel for qo_do
@@ -1876,7 +1884,12 @@ def _reduce_partial_out_lse(
                 lse=partial_remote_lse,
             )
             internode_output_seqlen: int = group_reduce_args.get(
-                "internode_output_seqlen", -1
+                "internode_output_seqlen", 0
+            )
+            group_cast_out_lse_rdma_bytes = self.compute_group_comm_bytes(
+                comm_tokens=internode_output_seqlen,
+                input=partial_remote_out,
+                lse=partial_remote_lse,
             )
 
             with nvtx.add_nvtx_event(
@@ -1891,7 +1904,7 @@ def _reduce_partial_out_lse(
                     f"input_lse.dtype={partial_remote_lse.dtype} | "
                     f"output_lse.shape={partial_local_lse.shape} | "
                     f"output_lse.dtype={partial_local_lse.dtype} | "
-                    f"{internode_output_seqlen=}"
+                    f"{group_cast_out_lse_rdma_bytes=}"
                 )
             ):
                 # launch group-reduce kernel
@@ -2022,7 +2035,11 @@ def _reduce_partial_dkv(
             input=partial_remote_dkv if self.concat_dkv else partial_remote_dkv[0],  # type: ignore
         )
         internode_output_seqlen: int = group_reduce_args.get(
-            "internode_output_seqlen", -1
+            "internode_output_seqlen", 0
+        )
+        group_cast_dkv_rdma_bytes = self.compute_group_comm_bytes(
+            comm_tokens=internode_output_seqlen * num_tensors_of_dkv,
+            input=partial_remote_dkv if self.concat_dkv else partial_remote_dkv[0],  # type: ignore
         )
         with nvtx.add_nvtx_event(
             (
@@ -2033,7 +2050,7 @@ def _reduce_partial_dkv(
                 f"{output_dkv_shape=} | "
                 f"{output_dkv_dtype=} | "
                 f"{num_tensors_of_dkv=} | "
-                f"{internode_output_seqlen=}"
+                f"{group_cast_dkv_rdma_bytes=}"
             )
         ):
             # launch group-reduce kernel
@@ -2116,7 +2133,11 @@ def _reduce_partial_dq(
                 input=partial_remote_dq,
             )
             internode_output_seqlen: int = group_reduce_args.get(
-                "internode_output_seqlen", -1
+                "internode_output_seqlen", 0
+            )
+            group_cast_dq_rdma_bytes = self.compute_group_comm_bytes(
+                comm_tokens=internode_output_seqlen,
+                input=partial_remote_dq,
             )
 
             with nvtx.add_nvtx_event(
@@ -2128,7 +2149,7 @@ def _reduce_partial_dq(
                     f"output_dq.shape={partial_local_dq.shape} | "
                     f"output_dq.dtype={partial_local_dq.dtype} | "
                     f"tensors_num_of_dq=1 | "
-                    f"{internode_output_seqlen=}"
+                    f"{group_cast_dq_rdma_bytes=}"
                 )
             ):
                 # launch group-reduce kernel