Better look

bar_plot.py

@@ -21,14 +21,26 @@ def reorder_data(per_scenario_data: dict) -> dict:
     return per_scenario_data
 
 
-def make_bar_kwargs(key: str) -> tuple[dict, list]:
+def infer_bar_label(config: dict) -> str:
+    """Format legend labels to be more readable."""
+    attn_implementation = {
+        "flash_attention_2": "Flash attention",
+        "sdpa": "SDPA",
+        "eager": "Eager",
+    }[config["attn_implementation"]]
+    compile = "compiled" if config["compilation"] 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]:
     # Prepare accumulators
     current_x = 0
     bar_kwargs = {"x": [], "height": [], "color": [], "label": []}
     errors_bars = []
     x_ticks = []
 
-    for device_name, device_data in load_data(keep_common_scenarios_only=False).items():
+    for device_name, device_data in per_device_data.items():
         per_scenario_data = device_data.get_bar_plot_data()
         per_scenario_data = reorder_data(per_scenario_data)
         device_xs = []
@@ -37,7 +49,7 @@ def make_bar_kwargs(key: str) -> tuple[dict, list]:
             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(scenario_name)
+            bar_kwargs["label"].append(infer_bar_label(scenario_data["config"]))
             errors_bars.append(np.std(scenario_data[key]))
             device_xs.append(current_x)
             current_x += 1
@@ -54,12 +66,19 @@ def create_matplotlib_bar_plot() -> None:
     fig, axs = plt.subplots(2, 1, figsize=(30, 16), sharex=True) 
     fig.patch.set_facecolor('#000000')
 
+    # Load and sanitize data
+    per_device_data = load_data()
+    batch_sizes = {name: device_data.get_main_batch_size() for name, device_data in per_device_data.items()}
+    if len(set(batch_sizes.values())) > 1:
+        fig.suptitle(f"Unmatched batch sizes: {batch_sizes}", color='white', fontsize=18, pad=20)
+        return None
+
     # TTFT Plot (left)
-    ttft_bars, ttft_errors, x_ticks = make_bar_kwargs("ttft")
+    ttft_bars, ttft_errors, x_ticks = make_bar_kwargs(per_device_data, "ttft")
     draw_bar_plot(axs[0], ttft_bars, ttft_errors, "Time to first token and inter-token latency (lower is better)", "TTFT (seconds)", x_ticks)
 
     # # ITL Plot (right)
-    itl_bars, itl_errors, x_ticks = make_bar_kwargs("itl")
+    itl_bars, itl_errors, x_ticks = make_bar_kwargs(per_device_data, "itl")
     draw_bar_plot(axs[1], itl_bars, itl_errors, None, "ITL (seconds)", x_ticks)
 
     # # E2E Plot (right)
@@ -68,8 +87,9 @@ def create_matplotlib_bar_plot() -> None:
     plt.tight_layout()
 
     # Add common legend with full text
-    legend_labels = ttft_bars["label"]  # Use full labels without truncation
-    legend_handles = [plt.Rectangle((0,0),1,1, color=color) for color in ttft_bars["color"]]
+    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]
 
     # Put a legend to the right of the current axis
     fig.legend(legend_handles, legend_labels, loc='lower center', ncol=4,
@@ -103,7 +123,7 @@ def draw_bar_plot(ax: plt.Axes, bar_kwargs: dict, errors: list, title: str, ylab
     # 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,
+        fmt='none', ecolor='white', alpha=0.8, elinewidth=1.5, capthick=1.5, capsize=4, zorder=4,
     )
     # Set labels and title
     ax.set_ylabel(ylabel, color='white', fontsize=16)

data.py

@@ -25,6 +25,16 @@ class ModelBenchmarkData:
         num_tokens = len(measures["t_tokens"]) - 1
         return delta_t / num_tokens
 
+    def get_main_batch_size(self) -> int:
+        batch_sizes = {}
+        for cfg_name, data in self.data.items():
+            for measure in data["measures"]:
+                bs = measure["batch_size"]
+                if bs not in batch_sizes:
+                    batch_sizes[bs] = 0
+                batch_sizes[bs] += 1
+        return max(batch_sizes, key=batch_sizes.get)
+
     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 = {}
@@ -52,7 +62,7 @@ class ModelBenchmarkData:
         return per_scenario_data
 
 
-def load_data(keep_common_scenarios_only: bool = True) -> 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"),

plot_utils.py

@@ -4,9 +4,11 @@ def hex_to_rgb(hex_color):
     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(rgb, hex_color, blend_strength):
-    other_rgb = hex_to_rgb(hex_color)
-    return tuple(map(lambda i: int(rgb[i] * blend_strength + other_rgb[i] * (1 - blend_strength)), range(3)))
+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)))
+    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)))
@@ -25,29 +27,25 @@ def rgb_to_hex(r, g, b):
 
 # Color assignment function
 def get_color_for_config(config: dict):
-    # Determine the main hue for the attention implementation
     attn_implementation, sdpa_backend = config["attn_implementation"], config["sdpa_backend"]
-    compilation = config["compilation"]
+    barycenter = 1 - (config["compilation"] + 2 * config["kernelize"]) / 3
+
+    # Eager 
     if attn_implementation == "eager":
-        main_hue = "#FF4B4BFF" if compilation else "#FF4141FF"
-    elif attn_implementation == "sdpa":
-        main_hue = {
-            None:                  "#4A90E2" if compilation else "#2E82E1FF",
-            "math":                "#408DDB" if compilation else "#227BD3FF",
-            "flash_attention":     "#35A34D" if compilation else "#219F3CFF",
-            "efficient_attention": "#605895" if compilation else "#423691FF",
-            "cudnn_attention":     "#774AE2" if compilation else "#5D27DCFF",
-        }[sdpa_backend]  # fmt: off
+        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":
+        color = blend_colors("#81FF9CFF", "#219F3CFF", barycenter)
+
+    # Flash attention
     elif attn_implementation == "flash_attention_2":
-        main_hue = "#FFD700" if compilation else "#FFBF00FF"
+        color = blend_colors("#FFDB70FF", "#DFD002FF", barycenter)
     else:
         raise ValueError(f"Unknown attention implementation: {attn_implementation}")
-    # Apply color modifications for compilation and kernelization
-    r, g, b = hex_to_rgb(main_hue)
-    if config["compilation"]:
-        delta = 0.2 + 0.2 * (len(config["compile_mode"]) - 7) / 8
-        r, g, b = increase_brightness(r, g, b, delta)
-    if config["kernelize"]:
-        r, g, b = decrease_brightness(r, g, b, 0.8)    
-    # Return the color as a hex string
-    return rgb_to_hex(r, g, b)
+    
+    return color