add gpu timer decorator (#207)

xfuser/core/__init__.py

@@ -4,10 +4,12 @@
     xFuserJointLongContextAttention,
     xFuserFluxLongContextAttention,
 )
+from .utils import gpu_timer_decorator
 
 __all__ = [
     "CacheManager",
     "xFuserLongContextAttention",
     "xFuserJointLongContextAttention",
     "xFuserFluxLongContextAttention",
+    "gpu_timer_decorator",
 ]

xfuser/core/utils/__init__.py

@@ -0,0 +1 @@
+from .timer import gpu_timer_decorator

xfuser/core/utils/timer.py

@@ -0,0 +1,19 @@
+import torch
+import time
+
+
+def gpu_timer_decorator(func):
+    def wrapper(*args, **kwargs):
+        torch.cuda.synchronize()
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        torch.cuda.synchronize()
+        end_time = time.time()
+
+        if torch.distributed.get_rank() == 0:
+            print(
+                f"{func.__name__} took {end_time - start_time} seconds to run on GPU."
+            )
+        return result
+
+    return wrapper