Function plot_relative_performance

benchmarking/plot-triton-bench.py:185–233

(
    bdf_rel_long: pd.DataFrame,
    show_providers: list[str],
    use_errorbar: bool = False,
)

def plot_relative_performance(
    bdf_rel_long: pd.DataFrame,
    show_providers: list[str],
    use_errorbar: bool = False,
) -> None:
    plot_df = bdf_rel_long.query("provider in @show_providers")
    palette = _provider_palette(show_providers)
    _, ax = plt.subplots(figsize=(10, 6))
    barplot_kwargs = {}
    if use_errorbar:
        barplot_kwargs["errorbar"] = minmax_skip_zero_range
        barplot_kwargs["err_kws"] = {"linewidth": 2.5, "color": "black"}
    sns.barplot(
        plot_df,
        x="n_hidden_states",
        y="relative-perf",
        hue="provider",
        hue_order=show_providers,
        palette=palette,
        ax=ax,
        **barplot_kwargs,
    )
    hatches = [PROVIDER_HATCHES.get(p, "") for p in show_providers]
    bar_containers = [
        c for c in ax.containers if isinstance(c, plt.matplotlib.container.BarContainer)
    ]
    for container, hatch in zip(bar_containers, hatches):
        for bar in container:
            bar.set_hatch(hatch)
    ax.grid(alpha=0.5, axis="y")
    fontsize = 18
    sns.move_legend(
        ax,
        "upper center",
        title="Method",
        bbox_to_anchor=(0.5, 1.3),
        ncol=len(show_providers),
        fontsize=fontsize,
        title_fontsize=fontsize,
    )
    for handle, hatch in zip(ax.get_legend().legend_handles, hatches):
        handle.set_hatch(hatch)
    ax.set_xlabel("Batch Size", fontsize=fontsize)
    ax.set_ylabel("Relative Performance", fontsize=fontsize)
    ax.set_xticks(ax.get_xticks(), labels=bdf_rel_long["n_hidden_states"].unique().astype(int))
    ax.yaxis.set_major_locator(plt.MultipleLocator(0.25))
    ax.tick_params(axis="both", labelsize=fontsize)
    ax.figure.tight_layout()
    return ax


def model_bytes(vocab_size: int, hidden_size: int, n_hidden_states: float) -> float: