feat(plugin/cuda): add missing UB init ops and fix CommOverlapP2P type conversion

transformer_engine/plugin/core/backends/vendor/cuda/cuda.py

@@ -1911,6 +1911,19 @@ def multi_tensor_compute_scale_inv_e8m0(
             chunk_size, noop_flag, tensor_lists, block_len
         )
 
+    # Device capability queries
+    def device_supports_multicast(self) -> bool:
+        tex = self._get_tex()
+        return tex.device_supports_multicast()
+
+    def ubuf_built_with_mpi(self) -> bool:
+        tex = self._get_tex()
+        return tex.ubuf_built_with_mpi()
+
+    def get_stream_priority_range(self) -> tuple:
+        tex = self._get_tex()
+        return tex.get_stream_priority_range()
+
     # Comm+GEMM Overlap
     def bulk_overlap_ag_with_external_gemm(
         self,
@@ -1992,6 +2005,7 @@ def create_comm_overlap_p2p(
         aggregate: bool = False,
     ) -> "CommOverlapP2P":
         tex = self._get_tex()
+        comm_type = tex.CommOverlapType(int(comm_type)) if comm_type is not None else None
         return tex.CommOverlapP2P(
             buffer_shape,
             buffer_dtype,

transformer_engine/plugin/core/backends/vendor/cuda/register_ops.py

@@ -1110,6 +1110,30 @@ def register_builtins(registry) -> None:
             priority=100,
         ),
         # Communication overlap operations
+        OpImpl(
+            op_name="device_supports_multicast",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.device_supports_multicast, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
+        OpImpl(
+            op_name="ubuf_built_with_mpi",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.ubuf_built_with_mpi, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
+        OpImpl(
+            op_name="get_stream_priority_range",
+            impl_id="vendor.cuda",
+            kind=BackendImplKind.VENDOR,
+            fn=_bind_is_available(backend.get_stream_priority_range, is_avail),
+            vendor="CUDA",
+            priority=100,
+        ),
         OpImpl(
             op_name="bulk_overlap_ag_with_external_gemm",
             impl_id="vendor.cuda",

transformer_engine/plugin/core/ops.py

@@ -1714,6 +1714,19 @@ def multi_tensor_compute_scale_inv_e8m0(
     ) -> None:
         raise NotImplementedError
 
+    # Device capability queries
+    def device_supports_multicast(self) -> bool:
+        """Check if current CUDA device supports multicast for comm+GEMM overlap."""
+        raise NotImplementedError
+
+    def ubuf_built_with_mpi(self) -> bool:
+        """Check if Userbuffers was built with MPI support."""
+        raise NotImplementedError
+
+    def get_stream_priority_range(self) -> tuple:
+        """Get the range of stream priorities (min_priority, max_priority)."""
+        raise NotImplementedError
+
     # Comm+GEMM Overlap
     def bulk_overlap_ag_with_external_gemm(
         self,