| 75 | return pd.concat(frames, ignore_index=True) |
| 76 | |
| 77 | |
| 78 | def plot_tp_scaling( |
| 79 | long: pd.DataFrame, |
| 80 | h_values: list[int], |
| 81 | providers: list[str], |
| 82 | ) -> plt.Figure: |
| 83 | sns.set_context("talk") |
| 84 | plot_df = long.query("n_hidden_states in @h_values and provider in @providers").copy() |
| 85 | plot_df["n_hidden_states"] = plot_df["n_hidden_states"].astype(int) |
| 86 | plot_df["time[us]"] = plot_df["time[ms]"] * 1000 |
| 87 | |
| 88 | palette = {p: PROVIDER_COLORS[p] for p in providers} |
| 89 | markers = {p: PROVIDER_MARKERS[p] for p in providers} |
| 90 | |
| 91 | fig, axes = plt.subplots(1, len(h_values), figsize=(5 * len(h_values), 4), sharey=False) |
| 92 | if len(h_values) == 1: |
| 93 | axes = [axes] |
| 94 | |
| 95 | unique_tps = sorted(plot_df["tp"].unique()) |
| 96 | |
| 97 | lineplot_kwargs = {"estimator": "min", "errorbar": None} |
| 98 | |
| 99 | for ax_idx, (ax, h) in enumerate(zip(axes, h_values)): |
| 100 | sub = plot_df.query("n_hidden_states == @h") |
| 101 | sns.lineplot( |
| 102 | sub, |
| 103 | x="tp", |
| 104 | y="time[us]", |
| 105 | hue="provider", |
| 106 | hue_order=providers, |
| 107 | style="provider", |
| 108 | style_order=providers, |
| 109 | markers=markers, |
| 110 | markersize=12, |
| 111 | dashes=False, |
| 112 | ax=ax, |
| 113 | palette=palette, |
| 114 | **lineplot_kwargs, |
| 115 | ) |
| 116 | |
| 117 | # Ideal 1/TP reference, anchored at FlashSampling TP=1 (min across runs). |
| 118 | fs_name = FLASHSAMPLING_RENAMES[L.fmms_triton] |
| 119 | fs_tp1 = sub.query("provider == @fs_name and tp == 1")["time[us]"].min() |
| 120 | ref_values = [fs_tp1 / tp for tp in unique_tps] |
| 121 | ax.plot( |
| 122 | unique_tps, |
| 123 | ref_values, |
| 124 | linestyle=":", |
| 125 | color=PROVIDER_COLORS[fs_name], |
| 126 | linewidth=2.5, |
| 127 | marker="*", |
| 128 | markersize=14, |
| 129 | label="Ideal 1/TP", |
| 130 | zorder=1, |
| 131 | ) |
| 132 | |
| 133 | ax.set_xscale("log") |
| 134 | ax.set_xticks(unique_tps, labels=[str(t) for t in unique_tps]) |