Classic roofline plot: achieved TFLOP/s vs arithmetic intensity (FLOP/byte).
(
bdf_long: pd.DataFrame,
vocab_size: int,
hidden_size: int,
peak_bw_gbs: float,
peak_compute_tflops: float,
)
| 343 | |
| 344 | |
| 345 | def plot_roofline( |
| 346 | bdf_long: pd.DataFrame, |
| 347 | vocab_size: int, |
| 348 | hidden_size: int, |
| 349 | peak_bw_gbs: float, |
| 350 | peak_compute_tflops: float, |
| 351 | ): |
| 352 | """Classic roofline plot: achieved TFLOP/s vs arithmetic intensity (FLOP/byte).""" |
| 353 | df = bdf_long.copy() |
| 354 | df["flops"] = df["n_hidden_states"].apply(lambda h: model_flops(vocab_size, hidden_size, h)) |
| 355 | df["bytes"] = df["n_hidden_states"].apply(lambda h: model_bytes(vocab_size, hidden_size, h)) |
| 356 | df["ai"] = df["flops"] / df["bytes"] # arithmetic intensity (FLOP/byte) |
| 357 | df["achieved_tflops"] = df["flops"] / (df["time[ms]"] / 1000) / 1e12 |
| 358 | |
| 359 | fig, ax = plt.subplots(figsize=(8, 6), layout="constrained") |
| 360 | fontsize = 18 |
| 361 | |
| 362 | # Roofline ceiling |
| 363 | ridge_ai = peak_compute_tflops / (peak_bw_gbs / 1000) # TFLOP/s / (TB/s) = FLOP/byte |
| 364 | ai_min = df["ai"].min() * 0.5 |
| 365 | ai_max = max(df["ai"].max() * 2, ridge_ai * 2) |
| 366 | ai_range = np.geomspace(ai_min, ai_max, 200) |
| 367 | mem_ceiling = peak_bw_gbs / 1000 * ai_range # TB/s * FLOP/byte = TFLOP/s |
| 368 | compute_ceiling = np.full_like(ai_range, peak_compute_tflops) |
| 369 | roofline = np.minimum(mem_ceiling, compute_ceiling) |
| 370 | ax.plot(ai_range, roofline, color="black", linewidth=2, label="_nolegend_", zorder=1) |
| 371 | |
| 372 | # Ridge point annotation |
| 373 | ax.axvline(ridge_ai, color="gray", linestyle=":", linewidth=0.8, alpha=0.6) |
| 374 | ax.annotate( |
| 375 | f"Ridge: AI={ridge_ai:.0f}", |
| 376 | xy=(ridge_ai, peak_compute_tflops), |
| 377 | xytext=(ridge_ai * 1.3, peak_compute_tflops * 0.7), |
| 378 | fontsize=fontsize - 4, |
| 379 | color="gray", |
| 380 | arrowprops=dict(arrowstyle="->", color="gray", lw=0.8), |
| 381 | ) |
| 382 | |
| 383 | # Data points per provider |
| 384 | providers = df["provider"].unique() |
| 385 | palette = _provider_palette(providers) |
| 386 | for idx, provider in enumerate(providers): |
| 387 | color = palette[provider] |
| 388 | pdf = df[df["provider"] == provider] |
| 389 | ax.plot( |
| 390 | pdf["ai"], |
| 391 | pdf["achieved_tflops"], |
| 392 | marker=PROVIDER_MARKERS.get(provider, "o"), |
| 393 | markersize=10, |
| 394 | label=provider, |
| 395 | color=color, |
| 396 | zorder=3, |
| 397 | ) |
| 398 | # Annotate each point with H value |
| 399 | # but only on the last provider |
| 400 | if idx == len(providers) - 1: |
| 401 | for _, row in pdf.iterrows(): |
| 402 | ax.annotate( |
no test coverage detected