Probably v1

app.py

@@ -6,7 +6,7 @@ from bar_plot import create_matplotlib_bar_plot
 
 
 # Configure matplotlib for better performance
-matplotlib.use('Agg')
+matplotlib.use("Agg")
 plt.ioff()
 
 
@@ -24,14 +24,25 @@ def refresh_plot():
     sidebar_text = "**Transformer CI Dashboard**<br>-<br>**AMD runs on MI325**<br>**NVIDIA runs on A10**<br><br>*This dashboard only tracks important models*<br>*(Data refreshed)*"
     return create_matplotlib_bar_plot(), sidebar_text
 
+
 # Create Gradio interface
-with gr.Blocks(title="Random Data Dashboard", css=load_css(), fill_height=True, fill_width=True) as demo:
+with gr.Blocks(
+    title="Random Data Dashboard", css=load_css(), fill_height=True, fill_width=True
+) as demo:
     with gr.Row():
         # Sidebar
         with gr.Column(scale=1, elem_classes=["sidebar"]):
             gr.Markdown("# 🤖 TCID", elem_classes=["sidebar-title"])
-            description = gr.Markdown("**Transformer CI Dashboard**<br>-<br>**AMD runs on MI325**<br>**NVIDIA runs on A10**<br><br>*This dashboard only tracks important models*", elem_classes=["sidebar-description"])
-            summary_btn = gr.Button("summary\n📊", variant="primary", size="lg", elem_classes=["summary-button"])
+            description = gr.Markdown(
+                "**Transformer CI Dashboard**<br>-<br>**AMD runs on MI325**<br>**NVIDIA runs on A10**<br><br>*This dashboard only tracks important models*",
+                elem_classes=["sidebar-description"],
+            )
+            summary_btn = gr.Button(
+                "summary\n📊",
+                variant="primary",
+                size="lg",
+                elem_classes=["summary-button"],
+            )
 
         # Main plot area
         with gr.Column(elem_classes=["main-content"]):

bar_plot.py

@@ -1,10 +1,11 @@
 import matplotlib.pyplot as plt
+import matplotlib.patches as mpatches
 import io
 import numpy as np
 import base64
 
 from plot_utils import get_color_for_config
-from data import load_data
+from data import load_data, ModelBenchmarkData
 
 
 def reorder_data(per_scenario_data: dict) -> dict:
@@ -13,8 +14,25 @@ def reorder_data(per_scenario_data: dict) -> dict:
     def sorting_fn(key: str) -> float:
         cfg = per_scenario_data[key]["config"]
         attn_implementation = cfg["attn_implementation"]
-        attn_implementation_prio = {"flash_attention_2": 0, "sdpa": 1, "eager": 2}[attn_implementation]
-        return attn_implementation_prio, cfg["sdpa_backend"], cfg["kernelize"], cfg["compilation"]
+        attn_impl_prio = {
+            "flash_attention_2": 0,
+            "sdpa": 1,
+            "eager": 2,
+            "flex_attention": 3,
+        }[attn_implementation]
+        sdpa_backend_prio = {
+            None: -1,
+            "flash_attention": 0,
+            "math": 1,
+            "efficient_attention": 2,
+            "cudnn_attention": 3,
+        }[cfg["sdpa_backend"]]
+        return (
+            attn_impl_prio,
+            sdpa_backend_prio,
+            cfg["kernelize"],
+            cfg["compile_mode"] is not None,
+        )
 
     keys.sort(key=sorting_fn)
     per_scenario_data = {k: per_scenario_data[k] for k in keys}
@@ -27,6 +45,8 @@ def infer_bar_label(config: dict) -> str:
         attn_implementation = "Eager"
     elif config["attn_implementation"] == "flash_attention_2":
         attn_implementation = "Flash attention"
+    elif config["attn_implementation"] == "flex_attention":
+        attn_implementation = "Flex attention"
     elif config["attn_implementation"] == "sdpa":
         attn_implementation = {
             "flash_attention": "SDPA (flash attention)",
@@ -37,15 +57,24 @@ def infer_bar_label(config: dict) -> str:
     else:
         attn_implementation = "Unknown"
 
-    compile = "compiled" if config["compilation"] else "no compile"
+    compile = "compiled" if config["compile_mode"] is not None else "no compile"
     kernels = "kernelized" if config["kernelize"] else "no kernels"
     return f"{attn_implementation}, {compile}, {kernels}"
 
 
-def make_bar_kwargs(per_device_data: dict, key: str) -> tuple[dict, list]:
+def infer_bar_hatch(config: dict) -> str:
+    if config["compile_mode"] is not None:
+        return "/"
+    else:
+        return ""
+
+
+def make_bar_kwargs(
+    per_device_data: dict[str, ModelBenchmarkData], key: str
+) -> tuple[dict, list]:
     # Prepare accumulators
     current_x = 0
-    bar_kwargs = {"x": [], "height": [], "color": [], "label": []}
+    bar_kwargs = {"x": [], "height": [], "color": [], "label": [], "hatch": []}
     errors_bars = []
     x_ticks = []
 
@@ -53,12 +82,13 @@ def make_bar_kwargs(per_device_data: dict, key: str) -> tuple[dict, list]:
         per_scenario_data = device_data.get_bar_plot_data()
         per_scenario_data = reorder_data(per_scenario_data)
         device_xs = []
-        
-        for scenario_name, scenario_data in per_scenario_data.items():    
+
+        for scenario_name, scenario_data in per_scenario_data.items():
             bar_kwargs["x"].append(current_x)
             bar_kwargs["height"].append(np.median(scenario_data[key]))
             bar_kwargs["color"].append(get_color_for_config(scenario_data["config"]))
             bar_kwargs["label"].append(infer_bar_label(scenario_data["config"]))
+            bar_kwargs["hatch"].append(infer_bar_hatch(scenario_data["config"]))
             errors_bars.append(np.std(scenario_data[key]))
             device_xs.append(current_x)
             current_x += 1
@@ -67,13 +97,14 @@ def make_bar_kwargs(per_device_data: dict, key: str) -> tuple[dict, list]:
         current_x += 1.5
     return bar_kwargs, errors_bars, x_ticks
 
+
 def create_matplotlib_bar_plot() -> None:
     """Create side-by-side matplotlib bar charts for TTFT and TPOT data."""
 
     # Create figure with dark theme - maximum size for full screen
-    plt.style.use('dark_background')
+    plt.style.use("dark_background")
     fig, axs = plt.subplots(2, 1, figsize=(20, 11), sharex=True)  # used to be 30, 16
-    fig.patch.set_facecolor('#000000')
+    fig.patch.set_facecolor("#000000")
 
     # Load data and ensure coherence
     per_device_data = load_data()
@@ -82,11 +113,16 @@ def create_matplotlib_bar_plot() -> None:
         bs, seqlen, n_tok = device_data.ensure_coherence()
         if batch_size is None:
             batch_size, sequence_length, num_tokens_to_generate = bs, seqlen, n_tok
-        elif (bs, seqlen, n_tok) != (batch_size, sequence_length, num_tokens_to_generate):
+        elif (bs, seqlen, n_tok) != (
+            batch_size,
+            sequence_length,
+            num_tokens_to_generate,
+        ):
             fig.suptitle(
                 f"Mismatch for batch size, sequence length and number of tokens to generate between configs: {bs} "
                 f"!= {batch_size}, {seqlen} != {sequence_length}, {n_tok} != {num_tokens_to_generate}",
-                color='white', fontsize=18
+                color="white",
+                fontsize=18,
             )
             return None
 
@@ -99,27 +135,58 @@ def create_matplotlib_bar_plot() -> None:
     draw_bar_plot(axs[1], itl_bars, itl_errors, "ITL (seconds)", x_ticks)
 
     # Title and tight layout
-    title = "\n".join([
-        "Time to first token and inter-token latency (lower is better)",
-        f"Batch size: {batch_size},  sequence length: {sequence_length},  new tokens: {num_tokens_to_generate}",
-    ])
-    fig.suptitle(title, color='white', fontsize=20, y=1.005, linespacing=1.5)
+    title = "\n".join(
+        [
+            "Time to first token and inter-token latency (lower is better)",
+            f"Batch size: {batch_size},  sequence length: {sequence_length},  new tokens: {num_tokens_to_generate}",
+        ]
+    )
+    fig.suptitle(title, color="white", fontsize=20, y=1.005, linespacing=1.5)
     plt.tight_layout()
 
     # Add common legend with full text
-    unique_bars = len(ttft_bars["label"]) // 2
-    legend_labels, legend_colors = ttft_bars["label"][:unique_bars], ttft_bars["color"][:unique_bars] 
-    legend_handles = [plt.Rectangle((0,0),1,1, color=color) for color in legend_colors]
+    legend_labels, legend_colors, legend_hatches = [], [], []
+    for label, color, hatch in zip(
+        ttft_bars["label"], ttft_bars["color"], ttft_bars["hatch"]
+    ):
+        if label not in legend_labels:
+            legend_labels.append(label)
+            legend_colors.append(color)
+            legend_hatches.append(hatch)
+
+    # Make sure all attn implementations are equally represented
+    # implementations = {}
+    # for label, color, hatch in zip(legend_labels, legend_colors, legend_hatches):
+    #     impl = label.split(",")[0]
+    #     implementations[impl] = implementations.get(impl, []) + [(label, color, hatch)]
+
+    # n_max = max(len(impls) for impls in implementations.values())
+    # for label_color_pairs in implementations.values():
+    #     for _ in range(len(label_color_pairs), n_max):
+    #         label_color_pairs.append(("", "#000000"))
+
+    # legend_labels, legend_colors = zip(*sum(implementations.values(), []))
+
+    legend_handles = [
+        mpatches.Patch(facecolor=color, hatch=hatch, label=label, edgecolor="white")
+        for color, hatch, label in zip(legend_colors, legend_hatches, legend_labels)
+    ]
 
     # Put a legend to the right of the current axis
-    fig.legend(legend_handles, legend_labels, loc='lower center', ncol=4,
-               bbox_to_anchor=(0.515, -0.11), facecolor='black', edgecolor='white',
-               labelcolor='white', fontsize=14)
+    fig.legend(
+        handles=legend_handles,
+        loc="lower center",
+        ncol=4,
+        bbox_to_anchor=(0.515, -0.11),
+        facecolor="black",
+        edgecolor="white",
+        labelcolor="white",
+        fontsize=14,
+    )
 
     # Save plot to bytes with high DPI for crisp text
     buffer = io.BytesIO()
-    plt.savefig(buffer, format='png', facecolor='#000000',
-                bbox_inches='tight', dpi=150)
+    plt.savefig(buffer, format="png", facecolor="#000000", bbox_inches="tight", dpi=150)
     buffer.seek(0)
 
     # Convert to base64 for HTML embedding
@@ -136,26 +203,34 @@ def create_matplotlib_bar_plot() -> None:
 
 
 def draw_bar_plot(
-    ax: plt.Axes, 
+    ax: plt.Axes,
     bar_kwargs: dict,
     errors: list,
     ylabel: str,
     xticks: list[tuple[float, str]],
     adapt_ylim: bool = False,
 ) -> None:
-    ax.set_facecolor('#000000')
-    ax.grid(True, alpha=0.3, color='white', axis='y', zorder=0)
+    ax.set_facecolor("#000000")
+    ax.grid(True, alpha=0.3, color="white", axis="y", zorder=0)
     # Draw bars
-    _ = ax.bar(**bar_kwargs, width=1.0, edgecolor='white', linewidth=1, zorder=3)
+    _ = ax.bar(**bar_kwargs, width=1.0, edgecolor="white", linewidth=1, zorder=3)
     # Add error bars
     ax.errorbar(
-        bar_kwargs["x"], bar_kwargs["height"], yerr=errors, 
-        fmt='none', ecolor='white', alpha=0.8, elinewidth=1.5, capthick=1.5, capsize=4, zorder=4,
+        bar_kwargs["x"],
+        bar_kwargs["height"],
+        yerr=errors,
+        fmt="none",
+        ecolor="white",
+        alpha=0.8,
+        elinewidth=1.5,
+        capthick=1.5,
+        capsize=4,
+        zorder=4,
     )
     # Set labels, ticks and grid
-    ax.set_ylabel(ylabel, color='white', fontsize=16)
+    ax.set_ylabel(ylabel, color="white", fontsize=16)
     ax.set_xticks([])
-    ax.tick_params(colors='white', labelsize=13)
+    ax.tick_params(colors="white", labelsize=13)
     ax.set_xticks([xt[0] for xt in xticks], [xt[1] for xt in xticks], fontsize=16)
     # Truncate axis to better fit the bars
     if adapt_ylim:
@@ -163,5 +238,5 @@ def draw_bar_plot(
         for h, e in zip(bar_kwargs["height"], errors):
             new_ymin = min(new_ymin, 0.98 * (h - e))
             new_ymax = max(new_ymax, 1.02 * (h + e))
-        ymin, ymax = ax.get_ylim() 
+        ymin, ymax = ax.get_ylim()
         ax.set_ylim(max(ymin, new_ymin), min(ymax, new_ymax))

data.py

@@ -1,20 +1,22 @@
 import json
+from copy import deepcopy
+
 import numpy as np
-from typing import Optional
+
 
 def make_id(config: dict, keys_to_ignore: list[str]) -> str:
     keys = sorted(set(config.keys()))
     return "_".join(str(config[k]) for k in keys if k not in keys_to_ignore)
 
-class ModelBenchmarkData:
 
+class ModelBenchmarkData:
     def __init__(self, json_path: str) -> None:
         with open(json_path, "r") as f:
             self.data: dict = json.load(f)
 
     def compute_ttft(self, measures: dict) -> list[float]:
         return [dts[0] for dts in measures["dt_tokens"]]
-    
+
     def compute_itl(self, measures: dict) -> list[float]:
         return [
             (dts[-1] - dts[0]) / (len(dts) - 1) if len(dts) > 2 else 0
@@ -34,7 +36,11 @@ class ModelBenchmarkData:
         all_hyperparams = set()
         for data in self.data.values():
             config = data["config"]
-            hyperparams = (config["batch_size"], config["sequence_length"], config["num_tokens_to_generate"])
+            hyperparams = (
+                config["batch_size"],
+                config["sequence_length"],
+                config["num_tokens_to_generate"],
+            )
             all_hyperparams.add(hyperparams)
         if len(all_hyperparams) > 1:
             raise ValueError(
@@ -42,7 +48,9 @@ class ModelBenchmarkData:
             )
         return all_hyperparams.pop()
 
-    def get_bar_plot_data(self, collapse_on_cache: bool = True, collapse_on_compile_mode: bool = True) -> dict:
+    def get_bar_plot_data(
+        self, collapse_on_cache: bool = True, collapse_on_compile_mode: bool = True
+    ) -> dict:
         # Gather data for each scenario
         per_scenario_data = {}
         for cfg_name, data in self.data.items():
@@ -57,25 +65,35 @@ class ModelBenchmarkData:
             collapsed_keys = {}
             for cfg_name, data in per_scenario_data.items():
                 keys_to_ignore = ["name"]
-                keys_to_ignore += (["use_cache"] if collapse_on_cache else [])
-                keys_to_ignore += (["compile_mode"] if collapse_on_compile_mode else [])
-                cfg_id = make_id(data["config"], keys_to_ignore)
+                keys_to_ignore += ["use_cache"] if collapse_on_cache else []
+                keys_to_ignore += ["compile_mode"] if collapse_on_compile_mode else []
+                duply_cfg = deepcopy(data["config"])
+                duply_cfg["compiled"] = duply_cfg["compile_mode"] is not None
+                cfg_id = make_id(duply_cfg, keys_to_ignore)
                 cfg_e2e = np.mean(data["e2e"])
                 other_name, other_e2e = collapsed_keys.get(cfg_id, (None, 1e16))
                 if cfg_e2e < other_e2e:
                     collapsed_keys[cfg_id] = (cfg_name, cfg_e2e)
-            per_scenario_data = {k: per_scenario_data[k] for k, _ in collapsed_keys.values()}
+            per_scenario_data = {
+                k: per_scenario_data[k] for k, _ in collapsed_keys.values()
+            }
 
         return per_scenario_data
 
 
-def load_data(keep_common_scenarios_only: bool = False) -> dict[str, ModelBenchmarkData]:
+def load_data(
+    keep_common_scenarios_only: bool = False,
+) -> dict[str, ModelBenchmarkData]:
     data = {
         "MI325": ModelBenchmarkData("mi325_data.json"),
         "H100": ModelBenchmarkData("h100_data.json"),
     }
     if keep_common_scenarios_only:
-        common_scenarios = set(data["MI325"].data.keys()) & set(data["H100"].data.keys())
+        common_scenarios = set(data["MI325"].data.keys()) & set(
+            data["H100"].data.keys()
+        )
         for device_name, device_data in data.items():
-            device_data.data = {k: v for k, v in device_data.data.items() if k in common_scenarios}
+            device_data.data = {
+                k: v for k, v in device_data.data.items() if k in common_scenarios
+            }
     return data

h100_data.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2f843d3f436d7919f67c071824fde3bc247b7e3a096a92d8abb191988d86a9d2
-size 2627476
+oid sha256:ee66b31725d29b9faaf38a437f4ca3ba8251f3ddc6eb9733650dac8b414bd73e
+size 1848790

mi325_data.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:628d50b1cbd1eca36f26d3e2c3e31d3062996ffbeb203f3c7ee04b604fec039e
-size 2674620
+oid sha256:e85e274fdf29798e4e1093df3beec82d1369fe306e720811670fe68176e9bc51
+size 1872352

plot_utils.py

@@ -1,25 +1,35 @@
 # Color manipulation functions
 def hex_to_rgb(hex_color):
-    hex_color = hex_color.lstrip('#')
+    hex_color = hex_color.lstrip("#")
     r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
     return r, g, b
 
+
 def blend_colors(color1, color2, blend_strength):
     rgb1 = hex_to_rgb(color1)
     rgb2 = hex_to_rgb(color2)
-    new_color = tuple(map(lambda i: int(rgb1[i] * blend_strength + rgb2[i] * (1 - blend_strength)), range(3)))
+    new_color = tuple(
+        map(
+            lambda i: int(rgb1[i] * blend_strength + rgb2[i] * (1 - blend_strength)),
+            range(3),
+        )
+    )
     return rgb_to_hex(*new_color)
 
+
 def increase_brightness(r, g, b, factor):
     return tuple(map(lambda x: int(x + (255 - x) * factor), (r, g, b)))
 
+
 def decrease_brightness(r, g, b, factor):
     return tuple(map(lambda x: int(x * factor), (r, g, b)))
 
+
 def increase_saturation(r, g, b, factor) -> tuple[int, int, int]:
     gray = 0.299 * r + 0.587 * g + 0.114 * b
     return tuple(map(lambda x: int(gray + (x - gray) * factor), (r, g, b)))
 
+
 def rgb_to_hex(r, g, b):
     r, g, b = map(lambda x: min(max(x, 0), 255), (r, g, b))
     return f"#{r:02x}{g:02x}{b:02x}"
@@ -27,25 +37,31 @@ def rgb_to_hex(r, g, b):
 
 # Color assignment function
 def get_color_for_config(config: dict):
-    attn_implementation, sdpa_backend = config["attn_implementation"], config["sdpa_backend"]
-    barycenter = 1 - (config["compilation"] + 2 * config["kernelize"]) / 3
+    attn_implementation, sdpa_backend = (
+        config["attn_implementation"],
+        config["sdpa_backend"],
+    )
+    compile_mode = config["compile_mode"] is not None
+    barycenter = 1 - (compile_mode + 2 * config["kernelize"]) / 3
 
-    # Eager 
+    # Eager
     if attn_implementation == "eager":
         color = blend_colors("#FA7F7FFF", "#FF2D2DFF", barycenter)
-
     # SDPA - math
     elif attn_implementation == "sdpa" and sdpa_backend == "math":
         color = blend_colors("#7AB8FFFF", "#277CD0FF", barycenter)
-
     # SDPA - flash attention
-    elif attn_implementation == "sdpa" and sdpa_backend == "flash_attention":
+    elif attn_implementation == "sdpa" and sdpa_backend in [None, "flash_attention"]:
         color = blend_colors("#81FF9CFF", "#219F3CFF", barycenter)
-
+    # SDPA - efficient attention
+    elif attn_implementation == "sdpa" and sdpa_backend == "efficient_attention":
+        color = blend_colors("#DB81FFFF", "#9C33B1FF", barycenter)
     # Flash attention
     elif attn_implementation == "flash_attention_2":
         color = blend_colors("#FFDB70FF", "#DFD002FF", barycenter)
+    # Flex attention
+    elif attn_implementation == "flex_attention":
+        color = blend_colors("#DB81FFFF", "#9C33B1FF", barycenter)
     else:
         raise ValueError(f"Unknown attention implementation: {attn_implementation}")
-    
     return color