refactor

examples/offline_profile.py

@@ -8,7 +8,7 @@
 import torch
 
 from vllm import LLM, SamplingParams
-from vllm.profiler import nm_profile
+from vllm.profiler import layerwise_profile
 
 BATCH_SIZE_DEFAULT = 1
 PROMPT_LEN_DEFAULT = 256
@@ -78,7 +78,7 @@ def run_profile(context: ProfileContext, csv_output: Optional[str],
               f"larger than max_num_batched_tokens ({max_num_batched_tokens}) "
               f"and therefore cannot be run in a single profile step, please "
               f"choose a smaller batch size or prompt length, or increase "
-              f"--max_num_batched_tokens")
+              f"--max-num-batched-tokens")
         sys.exit(-1)
     if batch_size >= max_num_seqs:
         print(
@@ -106,12 +106,12 @@ def run_profile(context: ProfileContext, csv_output: Optional[str],
             inputs={'prompt_token_ids': prompt_token_ids},
             params=sampling_params)
 
-    with nm_profile() as prefill_prof:
+    with layerwise_profile() as prefill_prof:
         llm.llm_engine.step()  # First step is prefill
 
     decode_results_list = []
     for x in range(args.output_len - 1):
-        with nm_profile() as decode_prof:
+        with layerwise_profile() as decode_prof:
             llm.llm_engine.step()
         decode_results_list.append(decode_prof.results)
 

vllm/profiler/__init__.py

@@ -1,5 +1,5 @@
-from .nm_profile import nm_profile
+from .layerwise_profile import layerwise_profile
 
 __all__ = [
-    "nm_profile",
+    "layerwise_profile",
 ]

vllm/profiler/nm_profile.py → vllm/profiler/layerwise_profile.py

@@ -63,7 +63,7 @@ class _StatsTreeNode:
 
 
 @dataclass
-class NMProfileResults(profile):
+class LayerwiseProfileResults(profile):
     _kineto_results: _ProfilerResult
     _kineto_event_correlation_map: Dict[int,
                                         List[_KinetoEvent]] = field(init=False)
@@ -331,7 +331,7 @@ def df_traversal(node: _StatsTreeNode, curr_json_list: List[Dict]):
         return root_dicts
 
 
-class nm_profile(profile):
+class layerwise_profile(profile):
 
     def __init__(self):
         super().__init__(
@@ -346,4 +346,4 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         super().__exit__(exc_type, exc_val, exc_tb)
-        self.results = NMProfileResults(self.profiler.kineto_results)
+        self.results = LayerwiseProfileResults(self.profiler.kineto_results)

neuralmagic/tools/profiler/print_table.py → vllm/profiler/print_layerwise_table.py

@@ -2,7 +2,7 @@
 import json
 from typing import Dict
 
-from vllm.profiler.nm_profile import ModelStatsEntry, SummaryStatsEntry
+from vllm.profiler.layerwise_profile import ModelStatsEntry, SummaryStatsEntry
 from vllm.profiler.utils import TablePrinter, indent_string
 
 

neuralmagic/tools/profiler/visualize_trace.py → vllm/profiler/visualize_layerwise_profile.py

@@ -1,3 +1,4 @@
+import os
 import argparse
 import copy
 import json
@@ -426,6 +427,9 @@ def make_plot_title_suffix(profile_json: dict) -> str:
 
     output_directory = args.output_directory if args.output_directory else Path(
         args.json_trace).parent
+
+    if not os.path.exists(output_directory):
+        os.makedirs(output_directory)
 
     main(Path(args.json_trace), output_directory, depth, args.plot_metric,
          make_names_unique, args.top_k, args.fold_json_node)