coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-11-09 22:59:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

Neuron: profile a prefill chunk too under DEEPINFER_NEURON_REGION_TIMING

Browse Source 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-23 06:38:31 +08:00
parent 32abff6662
commit a263b60514
1 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
deepinfer/worker/neuron_graph_runner.py
View File

@ -192,7 +192,9 @@ class NeuronGraphRunner(CudaGraphRunner):
torch.neuron.synchronize() torch.neuron.synchronize()
self._steps = getattr(self, "_steps", 0) + 1 self._steps = getattr(self, "_steps", 0) + 1
profiler = None 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 import cProfile
profiler = cProfile.Profile() profiler = cProfile.Profile()
@ -208,8 +210,8 @@ class NeuronGraphRunner(CudaGraphRunner):
profiler.disable() profiler.disable()
buf = io.StringIO() buf = io.StringIO()
pstats.Stats(profiler, stream=buf).sort_stats("tottime").print_stats(40) pstats.Stats(profiler, stream=buf).sort_stats("cumulative").print_stats(60)
logger.info(f"Neuron step profile (rank 0, step {self._steps}):\n{buf.getvalue()}") logger.info(f"Neuron step profile (rank 0, step {self._steps}, bs={batch_size}):\n{buf.getvalue()}")
timings = regions.take_timings() timings = regions.take_timings()
in_regions = sum(sum(v) for v in timings.values()) in_regions = sum(sum(v) for v in timings.values())
if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0: if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0: