Fix BeamSearchStrategy config validation, candidate schema, and PTX fingerprint logging

tests/test_beam_search_strategy.py

@@ -0,0 +1,148 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for BeamSearchStrategy config validation and candidate schema contract."""
+
+import logging
+import os
+import sys
+from unittest.mock import MagicMock
+
+import pytest
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+def _make_metrics(**kw):
+    from triton_kernel_agent.opt_worker_component.searching.history.models import ProgramMetrics
+    defaults = dict(time_ms=1.0)
+    defaults.update(kw)
+    return ProgramMetrics(**defaults)
+
+
+def _make_entry(**kw):
+    from triton_kernel_agent.opt_worker_component.searching.history.models import ProgramEntry
+    defaults = dict(
+        program_id="prog_0",
+        kernel_code="def k(): pass",
+        metrics=_make_metrics(),
+        problem_id="p0",
+    )
+    defaults.update(kw)
+    return ProgramEntry(**defaults)
+
+
+def _make_strategy(database=None, **kw):
+    from triton_kernel_agent.opt_worker_component.searching.strategy.beam_search import BeamSearchStrategy
+    defaults = dict(
+        num_top_kernels=4,
+        num_bottlenecks=2,
+        models=["claude-sonnet-4-6"],
+        samples_per_prompt=1,
+        num_expanding_parents=2,
+        database=database,
+    )
+    defaults.update(kw)
+    return BeamSearchStrategy(**defaults)
+
+
+class TestNumExpandingParentsValidation:
+    """Tests for num_expanding_parents config validation."""
+
+    def test_zero_parents_clamped_to_one(self, caplog):
+        with caplog.at_level(logging.WARNING, logger="BeamSearchStrategy"):
+            s = _make_strategy(num_expanding_parents=0)
+        assert s.num_expanding_parents == 1
+        assert any("clamping to 1" in r.message for r in caplog.records)
+
+    def test_negative_parents_clamped_to_one(self, caplog):
+        with caplog.at_level(logging.WARNING, logger="BeamSearchStrategy"):
+            s = _make_strategy(num_expanding_parents=-1)
+        assert s.num_expanding_parents == 1
+
+    def test_none_parents_not_clamped(self):
+        """num_expanding_parents=None is valid and means use all top kernels."""
+        s = _make_strategy(num_expanding_parents=None)
+        assert s.num_expanding_parents is None
+
+    def test_zero_parents_still_produces_candidates(self):
+        s = _make_strategy(num_expanding_parents=0)
+        s.initialize(_make_entry())
+        candidates = s.select_candidates(round_num=1)
+        assert len(candidates) > 0
+
+
+class TestCandidateSchemaContract:
+    """Tests for SearchStrategy Protocol schema compliance in beam search candidates."""
+
+    def test_candidates_contain_inspirations_key(self):
+        """Every candidate dict must include the 'inspirations' key per SearchStrategy Protocol."""
+        mock_db = MagicMock()
+        mock_db.sample_inspirations.return_value = []
+        s = _make_strategy(database=mock_db)
+        s.initialize(_make_entry())
+        for candidate in s.select_candidates(round_num=1):
+            assert "inspirations" in candidate
+            assert isinstance(candidate["inspirations"], list)
+
+    def test_candidates_satisfy_full_protocol_schema(self):
+        """All required Protocol keys must be present in every candidate."""
+        mock_db = MagicMock()
+        mock_db.sample_inspirations.return_value = []
+        s = _make_strategy(database=mock_db)
+        s.initialize(_make_entry())
+        required_keys = {"parent", "bottleneck_id", "inspirations"}
+        for candidate in s.select_candidates(round_num=1):
+            assert required_keys.issubset(candidate.keys())
+
+    def test_inspirations_excludes_parent_kernel(self):
+        mock_db = MagicMock()
+        mock_db.sample_inspirations.return_value = []
+        s = _make_strategy(database=mock_db)
+        entry = _make_entry(program_id="prog_parent")
+        s.initialize(entry)
+        s.select_candidates(round_num=1)
+        calls = mock_db.sample_inspirations.call_args_list
+        assert len(calls) > 0
+        for call in calls:
+            exclude = call.kwargs.get("exclude_ids") or (
+                call.args[1] if len(call.args) > 1 else None
+            )
+            assert exclude is not None and "prog_parent" in exclude
+
+    def test_no_database_returns_empty_inspirations(self):
+        """When database=None, inspirations should be empty list, not an error."""
+        s = _make_strategy(database=None)
+        s.initialize(_make_entry())
+        for candidate in s.select_candidates(round_num=1):
+            assert candidate["inspirations"] == []
+
+
+class TestWorkerCount:
+    """Tests for num_workers_needed property."""
+
+    def test_workers_needed_matches_fanout(self):
+        """num_workers_needed must equal parents × bottlenecks × models × samples."""
+        s = _make_strategy(
+            num_expanding_parents=2,
+            num_bottlenecks=3,
+            models=["a", "b"],
+            samples_per_prompt=2,
+        )
+        assert s.num_workers_needed == 2 * 3 * 2 * 2
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/test_ptx_fingerprint.py

@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for PTX-based kernel fingerprinting and deduplication."""
+
+import os
+
+import pytest
+
+from triton_kernel_agent.opt_worker_component.searching import ptx_fingerprint as _ptx
+normalize_ptx = _ptx.normalize_ptx
+ptx_hash_from_cache = _ptx.ptx_hash_from_cache
+
+BASE_PTX = """\
+.version 8.2
+.target sm_90
+// a comment
+.visible .entry kern()
+{
+    .reg .b32 %r<3>;
+    mov.u32 %r1, 5;
+    add.s32 %r2, %r1, 7;
+$L__BB0_1:
+    bra $L__BB0_1;
+}
+"""
+
+
+class TestNormalizePtx:
+    """Tests for PTX normalization invariants."""
+
+    def test_register_renumbering_invariant(self):
+        """Different register numbering should not change the fingerprint."""
+        renumbered = BASE_PTX.replace("%r1", "%r7").replace("%r2", "%r4")
+        assert normalize_ptx(BASE_PTX) == normalize_ptx(renumbered)
+
+    def test_comment_and_whitespace_invariant(self):
+        """Comments and extra whitespace should not affect the fingerprint."""
+        noisy = BASE_PTX.replace("// a comment", "// completely different comment\n\n   ")
+        assert normalize_ptx(BASE_PTX) == normalize_ptx(noisy)
+
+    def test_version_and_target_directives_ignored(self):
+        """PTX version and target directives should not affect the fingerprint."""
+        other_version = BASE_PTX.replace(".version 8.2", ".version 8.4").replace(
+            "sm_90", "sm_80"
+        )
+        assert normalize_ptx(BASE_PTX) == normalize_ptx(other_version)
+
+    def test_label_renaming_invariant(self):
+        """Different label names should not change the fingerprint."""
+        relabeled = BASE_PTX.replace("$L__BB0_1", "$L__BB0_9")
+        assert normalize_ptx(BASE_PTX) == normalize_ptx(relabeled)
+
+    def test_constant_difference_is_preserved(self):
+        """Kernels differing only in a constant value must produce different fingerprints.
+
+        This is the load-bearing invariant: canonicalizing numeric constants away
+        would cause incorrect kernel-equivalence merges in beam search deduplication.
+        """
+        changed_constant = BASE_PTX.replace("mov.u32 %r1, 5;", "mov.u32 %r1, 6;")
+        assert normalize_ptx(BASE_PTX) != normalize_ptx(changed_constant)
+
+    def test_register_classes_not_conflated(self):
+        """32-bit and 64-bit register classes must remain distinct after normalization."""
+        b32 = "mov.u32 %r10, 1;"
+        b64 = "mov.u64 %rd10, 1;"
+        assert normalize_ptx(b32) != normalize_ptx(b64)
+
+
+class TestFingerprintKernelDir:
+    """Tests for directory-level PTX fingerprinting."""
+
+    def test_missing_directory_returns_none(self, tmp_path):
+        """A nonexistent directory should return None."""
+
+        assert ptx_hash_from_cache(tmp_path / "does_not_exist") is None
+
+    def test_empty_directory_returns_none(self, tmp_path):
+        """An empty directory with no PTX files should return None."""
+
+        assert ptx_hash_from_cache(tmp_path) is None
+
+    def test_single_file_returns_fingerprint(self, tmp_path):
+        """A directory with one PTX file should return a non-None fingerprint."""
+
+        (tmp_path / "kernel.ptx").write_text(BASE_PTX)
+        result = ptx_hash_from_cache(tmp_path)
+        assert result is not None
+        assert isinstance(result, str)
+        assert len(result) > 0
+
+    def test_identical_content_produces_same_fingerprint(self, tmp_path):
+        """Two directories with identical PTX content should produce the same fingerprint."""
+
+        dir_a = tmp_path / "a"
+        dir_b = tmp_path / "b"
+        dir_a.mkdir()
+        dir_b.mkdir()
+        (dir_a / "kernel.ptx").write_text(BASE_PTX)
+        (dir_b / "kernel.ptx").write_text(BASE_PTX)
+        assert ptx_hash_from_cache(dir_a) == ptx_hash_from_cache(dir_b)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

triton_kernel_agent/opt_worker_component/searching/mutation/mutator.py

@@ -43,7 +43,7 @@ class SimpleMutator:
     def __init__(self, store: ProgramDatabase) -> None:
         self.store = store
 
-    def build_prompt(self, parent: AttemptRecord) -> str:
+    def build_prompt(self, parent: AttemptRecord, inspirations: list[AttemptRecord] | None = None) -> str:
         lines = [
             "# Optimize this Triton kernel\n",
             f"Current performance: {parent.time_ms:.4f}ms\n",
@@ -55,8 +55,15 @@ def build_prompt(self, parent: AttemptRecord) -> str:
             for a in history:
                 lines.append(f"- [{a.outcome.value}] {a.time_ms:.4f}ms\n")
 
-        lines.append("\n## Kernel:\n```python\n")
+        if inspirations:
+            lines.append("\n## High-performing reference kernels:\n")
+            for i, insp in enumerate(inspirations):
+                lines.append(f"\n### Reference {i + 1} ({insp.time_ms:.4f}ms):\n```python\n")
+                lines.append(insp.kernel_code)
+                lines.append("\n```\n")
+
+        lines.append("\n## Kernel to optimize:\n```python\n")
         lines.append(parent.kernel_code)
         lines.append("\n```\n")
 
-        return "".join(lines)
+        return "".join(lines)

triton_kernel_agent/opt_worker_component/searching/ptx_fingerprint.py

@@ -36,6 +36,9 @@
 import re
 from pathlib import Path
 
+import logging
+logger = logging.getLogger(__name__)
+
 # --- Normalization regex patterns -------------------------------------------
 
 # Strip ``//`` line comments and ``/* */`` block comments.
@@ -156,7 +159,14 @@ def ptx_hash_from_cache(cache_dir: Path) -> str | None:
     for rel, path in rel_sorted:
         try:
             normalized = normalize_ptx(path.read_text(errors="replace"))
-        except OSError:
+        except OSError as e:
+            logger.warning(
+                "Failed to read PTX file %s: %s. "
+                "Fingerprint will be computed from remaining files only "
+                "and may not represent the full kernel.",
+                path,
+                e,
+            )
             continue
         # Include the relative filename in the hash so two kernels with the
         # same PTX content under different function names still differ.

triton_kernel_agent/opt_worker_component/searching/strategy/beam_search.py

@@ -88,6 +88,14 @@ def __init__(
         self.top_kernels: list[ProgramEntry] = []
         self.models = models
         self.samples_per_prompt = max(1, samples_per_prompt)
+        if num_expanding_parents is not None and num_expanding_parents < 1:
+            self.logger.warning(
+                "num_expanding_parents=%d is invalid (must be >= 1); clamping to 1. "
+                "A value of 0 causes select_candidates() to return no candidates, "
+                "silently spawning zero workers and wasting the entire round.",
+                num_expanding_parents,
+            )
+            num_expanding_parents = 1
         self.num_expanding_parents = num_expanding_parents
         # Internal iteration list: [None] means "use runner default".
         self._expansion_models: list[str | None] = list(models) if models else [None]
@@ -155,6 +163,11 @@ def select_candidates(self, round_num: int) -> list[dict[str, Any]]:
                                 "kernel_rank": rank,
                                 "openai_model": model,
                                 "sample_idx": sample_idx,
+                                "inspirations": self.database.sample_inspirations(
+                                    n=2,
+                                    exclude_ids=[kernel.program_id],
+                                    problem_id=self.problem_id,
+                                ) if self.database else [],
                             }
                         )
         return candidates