Download code/models/tt_transformers/tests/test_device_perf.py from DeAIafdeling/qwen38-flash-next: direct link, hf CLI and curl.
- Browser
- Download file 15.9 kB
-
https://huggingface.co/DeAIafdeling/qwen38-flash-next/resolve/main/code/models/tt_transformers/tests/test_device_perf.py
- Command line
-
hf download hf://DeAIafdeling/qwen38-flash-next/code/models/tt_transformers/tests/test_device_perf.py
-
curl -L -o test_device_perf.py https://huggingface.co/DeAIafdeling/qwen38-flash-next/resolve/main/code/models/tt_transformers/tests/test_device_perf.py
15.9 kB
| # SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| import json | |
| import os | |
| import pandas as pd | |
| import pytest | |
| from loguru import logger | |
| from models.perf.benchmarking_utils import BenchmarkData, BenchmarkProfiler | |
| from models.perf.device_perf_utils import run_device_perf | |
| from models.tt_transformers.tests.test_utils import ( | |
| merge_device_rows, | |
| print_dict, | |
| process_measurements, | |
| split_compile_and_trace, | |
| verify_value_within_margin, | |
| ) | |
| from tools.tracy.common import PROFILER_DEFAULT_OP_SUPPORT_COUNT | |
| from tools.tracy.process_model_log import get_latest_ops_log_filename | |
| # This pytest flag is necessary to ensure that we do NOT open the device in the main process for device perf tests that run | |
| # the test inside a subprocess since UMD does not allow multiple subprocesses opening the device at the same time. | |
| # Add more models here as needed | |
| def test_device_perf_one_iter( | |
| num_layers, | |
| model_name, | |
| batch_size, | |
| data_parallel, | |
| max_seq_len, | |
| mode, | |
| num_runs, | |
| max_generated_tokens, | |
| export_measurements, | |
| ): | |
| cmd = f"pytest models/tt_transformers/demo/simple_text_demo.py -k 'device-perf and performance' --num_layers {num_layers} --data_parallel {data_parallel} --max_seq_len {max_seq_len} --max_generated_tokens {max_generated_tokens} --paged_attention 1 --batch_size {batch_size} --mode {mode} --use_prefetcher True" | |
| cols = ["DEVICE FW", "DEVICE KERNEL", "DEVICE BRISC KERNEL"] | |
| device_analysis_types = ["device_kernel_duration", "device_kernel_first_to_last_start"] | |
| subdir = f"ttt-device-perf-{mode}" | |
| profiler = BenchmarkProfiler() | |
| benchmark_data = BenchmarkData() | |
| profiler.start("run") | |
| profiler.start("decoder-perf-op-metrics") | |
| # Load perf targets | |
| perf_targets = {} | |
| try: | |
| with open(f"models/tt_transformers/tests/perf_targets/device_perf_{mode}.json", "r") as f: | |
| perf_targets = json.load(f) | |
| except FileNotFoundError: | |
| logger.error( | |
| f"Perf targets file not found, device perf test will proceed without performance target comparison" | |
| ) | |
| # Large sequence lengths dispatch many more programs than the default profiler buffer | |
| # (1333 slots) can hold, causing buffer overflow and missing device perf data. | |
| # Scale op_support_count proportionally with max_seq_len (empirically, seq=131072 needs ~10000). | |
| op_support_count = max( | |
| int(PROFILER_DEFAULT_OP_SUPPORT_COUNT * 1.333), # preserve previous minimum | |
| int(max_seq_len / 10), | |
| ) | |
| _ = run_device_perf( | |
| cmd, | |
| subdir, | |
| num_iterations=1, | |
| cols=cols, | |
| batch_size=batch_size, | |
| device_analysis_types=device_analysis_types, | |
| op_support_count=op_support_count, | |
| ) | |
| profiler.end("decoder-perf-op-metrics") | |
| profiler.end("run") | |
| # Parse the latest ops CSV and aggregate per-op metrics | |
| filename = get_latest_ops_log_filename(subdir) | |
| df = pd.read_csv(filename) | |
| df = df[df["OP TYPE"].isin(["tt_dnn_device"])] | |
| df = merge_device_rows(df) | |
| # Split compile and trace | |
| ( | |
| df_model_compilation, | |
| df_model_trace, | |
| df_first_layer_compilation, | |
| df_first_layer_trace, | |
| df_mid_layers_compilation, | |
| df_mid_layers_trace, | |
| df_model_tail_compilation, | |
| df_model_tail_trace, | |
| ) = split_compile_and_trace( | |
| df, | |
| mode=mode, | |
| num_runs=num_runs, | |
| num_layers=num_layers, | |
| ) | |
| ( | |
| kernel_agg_first_layer_compile, | |
| dispatch_agg_first_layer_compile, | |
| firstlast_agg_first_layer_compile, | |
| ) = process_measurements(df_first_layer_compilation, 1) | |
| ( | |
| kernel_agg_first_layer_trace, | |
| dispatch_agg_first_layer_trace, | |
| firstlast_agg_first_layer_trace, | |
| ) = process_measurements(df_first_layer_trace, 1) | |
| if num_layers > 1: | |
| ( | |
| kernel_agg_mid_layers_compile, | |
| dispatch_agg_mid_layers_compile, | |
| firstlast_agg_mid_layers_compile, | |
| ) = process_measurements( | |
| df_mid_layers_compilation, num_layers - 1 | |
| ) # we dont count the first layer | |
| ( | |
| kernel_agg_mid_layers_trace, | |
| dispatch_agg_mid_layers_trace, | |
| firstlast_agg_mid_layers_trace, | |
| ) = process_measurements(df_mid_layers_trace, num_layers - 1) | |
| if df_model_tail_compilation is not None: | |
| ( | |
| kernel_agg_model_tail_compile, | |
| dispatch_agg_model_tail_compile, | |
| firstlast_agg_model_tail_compile, | |
| ) = process_measurements(df_model_tail_compilation, 1) | |
| ( | |
| kernel_agg_model_tail_trace, | |
| dispatch_agg_model_tail_trace, | |
| firstlast_agg_model_tail_trace, | |
| ) = process_measurements(df_model_tail_trace, 1) | |
| # Print measurements | |
| print_dict(kernel_agg_first_layer_compile, "KERNEL AVERAGE DURATION FOR FIRST LAYER COMPILE") | |
| print_dict(kernel_agg_first_layer_trace, "KERNEL AVERAGE DURATION FOR FIRST LAYER TRACE") | |
| if num_layers > 1: | |
| print_dict(kernel_agg_mid_layers_compile, "KERNEL AVERAGE DURATION FOR MID LAYERS COMPILE") | |
| print_dict(kernel_agg_mid_layers_trace, "KERNEL AVERAGE DURATION FOR MID LAYERS TRACE") | |
| print_dict(dispatch_agg_mid_layers_trace, "DISPATCH AVERAGE DURATION FOR MID LAYERS TRACE") | |
| print_dict(firstlast_agg_mid_layers_trace, "FIRST TO LAST AVERAGE START TIME FOR MID LAYERS TRACE") | |
| if df_model_tail_compilation is not None: | |
| print_dict(kernel_agg_model_tail_compile, "KERNEL AVERAGE DURATION FOR MODEL TAIL COMPILE") | |
| print_dict(kernel_agg_model_tail_trace, "KERNEL AVERAGE DURATION FOR MODEL TAIL TRACE") | |
| print_dict(dispatch_agg_model_tail_trace, "DISPATCH AVERAGE DURATION FOR MODEL TAIL TRACE") | |
| print_dict(firstlast_agg_model_tail_trace, "FIRST TO LAST AVERAGE START TIME FOR MODEL TAIL TRACE") | |
| # Prefer trace for collectives, compile for others | |
| def is_collective(op_code: str) -> bool: | |
| return any(x in op_code for x in ("AllGather", "ReduceScatter", "AllReduce", "Matmul_RS")) | |
| # Prepare export structure and default margins | |
| perf_measurements_export: dict[str, dict[str, dict[str, float]]] = {} | |
| # Export metrics for an op group (first layer, mid layers, model tail) | |
| def export_group( | |
| group_name: str, | |
| kernel_agg_compile: dict, | |
| kernel_agg_trace: dict, | |
| dispatch_agg_trace: dict, | |
| firstlast_agg_trace: dict | None, | |
| ): | |
| all_passing = True | |
| perf_measurements_export[group_name] = {} | |
| op_codes = set(list(kernel_agg_compile["avg"].keys()) + list(kernel_agg_trace["avg"].keys())) | |
| for op_code in op_codes: | |
| # kernel avg | |
| k_avg_trace = kernel_agg_trace["avg"].get(op_code) | |
| k_min_trace = kernel_agg_trace["min"].get(op_code) | |
| k_max_trace = kernel_agg_trace["max"].get(op_code) | |
| k_avg_comp = kernel_agg_compile["avg"].get(op_code) | |
| k_min_comp = kernel_agg_compile["min"].get(op_code) | |
| k_max_comp = kernel_agg_compile["max"].get(op_code) | |
| if is_collective(op_code): | |
| k_avg = k_avg_trace if k_avg_trace is not None else k_avg_comp | |
| k_min = k_min_trace if k_min_trace is not None else k_min_comp | |
| k_max = k_max_trace if k_max_trace is not None else k_max_comp | |
| else: | |
| k_avg = k_avg_comp if k_avg_comp is not None else k_avg_trace | |
| k_min = k_min_comp if k_min_comp is not None else k_min_trace | |
| k_max = k_max_comp if k_max_comp is not None else k_max_trace | |
| if k_avg is not None: | |
| benchmark_data.add_measurement( | |
| profiler, 0, "decoder-perf-op-metrics", f"{op_code}-{group_name}-kernel-avg", float(k_avg) | |
| ) | |
| if k_min is not None: | |
| benchmark_data.add_measurement( | |
| profiler, 0, "decoder-perf-op-metrics", f"{op_code}-{group_name}-kernel-min", float(k_min) | |
| ) | |
| if k_max is not None: | |
| benchmark_data.add_measurement( | |
| profiler, 0, "decoder-perf-op-metrics", f"{op_code}-{group_name}-kernel-max", float(k_max) | |
| ) | |
| # Initialize export entry for this op | |
| export_entry: dict[str, float] = {} | |
| if k_avg is not None: | |
| export_entry["kernel_duration"] = float(k_avg) | |
| # Check that perf_targets, group_name, and op_code exist and keys exist | |
| if perf_targets and group_name in perf_targets and op_code in perf_targets[group_name]: | |
| passing = verify_value_within_margin( | |
| k_avg, | |
| perf_targets[group_name][op_code]["kernel_duration"], | |
| perf_targets[group_name][op_code]["kernel_duration_relative_margin"], | |
| op_code, | |
| "kernel", | |
| ) | |
| all_passing = all_passing and passing | |
| else: | |
| logger.warning(f"Warning: {op_code}-{group_name}-kernel not found in perf_targets") | |
| # dispatch from trace only | |
| d_avg = dispatch_agg_trace["avg"].get(op_code) | |
| d_min = dispatch_agg_trace["min"].get(op_code) | |
| d_max = dispatch_agg_trace["max"].get(op_code) | |
| if d_avg is not None: | |
| benchmark_data.add_measurement( | |
| profiler, 0, "decoder-perf-op-metrics", f"{op_code}-{group_name}-op_to_op-avg", float(d_avg) | |
| ) | |
| if d_min is not None: | |
| benchmark_data.add_measurement( | |
| profiler, 0, "decoder-perf-op-metrics", f"{op_code}-{group_name}-op_to_op-min", float(d_min) | |
| ) | |
| if d_max is not None: | |
| benchmark_data.add_measurement( | |
| profiler, 0, "decoder-perf-op-metrics", f"{op_code}-{group_name}-op_to_op-max", float(d_max) | |
| ) | |
| if d_avg is not None: | |
| export_entry["op_to_op"] = float(d_avg) | |
| if perf_targets and group_name in perf_targets and op_code in perf_targets[group_name]: | |
| passing = verify_value_within_margin( | |
| d_avg, | |
| perf_targets[group_name][op_code]["op_to_op"], | |
| perf_targets[group_name][op_code]["op_to_op_duration_relative_margin"], | |
| op_code, | |
| "op_to_op", | |
| ) | |
| all_passing = all_passing and passing | |
| else: | |
| logger.warning(f"Warning: {op_code}-{group_name}-op_to_op not found in perf_targets") | |
| # first_to_last from trace only (if provided) | |
| if firstlast_agg_trace is not None: | |
| fl_avg = firstlast_agg_trace["avg"].get(op_code) | |
| fl_min = firstlast_agg_trace["min"].get(op_code) | |
| fl_max = firstlast_agg_trace["max"].get(op_code) | |
| if fl_avg is not None: | |
| benchmark_data.add_measurement( | |
| profiler, | |
| 0, | |
| "decoder-perf-op-metrics", | |
| f"{op_code}-{group_name}-first_to_last-avg", | |
| float(fl_avg), | |
| ) | |
| if fl_min is not None: | |
| benchmark_data.add_measurement( | |
| profiler, | |
| 0, | |
| "decoder-perf-op-metrics", | |
| f"{op_code}-{group_name}-first_to_last-min", | |
| float(fl_min), | |
| ) | |
| if fl_max is not None: | |
| benchmark_data.add_measurement( | |
| profiler, | |
| 0, | |
| "decoder-perf-op-metrics", | |
| f"{op_code}-{group_name}-first_to_last-max", | |
| float(fl_max), | |
| ) | |
| if fl_avg is not None: | |
| export_entry["first_to_last_start"] = float(fl_avg) | |
| if perf_targets and group_name in perf_targets and op_code in perf_targets[group_name]: | |
| passing = verify_value_within_margin( | |
| fl_avg, | |
| perf_targets[group_name][op_code]["first_to_last_start"], | |
| perf_targets[group_name][op_code]["first_to_last_start_relative_margin"], | |
| op_code, | |
| "first_to_last_start", | |
| ) | |
| all_passing = all_passing and passing | |
| else: | |
| logger.warning(f"Warning: {op_code}-{group_name}-first_to_last not found in perf_targets") | |
| # Save the entry for this op if we collected any metrics | |
| if export_entry: | |
| perf_measurements_export[group_name][op_code] = export_entry | |
| return all_passing | |
| # Export per-op metrics for each group | |
| all_passing = True | |
| all_passing = all_passing and export_group( | |
| group_name=f"decoder-first", | |
| kernel_agg_compile=kernel_agg_first_layer_compile, | |
| kernel_agg_trace=kernel_agg_first_layer_trace, | |
| dispatch_agg_trace=dispatch_agg_first_layer_trace, | |
| firstlast_agg_trace=firstlast_agg_first_layer_trace, | |
| ) | |
| if num_layers > 1: | |
| all_passing = all_passing and export_group( | |
| group_name=f"decoder-mid", | |
| kernel_agg_compile=kernel_agg_mid_layers_compile, | |
| kernel_agg_trace=kernel_agg_mid_layers_trace, | |
| dispatch_agg_trace=dispatch_agg_mid_layers_trace, | |
| firstlast_agg_trace=firstlast_agg_mid_layers_trace, | |
| ) | |
| if df_model_tail_compilation is not None: | |
| all_passing = all_passing and export_group( | |
| group_name=f"model-tail", | |
| kernel_agg_compile=kernel_agg_model_tail_compile, | |
| kernel_agg_trace=kernel_agg_model_tail_trace, | |
| dispatch_agg_trace=dispatch_agg_model_tail_trace, | |
| firstlast_agg_trace=None, # align with decoder tail export (no first_to_last) | |
| ) | |
| # Write out current measurements in the perf-measurements JSON schema | |
| if export_measurements: | |
| try: | |
| output_dir = os.path.join(os.path.dirname(__file__), "perf_measurements") | |
| os.makedirs(output_dir, exist_ok=True) | |
| output_path = os.path.join( | |
| output_dir, | |
| f"device_perf_measurements_{model_name}_{mode}_bs{batch_size}_dp{data_parallel}_layers{num_layers}_seq{max_seq_len}.json", | |
| ) | |
| with open(output_path, "w", encoding="utf-8") as f: | |
| json.dump(perf_measurements_export, f, indent=4, sort_keys=True) | |
| logger.info(f"Wrote perf measurements export to {output_path}") | |
| except Exception as e: | |
| logger.error(f"Failed to write perf measurements export: {e}") | |
| # Save partial run | |
| benchmark_data.save_partial_run_json( | |
| profiler, | |
| run_type="ttnn_decoder_unit", | |
| ml_model_name=f"{model_name}-{mode}-{data_parallel}dp-{num_layers}layers-{max_seq_len}seq", | |
| ) | |
| # No strict assertions on perf; test succeeds if profiling and export ran | |
| assert True | |