From 9a12a5d4ccf1924827350e0c9bd2d997e888e2d8 Mon Sep 17 00:00:00 2001 From: jayden1711 Date: Sun, 19 Jul 2026 17:13:28 -0700 Subject: [PATCH] Fix beam search config validation, candidate schema, and PTX fingerprint logging - Guard num_expanding_parents against 0 and negative values by clamping to 1 with a warning; previously a value of 0 caused select_candidates() to return no candidates, silently wasting entire optimization rounds - Add None guard so num_expanding_parents=None (expand all) is preserved - Add missing 'inspirations' key to beam search candidate dicts to match the SearchStrategy Protocol contract already honored by GreedyStrategy; populates via database.sample_inspirations() with the parent excluded - Extend SimpleMutator.build_prompt to accept and render inspirations as reference kernels in the optimization prompt - Log a warning in ptx_hash_from_cache when a PTX file cannot be read instead of silently skipping it, so callers know the fingerprint may be incomplete - Add tests for BeamSearchStrategy config validation, candidate schema contract, and worker fanout math - Add tests for ptx_hash_from_cache normalization invariants, constant preservation, and directory edge cases --- tests/test_beam_search_strategy.py | 148 ++++++++++++++++++ tests/test_ptx_fingerprint.py | 117 ++++++++++++++ .../searching/mutation/mutator.py | 13 +- .../searching/ptx_fingerprint.py | 12 +- .../searching/strategy/beam_search.py | 13 ++ 5 files changed, 299 insertions(+), 4 deletions(-) create mode 100644 tests/test_beam_search_strategy.py create mode 100644 tests/test_ptx_fingerprint.py diff --git a/tests/test_beam_search_strategy.py b/tests/test_beam_search_strategy.py new file mode 100644 index 00000000..1445c9c4 --- /dev/null +++ b/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"]) \ No newline at end of file diff --git a/tests/test_ptx_fingerprint.py b/tests/test_ptx_fingerprint.py new file mode 100644 index 00000000..edd2b405 --- /dev/null +++ b/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"]) \ No newline at end of file diff --git a/triton_kernel_agent/opt_worker_component/searching/mutation/mutator.py b/triton_kernel_agent/opt_worker_component/searching/mutation/mutator.py index 42dfb36a..572c5689 100644 --- a/triton_kernel_agent/opt_worker_component/searching/mutation/mutator.py +++ b/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) \ No newline at end of file diff --git a/triton_kernel_agent/opt_worker_component/searching/ptx_fingerprint.py b/triton_kernel_agent/opt_worker_component/searching/ptx_fingerprint.py index 5da6bc6a..1b2cb815 100644 --- a/triton_kernel_agent/opt_worker_component/searching/ptx_fingerprint.py +++ b/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. diff --git a/triton_kernel_agent/opt_worker_component/searching/strategy/beam_search.py b/triton_kernel_agent/opt_worker_component/searching/strategy/beam_search.py index 39a034bf..eebd36e8 100644 --- a/triton_kernel_agent/opt_worker_component/searching/strategy/beam_search.py +++ b/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