d

Author: codelion

examples/circle_packing_with_artifacts/evaluator.py

@@ -295,9 +295,9 @@ def evaluate(program_path):
         # Add successful packing stats for good solutions
         if valid and target_ratio > 0.95:  # Near-optimal solutions
             artifacts["stdout"] = f"Excellent packing! Achieved {target_ratio:.1%} of target value"
-            artifacts["radius_stats"] = (
-                f"Min: {validation_details['min_radius']:.6f}, Max: {validation_details['max_radius']:.6f}, Avg: {validation_details['avg_radius']:.6f}"
-            )
+            artifacts[
+                "radius_stats"
+            ] = f"Min: {validation_details['min_radius']:.6f}, Max: {validation_details['max_radius']:.6f}, Avg: {validation_details['avg_radius']:.6f}"
 
         return EvaluationResult(
             metrics={
@@ -404,9 +404,9 @@ def evaluate_stage1(program_path):
 
             # Add validation issues if any
             if not valid:
-                artifacts["stderr"] = (
-                    f"Validation failed: {len(validation_details.get('boundary_violations', []))} boundary violations, {len(validation_details.get('overlaps', []))} overlaps"
-                )
+                artifacts[
+                    "stderr"
+                ] = f"Validation failed: {len(validation_details.get('boundary_violations', []))} boundary violations, {len(validation_details.get('overlaps', []))} overlaps"
                 artifacts["failure_stage"] = "stage1_geometric_validation"
                 if validation_details.get("boundary_violations"):
                     artifacts["boundary_issues"] = validation_details["boundary_violations"][

examples/mlx_metal_kernel_opt/best_program.py

@@ -24,22 +24,22 @@
 def qwen3_custom_gqa_attention(queries, keys, values, scale=1.0, mask=None):
     """
     Custom Metal kernel implementation for Qwen3 GQA attention.
-    
+
     Args:
-        queries: [B, num_heads=40, L, head_dim=128] 
+        queries: [B, num_heads=40, L, head_dim=128]
         keys: [B, num_kv_heads=8, L, head_dim=128]
         values: [B, num_kv_heads=8, L, head_dim=128]
         scale: Attention scaling factor (1/sqrt(head_dim))
         mask: Attention mask (None, "causal", or boolean tensor)
-        
+
     Returns:
         Attention output [B, num_heads=40, L, head_dim=128]
     """
-    
+
     B, num_heads, L, head_dim = queries.shape
     _, num_kv_heads, _, _ = keys.shape
     heads_per_kv = num_heads // num_kv_heads  # Should be 5 for Qwen3
-    
+
     # Handle mask conversion
     if mask == "causal" or mask is None:
         # Create causal mask for autoregressive attention
@@ -56,13 +56,13 @@ def qwen3_custom_gqa_attention(queries, keys, values, scale=1.0, mask=None):
     else:
         # Fallback for unsupported mask types
         return mx.fast.scaled_dot_product_attention(queries, keys, values, scale=scale, mask=mask)
-    
+
     # Expand mask to match batch and head dimensions if needed
     if mask_tensor.ndim == 2:
         mask_tensor = mx.broadcast_to(mask_tensor[None, None, :, :], (B, num_heads, L, L))
     elif mask_tensor.ndim == 3:
         mask_tensor = mx.broadcast_to(mask_tensor[:, None, :, :], (B, num_heads, L, L))
-    
+
     # EVOLVE-BLOCK-START
     # Custom Metal kernel source for Qwen3 GQA optimization
     # This kernel leverages the 40:8 head ratio and Apple Silicon architecture
@@ -169,23 +169,23 @@ def qwen3_custom_gqa_attention(queries, keys, values, scale=1.0, mask=None):
     }
     """
     # EVOLVE-BLOCK-END
-    
+
     try:
         # Prepare kernel inputs
         scale_tensor = mx.array([scale], dtype=queries.dtype)
         use_mask_tensor = mx.array([1 if use_mask else 0], dtype=mx.int32)
-        
+
         # Create and execute custom Metal kernel
         kernel = mx.fast.metal_kernel(
             name="qwen3_gqa_attention_kernel",
             input_names=["queries", "keys", "values", "mask", "scale", "use_mask"],
             output_names=["output"],
             source=kernel_source,
         )
-        
+
         # Optimize thread group size for Apple Silicon
         threadgroup_size = min(32, L)  # Adapt to sequence length
-        
+
         # Execute kernel
         outputs = kernel(
             inputs=[queries, keys, values, mask_tensor, scale_tensor, use_mask_tensor],
@@ -203,9 +203,9 @@ def qwen3_custom_gqa_attention(queries, keys, values, scale=1.0, mask=None):
                 ("HEADS_PER_KV", heads_per_kv),
             ],
         )
-        
+
         return outputs[0]
-        
+
     except Exception as e:
         # Fallback to standard MLX implementation if custom kernel fails
         print(f"⚠️ Custom GQA kernel failed: {e}, falling back to MLX SPDA")
@@ -215,7 +215,7 @@ def qwen3_custom_gqa_attention(queries, keys, values, scale=1.0, mask=None):
 class CustomGQAAttention(nn.Module):
     """
     Qwen3 attention module with custom Metal kernel optimization.
-    
+
     This module integrates the custom Metal kernel while maintaining
     compatibility with the standard MLX-LM interface.
     """
@@ -244,6 +244,7 @@ def __init__(self, args):
         # Standard MLX-LM RoPE
         try:
             from mlx_lm.models.rope_utils import initialize_rope
+
             self.rope = initialize_rope(
                 head_dim,
                 base=args.rope_theta,
@@ -254,7 +255,7 @@ def __init__(self, args):
         except ImportError:
             print("⚠️ Could not import mlx_lm rope_utils, using basic RoPE")
             self.rope = None
-        
+
         print(f"🔧 Initialized Custom Metal GQA Attention")
         print(f"   📊 Architecture: {n_heads}:{n_kv_heads} heads ({n_heads//n_kv_heads}:1 ratio)")
         print(f"   🎯 Head dimension: {head_dim}")
@@ -423,11 +424,11 @@ class MockArgs:
     output = metal_attn(x, mask=mask)
 
     print(f"✅ Metal GQA output shape: {output.shape}")
-    
+
     # Check for valid output
     has_nan = bool(mx.any(mx.isnan(output)))
     has_inf = bool(mx.any(mx.isinf(output)))
-    
+
     print(f"✅ Has NaN: {has_nan}, Has Inf: {has_inf}")
 
     # Check output statistics
@@ -443,10 +444,10 @@ class MockArgs:
     k = mx.random.normal((B, 8, L, D))  # 8 KV heads
     v = mx.random.normal((B, 8, L, D))
     scale = 1.0 / math.sqrt(D)
-    
+
     kernel_output = qwen3_custom_gqa_attention(q, k, v, scale=scale, mask="causal")
     print(f"✅ Direct kernel output shape: {kernel_output.shape}")
-    
+
     kernel_mean = float(mx.mean(kernel_output))
     kernel_std = float(mx.std(kernel_output))
     print(f"✅ Direct kernel stats - Mean: {kernel_mean:.6f}, Std: {kernel_std:.6f}")
@@ -470,7 +471,7 @@ class MockArgs:
     print("Ready for Metal Kernel Evolution")
     print("Evolution focus:")
     print("1. 🔧 Metal kernel source code optimization")
-    print("2. 💾 Memory access pattern improvements for Apple Silicon") 
+    print("2. 💾 Memory access pattern improvements for Apple Silicon")
     print("3. 🎯 GQA-specific optimizations for 40:8 head ratio")
     print("4. ⚡ Vectorization and SIMD optimization")
     print("5. 🚀 Thread group and grid configuration tuning")