| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import argparse |
| 5 | import copy |
| 6 | import json |
| 7 | import math |
| 8 | import os |
| 9 | from pathlib import Path |
| 10 | from typing import Any |
| 11 | |
| 12 | import matplotlib.pyplot as plt |
| 13 | import pandas as pd |
| 14 | |
| 15 | ## JSON parsing utils #### |
| 16 | |
| 17 | |
| 18 | def largest_dist_from_leaf(node: dict, depth: int = 0): |
| 19 | if len(node["children"]) == 0: |
| 20 | return depth |
| 21 | return max( |
| 22 | [largest_dist_from_leaf(child, depth=depth + 1) for child in node["children"]] |
| 23 | ) |
| 24 | |
| 25 | |
| 26 | def get_entries_at_depth( |
| 27 | depth: int, |
| 28 | entries_and_traces: list[tuple[Any, Any]], |
| 29 | node: dict, |
| 30 | curr_depth: int = 0, |
| 31 | trace=(), |
| 32 | ): |
| 33 | # assert that the query is at kernel or module level |
| 34 | assert depth == -1 or depth == -2 |
| 35 | |
| 36 | if curr_depth == 0 and largest_dist_from_leaf(node) <= (abs(depth) - 1): |
| 37 | # The tree is not tall enough! |
| 38 | entries_and_traces.append((node["entry"], trace)) |
| 39 | return |
| 40 | |
| 41 | if largest_dist_from_leaf(node) == (abs(depth) - 1): |
| 42 | entries_and_traces.append((node["entry"], trace)) |
| 43 | |
| 44 | trace = (node["entry"]["name"],) + trace |
| 45 | for child in node["children"]: |
| 46 | get_entries_at_depth( |
| 47 | depth, entries_and_traces, child, curr_depth=curr_depth + 1, trace=trace |
| 48 | ) |
| 49 | |
| 50 | |
| 51 | def fold_nodes(root: dict, nodes_to_fold: list[str]): |
| 52 | stack: list[dict] = [root] |
| 53 | while len(stack) != 0: |
| 54 | node = stack.pop() |
| 55 | if node["entry"]["name"] in nodes_to_fold: |
| 56 | node["children"] = [] |
| 57 | continue |
| 58 | for child in node["children"]: |
| 59 | stack.append(child) |
| 60 | return root |
| 61 | |
| 62 | |
| 63 | ## Operation name cleanup utils #### |
| 64 | |
| 65 | |
| 66 | def trim_string_back(string: str, width: int) -> str: |
| 67 | if len(string) > width: |
| 68 | offset = len(string) - width + 3 |
| 69 | string = string[:-offset] |
| 70 | if len(string) > 3: |
| 71 | string = string + "..." |
| 72 | return string |
| 73 | |
| 74 | |
| 75 | def shorten_plot_legend_strings(legend, max_char_len: int): |
| 76 | for t in legend.get_texts(): |
| 77 | t.set_text(trim_string_back(abbreviate_known_names(t.get_text()), max_char_len)) |
| 78 | |
| 79 | |
| 80 | def abbreviate_known_names(name: str) -> str: |
| 81 | abbreviations = { |
| 82 | "MergedColumnParallelLinear": "MCPLinear", |
| 83 | "QKVParallelLinear": "QKVPLinear", |
| 84 | "RowParallelLinear": "RPLinear", |
| 85 | "weight=": "w=", |
| 86 | "bfloat16": "bf16", |
| 87 | "float16": "f16", |
| 88 | } |
| 89 | for key, value in abbreviations.items(): |
| 90 | name = name.replace(key, value) |
| 91 | return name |
| 92 | |
| 93 | |
| 94 | def attempt_to_make_names_unique(entries_and_traces): |
| 95 | names, non_unique_names = (set(), set()) |
| 96 | |
| 97 | def all_the_same(items) -> bool: |
| 98 | return all(i == items[0] for i in items) |
| 99 | |
| 100 | for entry, _ in entries_and_traces: |
| 101 | if entry["name"] in names: |
| 102 | non_unique_names.add(entry["name"]) |
| 103 | else: |
| 104 | names.add(entry["name"]) |
| 105 | |
| 106 | for name in non_unique_names: |
| 107 | entries_and_traces_with_name = [ |
| 108 | (entry, trace) |
| 109 | for entry, trace in entries_and_traces |
| 110 | if entry["name"] == name |
| 111 | ] |
| 112 | |
| 113 | zipped_traces = list(zip(*[trace for _, trace in entries_and_traces_with_name])) |
| 114 | first_trace_difference = next( |
| 115 | ( |
| 116 | i |
| 117 | for i, trace_eles in enumerate(zipped_traces) |
| 118 | if not all_the_same(trace_eles) |
| 119 | ), |
| 120 | None, |
| 121 | ) |
| 122 | |
| 123 | if first_trace_difference is None: |
| 124 | # can't create a unique name, leave the names as they |
| 125 | # are they will get aggregated by the pivot_table call |
| 126 | continue |
| 127 | |
| 128 | for entry, trace in entries_and_traces_with_name: |
| 129 | entry["name"] = " <- ".join( |
| 130 | (entry["name"],) + trace[: first_trace_difference + 1] |
| 131 | ) |
| 132 | |
| 133 | |
| 134 | ## Operation grouping utils #### |
| 135 | """ |
| 136 | Group operations in the given dataframe by some high-level ops like, |
| 137 | - gemms |
| 138 | - attention |
| 139 | - rms_norm |
| 140 | etc. |
| 141 | """ |
| 142 | |
| 143 | |
| 144 | def group_trace_by_operations(trace_df: "pd.DataFrame") -> "pd.DataFrame": |
| 145 | def is_rms_norm(op_name: str): |
| 146 | if "rms_norm_kernel" in op_name: |
| 147 | return True |
| 148 | |
| 149 | def is_attention_block(op_name: str): |
| 150 | if "flash_fwd" in op_name or "reshape_and_cache_flash_kernel" in op_name: |
| 151 | return True |
| 152 | |
| 153 | def is_quant(op_name: str): |
| 154 | if "scaled_fp8_quant" in op_name or "scaled_int8_quant" in op_name: |
| 155 | return True |
| 156 | |
| 157 | # LoRA ops |
| 158 | def is_sgmv_shrink(op_name: str): |
| 159 | return "sgmv_shrink" in op_name |
| 160 | |
| 161 | def is_sgmv_expand(op_name: str): |
| 162 | return "sgmv_expand" in op_name |
| 163 | |
| 164 | def is_bgmv_shrink(op_name: str): |
| 165 | return "bgmv_shrink" in op_name |
| 166 | |
| 167 | def is_bgmv_expand(op_name: str): |
| 168 | return "bgmv_expand" in op_name |
| 169 | |
| 170 | def is_cutlass_gemm_op(op_name: str): |
| 171 | return ( |
| 172 | "void cutlass::Kernel" in op_name |
| 173 | or "void cutlass::device_kernel" in op_name |
| 174 | ) |
| 175 | |
| 176 | def is_gemm_op(op_name: str): |
| 177 | if is_quant(op_name): |
| 178 | return False |
| 179 | return ( |
| 180 | is_cutlass_gemm_op(op_name) |
| 181 | or "xmma_gemm" in op_name |
| 182 | or "gemv2T_kernel" in op_name |
| 183 | or "splitKreduce" in op_name |
| 184 | or "s16816gemm" in op_name |
| 185 | ) |
| 186 | |
| 187 | def is_elementwise_op(op_name: str): |
| 188 | return "elementwise_kernel" in op_name |
| 189 | |
| 190 | def is_mem_op(op_name: str): |
| 191 | return "memcpy" in op_name.lower() or "memset" in op_name.lower() |
| 192 | |
| 193 | def is_vocab_embedding_op(op_name: str): |
| 194 | return "vocabparallelembed" in op_name.lower() |
| 195 | |
| 196 | # nccl ops |
| 197 | def is_nccl_op(op_name: str): |
| 198 | return "nccl" in op_name.lower() |
| 199 | |
| 200 | def is_nccl_all_reduce(op_name: str): |
| 201 | return is_nccl_op(op_name) and ( |
| 202 | "all_reduce" in op_name.lower() or "allreduce" in op_name.lower() |
| 203 | ) |
| 204 | |
| 205 | def is_nccl_gather(op_name: str): |
| 206 | return is_nccl_op(op_name) and "gather" in op_name.lower() |
| 207 | |
| 208 | def is_nccl_broadcast(op_name: str): |
| 209 | return is_nccl_op(op_name) and "broadcast" in op_name.lower() |
| 210 | |
| 211 | # Reduce ops types |
| 212 | def is_cross_device_reduce_1stage(op_name: str): |
| 213 | return "cross_device_reduce_1stage" in op_name |
| 214 | |
| 215 | def is_cross_device_reduce_2stage(op_name: str): |
| 216 | return "cross_device_reduce_2stage" in op_name |
| 217 | |
| 218 | def is_custom_ar_all_reduce(op_name: str): |
| 219 | return "_C_custom_ar::all_reduce" in op_name |
| 220 | |
| 221 | def is_reduce_kernel(op_name: str): |
| 222 | return "reduce_kernel" in op_name |
| 223 | |
| 224 | headers = list(trace_df) |
| 225 | ops = copy.deepcopy(headers) |
| 226 | |
| 227 | attention_ops = list(filter(lambda x: is_attention_block(x), ops)) |
| 228 | ops = list(filter(lambda x: x not in attention_ops, ops)) |
| 229 | |
| 230 | quant_ops = list(filter(lambda x: is_quant(x), ops)) |
| 231 | ops = list(filter(lambda x: x not in quant_ops, ops)) |
| 232 | |
| 233 | sgmv_shrink_ops = list(filter(lambda x: is_sgmv_shrink(x), ops)) |
| 234 | ops = list(filter(lambda x: x not in sgmv_shrink_ops, ops)) |
| 235 | sgmv_expand_ops = list(filter(lambda x: is_sgmv_expand(x), ops)) |
| 236 | ops = list(filter(lambda x: x not in sgmv_expand_ops, ops)) |
| 237 | bgmv_shrink_ops = list(filter(lambda x: is_bgmv_shrink(x), ops)) |
| 238 | ops = list(filter(lambda x: x not in bgmv_shrink_ops, ops)) |
| 239 | bgmv_expand_ops = list(filter(lambda x: is_bgmv_expand(x), ops)) |
| 240 | ops = list(filter(lambda x: x not in bgmv_expand_ops, ops)) |
| 241 | |
| 242 | cutlass_gemm_ops = list(filter(lambda x: is_cutlass_gemm_op(x), ops)) |
| 243 | ops = list(filter(lambda x: x not in cutlass_gemm_ops, ops)) |
| 244 | |
| 245 | gemm_ops = list(filter(lambda x: is_gemm_op(x), ops)) |
| 246 | ops = list(filter(lambda x: x not in gemm_ops, ops)) |
| 247 | |
| 248 | rms_norm_ops = list(filter(lambda x: is_rms_norm(x), ops)) |
| 249 | ops = list(filter(lambda x: x not in rms_norm_ops, ops)) |
| 250 | |
| 251 | vocab_embed_ops = list(filter(lambda x: is_vocab_embedding_op(x), ops)) |
| 252 | ops = list(filter(lambda x: x not in vocab_embed_ops, ops)) |
| 253 | |
| 254 | mem_ops = list(filter(lambda x: is_mem_op(x), ops)) |
| 255 | ops = list(filter(lambda x: x not in mem_ops, ops)) |
| 256 | |
| 257 | elementwise_ops = list(filter(lambda x: is_elementwise_op(x), ops)) |
| 258 | ops = list(filter(lambda x: x not in elementwise_ops, ops)) |
| 259 | |
| 260 | nccl_all_reduce_ops = list(filter(lambda x: is_nccl_all_reduce(x), ops)) |
| 261 | ops = list(filter(lambda x: x not in nccl_all_reduce_ops, ops)) |
| 262 | |
| 263 | nccl_gather_ops = list(filter(lambda x: is_nccl_gather(x), ops)) |
| 264 | ops = list(filter(lambda x: x not in nccl_gather_ops, ops)) |
| 265 | |
| 266 | nccl_broadcast_ops = list(filter(lambda x: is_nccl_broadcast(x), ops)) |
| 267 | ops = list(filter(lambda x: x not in nccl_broadcast_ops, ops)) |
| 268 | |
| 269 | nccl_other_ops = list(filter(lambda x: is_nccl_op(x), ops)) |
| 270 | ops = list(filter(lambda x: x not in nccl_other_ops, ops)) |
| 271 | |
| 272 | cross_device_reduce_1stage_ops = list( |
| 273 | filter(lambda x: is_cross_device_reduce_1stage(x), ops) |
| 274 | ) |
| 275 | ops = list(filter(lambda x: x not in cross_device_reduce_1stage_ops, ops)) |
| 276 | |
| 277 | cross_device_reduce_2stage_ops = list( |
| 278 | filter(lambda x: is_cross_device_reduce_2stage(x), ops) |
| 279 | ) |
| 280 | ops = list(filter(lambda x: x not in cross_device_reduce_2stage_ops, ops)) |
| 281 | |
| 282 | custom_ar_all_reduce_ops = list(filter(lambda x: is_custom_ar_all_reduce(x), ops)) |
| 283 | ops = list(filter(lambda x: x not in custom_ar_all_reduce_ops, ops)) |
| 284 | |
| 285 | reduce_kernel_ops = list(filter(lambda x: is_reduce_kernel(x), ops)) |
| 286 | ops = list(filter(lambda x: x not in reduce_kernel_ops, ops)) |
| 287 | |
| 288 | if len(attention_ops): |
| 289 | trace_df["attention"] = trace_df[attention_ops].agg("sum", axis=1) |
| 290 | if len(quant_ops): |
| 291 | trace_df["quant_ops"] = trace_df[quant_ops].agg("sum", axis=1) |
| 292 | |
| 293 | if len(sgmv_shrink_ops): |
| 294 | trace_df["sgmv_shrink_ops"] = trace_df[sgmv_shrink_ops].agg("sum", axis=1) |
| 295 | if len(sgmv_expand_ops): |
| 296 | trace_df["sgmv_expand_ops"] = trace_df[sgmv_expand_ops].agg("sum", axis=1) |
| 297 | if len(bgmv_shrink_ops): |
| 298 | trace_df["bgmv_shrink_ops"] = trace_df[bgmv_shrink_ops].agg("sum", axis=1) |
| 299 | if len(bgmv_expand_ops): |
| 300 | trace_df["bgmv_expand_ops"] = trace_df[bgmv_expand_ops].agg("sum", axis=1) |
| 301 | |
| 302 | if len(cutlass_gemm_ops): |
| 303 | trace_df["cutlass_gemm_ops"] = trace_df[cutlass_gemm_ops].agg("sum", axis=1) |
| 304 | |
| 305 | if len(gemm_ops): |
| 306 | trace_df["gemm_ops"] = trace_df[gemm_ops].agg("sum", axis=1) |
| 307 | if len(rms_norm_ops): |
| 308 | trace_df["rms_norm_ops"] = trace_df[rms_norm_ops].agg("sum", axis=1) |
| 309 | if len(vocab_embed_ops): |
| 310 | trace_df["vocab_embed_ops"] = trace_df[vocab_embed_ops].agg("sum", axis=1) |
| 311 | if len(mem_ops): |
| 312 | trace_df["mem_ops"] = trace_df[mem_ops].agg("sum", axis=1) |
| 313 | if len(elementwise_ops): |
| 314 | trace_df["elementwise_ops"] = trace_df[elementwise_ops].agg("sum", axis=1) |
| 315 | |
| 316 | if len(nccl_all_reduce_ops): |
| 317 | trace_df["nccl_all_reduce_ops"] = trace_df[nccl_all_reduce_ops].agg( |
| 318 | "sum", axis=1 |
| 319 | ) |
| 320 | if len(nccl_gather_ops): |
| 321 | trace_df["nccl_gather_ops"] = trace_df[nccl_gather_ops].agg("sum", axis=1) |
| 322 | if len(nccl_broadcast_ops): |
| 323 | trace_df["nccl_broadcast_ops"] = trace_df[nccl_broadcast_ops].agg("sum", axis=1) |
| 324 | if len(nccl_other_ops): |
| 325 | trace_df["nccl_other_ops"] = trace_df[nccl_other_ops].agg("sum", axis=1) |
| 326 | |
| 327 | if len(cross_device_reduce_1stage_ops): |
| 328 | trace_df["cross_device_reduce_1stage_ops"] = trace_df[ |
| 329 | cross_device_reduce_1stage_ops |
| 330 | ].agg("sum", axis=1) |
| 331 | if len(cross_device_reduce_2stage_ops): |
| 332 | trace_df["cross_device_reduce_2stage_ops"] = trace_df[ |
| 333 | cross_device_reduce_2stage_ops |
| 334 | ].agg("sum", axis=1) |
| 335 | if len(custom_ar_all_reduce_ops): |
| 336 | trace_df["custom_ar_all_reduce_ops"] = trace_df[custom_ar_all_reduce_ops].agg( |
| 337 | "sum", axis=1 |
| 338 | ) |
| 339 | if len(reduce_kernel_ops): |
| 340 | trace_df["reduce_kernel_ops"] = trace_df[reduce_kernel_ops].agg("sum", axis=1) |
| 341 | |
| 342 | trace_df.drop( |
| 343 | attention_ops |
| 344 | + quant_ops |
| 345 | + sgmv_shrink_ops |
| 346 | + sgmv_expand_ops |
| 347 | + bgmv_shrink_ops |
| 348 | + bgmv_expand_ops |
| 349 | + cutlass_gemm_ops |
| 350 | + gemm_ops |
| 351 | + rms_norm_ops |
| 352 | + vocab_embed_ops |
| 353 | + mem_ops |
| 354 | + elementwise_ops |
| 355 | + nccl_all_reduce_ops |
| 356 | + nccl_gather_ops |
| 357 | + nccl_broadcast_ops |
| 358 | + nccl_other_ops |
| 359 | + cross_device_reduce_1stage_ops |
| 360 | + cross_device_reduce_2stage_ops |
| 361 | + custom_ar_all_reduce_ops |
| 362 | + reduce_kernel_ops, |
| 363 | axis=1, |
| 364 | inplace=True, |
| 365 | ) |
| 366 | return trace_df |
| 367 | |
| 368 | |
| 369 | ## Data plotting utils #### |
| 370 | |
| 371 | |
| 372 | def plot_trace_df( |
| 373 | traces_df: "pd.DataFrame", |
| 374 | plot_metric: str, |
| 375 | plot_title: str, |
| 376 | output: Path | None = None, |
| 377 | ): |
| 378 | def get_phase_description(traces_df: "pd.DataFrame", phase: str) -> str: |
| 379 | phase_df = traces_df.query(f'phase == "{phase}"') |
| 380 | descs = phase_df["phase_desc"].to_list() |
| 381 | assert all([desc == descs[0] for desc in descs]) |
| 382 | return descs[0] |
| 383 | |
| 384 | phases = traces_df["phase"].unique() |
| 385 | phase_descs = [get_phase_description(traces_df, p) for p in phases] |
| 386 | traces_df = traces_df.pivot_table( |
| 387 | index="phase", columns="name", values=plot_metric, aggfunc="sum" |
| 388 | ) |
| 389 | |
| 390 | traces_df = group_trace_by_operations(traces_df) |
| 391 | |
| 392 | # Make the figure |
| 393 | fig_size_x = max(5, len(phases)) |
| 394 | fig, ax = plt.subplots(1, figsize=(fig_size_x, 8), sharex=True) |
| 395 | |
| 396 | # Draw the stacked bars |
| 397 | ops = list(traces_df) |
| 398 | bottom = [0] * len(phases) |
| 399 | for op in ops: |
| 400 | values = [traces_df[op][phase] for phase in phases] |
| 401 | values = list(map(lambda x: 0.0 if math.isnan(x) else x, values)) |
| 402 | ax.bar(phase_descs, values, label=op, bottom=bottom) |
| 403 | bottom = [bottom[j] + values[j] for j in range(len(phases))] |
| 404 | |
| 405 | # Write the values as text on the bars |
| 406 | for bar in ax.patches: |
| 407 | if bar.get_height() != 0: |
| 408 | ax.text( |
| 409 | bar.get_x() + bar.get_width() / 2, |
| 410 | bar.get_height() / 2 + bar.get_y(), |
| 411 | f"{round(bar.get_height(), 2)}", |
| 412 | ha="center", |
| 413 | color="w", |
| 414 | weight="bold", |
| 415 | size=5, |
| 416 | ) |
| 417 | |
| 418 | # Setup legend |
| 419 | handles, labels = plt.gca().get_legend_handles_labels() |
| 420 | legend = fig.legend(handles, labels, loc="center left", bbox_to_anchor=(1, 1)) |
| 421 | shorten_plot_legend_strings(legend, 50) |
| 422 | |
| 423 | # Setup labels and title |
| 424 | plt.setp(ax.get_xticklabels(), rotation=90) |
| 425 | ax.set_ylabel(plot_metric) |
| 426 | plt.suptitle(plot_title) |
| 427 | |
| 428 | plt.savefig(output, bbox_inches="tight") |
| 429 | print("Created: ", output) |
| 430 | |
| 431 | |
| 432 | def main( |
| 433 | json_trace: Path, |
| 434 | output_directory: Path, |
| 435 | depth: int, # Fetch/Plot operations at this depth of the Json tree |
| 436 | plot_metric: str, |
| 437 | make_names_unique: bool, |
| 438 | top_k: int, |
| 439 | json_nodes_to_fold: list[str], |
| 440 | ): |
| 441 | def prepare_data(profile_json: dict, step_keys: list[str]) -> "pd.DataFrame": |
| 442 | def get_entries_and_traces(key: str): |
| 443 | entries_and_traces: list[tuple[Any, Any]] = [] |
| 444 | for root in profile_json[key]["summary_stats"]: |
| 445 | # Fold nodes in the traces as per user request. i.e. simply |
| 446 | # make the requested nodes leaf-nodes. |
| 447 | root = fold_nodes(root, json_nodes_to_fold) |
| 448 | get_entries_at_depth(depth, entries_and_traces, root) |
| 449 | return entries_and_traces |
| 450 | |
| 451 | def keep_only_top_entries( |
| 452 | df: "pd.DataFrame", metric: str, top_k: int = 9 |
| 453 | ) -> "pd.DataFrame": |
| 454 | df.loc[df.nsmallest(len(df) - top_k + 1, metric).index, ["name"]] = "others" |
| 455 | return df |
| 456 | |
| 457 | def get_phase_description(key: str) -> str: |
| 458 | num_running_seqs = profile_json[key]["metadata"]["num_running_seqs"] |
| 459 | if num_running_seqs is not None: |
| 460 | return f"{key}-seqs-{num_running_seqs}" |
| 461 | else: |
| 462 | return key |
| 463 | |
| 464 | # Get data for each key |
| 465 | traces = list(map(lambda x: get_entries_and_traces(x), step_keys)) |
| 466 | |
| 467 | # Attempt some cleanup |
| 468 | if make_names_unique: |
| 469 | for trace in traces: |
| 470 | attempt_to_make_names_unique(trace) |
| 471 | |
| 472 | # To pandas dataframe |
| 473 | trace_dfs = list( |
| 474 | map(lambda t: pd.DataFrame([entry for entry, _ in t]).fillna(0), traces) |
| 475 | ) |
| 476 | |
| 477 | # Respect top_k |
| 478 | if top_k: |
| 479 | trace_dfs = list( |
| 480 | map( |
| 481 | lambda trace_df: keep_only_top_entries( |
| 482 | trace_df, "cuda_time_us", top_k |
| 483 | ), |
| 484 | trace_dfs, |
| 485 | ) |
| 486 | ) |
| 487 | |
| 488 | # Fill in information about the step-keys |
| 489 | for trace_df, step_key in zip(trace_dfs, step_keys): |
| 490 | trace_df["phase"] = step_key |
| 491 | trace_df["phase_desc"] = get_phase_description(step_key) |
| 492 | |
| 493 | # Combine all data frames so they can be put in a single plot |
| 494 | traces_df = pd.concat(trace_dfs) |
| 495 | |
| 496 | # Add a derived metric `cuda_time_ms` |
| 497 | traces_df["cuda_time_ms"] = traces_df["cuda_time_us"] / 1000 |
| 498 | traces_df = traces_df.fillna(0) |
| 499 | |
| 500 | return traces_df |
| 501 | |
| 502 | def make_plot_title_suffix(profile_json: dict) -> str: |
| 503 | context = profile_json["context"] |
| 504 | sparsity = context.get("sparsity", None) |
| 505 | run_type = ( |
| 506 | f"Run {context['num_steps']} steps" |
| 507 | if context["num_steps"] |
| 508 | else ( |
| 509 | f"Complete {context['complete_num_requests_per_step']} per " |
| 510 | f"step; Run till completion" |
| 511 | ) |
| 512 | ) |
| 513 | return ( |
| 514 | f"{context['engine_args']['model']}\n" |
| 515 | f"Batch={context['batch_size']}, " |
| 516 | f"PromptLen={context['prompt_len']}, " |
| 517 | f"NumGpus={context['engine_args']['tensor_parallel_size']}" |
| 518 | f"{', Sparsity ' + sparsity if sparsity else ''}\n" |
| 519 | f"Run Type: {run_type}" |
| 520 | ) |
| 521 | |
| 522 | profile_json = None |
| 523 | with open(json_trace) as f: |
| 524 | profile_json = json.load(f) |
| 525 | assert profile_json is not None |
| 526 | |
| 527 | # Get all `llm.generate.step()` profile |
| 528 | step_traces = list(profile_json.keys()) |
| 529 | assert step_traces[0] == "context" |
| 530 | step_traces = step_traces[1:] # have only prefill and decodes |
| 531 | prefills = list(filter(lambda x: "prefill" in x, step_traces)) |
| 532 | all_decodes = list(filter(lambda x: "decode" in x, step_traces)) |
| 533 | assert len(prefills) + len(all_decodes) == len(step_traces) |
| 534 | assert len(prefills) == 1 |
| 535 | |
| 536 | decodes = all_decodes[:: args.step_plot_interval] |
| 537 | if decodes[-1] != all_decodes[-1]: |
| 538 | # Always have the last decode |
| 539 | decodes.append(all_decodes[-1]) |
| 540 | |
| 541 | prefill_traces = prepare_data(profile_json, prefills) |
| 542 | decode_traces = prepare_data(profile_json, decodes) |
| 543 | |
| 544 | plot_title_suffix = make_plot_title_suffix(profile_json) |
| 545 | |
| 546 | plot_trace_df( |
| 547 | prefill_traces, |
| 548 | plot_metric, |
| 549 | "prefill " + plot_title_suffix, |
| 550 | output_directory / Path("prefill.png"), |
| 551 | ) |
| 552 | plot_trace_df( |
| 553 | decode_traces, |
| 554 | plot_metric, |
| 555 | "decodes " + plot_title_suffix, |
| 556 | output_directory / Path("decode_steps.png"), |
| 557 | ) |
| 558 | |
| 559 | |
| 560 | if __name__ == "__main__": |
| 561 | parser = argparse.ArgumentParser() |
| 562 | |
| 563 | parser.add_argument( |
| 564 | "--json-trace", |
| 565 | type=str, |
| 566 | required=True, |
| 567 | help=( |
| 568 | "JSON trace file generated by scripts that use " |
| 569 | "vllm.profiler.layerwise_profile" |
| 570 | ), |
| 571 | ) |
| 572 | parser.add_argument( |
| 573 | "--output-directory", type=str, required=False, help="Directory to output plots" |
| 574 | ) |
| 575 | parser.add_argument( |
| 576 | "--level", type=str, default="module", choices=["module", "kernel"] |
| 577 | ) |
| 578 | parser.add_argument( |
| 579 | "--top-k", |
| 580 | type=int, |
| 581 | default=12, |
| 582 | help="Only graph the top `top_k` entries by time.", |
| 583 | ) |
| 584 | parser.add_argument( |
| 585 | "--fold-json-node", |
| 586 | nargs="+", |
| 587 | default=["Sampler", "LogitsProcessor"], |
| 588 | help="Do not plot the children of these nodes. Let, \ |
| 589 | the node represent the aggregate of all its \ |
| 590 | children", |
| 591 | ) |
| 592 | parser.add_argument( |
| 593 | "--plot-metric", |
| 594 | type=str, |
| 595 | default="cuda_time_ms", |
| 596 | help="Metric to plot. some options are cuda_time_ms, \ |
| 597 | pct_cuda_time", |
| 598 | ) |
| 599 | parser.add_argument( |
| 600 | "--step-plot-interval", |
| 601 | type=int, |
| 602 | default=4, |
| 603 | help="For every `step_plot_interval` steps, plot 1 step", |
| 604 | ) |
| 605 | |
| 606 | args = parser.parse_args() |
| 607 | |
| 608 | # Prepare/Extract relevant args |
| 609 | make_names_unique = False |
| 610 | if args.level == "module": |
| 611 | depth = -2 |
| 612 | make_names_unique = True |
| 613 | elif args.level == "kernel": |
| 614 | depth = -1 |
| 615 | else: |
| 616 | raise Exception(f"Unexpected level value ({args.level})") |
| 617 | |
| 618 | output_directory = ( |
| 619 | args.output_directory if args.output_directory else Path(args.json_trace).parent |
| 620 | ) |
| 621 | |
| 622 | if not os.path.exists(output_directory): |
| 623 | os.makedirs(output_directory) |
| 624 | |
| 625 | main( |
| 626 | Path(args.json_trace), |
| 627 | output_directory, |
| 628 | depth, |
| 629 | args.plot_metric, |
| 630 | make_names_unique, |
| 631 | args.top_k, |
| 632 | args.fold_json_node, |
| 633 | ) |