diff --git a/deepinfer/worker/neuron_graph_runner.py b/deepinfer/worker/neuron_graph_runner.py
index b0055c1..d59c100 100644
--- a/deepinfer/worker/neuron_graph_runner.py
+++ b/deepinfer/worker/neuron_graph_runner.py
@@ -192,7 +192,9 @@ class NeuronGraphRunner(CudaGraphRunner):
         torch.neuron.synchronize()
         self._steps = getattr(self, "_steps", 0) + 1
         profiler = None
-        if self._steps == 8 and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0):
+        self._bucket_steps = getattr(self, "_bucket_steps", 0) + (batch_size > 1)
+        profile = self._steps == 8 or (batch_size > 1 and self._bucket_steps == 2)  # a decode step, the 2nd prefill chunk
+        if profile and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0):
             import cProfile
 
             profiler = cProfile.Profile()
@@ -208,8 +210,8 @@ class NeuronGraphRunner(CudaGraphRunner):
 
             profiler.disable()
             buf = io.StringIO()
-            pstats.Stats(profiler, stream=buf).sort_stats("tottime").print_stats(40)
-            logger.info(f"Neuron step profile (rank 0, step {self._steps}):\n{buf.getvalue()}")
+            pstats.Stats(profiler, stream=buf).sort_stats("cumulative").print_stats(60)
+            logger.info(f"Neuron step profile (rank 0, step {self._steps}, bs={batch_size}):\n{buf.getvalue()}")
         timings = regions.take_timings()
         in_regions = sum(sum(v) for v in timings.values())
         if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0: