s

Author: codelion

examples/mlx_metal_kernel_opt/best_program.py

@@ -55,7 +55,9 @@ def qwen3_custom_gqa_attention(queries, keys, values, scale=1.0, mask=None):
             use_mask = True
     else:
         # Raise error for unsupported mask types - no fallback
-        raise ValueError(f"Unsupported mask type: {type(mask)}. Custom kernel requires None, 'causal', or mx.array mask.")
+        raise ValueError(
+            f"Unsupported mask type: {type(mask)}. Custom kernel requires None, 'causal', or mx.array mask."
+        )
 
     # Expand mask to match batch and head dimensions if needed
     if mask_tensor.ndim == 2:

examples/mlx_metal_kernel_opt/initial_program.py

@@ -55,7 +55,9 @@ def qwen3_custom_gqa_attention(queries, keys, values, scale=1.0, mask=None):
             use_mask = True
     else:
         # Raise error for unsupported mask types - no fallback
-        raise ValueError(f"Unsupported mask type: {type(mask)}. Custom kernel requires None, 'causal', or mx.array mask.")
+        raise ValueError(
+            f"Unsupported mask type: {type(mask)}. Custom kernel requires None, 'causal', or mx.array mask."
+        )
 
     # Expand mask to match batch and head dimensions if needed
     if mask_tensor.ndim == 2:

examples/mlx_metal_kernel_opt/run_benchmarks.py

@@ -106,16 +106,20 @@ def run_optimized_benchmark(args, original_dir):
     try:
         # Import the optimized attention implementation
         # First, try the OpenEvolve output directory (most likely location)
-        best_program_path = os.path.join(original_dir, "openevolve_output", "best", "best_program.py")
-        
+        best_program_path = os.path.join(
+            original_dir, "openevolve_output", "best", "best_program.py"
+        )
+
         # Fallback to root directory if not found in openevolve_output
         if not os.path.exists(best_program_path):
             best_program_path = os.path.join(original_dir, "best_program.py")
-        
+
         if not os.path.exists(best_program_path):
             print(f"❌ Error: Optimized program not found")
             print("Searched in the following locations:")
-            print(f"  1. {os.path.join(original_dir, 'openevolve_output', 'best', 'best_program.py')}")
+            print(
+                f"  1. {os.path.join(original_dir, 'openevolve_output', 'best', 'best_program.py')}"
+            )
             print(f"  2. {os.path.join(original_dir, 'best_program.py')}")
             print("Please ensure OpenEvolve has generated an optimized solution")
             print("Expected path: ./openevolve_output/best/best_program.py")
@@ -454,7 +458,9 @@ def print_comparison_summary(comparison_results):
 
     print(f"\n📊 ABSOLUTE PERFORMANCE:")
     print(f"  🔵 Standard MLX-LM:     {summary['avg_standard_decode_speed']:.1f} tokens/sec average")
-    print(f"  🟠 Metal Kernel Optimized: {summary['avg_optimized_decode_speed']:.1f} tokens/sec average")
+    print(
+        f"  🟠 Metal Kernel Optimized: {summary['avg_optimized_decode_speed']:.1f} tokens/sec average"
+    )
     print(
         f"  📈 Net Improvement:     {summary['avg_optimized_decode_speed'] - summary['avg_standard_decode_speed']:+.1f} tokens/sec"
     )