Add profiler for flux tp and cp example

docs/PROFILER.md

@@ -1,7 +1,6 @@
 # Torch Profiler Usage
 
-Reference: Adapted from `sglang/python/sglang/bench_one_batch.py` and `sglang/python/sglang/srt/managers/scheduler_profiler_mixin.py`
-
+Reference: Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/bench_one_batch.py .
 ## Quick Start
 
 ### Basic Usage

examples/parallelism/run_flux_cp.py

@@ -17,6 +17,7 @@
     maybe_init_distributed,
     maybe_destroy_distributed,
     MemoryTracker,
+    create_profiler_from_args,
 )
 import cache_dit
 
@@ -69,11 +70,14 @@
 
 
 def run_pipe(pipe: FluxPipeline):
+    steps = 28 if args.steps is None else args.steps
+    if args.profile and args.steps is None:
+        steps = 3
     image = pipe(
         prompt,
         height=height,
         width=width,
-        num_inference_steps=28 if args.steps is None else args.steps,
+        num_inference_steps=steps,
         generator=torch.Generator("cpu").manual_seed(0),
     ).images[0]
     return image
@@ -91,7 +95,14 @@ def run_pipe(pipe: FluxPipeline):
     memory_tracker.__enter__()
 
 start = time.time()
-image = run_pipe(pipe)
+if args.profile:
+    profiler = create_profiler_from_args(args, profile_name="flux_cp_inference")
+    with profiler:
+        image = run_pipe(pipe)
+    if rank == 0:
+        print(f"Profiler traces saved to: {profiler.output_dir}/{profiler.trace_path.name}")
+else:
+    image = run_pipe(pipe)
 end = time.time()
 
 if memory_tracker:

examples/parallelism/run_flux_tp.py

@@ -16,6 +16,7 @@
     maybe_init_distributed,
     maybe_destroy_distributed,
     MemoryTracker,
+    create_profiler_from_args,
 )
 import cache_dit
 
@@ -64,11 +65,14 @@
 
 
 def run_pipe(warmup: bool = False):
+    steps = 5 if warmup else (28 if args.steps is None else args.steps)
+    if args.profile and args.steps is None and not warmup:
+        steps = 3
     image = pipe(
         prompt,
         height=1024 if args.height is None else args.height,
         width=1024 if args.width is None else args.width,
-        num_inference_steps=5 if warmup else (28 if args.steps is None else args.steps),
+        num_inference_steps=steps,
         generator=torch.Generator("cpu").manual_seed(0),
     ).images[0]
     return image
@@ -86,7 +90,14 @@ def run_pipe(warmup: bool = False):
     memory_tracker.__enter__()
 
 start = time.time()
-image = run_pipe()
+if args.profile:
+    profiler = create_profiler_from_args(args, profile_name="flux_tp_inference")
+    with profiler:
+        image = run_pipe()
+    if rank == 0:
+        print(f"Profiler traces saved to: {profiler.output_dir}/{profiler.trace_path.name}")
+else:
+    image = run_pipe()
 end = time.time()
 
 if memory_tracker:

src/cache_dit/profiler.py

@@ -1,8 +1,7 @@
 """
 Torch Profiler for cache-dit.
 
-Reference: Adapted from sglang/python/sglang/bench_one_batch.py and
-           sglang/python/sglang/srt/managers/scheduler_profiler_mixin.py
+Reference: Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/bench_one_batch.py
 """
 
 import logging