| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """This generates gpu kernel analysis output from nsys rep. Will call nsys |
| 4 | stats -r cuda_gpu_kern_trace, get non-overlapped gpu cycles, then generate |
| 5 | csv and html output for analysis |
| 6 | """ |
| 7 | |
| 8 | import argparse |
| 9 | import logging |
| 10 | import os |
| 11 | |
| 12 | import regex as re |
| 13 | |
| 14 | logger = logging.getLogger(__name__) |
| 15 | |
| 16 | |
| 17 | # helper data class for annotating kernels |
| 18 | def load_engine_model(): |
| 19 | """Returns engine_model built from all json files in the current dir.""" |
| 20 | import glob |
| 21 | import json |
| 22 | |
| 23 | engine_model = {} |
| 24 | |
| 25 | json_files = glob.glob(os.path.join(os.path.dirname(__file__) or ".", "*.json")) |
| 26 | for fname in json_files: |
| 27 | with open(fname, encoding="utf-8") as f: |
| 28 | engine_model.update(json.load(f)) |
| 29 | return engine_model |
| 30 | |
| 31 | |
| 32 | class GPUTrace2Graph: |
| 33 | """Parses output of nsys report, generates csv and bar chart output.""" |
| 34 | |
| 35 | def __init__(self): |
| 36 | import pandas as pd # avoid importing till needed |
| 37 | |
| 38 | self.pd = pd |
| 39 | self.pd.options.mode.copy_on_write = True |
| 40 | |
| 41 | # helper functions for generating trace->summary csvs |
| 42 | def gen_nonoverlapped_sum_from_gputrace(self, in_file, out_file): |
| 43 | logger.info("loading %s", in_file) |
| 44 | df = self.pd.read_csv( |
| 45 | in_file, usecols=["Start (ns)", "Duration (ns)", "Device", "Strm", "Name"] |
| 46 | ) |
| 47 | df["End (ns)"] = df["Start (ns)"] + df["Duration (ns)"] |
| 48 | df = self.sum_non_overlapping_intervals(df) |
| 49 | # get ready to print table with elapsed times per kernel |
| 50 | df["Instances"] = 1 |
| 51 | df_sum = df.groupby("Name", as_index=False).agg( |
| 52 | {"Elapsed Time (ns)": "sum", "Duration (ns)": "sum", "Instances": "size"} |
| 53 | ) |
| 54 | |
| 55 | # generate csv |
| 56 | df_sum["Total Time (sec)"] = df_sum["Duration (ns)"] / 1e9 |
| 57 | df_sum["Elapsed Time (sec)"] = df_sum["Elapsed Time (ns)"] / 1e9 |
| 58 | df_sum = df_sum.sort_values(by="Elapsed Time (sec)", ascending=False) |
| 59 | df_sum[["Elapsed Time (sec)", "Total Time (sec)", "Instances", "Name"]].to_csv( |
| 60 | out_file, index=False |
| 61 | ) |
| 62 | |
| 63 | def sum_non_overlapping_intervals(self, df): |
| 64 | """Returns new sorted df with Elapsed Time (ns) column using |
| 65 | vectorized operations |
| 66 | """ |
| 67 | logger.info("sorting %s trace records by start time", str(df.shape)) |
| 68 | |
| 69 | # Sort by start time and reset index |
| 70 | df = df.sort_values(by="Start (ns)").reset_index(drop=True) |
| 71 | |
| 72 | # Initialize elapsed time as duration |
| 73 | df["Elapsed Time (ns)"] = df["Duration (ns)"] |
| 74 | |
| 75 | # Get numpy arrays for faster operations |
| 76 | starts = df["Start (ns)"].values |
| 77 | ends = df["End (ns)"].values |
| 78 | |
| 79 | # Keep track of current interval end |
| 80 | current_end = ends[0] |
| 81 | display_units = int(len(df) / 100) |
| 82 | # Update current_end for overlapping intervals |
| 83 | for i in range(1, len(df)): |
| 84 | if i % display_units == 0: |
| 85 | print(f"processing trace: {int(i / len(df) * 100)} %", end="\r") |
| 86 | if starts[i] <= current_end: |
| 87 | if ends[i] > current_end: |
| 88 | # Partial overlap |
| 89 | df.iloc[i, df.columns.get_loc("Elapsed Time (ns)")] = ( |
| 90 | ends[i] - current_end |
| 91 | ) |
| 92 | current_end = ends[i] |
| 93 | else: |
| 94 | # Complete overlap |
| 95 | df.iloc[i, df.columns.get_loc("Elapsed Time (ns)")] = 0 |
| 96 | else: |
| 97 | # No overlap |
| 98 | current_end = ends[i] |
| 99 | |
| 100 | return df |
| 101 | |
| 102 | # functions for generating html files |
| 103 | def make_html(self, df, output_dir, title): |
| 104 | """Make html graph from df.""" |
| 105 | import plotly.express as px |
| 106 | |
| 107 | if df.empty: |
| 108 | return |
| 109 | output_name = output_dir + "/result" |
| 110 | if not title: |
| 111 | title = "Model_Engine" |
| 112 | x = "Model_Engine" |
| 113 | y = "Elapsed Time (sec)" |
| 114 | color = "Category" |
| 115 | """ generate kernel mapping table """ |
| 116 | # Sort Model_Engine categories by last field after underscore |
| 117 | df["Model_Engine"] = self.pd.Categorical( |
| 118 | df["Model_Engine"], |
| 119 | sorted(df["Model_Engine"].unique(), key=lambda x: x.split("_")[-1]), |
| 120 | ) |
| 121 | df[["Model_Engine", color, "Instances", "Name", y]].sort_values( |
| 122 | by=color |
| 123 | ).to_csv(f"{output_name}.csv", index=False) |
| 124 | graph = px.histogram( |
| 125 | df.round(2), |
| 126 | x=x, |
| 127 | y=y, |
| 128 | title=(f"{y} for {title}"), |
| 129 | color=color, |
| 130 | text_auto=True, |
| 131 | ) |
| 132 | # wrap x axis labels |
| 133 | graph.update_xaxes(automargin=True) |
| 134 | graph.write_html(f"{output_name}.html") |
| 135 | """ |
| 136 | Generate data table with columns per Model_Engine into result.html |
| 137 | """ |
| 138 | pivot_df = df.pivot_table( |
| 139 | values="Elapsed Time (sec)", |
| 140 | index="Category", |
| 141 | columns="Model_Engine", |
| 142 | aggfunc="sum", |
| 143 | observed=False, |
| 144 | ).round(2) |
| 145 | # Add sum row at bottom |
| 146 | pivot_df.loc["total_elapsed_sec"] = pivot_df.sum() |
| 147 | pivot_df.fillna("").to_html("temp.html") |
| 148 | with ( |
| 149 | open(f"{output_name}.html", "a", encoding="utf-8") as outfile, |
| 150 | open("temp.html", encoding="utf-8") as infile, |
| 151 | ): |
| 152 | outfile.write(infile.read()) |
| 153 | os.remove("temp.html") |
| 154 | |
| 155 | print( |
| 156 | f"Finished generating: \n" |
| 157 | f" {output_name}.html for stack bar chart \n" |
| 158 | f" {output_name}.csv for Kernel-Category mapping" |
| 159 | ) |
| 160 | |
| 161 | def anno_gpu_kernname(self, df, mapping): |
| 162 | """Add "Category" column.""" |
| 163 | |
| 164 | def anno_gpu_kernname_helper(name): |
| 165 | for kern_name, val in mapping.items(): |
| 166 | if re.search(kern_name, name): |
| 167 | return val |
| 168 | |
| 169 | df["Category"] = df["Name"].apply(anno_gpu_kernname_helper) |
| 170 | |
| 171 | def make_nongpu_row(self, df, nongpu_sec): |
| 172 | """This will append non-gpu time entry at end of df.""" |
| 173 | nongpu_row = self.pd.DataFrame([df.iloc[-1]]) |
| 174 | nongpu_row["Category"] = nongpu_row["Name"] = "CPU(non-GPU)" |
| 175 | nongpu_row["Instances"] = 1 |
| 176 | nongpu_row["Elapsed Time (sec)"] = nongpu_sec |
| 177 | return nongpu_row |
| 178 | |
| 179 | def is_valid_file(self, base_file): |
| 180 | """Asserts if base_file is non-existent or is empty.""" |
| 181 | assert os.path.isfile(base_file) and os.path.getsize(base_file) > 0, ( |
| 182 | f"{base_file} doesn't exist or is empty" |
| 183 | ) |
| 184 | |
| 185 | def should_gen_file(self, new_file, base_file): |
| 186 | """Figure out if new file should be generated from base_file.""" |
| 187 | self.is_valid_file(base_file) |
| 188 | if ( |
| 189 | os.path.exists(new_file) |
| 190 | and (os.path.getmtime(new_file) > os.path.getmtime(base_file)) |
| 191 | and (os.path.getsize(base_file) > 0) |
| 192 | ): |
| 193 | logger.info("reusing %s", new_file) |
| 194 | return False |
| 195 | else: |
| 196 | logger.info("generating %s", new_file) |
| 197 | return True |
| 198 | |
| 199 | def gen_sum_file(self, file, nsys_cmd): |
| 200 | """Generates sum file from nsys trace with times per kernel and |
| 201 | returns the name of the sum file |
| 202 | """ |
| 203 | import subprocess |
| 204 | |
| 205 | file_dir = os.path.dirname(file) |
| 206 | file_name = os.path.basename(file) |
| 207 | |
| 208 | if not file_dir: |
| 209 | file_dir = "." |
| 210 | # Walk through trace and get the total non-overlapped time |
| 211 | nsys_stats_file = f"{file_dir}/{file_name}_cuda_gpu_trace.csv" |
| 212 | sum_file = f"{file_dir}/{file_name}_cuda_gpu_kernel_tracesum.csv" |
| 213 | if self.should_gen_file(nsys_stats_file, file): |
| 214 | cmd = [ |
| 215 | nsys_cmd, |
| 216 | "stats", |
| 217 | "-r", |
| 218 | "cuda_gpu_trace", |
| 219 | file, |
| 220 | "-o", |
| 221 | f"{file_dir}/{file_name}", |
| 222 | ] |
| 223 | cmd_str = " ".join(cmd) |
| 224 | logger.info("+ %s", cmd_str) |
| 225 | # estimate time based on calibrated 240M/min |
| 226 | file_size_mb = os.path.getsize(file) / 1e6 |
| 227 | logger.info( |
| 228 | "nsys stats for %.2f MB file expected to take %.2f min", |
| 229 | file_size_mb, |
| 230 | file_size_mb / 240, |
| 231 | ) |
| 232 | try: |
| 233 | subprocess.run(cmd, check=True) |
| 234 | except Exception: |
| 235 | logger.error("%s failed; Use --nsys_cmd to specify nsys path", cmd_str) |
| 236 | exit(1) |
| 237 | logger.info("generating non-overalapped sum %s", sum_file) |
| 238 | self.gen_nonoverlapped_sum_from_gputrace(nsys_stats_file, sum_file) |
| 239 | self.is_valid_file(sum_file) |
| 240 | logger.info("Finished generating %s", sum_file) |
| 241 | return sum_file |
| 242 | |
| 243 | def gen_graph(self, in_file, out_dir, title, nsys_cmd, engine_model): |
| 244 | """Generates graph and csv file from in_file into out_dir.""" |
| 245 | # Initialize an empty DataFrame to store combined data |
| 246 | combined_df = self.pd.DataFrame() |
| 247 | for idx, (file, engine, model, total_sec) in enumerate(in_file): |
| 248 | file_dir = os.path.dirname(file) |
| 249 | file_name = os.path.basename(file) |
| 250 | if not file_dir: |
| 251 | file_dir = "." |
| 252 | sum_file = self.gen_sum_file(file, nsys_cmd) |
| 253 | # read kernel summary file |
| 254 | df = self.pd.read_csv(sum_file) |
| 255 | # annotate kernel to their categories |
| 256 | assert engine_model.get(engine), f"engine {engine} unknown" |
| 257 | assert engine_model[engine].get(model), f"model {model} unknown" |
| 258 | # remove nsys-rep from file_name for shorter x-label |
| 259 | file_name = file_name.replace(".nsys-rep", "") |
| 260 | df["Model_Engine"] = f"{model}_{engine}_{file_name}_{idx}" |
| 261 | self.anno_gpu_kernname(df, engine_model[engine][model]) |
| 262 | # patch in non-gpu time |
| 263 | gpu_sec = round(df["Elapsed Time (sec)"].sum(), 1) |
| 264 | total_sec = round(float(total_sec), 1) |
| 265 | if total_sec < gpu_sec: |
| 266 | logger.warning( |
| 267 | "Elapsed sec %.2f < GPU sec %.2f resetting Elapsed sec ", |
| 268 | total_sec, |
| 269 | gpu_sec, |
| 270 | ) |
| 271 | total_sec = gpu_sec |
| 272 | nongpu_row = self.make_nongpu_row(df, total_sec - gpu_sec) |
| 273 | df = self.pd.concat([df, nongpu_row], ignore_index=True) |
| 274 | combined_df = self.pd.concat([combined_df, df], ignore_index=True) |
| 275 | if out_dir is None: |
| 276 | out_dir = "." |
| 277 | else: |
| 278 | os.makedirs(out_dir, exist_ok=True) |
| 279 | # generate html file |
| 280 | self.make_html(combined_df, out_dir, title) |
| 281 | |
| 282 | |
| 283 | def parse_tuple(s): |
| 284 | return tuple(s.split(",")) |
| 285 | |
| 286 | |
| 287 | def main(): |
| 288 | logging.basicConfig( |
| 289 | format=("%(asctime)s - %(levelname)s - %(message)s"), level=logging.INFO |
| 290 | ) |
| 291 | parser = argparse.ArgumentParser( |
| 292 | description=( |
| 293 | "Process nsys rep and generate kernel non-overlapped cycles. \n" |
| 294 | "Example:\n" |
| 295 | "gputrc2graph.py --in_file d1.nsys-rep,vllm,llama,100 \n" |
| 296 | "d2.nsys-rep,vllm,gpt-oss,102 " |
| 297 | '--out_dir results/ --title "Model=gpt-oss vLLM chart"' |
| 298 | ), |
| 299 | formatter_class=argparse.RawDescriptionHelpFormatter, |
| 300 | ) |
| 301 | |
| 302 | # load supported engine_model |
| 303 | engine_model_supported = load_engine_model() |
| 304 | # Get a string representation of supported engine/model combinations |
| 305 | engine_model_supported_str = ", ".join( |
| 306 | f"{engine}:[{', '.join(models.keys())}]" |
| 307 | for engine, models in engine_model_supported.items() |
| 308 | ) |
| 309 | parser.add_argument( |
| 310 | "--in_file", |
| 311 | type=parse_tuple, |
| 312 | nargs="+", |
| 313 | help=( |
| 314 | "list of (nsys-rep, engine, model, elapsed_nonprofiled_sec) " |
| 315 | "separated by space. Elapsed_nonprofiled_sec is runtime without " |
| 316 | "profiling used to calculate non-gpu time. Specify 0 to use " |
| 317 | "elapsed time from nsys-rep but that might inflate non-gpu time. " |
| 318 | f"Available engine:[model] are: {engine_model_supported_str} " |
| 319 | f"Example: --infile d1.nsys-rep,vllm,llama,100 " |
| 320 | "d2.nsys-rep,vllm,gpt-oss,102" |
| 321 | ), |
| 322 | required=True, |
| 323 | ) |
| 324 | parser.add_argument("--out_dir", help=("output dir for result.csv/html")) |
| 325 | parser.add_argument("--title", help=("title for html chart")) |
| 326 | parser.add_argument( |
| 327 | "--nsys_cmd", |
| 328 | help=("nsys cmd, e.g. /usr/bin/nsys, Default: nsys"), |
| 329 | default="nsys", |
| 330 | ) |
| 331 | args = parser.parse_args() |
| 332 | gputrace = GPUTrace2Graph() |
| 333 | gputrace.gen_graph( |
| 334 | args.in_file, args.out_dir, args.title, args.nsys_cmd, engine_model_supported |
| 335 | ) |
| 336 | |
| 337 | |
| 338 | if __name__ == "__main__": |
| 339 | main() |