From b47aa2c24643b9166527f47298ba94a605639af9 Mon Sep 17 00:00:00 2001 From: Guanghua Li Date: Thu, 25 Jun 2026 22:52:18 +0000 Subject: [PATCH 1/2] beam_search: implement worker, orchestrator, integration tests, and documentation --- MaxKernel/auto_agent/beam_worker.py | 28 + .../subagents/beam_worker_pipeline.py | 135 ++++ .../tests/test_beam_worker_pipeline.py | 211 ++++++ MaxKernel/beam_search/__init__.py | 0 .../docs/beam_search_expansion_design_doc.md | 139 ++++ .../summary_12p_RMSNorm_aa_r3_b4_0625_0645.md | 117 ++++ .../summary_12p_RMSNorm_bw_r3_b4_0625_0216.md | 118 ++++ .../docs/worker_integration_plan.md | 198 ++++++ MaxKernel/beam_search/orchestrator.py | 603 ++++++++++++++++++ MaxKernel/beam_search/run_beam_search.py | 95 +++ MaxKernel/beam_search/test_orchestrator.py | 389 +++++++++++ MaxKernel/beam_search/utils.py | 73 +++ MaxKernel/run_beam_search.sh | 29 + 13 files changed, 2135 insertions(+) create mode 100644 MaxKernel/auto_agent/beam_worker.py create mode 100644 MaxKernel/auto_agent/subagents/beam_worker_pipeline.py create mode 100644 MaxKernel/auto_agent/tests/test_beam_worker_pipeline.py create mode 100644 MaxKernel/beam_search/__init__.py create mode 100644 MaxKernel/beam_search/docs/beam_search_expansion_design_doc.md create mode 100644 MaxKernel/beam_search/docs/summary_12p_RMSNorm_aa_r3_b4_0625_0645.md create mode 100644 MaxKernel/beam_search/docs/summary_12p_RMSNorm_bw_r3_b4_0625_0216.md create mode 100644 MaxKernel/beam_search/docs/worker_integration_plan.md create mode 100644 MaxKernel/beam_search/orchestrator.py create mode 100644 MaxKernel/beam_search/run_beam_search.py create mode 100644 MaxKernel/beam_search/test_orchestrator.py create mode 100644 MaxKernel/beam_search/utils.py create mode 100755 MaxKernel/run_beam_search.sh diff --git a/MaxKernel/auto_agent/beam_worker.py b/MaxKernel/auto_agent/beam_worker.py new file mode 100644 index 0000000..bf098e2 --- /dev/null +++ b/MaxKernel/auto_agent/beam_worker.py @@ -0,0 +1,28 @@ +"""Beam Search Worker Agent registration. + +This module instantiates the correctness-only pipeline agent +used to generate and verify kernel candidates for Beam Search. +""" + +from auto_agent.subagents.beam_worker_pipeline import BeamWorkerPipeline +from auto_agent.subagents.kernel_writing import ( + implement_kernel_agent, + plan_kernel_agent, + validate_kernel_compilation_agent, +) +from auto_agent.subagents.testing import ( + unified_test_agent, + validated_test_generation_agent, +) + +beam_worker_agent = BeamWorkerPipeline( + name="BeamWorkerPipeline", + plan_agent=plan_kernel_agent, + implement_agent=implement_kernel_agent, + validate_agent=validate_kernel_compilation_agent, + test_gen_agent=validated_test_generation_agent, + test_run_agent=unified_test_agent, + max_iterations=5, +) + +__all__ = ["beam_worker_agent"] diff --git a/MaxKernel/auto_agent/subagents/beam_worker_pipeline.py b/MaxKernel/auto_agent/subagents/beam_worker_pipeline.py new file mode 100644 index 0000000..a73ebc9 --- /dev/null +++ b/MaxKernel/auto_agent/subagents/beam_worker_pipeline.py @@ -0,0 +1,135 @@ +"""Specialized correctness-only pipeline for Beam Search workers.""" + +import logging +import os +from typing import AsyncGenerator + +from auto_agent.subagents.pipeline_agent import AutonomousPipelineAgent +from google.adk.events import Event, EventActions +from google.adk.agents.invocation_context import InvocationContext + +class BeamWorkerPipeline(AutonomousPipelineAgent): + """Subclass of AutonomousPipelineAgent that exits early on correctness success.""" + + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + iteration = 0 + yield self._initialize_state(ctx) + + while iteration < self.max_iterations: + logging.info( + f"[{self.name}] Starting pipeline iteration {iteration + 1}/{self.max_iterations}" + ) + + # Step 1: Plan + logging.info(f"[{self.name}] Running PlanKernelAgent...") + async for event in self.plan_agent.run_async(ctx): + yield event + + # Step 2: Implement + logging.info(f"[{self.name}] Running ImplementKernelAgent...") + async for event in self.implement_agent.run_async(ctx): + yield event + + # Step 3: Validate (Compilation) + logging.info(f"[{self.name}] Running ValidateKernelCompilationAgent...") + async for event in self.validate_agent.run_async(ctx): + yield event + + # Check if compilation succeeded + compilation_status = ctx.session.state.get("kernel_compilation_status", {}) + if not compilation_status.get("success", False): + logging.error( + f"[{self.name}] Compilation failed. Looping back to planning." + ) + self._save_iteration_files( + ctx, iteration, keys_to_save=["optimized_kernel_path"] + ) + iteration += 1 + continue + + # Step 4: Test Gen + logging.info(f"[{self.name}] Running ValidatedTestGenerationAgent...") + async for event in self.test_gen_agent.run_async(ctx): + yield event + + # Check if test generation succeeded + validation_status = ctx.session.state.get("validation_loop_status", {}) + if not validation_status.get("success", False): + logging.error( + f"[{self.name}] Test generation/validation failed. Looping back to planning." + ) + self._save_iteration_files( + ctx, + iteration, + keys_to_save=["optimized_kernel_path", "test_file_path"], + ) + iteration += 1 + continue + + # Step 5: Test Run + logging.info(f"[{self.name}] Running UnifiedTestAgent...") + async for event in self.test_run_agent.run_async(ctx): + yield event + + # Check if correctness tests passed (Early Exit check) + test_results = ctx.session.state.get("test_results", {}) + if not test_results.get("success", False): + logging.error(f"[{self.name}] Tests failed. Looping back to planning.") + self._save_iteration_files( + ctx, + iteration, + keys_to_save=["optimized_kernel_path", "test_file_path"] + ) + iteration += 1 + continue + + kernel_path = ctx.session.state.get("optimized_kernel_path") + # Snapshot the successful implementation + kernel_code = "" + if kernel_path and os.path.exists(kernel_path): + try: + with open(kernel_path, "r") as f: + kernel_code = f.read() + except Exception as e: + logging.error( + f"[{self.name}] Failed to read kernel file for snapshot: {e}" + ) + + # Extract latency + latency = self._extract_latency(ctx) + + yield Event( + author=self.name, + actions=EventActions( + state_delta={ + "worker_status": "Success", + "kernel_code": kernel_code, + "compilation_status": ctx.session.state.get( + "kernel_compilation_status", {} + ), + "test_status": ctx.session.state.get("test_results", {}), + "latency_ms": latency, + "pipeline_status": "Completed", + "pipeline_iteration": iteration, + "best_iteration": iteration, + } + ), + ) + self._save_iteration_files(ctx, iteration) + return # Terminate worker early since code is correct + + # If we exit the loop, it means we failed to reach correctness + logging.warning( + f"[{self.name}] Failed to generate correct code within iteration limit." + ) + yield Event( + author=self.name, + actions=EventActions( + state_delta={ + "worker_status": "Failed", + "pipeline_status": "Failed", + } + ), + ) diff --git a/MaxKernel/auto_agent/tests/test_beam_worker_pipeline.py b/MaxKernel/auto_agent/tests/test_beam_worker_pipeline.py new file mode 100644 index 0000000..c434f94 --- /dev/null +++ b/MaxKernel/auto_agent/tests/test_beam_worker_pipeline.py @@ -0,0 +1,211 @@ +"""Unit tests for BeamWorkerPipeline execution logic.""" + +import asyncio +from typing import AsyncGenerator +from unittest.mock import MagicMock, patch, mock_open +import pytest + +from auto_agent.subagents.beam_worker_pipeline import BeamWorkerPipeline +from google.adk.agents.base_agent import BaseAgent +from google.adk.events import Event +from google.adk.agents.invocation_context import InvocationContext + + +# Define a dummy agent subclass that inherits from BaseAgent to satisfy Pydantic validations +class DummySubAgent(BaseAgent): + # Custom mock callback runner + mock_run: MagicMock = None + + async def _run_async_impl(self, ctx) -> AsyncGenerator[Event, None]: + if self.mock_run: + async for event in self.mock_run(ctx): + yield event + else: + yield Event(author=self.name) + + +# Helper to construct basic yield generator +async def default_mock_run(ctx): + yield Event(author="mock_agent") + + +@pytest.fixture +def mock_subagents(): + """Construct valid subagent instances using the DummySubAgent subclass.""" + subagents = { + "plan": DummySubAgent(name="PlanAgent"), + "implement": DummySubAgent(name="ImplementAgent"), + "validate": DummySubAgent(name="ValidateAgent"), + "test_gen": DummySubAgent(name="TestGenAgent"), + "test_run": DummySubAgent(name="TestRunAgent") + } + + # Set default side effects + for agent in subagents.values(): + agent.mock_run = MagicMock(side_effect=default_mock_run) + + return subagents + + +@pytest.fixture +def mock_context(): + """Create a mock ADK InvocationContext.""" + from unittest.mock import AsyncMock + + ctx = MagicMock(spec=InvocationContext) + + # Ensure copies of context return the same configured mock context + ctx.model_copy.return_value = ctx + ctx.end_invocation = False + + ctx.session = MagicMock() + + ctx.session.id = "mock_session_id" + ctx.session.state = { + "history": [], + "kernel_compilation_status": {}, + "validation_loop_status": {}, + "test_results": {}, + "autotune_results": {} + } + + # Mock async plugin manager callbacks to be awaitable + ctx.plugin_manager = MagicMock() + ctx.plugin_manager.run_before_agent_callback = AsyncMock(return_value=None) + ctx.plugin_manager.run_after_agent_callback = AsyncMock(return_value=None) + + return ctx + + + +@pytest.fixture +def worker_pipeline(mock_subagents): + """Instantiate BeamWorkerPipeline with mocked subagents.""" + pipeline = BeamWorkerPipeline( + name="TestBeamWorker", + plan_agent=mock_subagents["plan"], + implement_agent=mock_subagents["implement"], + validate_agent=mock_subagents["validate"], + test_gen_agent=mock_subagents["test_gen"], + test_run_agent=mock_subagents["test_run"], + max_iterations=3 + ) + return pipeline + + +@pytest.mark.asyncio +@patch("os.makedirs") +@patch("os.path.exists", return_value=True) +@patch("shutil.copy2") +@patch("builtins.open", new_callable=mock_open, read_data="def optimized_kernel(): pass") +async def test_early_exit_on_success(mock_file_open, mock_copy, mock_exists, mock_makedirs, worker_pipeline, mock_context): + """Test Case 1: Verifies worker pipeline exits early and yields Success event on correctness success.""" + # Setup: Set compilation, test generation, and test execution successes to True + mock_context.session.state["kernel_compilation_status"] = {"success": True} + mock_context.session.state["validation_loop_status"] = {"success": True} + mock_context.session.state["test_results"] = { + "success": True, + "output": "PERF_METRICS: 8.52 ms" + } + mock_context.session.state["optimized_kernel_path"] = "/dummy/optimized_kernel.py" + + events = [] + async for event in worker_pipeline._run_async_impl(mock_context): + events.append(event) + + # Check that each subagent's mock callback was invoked once + assert worker_pipeline.plan_agent.mock_run.call_count == 1 + assert worker_pipeline.implement_agent.mock_run.call_count == 1 + assert worker_pipeline.validate_agent.mock_run.call_count == 1 + assert worker_pipeline.test_gen_agent.mock_run.call_count == 1 + assert worker_pipeline.test_run_agent.mock_run.call_count == 1 + + # Check final emitted success state delta + delta_event = events[-1] + assert delta_event.actions.state_delta["worker_status"] == "Success" + assert delta_event.actions.state_delta["pipeline_status"] == "Completed" + assert delta_event.actions.state_delta["latency_ms"] == 8.52 + assert delta_event.actions.state_delta["kernel_code"] == "def optimized_kernel(): pass" + + +@pytest.mark.asyncio +@patch("os.makedirs") +@patch("os.path.exists", return_value=True) +@patch("shutil.copy2") +@patch("builtins.open", new_callable=mock_open, read_data="def code(): pass") +async def test_iteration_loop_retries(mock_file_open, mock_copy, mock_exists, mock_makedirs, worker_pipeline, mock_context): + """Test Case 2: Verifies pipeline loops back to planning and retries on failures.""" + states_history = [ + # Iteration 0 (Compilation failure) + { + "kernel_compilation_status": {"success": False}, + "validation_loop_status": {}, + "test_results": {} + }, + # Iteration 1 (Test generation failure) + { + "kernel_compilation_status": {"success": True}, + "validation_loop_status": {"success": False}, + "test_results": {} + }, + # Iteration 2 (Correctness Success) + { + "kernel_compilation_status": {"success": True}, + "validation_loop_status": {"success": True}, + "test_results": {"success": True, "output": "PERF_METRICS: 9.1 ms"} + } + ] + + # Mock run_async of plan_agent to dynamically adjust state context + call_count = 0 + async def plan_side_effect(ctx): + nonlocal call_count + # Inject state variables for this iteration + ctx.session.state["kernel_compilation_status"] = states_history[call_count]["kernel_compilation_status"] + ctx.session.state["validation_loop_status"] = states_history[call_count]["validation_loop_status"] + ctx.session.state["test_results"] = states_history[call_count]["test_results"] + ctx.session.state["optimized_kernel_path"] = "/dummy/optimized_kernel.py" + call_count += 1 + yield Event(author="PlanAgent") + + worker_pipeline.plan_agent.mock_run = MagicMock(side_effect=plan_side_effect) + + events = [] + async for event in worker_pipeline._run_async_impl(mock_context): + events.append(event) + + # Asserts it ran exactly 3 iterations (0, 1, 2) + assert call_count == 3 + assert worker_pipeline.plan_agent.mock_run.call_count == 3 + + # Check final status + delta_event = events[-1] + assert delta_event.actions.state_delta["worker_status"] == "Success" + assert delta_event.actions.state_delta["pipeline_status"] == "Completed" + assert delta_event.actions.state_delta["latency_ms"] == 9.1 + + +@pytest.mark.asyncio +@patch("os.makedirs") +@patch("os.path.exists", return_value=True) +@patch("shutil.copy2") +@patch("builtins.open", new_callable=mock_open, read_data="def code(): pass") +async def test_exhaust_iterations_failure(mock_file_open, mock_copy, mock_exists, mock_makedirs, worker_pipeline, mock_context): + """Test Case 3: Verifies failure event emitted when iteration limit is hit without correctness success.""" + # Set max_iterations = 2 + worker_pipeline.max_iterations = 2 + + # Set compilation and test generation success, but correctness tests always fail + mock_context.session.state["kernel_compilation_status"] = {"success": True} + mock_context.session.state["validation_loop_status"] = {"success": True} + mock_context.session.state["test_results"] = {"success": False} + mock_context.session.state["optimized_kernel_path"] = "/dummy/optimized_kernel.py" + + events = [] + async for event in worker_pipeline._run_async_impl(mock_context): + events.append(event) + + # Check final emitted event is failure + delta_event = events[-1] + assert delta_event.actions.state_delta["worker_status"] == "Failed" + assert delta_event.actions.state_delta["pipeline_status"] == "Failed" diff --git a/MaxKernel/beam_search/__init__.py b/MaxKernel/beam_search/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/MaxKernel/beam_search/docs/beam_search_expansion_design_doc.md b/MaxKernel/beam_search/docs/beam_search_expansion_design_doc.md new file mode 100644 index 0000000..eddef11 --- /dev/null +++ b/MaxKernel/beam_search/docs/beam_search_expansion_design_doc.md @@ -0,0 +1,139 @@ +# Design Doc: Dynamic Beam Search Expansion with Stochastic Focus Menu + +This document outlines the design and implementation plan to expand the Beam Search orchestrator (`beam_search/orchestrator.py`) in `MaxKernel` by adapting the stochastic menu dropout strategy and TPU-specific focus strategies from the `autocomp` repository. + +--- + +## 1. Objectives +- **Dynamic Width Scaling**: Scale the number of parallel worker sessions dynamically based on the configured `beam_size`. +- **TPU/Pallas Optimization Menu**: Import all 35 TPU-specific optimization strategies defined in `autocomp`. +- **Stochastic Menu Dropout**: Apply a random dropout filter (probability `p = 0.5`) to the global menu for each worker, generating a unique checklist in each worker's prompt to ensure diverse search directions. +- **Parent Regression budget gate**: Track the parent latency for each candidate and prune candidates that regress in performance compared to their parent (using a configurable `keep_factor`). +- **Semantic Code Deduplication**: Use AST-based normalization inside a separate utilities module to identify and prune duplicate code candidates that are functionally equivalent despite differing comments, variable names, or formatting. + +--- + +## 2. Complete TPU/Pallas Optimization Menu +We will declare the following global list of optimization strategies from `autocomp/agent_builder/.built/tpu-v5e/optimization_menu.yaml` in our code: + +```python +TPU_PALLAS_OPTIMIZATION_STRATEGIES = [ + "Reduce data movement", + "Overlap data movement and compute", + "Cache reused data in local memory instead of reloading from main memory", + "Loop tiling", + "Loop reordering and restructuring", + "Loop unrolling", + "Fuse operations", + "Use lower precision", + "Double buffering", + "Software pipelining", + "Hoist redundant operations out of loops", + "Eliminate redundant computation", + "Simplify or remove unnecessary code", + "Try new parameter values", + "Rewrite the algorithm to reduce total work", + "Place reduction axis last in grid to enable in-place SRAM accumulation without HBM round-trips", + "Align block dimensions to 8x128 tile boundaries to avoid wasted padding and register spills", + "Use scratch_shapes=[pltpu.VMEM(...)] for persistent high-precision accumulators during reduction loops", + "Maximize block sizes up to ~16 MB VMEM capacity to increase arithmetic intensity per pipeline step", + "Use scalar prefetch via PrefetchScalarGridSpec to load indices/metadata into SMEM without stalling vector core", + "Upcast bf16/int8 to float32 before elementwise ops, downcast only on final output write", + "Fuse transpose into lax.dot_general contraction dimensions instead of materializing transposed operands", + "Arrange grid iteration order so consecutive invocations reuse already-resident input slices", + "Increase pipeline buffer count beyond double buffering to hide memory latency for bandwidth-bound kernels", + "Generate random numbers inside kernel via hardware PRNG with key in SMEM instead of passing precomputed arrays", + "Avoid singleton dimensions in last two array axes to prevent full-tile waste per element", + "Reduce along second-to-last dimension rather than last dimension when possible", + "Prefer add/multiply over exp/tanh/division; restructure math to minimize expensive elementwise ops", + "Tune block sizes jointly — systematically vary BM, BN, BK together under the VMEM budget constraint (16 MiB including double-buffering)", + "Compute arithmetic intensity accounting for tiling amplification to predict compute-bound vs memory-bound regime", + "Minimize control flow inside kernels; consolidate into single basic blocks to avoid unrolling overhead", + "Pass all data as explicit kernel inputs with BlockSpec instead of closing over constants", + "Use pltpu.VMEM scoped scratch buffers for temporary storage within kernel lifetime", + "Balance block size against pipeline depth to amortize startup/drain bubble cost over enough iterations", + "Explicitly initialize accumulator buffers to zero on first reduction iteration since SRAM starts undefined" +] +``` + +--- + +## 3. High-Level Flow + +```mermaid +graph TD + P_Beam[Beam of N Candidates] --> P_Alloc[Dynamically Allocate N Workers] + + P_Alloc -->|Worker 1: Random Subset A| W_1[Worker Session 1] + P_Alloc -->|Worker 2: Random Subset B| W_2[Worker Session 2] + P_Alloc -->|Worker N: Random Subset Z| W_N[Worker Session N] + + W_1 --> P_Eval[Grouped Evaluation] + W_2 --> P_Eval + W_N --> P_Eval + + P_Eval --> P_Gate[Apply Parent Regression Gate & Deduplicate] + P_Gate --> P_Sort[Merge, Sort & Prune to Beam Size] +``` + +--- + +## 4. Detailed Design + +### 4.1. Parameter Adjustments +We update `AgenticSearchOrchestrator.__init__` and `run_search` to support: +- `beam_size`: (int) The size of the beam (default `2`). +- `dropout_menu_options`: (float) The probability of keeping each menu option for a worker's prompt (default `0.5`). +- `keep_factor`: (float) Regression threshold allowed against the parent candidate's score (default `1.0` for strict improvement only). + +### 4.2. Candidate-Worker Mapping & Beam Size Progression +When spawning `beam_size` parallel workers: +1. Sort the active `beam` list by latency. +2. For worker `idx` in `range(beam_size)`: + - Retrieve the parent candidate using: + `parent_candidate = beam[idx % len(beam)]` + - *Beam Size Progression:* At the start of the search (Round 1), the beam contains only 1 baseline candidate. Therefore, all workers branch from it. From Round 2 onwards, once the beam fills up with up to `beam_size` unique candidates, the workers will map to distinct parents in the beam (e.g. Worker 0 to `beam[0]`, Worker 1 to `beam[1]`, etc.). If the total number of surviving candidates is less than `beam_size`, some candidates will naturally be selected as parents by multiple workers (branching). + - *Parent Code Passing:* Inside `run_worker_session`, the orchestrator writes the selected parent candidate's code into the worker's session directory as `base_kernel.py`. + - *Parent State Injection*: We store the parent's latency inside the worker session state: + `session.state["parent_latency_ms"] = parent_candidate["latency_ms"]` + +### 4.3. Stochastic Menu Filtering +For each spawned worker session: +1. Filter the global `TPU_PALLAS_OPTIMIZATION_STRATEGIES` list: + ```python + import random + selected_opts = [ + opt for opt in TPU_PALLAS_OPTIMIZATION_STRATEGIES + if random.random() < dropout_menu_options + ] + if not selected_opts: + selected_opts = [random.choice(TPU_PALLAS_OPTIMIZATION_STRATEGIES)] + ``` +2. Inline the selected strategies into the worker's user prompt: + ```python + focus_text = "\n".join(f"- {opt}" for opt in selected_opts) + user_message = ( + "Optimize the JAX/Pallas kernel. " + f"Focus your optimization effort on applying the following strategies:\n{focus_text}" + ) + ``` + +### 4.4. Semantic Code Deduplication (AST Normalizer) +To keep the codebase modular, the AST normalization logic will be implemented in a separate module [beam_search/utils.py](file:///usr/local/google/home/ligh/github/accelerator-agents/MaxKernel/beam_search/utils.py): +- `normalize_ast_code(code: str) -> str`: + 1. Parses the source code using Python's built-in `ast` module. + 2. Traverses the AST tree and renames all function parameters and local variables sequentially to standard placeholders (`v0`, `v1`, `v2`, etc.). + 3. Renames user-defined sub-functions (excluding the main `solution` or library calls). + 4. Unparses the AST back to standard format (stripping docstrings and comments) via `ast.unparse()`. + +### 4.5. Parent Regression Gate & Pruning +At the end of each round: +1. For each newly generated candidate `cand` evaluated from the workers: + - Identify its parent candidate from the previous beam. + - If the candidate's latency is worse than its parent's latency multiplied by the `keep_factor`: + `cand["latency_ms"] >= parent_candidate["latency_ms"] * keep_factor` + - Discard the candidate from the list of results for that round. + - *Exemption justification*: Incumbent candidates from the previous beam are exempt from this check to prevent retroactive eviction if `keep_factor` is tightened dynamically in later rounds. +2. Run the AST Normalizer from the utility module on all candidate source codes to identify duplicates: + - For any duplicates, keep only the one with the lowest latency and discard the rest. +3. Merge the remaining valid new candidates into the beam, sort by latency, and prune to `beam_size`. diff --git a/MaxKernel/beam_search/docs/summary_12p_RMSNorm_aa_r3_b4_0625_0645.md b/MaxKernel/beam_search/docs/summary_12p_RMSNorm_aa_r3_b4_0625_0645.md new file mode 100644 index 0000000..5784188 --- /dev/null +++ b/MaxKernel/beam_search/docs/summary_12p_RMSNorm_aa_r3_b4_0625_0645.md @@ -0,0 +1,117 @@ +# Dynamic Beam Search Run Summary Report (AutoAgent) + +**Command Line:** +```bash +python3 run_beam_search.py --use_beam_worker false --rounds 3 --beam_size 4 --task_id 12p_RMSNorm +``` + +This document summarizes the execution and performance analysis of the local Dynamic Beam Search orchestrator run optimizing the **`12p_RMSNorm`** JAX/Pallas kernel using **`auto_agent`** mode. + +--- + +## 1. Run Metadata +* **Target Kernel**: `12p_RMSNorm` (RMS Normalization block) +* **Search Configuration**: + * Rounds: `3` + * Beam Size ($B$): `4` + * Keep Factor: `1.0` (Regression Budget Gate enabled) + * Worker Mode: `auto_agent` (using `AutonomousPipelineAgent` with compilation validation, autotuning, and profiling) +* **Orchestrator Mode**: Mock compiler evaluation mode (`mock_mode=True`) using `FakeKernelEvaluator` for fallbacks, with autotune results bypass enabled. +* **Baseline Latency**: **`11.400 ms`** + +--- + +## 2. Search Progression Timeline + +### Round 1 +* **Baseline Code**: Reference RMSNorm kernel ($11.400\text{ ms}$). +* **Worker Execution**: + * **`Round_1_Worker_0`** (`beam_r1_w0`): Succeeded. Autotuned latency parsed: **`8.030 ms`** (Bypassed evaluator). + * **`Round_1_Worker_1`** (`beam_r1_w1`): Succeeded (fallback to LLM selection, skipped autotuning). Evaluator fallback latency: **`10.830 ms`**. + * **`Round_1_Worker_2`** (`beam_r1_w2`): Succeeded. Autotuned latency parsed: **`8.000 ms`** (Bypassed evaluator). + * **`Round_1_Worker_3`** (`beam_r1_w3`): Failed correctness tests in Mode 1. Exited. +* **Incumbent Preservation**: Because `Round_1_Worker_3` failed correctness, the original baseline candidate was preserved as Rank 4. + +#### Round 1 Beam State: +``` +================================================== + ROUND 1 BEAM STATE +================================================== + Rank 1: Latency=8.000 ms | Path=beam_search/output/12p_RMSNorm_aa_r3_b4_0625_0645/beam_r1_w2_1782369908/optimized_kernel.py + Rank 2: Latency=8.030 ms | Path=beam_search/output/12p_RMSNorm_aa_r3_b4_0625_0645/beam_r1_w0_1782369908/optimized_kernel.py + Rank 3: Latency=10.830 ms | Path=beam_search/output/12p_RMSNorm_aa_r3_b4_0625_0645/beam_r1_w1_1782369908/optimized_kernel.py + Rank 4: Latency=11.400 ms | Path=beam_search/output/12p_RMSNorm/baseline.py (Incumbent Preserved) +================================================== +``` + +--- + +### Round 2 +* **Worker Launch Paths**: + * `Round_2_Worker_0` spawned from Round 1 Rank 1 baseline ($8.000\text{ ms}$). + * `Round_2_Worker_1` spawned from Round 1 Rank 2 baseline ($8.030\text{ ms}$). + * `Round_2_Worker_2` spawned from Round 1 Rank 3 baseline ($10.830\text{ ms}$). + * `Round_2_Worker_3` spawned from Round 1 Rank 4 baseline ($11.400\text{ ms}$). +* **Worker Execution**: + * **`Round_2_Worker_0`** (`beam_r2_w0`): Succeeded. Autotuned latency parsed: **`8.030 ms`** (Pruned since it didn't beat the parent limit). + * **`Round_2_Worker_1`** (`beam_r2_w1`): Succeeded. Autotuned latency parsed: **`6.100 ms`** (Rank 1). + * **`Round_2_Worker_2`** (`beam_r2_w2`): Succeeded. Autotuned latency parsed: **`8.020 ms`** (Rank 3). + * **`Round_2_Worker_3`** (`beam_r2_w3`): Succeeded after self-healing import naming errors in tests. Evaluator fallback latency: **`10.288 ms`** (Pruned). +* **Incumbent Preservation**: Round 2 Worker 0 and Worker 3 were pruned. The orchestrator preserved Rank 2 (`beam_r1_w2` @ 8.0 ms) and Rank 4 (`beam_r1_w0` @ 8.030 ms) from Round 1. + +#### Round 2 Beam State: +``` +================================================== + ROUND 2 BEAM STATE +================================================== + Rank 1: Latency=6.100 ms | Path=beam_search/output/12p_RMSNorm_aa_r3_b4_0625_0645/beam_r2_w1_1782371027/optimized_kernel.py + Rank 2: Latency=8.000 ms | Path=beam_search/output/12p_RMSNorm_aa_r3_b4_0625_0645/beam_r1_w2_1782369908/optimized_kernel.py (Incumbent Preserved) + Rank 3: Latency=8.020 ms | Path=beam_search/output/12p_RMSNorm_aa_r3_b4_0625_0645/beam_r2_w2_1782371027/optimized_kernel.py + Rank 4: Latency=8.030 ms | Path=beam_search/output/12p_RMSNorm_aa_r3_b4_0625_0645/beam_r1_w0_1782369908/optimized_kernel.py (Incumbent Preserved) +================================================== +``` + +--- + +### Round 3 +* **Worker Launch Paths**: + * `Round_3_Worker_0` spawned from Round 2 Rank 1 baseline ($6.100\text{ ms}$). + * `Round_3_Worker_1` spawned from Round 2 Rank 2 baseline ($8.000\text{ ms}$). + * `Round_3_Worker_2` spawned from Round 2 Rank 3 baseline ($8.020\text{ ms}$). + * `Round_3_Worker_3` spawned from Round 2 Rank 4 baseline ($8.030\text{ ms}$). +* **Worker Execution**: + * **`Round_3_Worker_0`** (`beam_r3_w0`): Succeeded. Autotuned latency parsed: **`6.010 ms`** (Rank 1). + * **`Round_3_Worker_1`** (`beam_r3_w1`): Succeeded. Evaluator fallback latency: **`9.774 ms`** (Pruned). + * **`Round_3_Worker_2`** (`beam_r3_w2`): Succeeded. Autotuned latency parsed: **`6.200 ms`** (Rank 4). + * **`Round_3_Worker_3`** (`beam_r3_w3`): Succeeded. Autotuned latency parsed: **`6.075 ms`** (Rank 2). +* **Incumbent Preservation**: Round 3 Worker 1 was pruned. The orchestrator preserved the best candidate from Round 2 (`beam_r2_w1` @ 6.100 ms) as Rank 3, which beat `beam_r3_w2` (6.200 ms). + +#### Round 3 Beam State (Final State): +``` +================================================== + ROUND 3 BEAM STATE +================================================== + Rank 1: Latency=6.010 ms | Path=beam_search/output/12p_RMSNorm_aa_r3_b4_0625_0645/beam_r3_w0_1782373360/optimized_kernel.py + Rank 2: Latency=6.075 ms | Path=beam_search/output/12p_RMSNorm_aa_r3_b4_0625_0645/beam_r3_w3_1782373360/optimized_kernel.py + Rank 3: Latency=6.100 ms | Path=beam_search/output/12p_RMSNorm_aa_r3_b4_0625_0645/beam_r2_w1_1782371027/optimized_kernel.py (Incumbent Preserved) + Rank 4: Latency=6.200 ms | Path=beam_search/output/12p_RMSNorm_aa_r3_b4_0625_0645/beam_r3_w2_1782373360/optimized_kernel.py +================================================== +``` + +--- + +## 3. Final Search Summary & Performance + +* **Best Kernel Candidate**: Produced by `Round_3_Worker_0` at: + [optimized_kernel.py](file:///usr/local/google/home/ligh/github/accelerator-agents/MaxKernel/beam_search/output/12p_RMSNorm_aa_r3_b4_0625_0645/beam_r3_w0_1782373360/optimized_kernel.py) +* **Baseline Latency**: `11.400 ms` +* **Best Optimized Latency**: **`6.010 ms`** +* **Speedup**: **`1.90x` speedup** (a **`47.3%` latency reduction**). + +--- + +## 4. Evaluator Bypassing and Orchestrator Mechanics Verified + +1. **Direct Metrics Extraction (Bypass)**: Successfully verified the core change: the orchestrator now parses the worker's `autotune_results.json` directly to retrieve `best_time_ms` instead of re-running the mock compiler profile process. This prevents the timing degradation that occurred when running sequential profiling steps. +2. **Robust Fallback Mechanics**: In cases where `autotune_results.json` was missing (e.g., `beam_r1_w1` and `beam_r3_w1` where autotuning was skipped), the orchestrator successfully fell back to evaluating the candidate using `FakeKernelEvaluator`. +3. **Cross-Round Preservations**: Verified that the incumbent preservation rules dynamically handle a mix of autotuned and evaluated baseline candidates. The preservation of `beam_r2_w1` (6.10 ms) in Rank 3 of Round 3 demonstrates correct comparative logic across candidate types. diff --git a/MaxKernel/beam_search/docs/summary_12p_RMSNorm_bw_r3_b4_0625_0216.md b/MaxKernel/beam_search/docs/summary_12p_RMSNorm_bw_r3_b4_0625_0216.md new file mode 100644 index 0000000..5417489 --- /dev/null +++ b/MaxKernel/beam_search/docs/summary_12p_RMSNorm_bw_r3_b4_0625_0216.md @@ -0,0 +1,118 @@ +# Dynamic Beam Search Run Summary Report + +**Command Line:** +```bash +./run_beam_search.sh --task_id 12p_RMSNorm --rounds 3 --beam_size 4 --use_beam_worker true +``` + +This document summarizes the execution and performance analysis of the local Dynamic Beam Search orchestrator run optimizing the **`12p_RMSNorm`** JAX/Pallas kernel. + +--- + +## 1. Run Metadata +* **Target Kernel**: `12p_RMSNorm` (RMS Normalization block) +* **Search Configuration**: + * Rounds: `3` + * Beam Size ($B$): `4` + * Keep Factor: `1.0` (Regression Budget Gate enabled) + * Mode: `correctness-only` (exits early upon passing correctness verification) +* **Orchestrator Mode**: Mock compiler evaluation mode (`mock_mode=True`) using `FakeKernelEvaluator` +* **Baseline Latency**: **`11.400 ms`** + +--- + +## 2. Search Progression Timeline + +### Round 1 +* **Baseline Code**: Reference RMSNorm kernel ($11.400\text{ ms}$). +* **Worker Execution**: + * **`Round_1_Worker_0`** (`beam_r1_w0`): Succeeded in iteration 2 after fixing correctness scripts. Latency: **`10.288 ms`** + * **`Round_1_Worker_1`** (`beam_r1_w1`): Succeeded in iteration 2. Latency: **`9.774 ms`** + * **`Round_1_Worker_2`** (`beam_r1_w2`): Succeeded in iteration 1. Latency: **`9.285 ms`** + * **`Round_1_Worker_3`** (`beam_r1_w3`): Succeeded in iteration 1. Latency: **`8.821 ms`** +* **AST Deduplication**: All 4 candidates were semantically unique. No pruning required. +* **Evaluation Fallback**: Fallback to sequential evaluation succeeded. + +#### Round 1 Beam State: +``` +================================================== + ROUND 1 BEAM STATE +================================================== + Rank 1: Latency=8.821 ms | Path=beam_search/output/12p_RMSNorm_bw_r3_b4_0625_0216/beam_r1_w3_1782353779/optimized_kernel.py + Rank 2: Latency=9.285 ms | Path=beam_search/output/12p_RMSNorm_bw_r3_b4_0625_0216/beam_r1_w2_1782353779/optimized_kernel.py + Rank 3: Latency=9.774 ms | Path=beam_search/output/12p_RMSNorm_bw_r3_b4_0625_0216/beam_r1_w1_1782353779/optimized_kernel.py + Rank 4: Latency=10.288 ms | Path=beam_search/output/12p_RMSNorm_bw_r3_b4_0625_0216/beam_r1_w0_1782353779/optimized_kernel.py +================================================== +``` + +--- + +### Round 2 +* **Worker Launch Paths**: + * `Round_2_Worker_0` spawned from Round 1 Rank 1 baseline ($8.821\text{ ms}$). + * `Round_2_Worker_1` spawned from Round 1 Rank 2 baseline ($9.285\text{ ms}$). + * `Round_2_Worker_2` spawned from Round 1 Rank 3 baseline ($9.774\text{ ms}$). + * `Round_2_Worker_3` spawned from Round 1 Rank 4 baseline ($10.288\text{ ms}$). +* **Worker Execution**: + * **`Round_2_Worker_0`** (`beam_r2_w0`): Succeeded. Latency: **`7.961 ms`** + * **`Round_2_Worker_1`** (`beam_r2_w1`): Succeeded. Latency: **`7.563 ms`** + * **`Round_2_Worker_2`** (`beam_r2_w2`): Encountered compilation failures. Self-healed on validation attempt 2. Latency: **`7.185 ms`** + * **`Round_2_Worker_3`** (`beam_r2_w3`): Succeeded. Latency: **`6.826 ms`** + +#### Round 2 Beam State: +``` +================================================== + ROUND 2 BEAM STATE +================================================== + Rank 1: Latency=6.826 ms | Path=beam_search/output/12p_RMSNorm_bw_r3_b4_0625_0216/beam_r2_w3_1782354715/optimized_kernel.py + Rank 2: Latency=7.185 ms | Path=beam_search/output/12p_RMSNorm_bw_r3_b4_0625_0216/beam_r2_w2_1782354715/optimized_kernel.py + Rank 3: Latency=7.563 ms | Path=beam_search/output/12p_RMSNorm_bw_r3_b4_0625_0216/beam_r2_w1_1782354715/optimized_kernel.py + Rank 4: Latency=7.961 ms | Path=beam_search/output/12p_RMSNorm_bw_r3_b4_0625_0216/beam_r2_w0_1782354715/optimized_kernel.py +================================================== +``` + +--- + +### Round 3 +* **Worker Launch Paths**: + * `Round_3_Worker_0` spawned from Round 2 Rank 1 baseline ($6.826\text{ ms}$). + * `Round_3_Worker_1` spawned from Round 2 Rank 2 baseline ($7.185\text{ ms}$). + * `Round_3_Worker_2` spawned from Round 2 Rank 3 baseline ($7.563\text{ ms}$). + * `Round_3_Worker_3` spawned from Round 2 Rank 4 baseline ($7.961\text{ ms}$). +* **Worker Execution**: + * **`Round_3_Worker_0`** (`beam_r3_w0`): Succeeded. Latency: **`6.160 ms`** + * **`Round_3_Worker_1`** (`beam_r3_w1`): Succeeded. Latency: **`5.852 ms`** + * **`Round_3_Worker_2`** (`beam_r3_w2`): Failed to pass correctness checks. Exited after max retries. + * **`Round_3_Worker_3`** (`beam_r3_w3`): Succeeded after self-healing compilation validation errors. Latency: **`5.559 ms`** + +* **Incumbent Preservation**: Round 3 Worker 2 failed. The remaining active candidates from Round 3 were Rank 1 ($5.559\text{ ms}$), Rank 2 ($5.852\text{ ms}$), and Rank 3 ($6.160\text{ ms}$). Because the beam size is $B=4$, the orchestrator engaged the **Incumbent Exception Rule**, comparing the Round 2 beam's best candidate (`beam_r2_w3` at $6.826\text{ ms}$) against Round 3 candidates. Since $6.826\text{ ms}$ is faster than having an empty slot or invalid failure, it was successfully preserved as Rank 4. + +#### Round 3 Beam State (Final State): +``` +================================================== + ROUND 3 BEAM STATE +================================================== + Rank 1: Latency=5.559 ms | Path=beam_search/output/12p_RMSNorm_bw_r3_b4_0625_0216/beam_r3_w3_1782357069/optimized_kernel.py + Rank 2: Latency=5.852 ms | Path=beam_search/output/12p_RMSNorm_bw_r3_b4_0625_0216/beam_r3_w1_1782357069/optimized_kernel.py + Rank 3: Latency=6.160 ms | Path=beam_search/output/12p_RMSNorm_bw_r3_b4_0625_0216/beam_r3_w0_1782357069/optimized_kernel.py + Rank 4: Latency=6.826 ms | Path=beam_search/output/12p_RMSNorm_bw_r3_b4_0625_0216/beam_r2_w3_1782354715/optimized_kernel.py (Incumbent Preserved) +================================================== +``` + +--- + +## 3. Final Search Summary & Performance + +* **Best Kernel Candidate**: Produced by `Round_3_Worker_3` at: + [optimized_kernel.py](file:///usr/local/google/home/ligh/github/accelerator-agents/MaxKernel/beam_search/output/12p_RMSNorm_bw_r3_b4_0625_0216/beam_r3_w3_1782357069/optimized_kernel.py) +* **Baseline Latency**: `11.400 ms` +* **Best Optimized Latency**: **`5.559 ms`** +* **Speedup**: **`2.05x` speedup** (a **`51.2%` latency reduction**). + +--- + +## 4. Scheduler and Evaluator Mechanics Verified + +1. **Correctness Self-Healing**: Successfully witnessed the self-healing validator repair compilation issues and syntax mismatches in Round 2 (Worker 2) and Round 3 (Worker 3) dynamically, restoring execution pipelines without manual intervention. +2. **Incumbent Exception Safeguard**: Successfully verified that the orchestrator preserves the best candidate of previous rounds when subsequent rounds yield failures, keeping the beam filled with valid candidates. +3. **FastAPI Queue Execution**: Verified that parallel workers coordinate JAX timing requests cleanly on a serialized lock via FastAPI, ensuring jitter-free latency metrics. diff --git a/MaxKernel/beam_search/docs/worker_integration_plan.md b/MaxKernel/beam_search/docs/worker_integration_plan.md new file mode 100644 index 0000000..a7d1620 --- /dev/null +++ b/MaxKernel/beam_search/docs/worker_integration_plan.md @@ -0,0 +1,198 @@ +# Implementation & Integration Plan: MaxKernel Beam Search + +This document outlines the architectural design, refactoring strategy, and implementation phases to integrate the Beam Search optimization engine (from AutoComp) into the **MaxKernel** repository. + +--- + +## 1. Goal & Context +Our objective is to combine AutoComp's high-dimensional beam search engine with MaxKernel's modular ADK-based agent orchestration. + +To achieve this efficiently, we split the responsibilities: +1. **Orchestrator Layer (`beam_search/`)**: Manages the beam repository, spawns worker agents in parallel, groups candidate implementations, and dispatches them to the evaluation backend for performance profiling. +2. **Worker Layer (`auto_agent/` - Subagent)**: Autonomously ensures code correctness (local compilation and unit tests) using a self-correction loop. It exits early as soon as the code passes correctness checks, delegating latency profiling to the orchestrator. + +--- + +## 2. Target Directory Structure +The orchestrator and worker agent will exist as sibling packages at the root of the `MaxKernel/` repository to maintain a clean separation of concerns: + +``` +MaxKernel/ +├── beam_search/ # Top-level Search Orchestrator (Main entry point) +│ ├── __init__.py +│ ├── orchestrator.py # Coordinates the Beam Search loop & rounds +│ ├── tools.py # Orchestrator tools (e.g. grouped harness generator) +│ └── docs/ # Documentation specific to Beam Search +│ ├── worker_integration_plan.md # This design document +│ └── worker_integration_checklist.md # Step-by-step progress checklist +│ +├── auto_agent/ # Autonomous Worker Agent (Spawned as subagents) +│ ├── agent.py # Entry point for standard single-agent runs +│ ├── beam_worker.py # Entry point for correctness-only worker runs +│ ├── subagents/ # Planner, Implementer, Validator, Pipelines +│ │ ├── beam_worker_pipeline.py # Shorter correctness-only pipeline for search +│ │ └── pipeline_agent.py # Production end-to-end tuning pipeline +│ └── tools/ # Compiler & local test checks used by workers +│ +├── evaluation/ # Backend profiling files +│ ├── jax_kernel_evaluator.py # Production TPUVM Pallas evaluator +│ ├── fake_kernel_evaluator.py # Mock/Fake evaluator for local verification +│ └── ... +``` + +--- + +## 3. Architecture: Why `BaseAgent` and not `LlmAgent`? +A key detail of MaxKernel's design is that coordinator agents (like `AutonomousPipelineAgent` and the proposed `BeamWorkerPipeline`) subclass ADK's `BaseAgent` directly, rather than `LlmAgent`. + +* **`LlmAgent` (Leaf/Reasoning Agents)**: Designed for agents whose execution loop is driven by prompting an LLM (using system instructions, chat history, and tool calls). Examples: `plan_kernel_agent`, `implement_kernel_agent`. +* **`BaseAgent` (Coordinator Agents)**: The abstract foundation for any agent that implements custom execution logic via Python code (`_run_async_impl`). +* **The Coordinator Role**: The pipeline agent acts as a deterministic state machine that chains subagents together. It does not need its own LLM reasoning loop to decide its next step; the execution sequence is defined programmatically in Python. Therefore, it subclasses `BaseAgent` to run Python control flow while delegating the LLM-heavy "thinking" steps to its `LlmAgent` subagents. + +--- + +## 4. Refactoring Design Options for the Worker Pipeline + +We evaluated two options for implementing the correctness-only worker in MaxKernel. + +### Option A: Subclass `AutonomousPipelineAgent` (Recommended) +We subclass `AutonomousPipelineAgent` and override the `_run_async_impl` method to shorten the loop. + +#### **Implementation Sketch (`auto_agent/subagents/beam_worker_pipeline.py`):** +```python +from auto_agent.subagents.pipeline_agent import AutonomousPipelineAgent +from google.adk.events import Event, EventActions +from google.adk.agents.invocation_context import InvocationContext +from typing import AsyncGenerator +import logging + +class BeamWorkerPipeline(AutonomousPipelineAgent): + """Subclass of AutonomousPipelineAgent that exits early on correctness success.""" + + async def _run_async_impl( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: + iteration = 0 + yield self._initialize_state(ctx) + + while iteration < self.max_iterations: + logging.info(f"[{self.name}] Iteration {iteration + 1}/{self.max_iterations}") + + # Reuse parent's subagent steps + async for event in self.plan_agent.run_async(ctx): yield event + async for event in self.implement_agent.run_async(ctx): yield event + async for event in self.validate_agent.run_async(ctx): yield event + + if not ctx.session.state.get("kernel_compilation_status", {}).get("success", False): + iteration += 1 + continue + + async for event in self.test_gen_agent.run_async(ctx): yield event + if not ctx.session.state.get("validation_loop_status", {}).get("success", False): + iteration += 1 + continue + + async for event in self.test_run_agent.run_async(ctx): yield event + + # Check correctness success and EXIT EARLY + test_results = ctx.session.state.get("test_results", {}) + if test_results.get("success", False): + optimized_code_path = ctx.session.state.get("optimized_kernel_path") + logging.info(f"[{self.name}] Correctness achieved: {optimized_code_path}") + yield Event( + author=self.name, + actions=EventActions( + state_delta={ + "worker_status": "Success", + "final_correct_code_path": optimized_code_path + } + ) + ) + return # Stop the pipeline loop immediately + + logging.error(f"[{self.name}] Correctness tests failed. Looping back.") + iteration += 1 + + yield Event( + author=self.name, + actions=EventActions(state_delta={"worker_status": "Failed"}) + ) +``` + +#### **Pros & Cons:** +* **Pros**: + * Clean separation: The execution loop logic for "tuning search" and "correctness validation" remain separate. + * No change to existing production pipeline code, minimizing regression risk. +* **Cons**: Requires maintaining the overridden async loop. + +--- + +### Option B: Optional Subagents in `AutonomousPipelineAgent` +We modify the base `AutonomousPipelineAgent` to accept `None` for `autotune_agent` and `profile_agent`, and adjust the control flow logic internally. + +#### **Changes to `AutonomousPipelineAgent`:** +1. Make subagents optional in constructor: + ```python + autotune_agent: BaseAgent | None = None + profile_agent: BaseAgent | None = None + ``` +2. Update the loop to check for `None` before executing steps 6 (Autotune) and 7 (Profile). +3. Add an **early exit check** after correctness test validation if tuning is disabled. + +#### **Pros & Cons:** +* **Pros**: + * No subclassing required. Single class manages all variants. +* **Cons**: + * Modifies core MaxKernel pipeline control logic, which requires careful testing. + * Introduces conditional checks inside a previously clean state machine loop. + +### Recommendation +We recommend **Option A (Subclassing)**. The control flow of `AutonomousPipelineAgent` is highly tuned for iterative improvement (running multiple iterations to find the *best* candidate even if the first one is correct). A Beam Search worker wants to exit *immediately* on the first correct compilation to save time, as the outer orchestrator handles performance selection. Subclassing allows us to cleanly express this difference in loop invariants without adding complex conditional checks inside the production pipeline. + +--- + +## 5. Detailed Implementation Phases + +### Phase 1: Implement the Correctness Worker in `auto_agent` +We will reuse MaxKernel's existing subagents by implementing a specialized correctness pipeline. + +1. **Create `BeamWorkerPipeline`**: + * Implement `MaxKernel/auto_agent/subagents/beam_worker_pipeline.py` subclassing `AutonomousPipelineAgent` (Option A). + * Add an early-exit check: as soon as `test_results.get("success")` is `True`, yield a success event with the `optimized_kernel_path` and terminate the generator loop. +2. **Expose the Worker Entrypoint**: + * Create `MaxKernel/auto_agent/beam_worker.py`. + * Instantiate `BeamWorkerPipeline` using the existing production instances: `plan_kernel_agent`, `implement_kernel_agent`, `validate_kernel_compilation_agent`, `validated_test_generation_agent`, and `unified_test_agent`. + +### Phase 2: Add Fake Evaluation Backend +To support local testing of the orchestrator loop without requiring physical TPU/Trainium hardware, we will port the fake backend as a standalone Python file. + +1. **Implement `FakeKernelEvaluator`**: + * Create `MaxKernel/evaluation/fake_kernel_evaluator.py`. + * The evaluator should accept multiple inlined code candidates, simulate compilation, and return randomized (but decreasing across rounds) execution latencies in milliseconds. + +### Phase 3: Implement the Beam Search Orchestrator +Implement the top-layer coordinator that manages the exploration loop. + +1. **Create Orchestrator Logic**: + * Create `MaxKernel/beam_search/orchestrator.py`. + * Implement `AgenticSearchOrchestrator` which: + * Accepts a starting baseline Pallas kernel and a target test harness. + * Spawns multiple `beam_worker_agent` instances in parallel (using `asyncio.gather`) with distinct optimization prompt directives (e.g. tiling-focus vs. memory-bound focus). + * Gathers the correct code paths from the completed worker sessions. + * Renames the entry point functions (e.g., `solution_0`, `solution_1`) and inlines them into a single grouped test harness file. + * Dispatches the grouped harness to the evaluation backend (`FakeKernelEvaluator` or real TPU backend `JAXKernelEvaluator`) for latency profiling. + * Ranks and filters the candidates to select the top candidates for the next round. +2. **Configure Mock Mode**: + * In `beam_search/tools.py` (or directly in the orchestrator config), read the environment variable `MOCK_COMPILER=True` or a command-line flag `--mock` to toggle between calling the `FakeKernelEvaluator` (simulated runs) and the real hardware test runner. + +### Phase 4: Local Verification +Verify the integrated orchestrator-worker loop. + +1. **Create local runner script**: + * Implement `MaxKernel/run_beam_search.sh` or a Python helper script to trigger `beam_search/orchestrator.py` with mock mode enabled. +2. **Run a test round**: + * Verify that: + * The orchestrator correctly spawns multiple worker pipelines. + * The workers execute the MaxKernel implementation agents and compile locally. + * On correctness success, workers exit early. + * The orchestrator groups the outputs, calls the fake backend, and prints the latency summary. diff --git a/MaxKernel/beam_search/orchestrator.py b/MaxKernel/beam_search/orchestrator.py new file mode 100644 index 0000000..b88c75b --- /dev/null +++ b/MaxKernel/beam_search/orchestrator.py @@ -0,0 +1,603 @@ +"""Top-level Beam Search orchestrator for MaxKernel.""" + +import asyncio +import json +import logging +import shutil +import os +import random +import re +import time +from pathlib import Path +from typing import List, Optional +from dotenv import load_dotenv + +# ADK imports +from google.adk import Runner +from google.adk.agents import RunConfig +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.genai import types + +# MaxKernel imports +from auto_agent.agent import root_agent +from auto_agent.beam_worker import beam_worker_agent +from auto_agent.config import WORKDIR +from beam_search.utils import normalize_ast_code +from evaluation.fake_kernel_evaluator import FakeKernelEvaluator +from evaluation.jax_kernel_evaluator import JAXKernelEvaluator + +logger = logging.getLogger(__name__) + +# Complete TPU/Pallas optimization strategy list from autocomp +TPU_PALLAS_OPTIMIZATION_STRATEGIES = [ + "Reduce data movement", + "Overlap data movement and compute", + "Cache reused data in local memory instead of reloading from main memory", + "Loop tiling", + "Loop reordering and restructuring", + "Loop unrolling", + "Fuse operations", + "Use lower precision", + "Double buffering", + "Software pipelining", + "Hoist redundant operations out of loops", + "Eliminate redundant computation", + "Simplify or remove unnecessary code", + "Try new parameter values", + "Rewrite the algorithm to reduce total work", + "Place reduction axis last in grid to enable in-place SRAM accumulation without HBM round-trips", + "Align block dimensions to 8x128 tile boundaries to avoid wasted padding and register spills", + "Use scratch_shapes=[pltpu.VMEM(...)] for persistent high-precision accumulators during reduction loops", + "Maximize block sizes up to ~16 MB VMEM capacity to increase arithmetic intensity per pipeline step", + "Use scalar prefetch via PrefetchScalarGridSpec to load indices/metadata into SMEM without stalling vector core", + "Upcast bf16/int8 to float32 before elementwise ops, downcast only on final output write", + "Fuse transpose into lax.dot_general contraction dimensions instead of materializing transposed operands", + "Arrange grid iteration order so consecutive invocations reuse already-resident input slices", + "Increase pipeline buffer count beyond double buffering to hide memory latency for bandwidth-bound kernels", + "Generate random numbers inside kernel via hardware PRNG with key in SMEM instead of passing precomputed arrays", + "Avoid singleton dimensions in last two array axes to prevent full-tile waste per element", + "Reduce along second-to-last dimension rather than last dimension when possible", + "Prefer add/multiply over exp/tanh/division; restructure math to minimize expensive elementwise ops", + "Tune block sizes jointly — systematically vary BM, BN, BK together under the VMEM budget constraint (16 MiB including double-buffering)", + "Compute arithmetic intensity accounting for tiling amplification to predict compute-bound vs memory-bound regime", + "Minimize control flow inside kernels; consolidate into single basic blocks to avoid unrolling overhead", + "Pass all data as explicit kernel inputs with BlockSpec instead of closing over constants", + "Use pltpu.VMEM scoped scratch buffers for temporary storage within kernel lifetime", + "Balance block size against pipeline depth to amortize startup/drain bubble cost over enough iterations", + "Explicitly initialize accumulator buffers to zero on first reduction iteration since SRAM starts undefined" +] + + +# Load API key configuration +local_env = Path(__file__).resolve().parents[1] / ".env" +if local_env.exists(): + load_dotenv(local_env) + + +class AgenticSearchOrchestrator: + """Orchestration layer that manages candidate exploration loops.""" + + def __init__( + self, + baseline_code_path: str, + reference_code_path: str, + task_yaml_path: str, + output_dir: Path, + mock_mode: Optional[bool] = None, + use_beam_worker: bool = True, # True = BeamWorkerPipeline, False = AutonomousPipelineAgent + ): + self.baseline_code_path = baseline_code_path + self.reference_code_path = reference_code_path + self.task_yaml_path = task_yaml_path + self.output_dir = output_dir + self.use_beam_worker = use_beam_worker + + # Configure evaluator backend (defaulting to MOCK_COMPILER env var) + if mock_mode is None: + mock_mode = os.environ.get("MOCK_COMPILER", "true").lower() == "true" + self.mock_mode = mock_mode + + if self.mock_mode: + logger.info("[Orchestrator] Running in MOCK mode. Using FakeKernelEvaluator.") + self.evaluator = FakeKernelEvaluator() + else: + logger.info("[Orchestrator] Running in PRODUCTION mode. Using JAXKernelEvaluator.") + self.evaluator = JAXKernelEvaluator(local=True) + + self.output_dir.mkdir(parents=True, exist_ok=True) + + # Rebind WORKDIR in config and all modules that imported it to ensure consistency + workdir_str = str(self.output_dir) + os.environ["WORKDIR"] = workdir_str + + import auto_agent.config + auto_agent.config.WORKDIR = workdir_str + + import auto_agent.subagents.pipeline_agent + auto_agent.subagents.pipeline_agent.WORKDIR = workdir_str + + import auto_agent.callbacks + auto_agent.callbacks.WORKDIR = workdir_str + + import auto_agent.tools.file_tools + auto_agent.tools.file_tools.WORKDIR = workdir_str + + # Update local reference in orchestrator module + global WORKDIR + WORKDIR = workdir_str + + # Read baseline source code + with open(self.baseline_code_path, "r") as f: + self.baseline_code = f.read() + + def _prepare_grouped_harness(self, candidate_paths: List[str]) -> str: + """Inlines all candidate implementations into a single file.""" + inlined_solutions = [] + for idx, path in enumerate(candidate_paths): + if not os.path.exists(path): + continue + try: + with open(path, "r") as f: + code = f.read() + # Rename function computation() to solution_{idx}() to prevent symbol collision and match mock evaluator expectations + renamed_code = re.sub(r"def computation\(", f"def solution_{idx}(", code) + inlined_solutions.append( + f"# --- Candidate {idx} ---\n{renamed_code}" + ) + except Exception as e: + logger.error(f"Failed to inline candidate {path}: {e}") + + return "\n\n".join(inlined_solutions) + + async def run_worker_session( + self, + worker_name: str, + session_id: str, + prompt_focus: str, + base_code: str, + parent_latency: float, + ) -> Optional[dict]: + """Spawns an isolated worker session and extracts its optimized kernel path.""" + logger.info(f"[Orchestrator] Spawning {worker_name} (session_id={session_id})...") + + # Select the agent class based on mode flag + if self.use_beam_worker: + agent_to_run = beam_worker_agent + logger.info(f"[{worker_name}] Mode: BeamWorkerPipeline (correctness-only).") + else: + agent_to_run = root_agent + # Allow up to 5 iterations for correctness repairs, but stop on first valid profiled candidate + agent_to_run.max_iterations = 5 + agent_to_run.stop_on_first_valid = True + logger.info( + f"[{worker_name}] Mode: AutonomousPipelineAgent (stop_on_first_valid=True, max_iterations=5)." + ) + + session_service = InMemorySessionService() + runner = Runner( + app_name=worker_name, + agent=agent_to_run, + session_service=session_service + ) + + user_message = f"Optimize the JAX/Pallas kernel. Focus: {prompt_focus}" + content = types.Content( + role="user", + parts=[types.Part.from_text(text=user_message)] + ) + + # Create the session and pre-inject the base kernel path so the worker can read it + session = await session_service.create_session( + app_name=worker_name, + user_id="orchestrator", + session_id=session_id + ) + + # Write the base candidate code to the expected session work directory + session_dir = os.path.join(WORKDIR, session_id) + os.makedirs(session_dir, exist_ok=True) + session_base_path = os.path.join(session_dir, "base_kernel.py") + with open(session_base_path, "w") as f: + f.write(base_code) + + # Pre-populate state + session.state["base_kernel_path"] = session_base_path + session.state["parent_latency_ms"] = parent_latency + + # Save back to database + session_service.sessions.setdefault(worker_name, {}).setdefault("orchestrator", {})[session_id] = session + + correct_code_path = None + + # Execute the worker pipeline + async for event in runner.run_async( + user_id="orchestrator", + session_id=session_id, + new_message=content, + run_config=RunConfig() + ): + if event.actions and event.actions.state_delta: + delta = event.actions.state_delta + if "final_correct_code_path" in delta: + correct_code_path = delta["final_correct_code_path"] + logger.info(f"[{worker_name}] Reported correctness achieved at: {correct_code_path}") + + # Inspect final session state + final_session = await session_service.get_session( + app_name=worker_name, + user_id="orchestrator", + session_id=session_id + ) + + if not final_session: + logger.warning(f"[{worker_name}] No final session found.") + return None + + # Extract optimized code path based on Mode + if not self.use_beam_worker: + # Mode 1 (AutonomousPipelineAgent): Check if a best iteration was selected + best_iter = final_session.state.get("best_iteration", -1) + if best_iter != -1: + opt_path = final_session.state.get("optimized_kernel_path") + logger.info(f"[{worker_name}] Output tuned code from iteration {best_iter}: {opt_path}") + return { + "path": opt_path, + "parent_latency_ms": parent_latency + } + logger.warning(f"[{worker_name}] Pipeline failed correctness tests in Mode 1.") + return None + else: + # Mode 2 (BeamWorkerPipeline): Check worker success status + if final_session.state.get("worker_status") == "Success": + opt_path = correct_code_path or final_session.state.get("optimized_kernel_path") + logger.info(f"[{worker_name}] Output correct code: {opt_path}") + return { + "path": opt_path, + "parent_latency_ms": parent_latency + } + logger.warning(f"[{worker_name}] Pipeline failed correctness tests in Mode 2.") + return None + + def _evaluate_candidate_performance( + self, path: str, parent_latency: float, fallback_analysis: str = "" + ) -> dict: + """Evaluates a single candidate and parses its performance metrics.""" + opt_dir = os.path.dirname(path) + + # Try to load autotune results first (Mode 1 auto_agent output) + autotune_json = os.path.join(opt_dir, "autotune_results.json") + best_latency = None + if os.path.exists(autotune_json): + try: + with open(autotune_json, "r") as f: + data = json.load(f) + if "best_time_ms" in data: + best_latency = float(data["best_time_ms"]) + logger.info(f"[Orchestrator] Found autotune latency {best_latency} ms in {autotune_json}") + except Exception as e: + logger.warning(f"Failed to read autotune results from {autotune_json}: {e}") + + # Fallback to evaluation_metrics.json if autotune was missing but metrics exist + metrics_json = os.path.join(opt_dir, "evaluation_metrics.json") + hbm_util = 0.0 + comp_util = 0.0 + density = 0.0 + analysis = fallback_analysis or "Evaluated performance" + + if best_latency is None and os.path.exists(metrics_json): + try: + with open(metrics_json, "r") as f: + meta = json.load(f) + perf_analysis = meta.get("performance_analysis", "") + match = re.search(r"latency:\s*([\d.]+)\s*ms", perf_analysis, re.IGNORECASE) + if match: + best_latency = float(match.group(1)) + logger.info(f"[Orchestrator] Parsed latency {best_latency} ms from metrics analysis text") + except Exception as e: + logger.warning(f"Failed to parse latency from existing metrics file: {e}") + + if best_latency is not None: + logger.info(f"[Orchestrator] Bypassing evaluator for candidate: {path}. Using parsed latency: {best_latency} ms") + latency_ms = best_latency + if self.mock_mode: + hbm_util = 10.0 + comp_util = 2.0 + density = 0.5 + analysis = f"Bypassed evaluator (mock metrics). Latency: {best_latency:.3f} ms" + else: + if os.path.exists(metrics_json): + try: + with open(metrics_json, "r") as f: + meta = json.load(f) + hbm_util = meta.get("estimated_hbm_utilization", 0.0) + comp_util = meta.get("estimated_compute_utilization", 0.0) + density = meta.get("estimated_computation_density_flops_per_byte", 0.0) + analysis = meta.get("performance_analysis", "Bypassed evaluator") + except Exception as e: + logger.warning(f"Failed to read metrics from {metrics_json}: {e}") + else: + logger.info(f"[Orchestrator] No cached metrics found. Evaluating candidate: {path}") + eval_res = self.evaluator.evaluate( + reference_code_path=self.reference_code_path, + optimized_code_path=path, + task_yaml_path=self.task_yaml_path + ) + latency_ms = eval_res.optimized_time_ms + if os.path.exists(metrics_json): + try: + with open(metrics_json, "r") as f: + meta = json.load(f) + hbm_util = meta.get("estimated_hbm_utilization", 0.0) + comp_util = meta.get("estimated_compute_utilization", 0.0) + density = meta.get("estimated_computation_density_flops_per_byte", 0.0) + analysis = meta.get("performance_analysis", "") + except Exception as e: + logger.warning(f"Failed to read evaluation metrics from {metrics_json}: {e}") + + with open(path, "r") as f: + code_content = f.read() + + return { + "path": path, + "latency_ms": latency_ms, + "parent_latency_ms": parent_latency, + "code": code_content, + "hbm_utilization_pct": hbm_util, + "compute_utilization_pct": comp_util, + "computation_density_flops_per_byte": density, + "analysis": analysis + } + + def _run_sequential_fallback_evaluation( + self, valid_results: List[dict], fallback_analysis: str = "Sequential fallback prediction" + ) -> List[dict]: + """Evaluates candidates sequentially when batched evaluation fails or is skipped.""" + round_candidates = [] + for res_entry in valid_results: + cand = self._evaluate_candidate_performance( + path=res_entry["path"], + parent_latency=res_entry["parent_latency_ms"], + fallback_analysis=fallback_analysis + ) + round_candidates.append(cand) + return round_candidates + + def _apply_parent_regression_gate( + self, candidates: List[dict], keep_factor: float + ) -> List[dict]: + """Filters out candidates whose latency regresses beyond the parent latency threshold.""" + valid_candidates = [] + for cand in candidates: + parent_limit = cand["parent_latency_ms"] * keep_factor + if cand["latency_ms"] >= parent_limit: + logger.info( + f"[Orchestrator] Pruning candidate {cand['path']} (latency {cand['latency_ms']:.3f} ms " + f">= parent limit {parent_limit:.3f} ms)" + ) + continue + valid_candidates.append(cand) + return valid_candidates + + def _deduplicate_candidates_ast(self, candidates: List[dict]) -> List[dict]: + """Deduplicates candidates based on their normalized AST representations.""" + seen_ast_hashes = set() + deduped_candidates = [] + for cand in candidates: + normalized_code = normalize_ast_code(cand["code"]) + ast_hash = "".join(normalized_code.split()) + if ast_hash not in seen_ast_hashes: + seen_ast_hashes.add(ast_hash) + deduped_candidates.append(cand) + else: + logger.info( + f"[Orchestrator] Pruning duplicate candidate {cand['path']} (AST duplicate of earlier candidate)" + ) + return deduped_candidates + + def _log_round_beam_state(self, round_idx: int, beam: List[dict]) -> None: + """Logs the beam ranking status at the end of a round.""" + print(f"\n==================================================") + print(f" ROUND {round_idx} BEAM STATE") + print(f"==================================================") + for rank, cand in enumerate(beam): + print(f" Rank {rank+1}: Latency={cand['latency_ms']:.3f} ms | Path={cand['path']}") + if self.mock_mode: + print( + f" HBM Util={cand['hbm_utilization_pct']}% | " + f"Comp Util={cand['compute_utilization_pct']}% | " + f"density={cand['computation_density_flops_per_byte']:.2f}" + ) + print(f"==================================================\n") + + def _prepare_worker_tasks( + self, round_idx: int, beam: List[dict], beam_size: int, dropout_menu_options: float + ) -> List: + """Prepares run_worker_session tasks for all worker beams in a round.""" + tasks = [] + for idx in range(beam_size): + # Select parent candidate (branching if len(beam) < beam_size) + parent_candidate = beam[idx % len(beam)] + + # Apply stochastic dropout to optimization strategies menu + selected_opts = [ + opt for opt in TPU_PALLAS_OPTIMIZATION_STRATEGIES + if random.random() < dropout_menu_options + ] + # Fallback: ensure at least one strategy focus is passed + if not selected_opts: + selected_opts = [random.choice(TPU_PALLAS_OPTIMIZATION_STRATEGIES)] + + focus_text = "\n".join(f"- {opt}" for opt in selected_opts) + prompt_focus = ( + "Focus your optimization effort on applying the following strategies:\n" + f"{focus_text}" + ) + + session_id = f"beam_r{round_idx}_w{idx}_{int(time.time())}" + tasks.append( + self.run_worker_session( + worker_name=f"Round_{round_idx}_Worker_{idx}", + session_id=session_id, + prompt_focus=prompt_focus, + base_code=parent_candidate["code"], + parent_latency=parent_candidate["latency_ms"] + ) + ) + return tasks + + def _evaluate_round_candidates( + self, round_idx: int, valid_results: List[dict] + ) -> List[dict]: + """Evaluates the round candidates using either grouped harness or sequential fallback.""" + valid_paths = [r["path"] for r in valid_results] + round_candidates = [] + + if not self.use_beam_worker: + # Mode 1: Evaluate candidates individually (since they have already run local autotuning) + logger.info(f"[Round {round_idx}] Mode 1: Evaluating candidates individually...") + round_candidates = self._run_sequential_fallback_evaluation( + valid_results, fallback_analysis="Individual autotuned run" + ) + + else: + # Mode 2: Dispatch a GROUPED harness containing all candidates + logger.info(f"[Round {round_idx}] Mode 2: Compiling grouped harness for {len(valid_paths)} candidates...") + grouped_code = self._prepare_grouped_harness(valid_paths) + grouped_path = self.output_dir / f"grouped_harness_r{round_idx}.py" + with open(grouped_path, "w") as f: + f.write(grouped_code) + + # Dispatch grouped harness to the evaluation backend + logger.info(f"[Round {round_idx}] Dispatching grouped harness to evaluation backend...") + try: + self.evaluator.evaluate( + reference_code_path=self.reference_code_path, + optimized_code_path=str(grouped_path), + task_yaml_path=self.task_yaml_path + ) + except Exception as e: + logger.warning(f"[Round {round_idx}] Grouped evaluation failed with error: {e}") + + # Parse metrics from JSON file + metrics_json = os.path.join(os.path.dirname(grouped_path), "evaluation_metrics.json") + valid_paths_fallback = False + if os.path.exists(metrics_json): + try: + with open(metrics_json, "r") as f: + meta = json.load(f) + + candidates_data = meta.get("candidates", []) + for idx, data in enumerate(candidates_data): + cand_id = data.get("candidate_id", idx) + if cand_id >= len(valid_results): + continue + res_entry = valid_results[cand_id] + path = res_entry["path"] + with open(path, "r") as f: + code_content = f.read() + + round_candidates.append({ + "path": path, + "latency_ms": data.get("latency_ms", self.evaluator.reference_time_ms), + "parent_latency_ms": res_entry["parent_latency_ms"], + "code": code_content, + "hbm_utilization_pct": data.get("estimated_hbm_utilization", 0.0), + "compute_utilization_pct": data.get("estimated_compute_utilization", 0.0), + "computation_density_flops_per_byte": data.get("estimated_computation_density", 0.0), + "analysis": data.get("analysis", "") + }) + except Exception as e: + logger.error(f"[Round {round_idx}] Failed to parse grouped evaluation metrics: {e}") + valid_paths_fallback = True + else: + valid_paths_fallback = True + + # Fallback to sequential evaluation if file parsing was unsuccessful + if valid_paths_fallback or not round_candidates: + logger.warning(f"[Round {round_idx}] Grouped harness metrics missing. Falling back to sequential evaluation.") + round_candidates = self._run_sequential_fallback_evaluation(valid_results) + + return round_candidates + + async def run_search( + self, + num_rounds: int = 3, + beam_size: int = 2, + dropout_menu_options: float = 0.5, + keep_factor: float = 1.0, + ) -> dict: + """Executes the complete multi-round Beam Search optimization process.""" + t0 = time.perf_counter() + logger.info( + f"[Orchestrator] Starting Beam Search: {num_rounds} rounds, Beam Size = {beam_size}, " + f"dropout_menu_options={dropout_menu_options}, keep_factor={keep_factor}, " + f"Mode={'BeamWorker' if self.use_beam_worker else 'AutonomousAgent'}" + ) + + # 1. Establish initial baseline performance + logger.info("[Orchestrator] Evaluating baseline kernel...") + base_cand = self._evaluate_candidate_performance( + path=self.baseline_code_path, + parent_latency=0.0, + fallback_analysis="Baseline reference" + ) + # Override baseline metadata defaults + base_cand["code"] = self.baseline_code + base_cand["parent_latency_ms"] = base_cand["latency_ms"] + beam = [base_cand] + + logger.info(f"[Orchestrator] Baseline Latency: {base_cand['latency_ms']:.3f} ms") + + for round_idx in range(1, num_rounds + 1): + logger.info(f"\n=== STARTING SEARCH ROUND {round_idx}/{num_rounds} ===") + + # 2. Distribute candidates from the beam to workers and run them + tasks = self._prepare_worker_tasks(round_idx, beam, beam_size, dropout_menu_options) + worker_results = await asyncio.gather(*tasks) + valid_results = [r for r in worker_results if r is not None] + + logger.info(f"[Round {round_idx}] Collected {len(valid_results)} correct candidate outputs.") + if not valid_results: + logger.warning(f"[Round {round_idx}] No valid candidates generated this round. Skipping evaluation.") + continue + + # 3. Profile mutations + round_candidates = self._evaluate_round_candidates(round_idx, valid_results) + + # 5. Apply Parent Regression Budget Gate (Incumbents are exempt) + valid_round_candidates = self._apply_parent_regression_gate(round_candidates, keep_factor) + + # Merge survivors with incumbents + all_candidates = beam + valid_round_candidates + + # 6. Deduplicate candidates using AST-based normalization + deduped_candidates = self._deduplicate_candidates_ast(all_candidates) + + # Sort by latency (lower is better) + deduped_candidates.sort(key=lambda x: x["latency_ms"]) + + # 7. Prune beam to beam_size + beam = deduped_candidates[:beam_size] + + self._log_round_beam_state(round_idx, beam) + + # Return final best candidate + best_candidate = beam[0] + elapsed = time.perf_counter() - t0 + logger.info(f"[Orchestrator] Beam Search finished. Duration: {elapsed:.2f} s. Best Latency: {best_candidate['latency_ms']:.3f} ms") + + # Save the final winning optimized code to output_dir + best_output_path = self.output_dir / "best_optimized_kernel.py" + try: + shutil.copy(best_candidate["path"], best_output_path) + logger.info(f"[Orchestrator] Saved winning optimized kernel to {best_output_path}") + best_path_to_return = str(best_output_path) + except Exception as e: + logger.error(f"[Orchestrator] Failed to copy winning kernel to {best_output_path}: {e}") + best_path_to_return = best_candidate["path"] + + return { + "status": "success", + "best_latency_ms": best_candidate["latency_ms"], + "best_code_path": best_path_to_return + } diff --git a/MaxKernel/beam_search/run_beam_search.py b/MaxKernel/beam_search/run_beam_search.py new file mode 100644 index 0000000..03cc755 --- /dev/null +++ b/MaxKernel/beam_search/run_beam_search.py @@ -0,0 +1,95 @@ +"""Local dry-run script to verify Beam Search Orchestrator functionality.""" + +import asyncio +import logging +import sys +from pathlib import Path + +# Add MaxKernel root to path +sys.path.append(str(Path(__file__).resolve().parents[1])) + +from beam_search.orchestrator import AgenticSearchOrchestrator + +# Configure logging to show pipeline steps clearly in console +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + handlers=[logging.StreamHandler(sys.stdout)], + force=True +) +logger = logging.getLogger("run_beam_search") + + +async def main(): + import argparse + import time + + # Parse arguments + parser = argparse.ArgumentParser(description="Run Beam Search Orchestrator") + parser.add_argument("--run_id", type=str, default=None, help="Descriptive Run ID") + parser.add_argument("--task_id", type=str, default="12p_RMSNorm", help="Task/Kernel ID (e.g. 12p_RMSNorm)") + parser.add_argument("--rounds", type=int, default=2, help="Number of search rounds") + parser.add_argument("--beam_size", type=int, default=2, help="Beam size") + parser.add_argument( + "--use_beam_worker", + type=str, + default="true", + help="Use BeamWorkerPipeline (correctness-only): true/false", + ) + args = parser.parse_args() + + # Convert string boolean + use_beam_worker = args.use_beam_worker.lower() == "true" + + # Paths to resources + import os + project_root = Path(__file__).resolve().parents[2] + baseline_path = str(project_root / f"JAXBench/benchmark/{args.task_id}/baseline.py") + reference_path = str(project_root / f"JAXBench/benchmark/{args.task_id}/baseline.py") + task_yaml = str(project_root / f"MaxKernel/evaluation/jaxbench_adapted_dataset/{args.task_id}/kernel_task.yaml") + output_dir = project_root / "MaxKernel/beam_search/output" + + # Validate file existence + if not os.path.exists(baseline_path): + logger.error(f"Baseline code not found at: {baseline_path}") + sys.exit(1) + if not os.path.exists(task_yaml): + logger.error(f"Task YAML configuration not found at: {task_yaml}") + sys.exit(1) + + # Build run_id from naming convention if not specified by user + if args.run_id: + run_id = args.run_id + else: + kernel_name = Path(task_yaml).parent.name + worker_short = "bw" if use_beam_worker else "aa" + timestamp_short = time.strftime("%m%d_%H%M") + run_id = f"{kernel_name}_{worker_short}_r{args.rounds}_b{args.beam_size}_{timestamp_short}" + + run_output_dir = output_dir / run_id + + logger.info("="*80) + logger.info(f" STARTING RUN: {run_id}") + logger.info(f" BEAM SIZE = {args.beam_size}, ROUNDS = {args.rounds}") + logger.info("="*80) + + orchestrator_mode2 = AgenticSearchOrchestrator( + baseline_code_path=baseline_path, + reference_code_path=reference_path, + task_yaml_path=task_yaml, + output_dir=run_output_dir, + mock_mode=True, + use_beam_worker=use_beam_worker + ) + + res_mode2 = await orchestrator_mode2.run_search( + num_rounds=args.rounds, + beam_size=args.beam_size, + dropout_menu_options=0.5, + keep_factor=1.0 + ) + logger.info(f"Search Completed. Result: {res_mode2}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/MaxKernel/beam_search/test_orchestrator.py b/MaxKernel/beam_search/test_orchestrator.py new file mode 100644 index 0000000..e4114ec --- /dev/null +++ b/MaxKernel/beam_search/test_orchestrator.py @@ -0,0 +1,389 @@ +"""Unit tests for AgenticSearchOrchestrator helper functions.""" + +import ast +import json +import os +import unittest +from unittest.mock import MagicMock, patch, mock_open +import pytest + +from beam_search.orchestrator import AgenticSearchOrchestrator + +from pathlib import Path + +# Setup dummy paths for constructor instantiation +DUMMY_BASELINE_PATH = "dummy_baseline.py" +DUMMY_REFERENCE_PATH = "dummy_reference.py" +DUMMY_TASK_YAML_PATH = "dummy_task.yaml" +DUMMY_OUTPUT_DIR = Path("dummy_output") + + + +@pytest.fixture +def mock_orchestrator(): + """Fixture to instantiate AgenticSearchOrchestrator with mocked dependencies.""" + with patch("builtins.open", mock_open(read_data="def computation(x): return x")): + with patch("os.path.exists", return_value=True): + with patch("pathlib.Path.mkdir"): + orchestrator = AgenticSearchOrchestrator( + baseline_code_path=DUMMY_BASELINE_PATH, + reference_code_path=DUMMY_REFERENCE_PATH, + task_yaml_path=DUMMY_TASK_YAML_PATH, + output_dir=DUMMY_OUTPUT_DIR, + mock_mode=True, # Use fake evaluator internally + ) + # Mock the internal evaluator completely for safety + orchestrator.evaluator = MagicMock() + return orchestrator + + +def test_apply_parent_regression_gate(mock_orchestrator): + """Test Case 1: Verifies that regression gate correctly filters based on parent limits.""" + candidates = [ + {"path": "c0", "latency_ms": 10.0, "parent_latency_ms": 12.0}, # Improved + {"path": "c1", "latency_ms": 13.0, "parent_latency_ms": 12.0}, # Regressed + {"path": "c2", "latency_ms": 9.2, "parent_latency_ms": 9.5}, # Improved + {"path": "c3", "latency_ms": 10.0, "parent_latency_ms": 9.5}, # Regressed + ] + + # Scenario 1: keep_factor = 1.0 (strict) + survivors_strict = mock_orchestrator._apply_parent_regression_gate(candidates, keep_factor=1.0) + assert len(survivors_strict) == 2 + assert survivors_strict[0]["path"] == "c0" + assert survivors_strict[1]["path"] == "c2" + + # Scenario 2: keep_factor = 1.1 (10% margin, limit for c1 is 13.2 ms, limit for c3 is 10.45 ms) + survivors_loose = mock_orchestrator._apply_parent_regression_gate(candidates, keep_factor=1.1) + assert len(survivors_loose) == 4 + + +def test_deduplicate_candidates_ast(mock_orchestrator): + """Test Case 2: Verifies AST deduplication preserves unique candidates.""" + candidates = [ + {"path": "c0", "code": "def computation(x):\n return x + 1", "latency_ms": 10.0}, + {"path": "c1", "code": "def computation(x):\n # Some comment\n return x + 1", "latency_ms": 11.0}, # AST duplicate of c0 + {"path": "c2", "code": "def computation(x):\n return x + 2", "latency_ms": 9.5}, # Unique + ] + + deduped = mock_orchestrator._deduplicate_candidates_ast(candidates) + assert len(deduped) == 2 + assert deduped[0]["path"] == "c0" + assert deduped[1]["path"] == "c2" + + +@patch("os.path.exists") +@patch("builtins.open") +def test_evaluate_candidate_performance_normal(mock_open, mock_exists, mock_orchestrator): + """Test Case 3: Verifies evaluator is called normally when no timing files exist.""" + mock_orchestrator.evaluator.evaluate = MagicMock() + mock_orchestrator.evaluator.reference_time_ms = 11.4 + mock_orchestrator.evaluator.evaluate.return_value = MagicMock(optimized_time_ms=8.5) + + # Mock file existence: only evaluation_metrics.json exists (written by evaluator) + # but NOT autotune_results.json (prior to eval) + def side_effect_exists(path): + if "autotune_results.json" in path: + return False + if "evaluation_metrics.json" in path: + return True + return True + mock_exists.side_effect = side_effect_exists + + metrics_content = json.dumps({ + "estimated_hbm_utilization": 22.5, + "estimated_compute_utilization": 14.8, + "estimated_computation_density_flops_per_byte": 1.25, + "performance_analysis": "Memory access pattern optimized" + }) + + def mock_open_helper(filename, *args, **kwargs): + content = "" + if "evaluation_metrics.json" in filename: + content = metrics_content + elif filename.endswith(".py"): + content = "def optimized_code(): pass" + m = MagicMock() + m.__enter__.return_value.read.return_value = content + return m + mock_open.side_effect = mock_open_helper + + res = mock_orchestrator._evaluate_candidate_performance( + path="dummy_candidate_path.py", + parent_latency=9.5, + fallback_analysis="Unit test run" + ) + + # Verify evaluate was called + mock_orchestrator.evaluator.evaluate.assert_called_once() + assert res["latency_ms"] == 8.5 + assert res["parent_latency_ms"] == 9.5 + assert res["hbm_utilization_pct"] == 22.5 + assert res["compute_utilization_pct"] == 14.8 + assert res["computation_density_flops_per_byte"] == 1.25 + assert res["analysis"] == "Memory access pattern optimized" + assert res["code"] == "def optimized_code(): pass" + + +@patch("os.path.exists") +@patch("builtins.open") +def test_evaluate_candidate_performance_bypassed_autotune(mock_open, mock_exists, mock_orchestrator): + """Test Case 3b: Verifies evaluator is bypassed when autotune_results.json exists with best_time_ms.""" + mock_orchestrator.evaluator.evaluate = MagicMock() + + # Mock file existence: autotune_results.json exists + def side_effect_exists(path): + if "autotune_results.json" in path: + return True + return False + mock_exists.side_effect = side_effect_exists + + autotune_content = json.dumps({ + "best_time_ms": 6.02, + "best_config": {"BLOCK_M": 64} + }) + + def mock_open_helper(filename, *args, **kwargs): + content = "" + if "autotune_results.json" in filename: + content = autotune_content + elif filename.endswith(".py"): + content = "def optimized_code(): pass" + m = MagicMock() + m.__enter__.return_value.read.return_value = content + return m + mock_open.side_effect = mock_open_helper + + res = mock_orchestrator._evaluate_candidate_performance( + path="dummy_candidate_path.py", + parent_latency=9.5, + fallback_analysis="Unit test run" + ) + + # Verify evaluate was NOT called + mock_orchestrator.evaluator.evaluate.assert_not_called() + assert res["latency_ms"] == 6.02 + assert res["parent_latency_ms"] == 9.5 + assert res["hbm_utilization_pct"] == 10.0 # Mocked metrics + assert res["compute_utilization_pct"] == 2.0 + assert res["computation_density_flops_per_byte"] == 0.5 + assert "Latency: 6.020 ms" in res["analysis"] + assert res["code"] == "def optimized_code(): pass" + + +@patch("os.path.exists") +@patch("builtins.open") +def test_evaluate_candidate_performance_bypassed_metrics(mock_open, mock_exists, mock_orchestrator): + """Test Case 3c: Verifies evaluator is bypassed when evaluation_metrics.json contains parsed latency.""" + mock_orchestrator.evaluator.evaluate = MagicMock() + + # Mock file existence: autotune_results.json does NOT exist, but evaluation_metrics.json exists + def side_effect_exists(path): + if "autotune_results.json" in path: + return False + if "evaluation_metrics.json" in path: + return True + return False + mock_exists.side_effect = side_effect_exists + + metrics_content = json.dumps({ + "estimated_hbm_utilization": 12.0, + "estimated_compute_utilization": 3.0, + "estimated_computation_density_flops_per_byte": 0.6, + "performance_analysis": "Fallback decay latency: 7.185 ms" + }) + + def mock_open_helper(filename, *args, **kwargs): + content = "" + if "evaluation_metrics.json" in filename: + content = metrics_content + elif filename.endswith(".py"): + content = "def optimized_code(): pass" + m = MagicMock() + m.__enter__.return_value.read.return_value = content + return m + mock_open.side_effect = mock_open_helper + + res = mock_orchestrator._evaluate_candidate_performance( + path="dummy_candidate_path.py", + parent_latency=9.5, + fallback_analysis="Unit test run" + ) + + # Verify evaluate was NOT called + mock_orchestrator.evaluator.evaluate.assert_not_called() + assert res["latency_ms"] == 7.185 + assert res["parent_latency_ms"] == 9.5 + assert res["hbm_utilization_pct"] == 10.0 # Mocked defaults for mock mode bypassing + assert res["compute_utilization_pct"] == 2.0 + assert res["computation_density_flops_per_byte"] == 0.5 + assert "Latency: 7.185 ms" in res["analysis"] + assert res["code"] == "def optimized_code(): pass" + + +@patch("os.path.exists") +@patch("builtins.open") +def test_evaluate_candidate_performance_bypassed_prod_metrics(mock_open, mock_exists): + """Test Case 3d: Verifies evaluator is bypassed in production mode and real metrics are loaded.""" + with patch("builtins.open", mock_open(read_data="def computation(x): return x")): + with patch("os.path.exists", return_value=True): + with patch("pathlib.Path.mkdir"): + orchestrator = AgenticSearchOrchestrator( + baseline_code_path=DUMMY_BASELINE_PATH, + reference_code_path=DUMMY_REFERENCE_PATH, + task_yaml_path=DUMMY_TASK_YAML_PATH, + output_dir=DUMMY_OUTPUT_DIR, + mock_mode=False, # Production mode + ) + orchestrator.evaluator = MagicMock() + + # Mock file existence: autotune_results.json exists, and evaluation_metrics.json exists + def side_effect_exists(path): + if "autotune_results.json" in path: + return True + if "evaluation_metrics.json" in path: + return True + return False + mock_exists.side_effect = side_effect_exists + + autotune_content = json.dumps({ + "best_time_ms": 5.4, + "best_config": {"BLOCK_M": 128} + }) + metrics_content = json.dumps({ + "estimated_hbm_utilization": 45.0, + "estimated_compute_utilization": 15.0, + "estimated_computation_density_flops_per_byte": 2.5, + "performance_analysis": "Real GPU profiling results" + }) + + def mock_open_helper(filename, *args, **kwargs): + content = "" + if "autotune_results.json" in filename: + content = autotune_content + elif "evaluation_metrics.json" in filename: + content = metrics_content + elif filename.endswith(".py"): + content = "def optimized_code(): pass" + m = MagicMock() + m.__enter__.return_value.read.return_value = content + return m + mock_open.side_effect = mock_open_helper + + res = orchestrator._evaluate_candidate_performance( + path="dummy_candidate_path.py", + parent_latency=9.5, + fallback_analysis="Unit test run" + ) + + # Verify evaluate was NOT called + orchestrator.evaluator.evaluate.assert_not_called() + assert res["latency_ms"] == 5.4 + assert res["parent_latency_ms"] == 9.5 + assert res["hbm_utilization_pct"] == 45.0 + assert res["compute_utilization_pct"] == 15.0 + assert res["computation_density_flops_per_byte"] == 2.5 + assert res["analysis"] == "Real GPU profiling results" + assert res["code"] == "def optimized_code(): pass" + + +def test_run_sequential_fallback_evaluation(mock_orchestrator): + """Test Case 4: Verifies sequential evaluations over a list of results.""" + valid_results = [ + {"path": "p0", "parent_latency_ms": 11.4}, + {"path": "p1", "parent_latency_ms": 10.0} + ] + + mock_metrics = { + "latency_ms": 9.2, + "parent_latency_ms": 11.4, + "code": "def code(): pass", + "hbm_utilization_pct": 10.0, + "compute_utilization_pct": 5.0, + "computation_density_flops_per_byte": 0.5, + "analysis": "fallback" + } + + with patch.object(mock_orchestrator, "_evaluate_candidate_performance", return_value=mock_metrics) as mock_eval: + res = mock_orchestrator._run_sequential_fallback_evaluation(valid_results) + assert len(res) == 2 + assert res[0]["latency_ms"] == 9.2 + assert mock_eval.call_count == 2 + + +def test_prepare_worker_tasks(mock_orchestrator): + """Test Case 5: Verifies tasks prepare correctly and round-robin allocations are mapped correctly.""" + beam = [ + {"code": "optimized_code_a", "latency_ms": 9.5}, # Rank 1 + {"code": "optimized_code_b", "latency_ms": 10.2} # Rank 2 + ] + + mock_orchestrator.run_worker_session = MagicMock(return_value="mock_coroutine") + + tasks = mock_orchestrator._prepare_worker_tasks( + round_idx=2, + beam=beam, + beam_size=4, + dropout_menu_options=0.5 + ) + + assert len(tasks) == 4 + assert tasks == ["mock_coroutine"] * 4 + + # Extract arguments passed to run_worker_session mock calls + call_args = mock_orchestrator.run_worker_session.call_args_list + assert len(call_args) == 4 + + # Verify round-robin parent candidate allocation + # Task 0 -> Candidate A (idx % 2 == 0) + assert call_args[0].kwargs["base_code"] == "optimized_code_a" + assert call_args[0].kwargs["parent_latency"] == 9.5 + # Task 1 -> Candidate B (idx % 2 == 1) + assert call_args[1].kwargs["base_code"] == "optimized_code_b" + assert call_args[1].kwargs["parent_latency"] == 10.2 + # Task 2 -> Candidate A (idx % 2 == 0) + assert call_args[2].kwargs["base_code"] == "optimized_code_a" + assert call_args[2].kwargs["parent_latency"] == 9.5 + # Task 3 -> Candidate B (idx % 2 == 1) + assert call_args[3].kwargs["base_code"] == "optimized_code_b" + assert call_args[3].kwargs["parent_latency"] == 10.2 + + +def test_beam_sorting_and_pruning(mock_orchestrator): + """Test Case 6: Verifies beam candidates are sorted (lowest latency first) and pruned to beam_size.""" + candidates = [ + {"path": "p0", "latency_ms": 12.1}, + {"path": "p1", "latency_ms": 9.8}, + {"path": "p2", "latency_ms": 11.5}, + {"path": "p3", "latency_ms": 8.4}, + {"path": "p4", "latency_ms": 10.0}, + ] + + # Sort by latency + candidates.sort(key=lambda x: x["latency_ms"]) + # Prune to beam_size = 3 + beam = candidates[:3] + + assert len(beam) == 3 + assert beam[0]["latency_ms"] == 8.4 + assert beam[1]["latency_ms"] == 9.8 + assert beam[2]["latency_ms"] == 10.0 + + +def test_incumbent_preservation_on_total_regression(mock_orchestrator): + """Test Case 7: Verifies that incumbents are preserved and carried forward when new round candidates regress.""" + incumbents = [{"path": "p0", "latency_ms": 9.5, "code": "code_p0"}] + + # New candidates all regress compared to parent + new_candidates = [ + {"path": "c0", "latency_ms": 11.5, "parent_latency_ms": 9.5} + ] + + # Apply parent regression gate (strict 1.0) + survivors = mock_orchestrator._apply_parent_regression_gate(new_candidates, keep_factor=1.0) + assert len(survivors) == 0 # c0 is pruned + + # Merging survivors with incumbents + all_candidates = incumbents + survivors + assert len(all_candidates) == 1 + assert all_candidates[0]["path"] == "p0" + assert all_candidates[0]["latency_ms"] == 9.5 diff --git a/MaxKernel/beam_search/utils.py b/MaxKernel/beam_search/utils.py new file mode 100644 index 0000000..ce78552 --- /dev/null +++ b/MaxKernel/beam_search/utils.py @@ -0,0 +1,73 @@ +"""Utility functions for Beam Search Orchestrator.""" + +import ast +import logging + +logger = logging.getLogger(__name__) + + +class PallasASTNormalizer(ast.NodeTransformer): + """AST NodeTransformer that normalizes variable names and strips metadata.""" + + def __init__(self): + super().__init__() + self.var_map = {} + self.var_counter = 0 + # Reserved names that should not be renamed (modules, main entrypoint, key JAX APIs) + self.reserved_names = { + "jax", + "jnp", + "lax", + "pl", + "pltpu", + "solution", + "mean", + "square", + "rsqrt", + "astype", + "dtype", + "float32", + "bfloat16", + } + + def _get_replacement_name(self, name: str) -> str: + if name in self.reserved_names: + return name + if name not in self.var_map: + self.var_map[name] = f"v{self.var_counter}" + self.var_counter += 1 + return self.var_map[name] + + def visit_Name(self, node): + node.id = self._get_replacement_name(node.id) + return self.generic_visit(node) + + def visit_arg(self, node): + node.arg = self._get_replacement_name(node.arg) + return self.generic_visit(node) + + def visit_FunctionDef(self, node): + # Do not rename the primary solution entrypoint function + if node.name != "solution": + node.name = self._get_replacement_name(node.name) + # Remove docstring expressions if they exist at the start of the function body + if ast.get_docstring(node): + if isinstance(node.body[0], ast.Expr) and isinstance( + node.body[0].value, ast.Constant + ): + node.body.pop(0) + return self.generic_visit(node) + + +def normalize_ast_code(source_code: str) -> str: + """Parses, normalizes, and unparses Python code to compare structural equivalence.""" + try: + tree = ast.parse(source_code) + normalizer = PallasASTNormalizer() + normalized_tree = normalizer.visit(tree) + ast.fix_missing_locations(normalized_tree) + return ast.unparse(normalized_tree) + except Exception as e: + logger.error(f"AST Normalization failed: {e}") + # Return basic whitespace-stripped string fallback on error + return "".join(source_code.split()) diff --git a/MaxKernel/run_beam_search.sh b/MaxKernel/run_beam_search.sh new file mode 100755 index 0000000..21d4845 --- /dev/null +++ b/MaxKernel/run_beam_search.sh @@ -0,0 +1,29 @@ +#!/bin/bash +# Local wrapper to execute Beam Search orchestrator dry-runs + +# Exit immediately if a command exits with a non-zero status +set -e + +# Move to the MaxKernel project root directory +CDPATH= cd -- "$(dirname -- "$0")" + +# Activate the python virtual environment from the autocomp clone +VENV_PATH="/usr/local/google/home/ligh/github/accelerator-agents/MaxKernel/.venv" +if [ -d "$VENV_PATH" ]; then + echo "[Shell] Activating virtual environment: $VENV_PATH" + source "$VENV_PATH/bin/activate" +else + echo "[Error] Virtual environment not found at $VENV_PATH" + exit 1 +fi + +# Set mock compiler environment variables +export MOCK_COMPILER=true +export GOOGLE_CLOUD_PROJECT="tpu-kernel-assist-sandbox" +export RAG_CORPUS="projects/tpu-kernel-assist-sandbox/locations/us-west1/ragCorpora/7991637538768945152" + + +echo "[Shell] Launching Beam Search local verification..." +python3 -u beam_search/run_beam_search.py "$@" + +echo "[Shell] Beam Search verification finished." From 628d3df3a44dd09eafdf366da55ed6d3af8a99ab Mon Sep 17 00:00:00 2001 From: Guanghua Li Date: Fri, 26 Jun 2026 16:35:27 +0000 Subject: [PATCH 2/2] Add a comment about an caveat in grouped harness --- MaxKernel/beam_search/orchestrator.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/MaxKernel/beam_search/orchestrator.py b/MaxKernel/beam_search/orchestrator.py index b88c75b..2928a92 100644 --- a/MaxKernel/beam_search/orchestrator.py +++ b/MaxKernel/beam_search/orchestrator.py @@ -130,6 +130,10 @@ def __init__( with open(self.baseline_code_path, "r") as f: self.baseline_code = f.read() + # BE AWARE: the current implementation is broken. It simply concatenates all candidate codes into one python file, thus the + # output python contains repeated main functions and won't execute all candidates. + # + # TODO(ligh-svg): Fix the grouped kernel harness to allow executing all candidates. def _prepare_grouped_harness(self, candidate_paths: List[str]) -> str: """Inlines all candidate implementations into a single file.""" inlined_solutions = []