NM Profiler : Update visualize_trace.py

[varun-sundar-rabindranath]

Update visualize trace utility.

• Add ability to plot decode steps
• Expose the metric to plot as a command line argument
• Generate 2 plots - 1 for prefill and another for decode-steps and store them in the user given output directory
• Remove the ignore_sampler arg, and instead add a fold_json_node arg - This argument collapses the specified JSON tree so the plot has less clutter.

Usage:

• python3 neuralmagic/tools/profiler/visualize_trace.py --json-trace profiler_fp8_trace.json --output-directory ./kernel --level kernel This command produce 2 output files : kernel/prefill.png and kernel/decode_steps.png which are stacked-bar graph plots. In these plots the operations are grouped together by high-level concepts such as gemms, attention, rms-norm etc.
  

• python3 neuralmagic/tools/profiler/visualize_trace.py --json-trace profiler_fp8_trace.json --output-directory ./module --level module --plot-metric pct_cuda_time This command also produces 2 output files : module/prefill.png and module/decode_steps.png which are stacked-bar graph plots. In these plots the bars sum up to a 100 as the requested plot metric is pct_cuda_time

[LucasWilkinson]

They Y-label in the --plot-metric pct_cuda_time graph appears to be wrong above, but in the code it seems to be getting set correctly? I assume this has been fixed?

[varun-sundar-rabindranath]

They Y-label in the --plot-metric pct_cuda_time graph appears to be wrong above, but in the code it seems to be getting set correctly? I assume this has been fixed?

Hey Lucas. Yes, I noticed that and fixed it. Sorry, should have mentioned it somewhere.

[varun-sundar-rabindranath]

Migrated all changes including all of the layer-by-layer profiling code to https://github.com/neuralmagic/vllm/pull/3