diff --git a/.env.template b/.env.template index 4ec053a26d..86e32db13d 100644 --- a/.env.template +++ b/.env.template @@ -226,10 +226,11 @@ KNOWLEDGE_STORE_MODE=local # applies to dense TunableOp recording, while block-FP8 uses standard Roofline. # HYPERLOOM_GEMM_SHAPE_CAPTURE=1 # HYPERLOOM_GEMM_SHAPE_CAPTURE_TIMEOUT_SEC=1800 -# Force the per-optimization wall-clock budget in minutes (default 60). The env +# Force the per-optimization wall-clock budget in minutes (default 90). The env # wins over the payload value, which is LLM-authored, so an operator raising the -# budget is not silently overridden. -# KERNEL_OPT_BACKEND_BUDGET_MIN=60 +# budget is not silently overridden. forge-loop reserves half the window for +# finalize, so this buys roughly half as much iteration as it reads. +# KERNEL_OPT_BACKEND_BUDGET_MIN=90 # --- Collective optimization lane (optional) ------------------------------- # The Coordinator drives this lane itself at KERNEL entry; it is never an agent diff --git a/docs/conceptual/kernel-execution-path.md b/docs/conceptual/kernel-execution-path.md index 2cfd6323e3..3fa26f9ec0 100644 --- a/docs/conceptual/kernel-execution-path.md +++ b/docs/conceptual/kernel-execution-path.md @@ -224,7 +224,7 @@ Optional: | `TRACELENS_INTERNAL_ROOT` | TraceLens internal extension; unset = open-source-only | | `KERNEL_OPT_MAX_PARALLEL` | Override the 8-concurrent-kernel default | | `INFERENCE_OPTIMIZER_KERNEL_OPT_MAX_PARTIAL` | Override partial-attempt retry cap (default 2) | -| `KERNEL_OPT_BACKEND_BUDGET_MIN` | Force the per-optimization wall-clock budget in minutes (default 60); wins over the LLM-authored payload value | +| `KERNEL_OPT_BACKEND_BUDGET_MIN` | Force the per-optimization wall-clock budget in minutes (default 90); wins over the LLM-authored payload value | Fusion lane: diff --git a/docs/reference/environment-variables.md b/docs/reference/environment-variables.md index e978a1e318..241508fd6f 100644 --- a/docs/reference/environment-variables.md +++ b/docs/reference/environment-variables.md @@ -189,7 +189,7 @@ The following variables control the kernel optimization backend ladder. | `HYPERLOOM_GEMM_SHAPE_CAPTURE` | `1` | Enables automatic runtime GEMM-shape capture for eligible single-node dense vLLM Forge tuning when no explicit shape input is available. Block-FP8 first reuses shapes from the TraceLens-selected steady-state trace of a successful Roofline with exactly matching model, workload, server arguments, environment, and backend controls. Missing or stale evidence triggers the same standard Roofline/ProfileExecutor/TraceLens steady-state pipeline as a fallback. Set to `0` to preserve the no-capture path. | | `HYPERLOOM_GEMM_SHAPE_CAPTURE_TIMEOUT_SEC` | `1800` | Timeout in seconds for the dense vLLM TunableOp recording benchmark. Block-FP8 fallback uses the standard Roofline/ProfileExecutor timeout. Values below `60` are clamped to `60`. | | `INFERENCE_OPTIMIZER`
`_KERNEL_OPT_MAX_PARTIAL` | Unset | Cap on how many `PARTIAL` kernel-opt verdicts an action can yield before it short-circuits to `NEEDS_REVIEW`. Useful for keeping budget contained when GEAK is consistently timing out. | -| `KERNEL_OPT_BACKEND_BUDGET_MIN` | `60` | Wall-clock budget in minutes for one optimization, mirrored by the `kernel_optimization.py` wrapper. The env deliberately wins over the payload `budget_minutes`, which is LLM-authored from a prompt template, so an operator raising the budget is not silently overridden. forge-loop reserves half the window for finalize, so `60` leaves roughly 30 minutes of real iteration. | +| `KERNEL_OPT_BACKEND_BUDGET_MIN` | `90` | Wall-clock budget in minutes for one optimization, mirrored by the `kernel_optimization.py` wrapper. The env deliberately wins over the payload `budget_minutes`, which is LLM-authored from a prompt template, so an operator raising the budget is not silently overridden. forge-loop reserves half the window for finalize, so `90` leaves roughly 45 minutes of real iteration. | | `AITER_LOG_TUNED_CONFIG` | `1` (set for every serving run) | Makes aiter log each tuned-config lookup it *hits*, not only the ones it misses. Two checks have no input without it: the GEMM demand list, which learns the shapes the runtime actually asks for (config-derived shapes covered 0.4% of them), and the apply verdict, which cannot tell "the tuned table was never read" from "it was read and did not help". A scan of 60 production logs found it set in none of them, so it is now injected by default. An operator value wins — set `0` to turn hit logging off, at the cost of both checks going inconclusive. Every miss already prints a line regardless of this setting; hit logging adds roughly one line per lookup that succeeds. | | `HYPERLOOM_GEMM_PAIRED_PAIRS` | `0` (off) | How many interleaved baseline/tuned pairs to re-measure before a GEMM tuning KEEP is reported as confirmed. One end-to-end measurement cannot separate a gain from drift on this fleet: three rounds of a single unchanged configuration spanned 58%, and one controlled repeat moved 16%. Each pair costs two extra benchmark rounds. When `0`, the gain is still promoted — it is the best number available — but recorded as an unpaired block comparison rather than presented as a paired one. | diff --git a/src/hyperloom/agents/kernel/tests/test_bypass_report.py b/src/hyperloom/agents/kernel/tests/test_bypass_report.py index 7bec704b8f..8cbb89716c 100644 --- a/src/hyperloom/agents/kernel/tests/test_bypass_report.py +++ b/src/hyperloom/agents/kernel/tests/test_bypass_report.py @@ -663,10 +663,11 @@ def _boom(*a, **k): # pragma: no cover - must not be called assert cand["shape_provenance"] == "torch_trace" -def test_routable_candidate_carries_shapes_for_orchestrator_gate(): - # The orchestrator shape gate reads candidate["shapes"] and rejects dispatch - # with "empty_kernel_shape" when empty, so a routable candidate with real - # trace-captured dims must expose a non-empty "shapes" list. +def test_routable_candidate_carries_shapes_for_dispatch(): + # Dispatch reads candidate["shapes"] to pin the harness to the serving dims, + # so a routable candidate whose trace DID record them must expose them in + # the downstream contract form rather than leaving the backend to recover + # dims that were measured all along. kernels = [ { "name": "triton_silu", @@ -682,7 +683,7 @@ def test_routable_candidate_carries_shapes_for_orchestrator_gate(): cand = report.build_candidates(_analyze(kernels), framework="vllm", target_platform="MI300X")["hot_kernels"][0] shapes = cand.get("shapes") assert isinstance(shapes, list) and shapes, ( - "routable candidate must expose a non-empty 'shapes' for the orchestrator gate" + "a routable candidate with trace-captured dims must expose a non-empty 'shapes'" ) assert cand["shape_provenance"] in {"torch_trace", "tuning_csv"} # shapes mirrors input_shapes in the harness-consumable contract form. @@ -699,15 +700,15 @@ def test_trace_shape_entries_contract_format(): assert out == [{"call_num": 5, "shape": "(4,1024) bf16
(1024,) fp32"}] # unmapped dtype -> bare shape (no suffix). assert report._trace_shape_entries([[8, 8]], ["weird"], 1) == [{"call_num": 1, "shape": "(8,8)"}] - # no renderable operand -> empty (gate will reject as empty_kernel_shape). + # no renderable operand -> empty, which the backend recovers dims for. assert report._trace_shape_entries([[]], ["float"], 1) == [] assert report._trace_shape_entries([], [], 1) == [] def test_unresolved_shape_candidate_has_empty_shapes(): # A kernel with no captured dims stays shape-less: "shapes" is an empty list - # (present, not absent) and the gate will correctly reject it as - # empty_kernel_shape. + # (present, not absent) and the provenance says why, so the dispatch can + # tell "no dims were recorded" from "these are the dims". kernels = [{"name": "mystery_kernel", "op_name": "aten::mystery", "gpu_time_us": 100.0, "count": 1}] cand = report.build_candidates(_analyze(kernels), framework="vllm", target_platform="MI300X")["hot_kernels"][0] assert cand.get("shapes") == [] diff --git a/src/hyperloom/agents/kernel/tests/test_candidate_review_agent.py b/src/hyperloom/agents/kernel/tests/test_candidate_review_agent.py new file mode 100644 index 0000000000..74f6701d29 --- /dev/null +++ b/src/hyperloom/agents/kernel/tests/test_candidate_review_agent.py @@ -0,0 +1,985 @@ +############################################################################### +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT +# +# See LICENSE for license information. +############################################################################### + +"""Guards on the agent review of the deterministic kernel-candidate table. + +The review exists because the deterministic tiers fail by being confidently +wrong, and it is given a tool-enabled session to check their work. That freedom +is what these tests bound: the session proposes, and everything it proposes is +either verified or dropped before it reaches a candidate row. + +Three properties carry the weight, and each has a concrete failure behind it: +a measured field overwritten by a model would corrupt the impact ranking and +the tuning harness that are computed from it; an invented path would hand a +backend the wrong file to rewrite, which is the failure the whole pipeline +exists to prevent; and a session that edited the framework tree would leave the +benchmark that follows measuring an unrecorded change. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tools")) + +import _candidate_review_agent as cra # noqa: E402 + + +def _candidate(**overrides) -> dict: + """A finalized candidate row shaped the way the pipeline emits one.""" + row = { + "kernel_id": "k001", + "name": "_gqa_sparse_decode_kernel", + "gpu_pct": 9.47, + "duration_us": 1234.0, + "call_count": 1710, + "shapes": ["(8192,1024) bf16"], + "raw_arg_spec": {"0": "tensor"}, + "source_file": "", + "source_resolution_method": "name_grep", + } + row.update(overrides) + return row + + +@pytest.fixture +def tree(tmp_path: Path): + """A framework root holding one real kernel source.""" + root = tmp_path / "vllm" + (root / "ops").mkdir(parents=True) + defines = root / "ops" / "sparse_attn.py" + defines.write_text("def _gqa_sparse_decode_kernel(): pass\n", encoding="utf-8") + return root, defines + + +# --- measured fields are evidence, not suggestions -------------------------- + + +class TestImmutableFields: + def test_trace_measurements_are_declared_immutable(self): + """Event durations off the trace; also the dispatch floor's input.""" + for field in ("gpu_pct", "duration_us", "call_count"): + assert field in cra.IMMUTABLE_FIELDS + + def test_join_keys_are_declared_immutable(self): + """Revising these detaches the row from its ledger and from the CSVs.""" + for field in ("kernel_id", "name", "device_kernel_name"): + assert field in cra.IMMUTABLE_FIELDS + + def test_judgement_fields_stay_revisable(self): + """Locking these would leave the review nothing to correct.""" + for field in ( + "source_file", + "reusable_native_kernel", + "skip_reason", + "benchmark_files", + "shapes", + "input_dtypes", + ): + assert field not in cra.IMMUTABLE_FIELDS + + def test_alternate_shape_representations_are_derived_not_revisable(self): + """Accepting these alongside shapes is how the three drift apart.""" + for field in ("input_shapes", "invocation_cases", "raw_arg_spec"): + assert field in cra.DERIVED_SHAPE_FIELDS + assert field not in cra.IMMUTABLE_FIELDS + + def test_a_revision_cannot_overwrite_what_the_trace_measured(self, tree): + """The impact ranking and the closing gain figure are computed from these. + + A plausible-looking edit here is indistinguishable from data, so it is + dropped and reported rather than trusted. + """ + root, defines = tree + row = _candidate() + notes = cra.apply_revisions( + [row], + [ + { + "kernel_id": "k001", + "action": "rewrite", + "source_file": str(defines), + "gpu_pct": 99.0, + "duration_us": 1.0, + "name": "something_else", + } + ], + framework_roots=(str(root),), + ) + assert row["gpu_pct"] == 9.47 + assert row["duration_us"] == 1234.0 + assert row["name"] == "_gqa_sparse_decode_kernel" + assert any("ignored measured field" in note for note in notes) + # The judgement half of the same revision still lands. + assert row["source_file"] == str(defines) + + def test_a_veto_without_a_reason_is_the_input_echoed_back(self, tree): + """Every unresolved row in the table carries ``reusable_native_kernel`` + false, and a session correcting such a row has been observed returning + that value while its own prose argued the file is editable. + + Honouring it refused four kernels the review had just located, 16% of + GPU time, in the run that found this. Nothing separates the echo from an + intended refusal except the reason the prompt asks for alongside it. + """ + root, defines = tree + row = _candidate(source_file="", reusable_native_kernel=False) + notes = cra.apply_revisions( + [row], + [ + { + "kernel_id": "k001", + "action": "rewrite", + "source_file": str(defines), + "reusable_native_kernel": False, + "skip_reason": "", + "reason": "opened it; defines the kernel at line 73", + } + ], + framework_roots=(str(root),), + ) + assert "review_reusable_hint" not in row + assert any("veto ignored" in note for note in notes) + # The half of the revision that was meant still lands. + assert row["source_file"] == str(defines) + + def test_a_veto_with_a_reason_is_still_honoured(self, tree): + """The refusal itself stays available; only the silent one is dropped.""" + root, defines = tree + row = _candidate(source_file="", reusable_native_kernel=False) + cra.apply_revisions( + [row], + [ + { + "kernel_id": "k001", + "action": "rewrite", + "source_file": str(defines), + "reusable_native_kernel": False, + "skip_reason": "vendor template; only its launch config is tunable", + } + ], + framework_roots=(str(root),), + ) + assert row["review_reusable_hint"] is False + assert row["review_skip_reason"] == "vendor template; only its launch config is tunable" + + def test_a_derived_representation_is_dropped_with_a_note(self, tree): + """Silently ignoring a field the prompt discusses is how drift hides.""" + root, defines = tree + row = _candidate(invocation_cases=[{"operation": "real"}]) + notes = cra.apply_revisions( + [row], + [ + { + "kernel_id": "k001", + "action": "rewrite", + "source_file": str(defines), + "invocation_cases": [{"operation": "invented"}], + "raw_arg_spec": {"0": "invented"}, + } + ], + framework_roots=(str(root),), + ) + assert row["invocation_cases"] == [{"operation": "real"}] + assert any("ignored derived field" in note for note in notes) + + +# --- operand dims the trace never recorded --------------------------------- + + +class TestShapeProposals: + """A graph replay records no arguments, so the hottest kernels of a + captured model arrive with no shape. Left empty, the tuning backend picks + its own without any view of the serving configuration. + """ + + def test_dims_are_staged_for_the_deterministic_pass_not_written(self, tree): + """Same split as the routability hint: stamping stays the only writer.""" + root, defines = tree + row = _candidate(shapes=[], source_file=str(defines)) + cra.apply_revisions( + [row], + [ + { + "kernel_id": "k001", + "action": "keep", + "shapes": ["(8192,6144) bf16", "(6144,1536) fp4"], + "input_dtypes": ["bf16", "fp4"], + "shape_provenance": cra.REVIEW_BACKFILL_PROVENANCE, + } + ], + framework_roots=(str(root),), + ) + assert row["shapes"] == [] + assert row["review_shapes"] == ["(8192,6144) bf16", "(6144,1536) fp4"] + assert row["review_input_dtypes"] == ["bf16", "fp4"] + assert row["review_shape_provenance"] == cra.REVIEW_BACKFILL_PROVENANCE + + def test_a_confirmed_path_still_carries_its_dims(self, tree): + """The rows most needing dims are the ones already resolved correctly. + + A rewrite naming the path the row already holds is not a correction, but + dropping the whole revision there would discard the shapes proposed with + it -- which is every kernel the deterministic tiers got right. + """ + root, defines = tree + row = _candidate(shapes=[], source_file=str(defines)) + cra.apply_revisions( + [row], + [ + { + "kernel_id": "k001", + "action": "rewrite", + "source_file": str(defines), + "shapes": ["(64,9216) bf16"], + } + ], + framework_roots=(str(root),), + ) + assert row["review_shapes"] == ["(64,9216) bf16"] + # The path did not move, so nothing was recorded as a correction. + assert "previous_source_file" not in row + assert row["source_resolution_method"] == "name_grep" + + def test_a_derivation_cannot_be_claimed_as_a_measurement(self, tree): + """Provenance is the only thing separating a recovered shape from a + computed one when a tuned kernel later fails to move throughput. + """ + root, defines = tree + row = _candidate(shapes=[], source_file=str(defines)) + cra.apply_revisions( + [row], + [ + { + "kernel_id": "k001", + "action": "keep", + "shapes": ["(1,1) fp32"], + "shape_provenance": "torch_trace", + } + ], + framework_roots=(str(root),), + ) + assert row["review_shape_provenance"] == cra.REVIEW_DERIVED_PROVENANCE + + def test_an_unlabelled_derivation_is_not_promoted(self, tree): + root, defines = tree + row = _candidate(shapes=[], source_file=str(defines)) + cra.apply_revisions( + [row], + [{"kernel_id": "k001", "action": "keep", "shapes": ["(8,8) bf16"]}], + framework_roots=(str(root),), + ) + assert row["review_shape_provenance"] == cra.REVIEW_DERIVED_PROVENANCE + + def test_an_empty_proposal_is_reported_rather_than_staged(self, tree): + """Clearing dims is not a correction the review has any use for.""" + root, defines = tree + row = _candidate(shapes=[], source_file=str(defines)) + notes = cra.apply_revisions( + [row], + [{"kernel_id": "k001", "action": "keep", "shapes": []}], + framework_roots=(str(root),), + ) + assert "review_shapes" not in row + assert any("empty shapes proposal" in note for note in notes) + + def test_an_unmentioned_row_keeps_its_dims(self, tree): + root, _ = tree + row = _candidate() + before = dict(row) + cra.apply_revisions([row], [], framework_roots=(str(root),)) + assert row == before + + def test_review_provenance_is_dispatchable(self): + """Dims the gate rejects are worse than none: an empty shape has an + override, an untrusted provenance does not. + """ + from hyperloom.common.kernel_shape_contract import DISPATCHABLE_SHAPE_PROVENANCE + + for provenance in cra.REVIEW_SHAPE_PROVENANCE: + assert provenance in DISPATCHABLE_SHAPE_PROVENANCE + + +# --- a path is taken only when it can be verified --------------------------- + + +class TestApplyRevisions: + def test_rewrite_to_a_verified_path_records_what_it_replaced(self, tree): + root, defines = tree + row = _candidate(source_file="/gone/wrong.py", source_resolution_method="name_grep") + notes = cra.apply_revisions( + [row], + [{"kernel_id": "k001", "action": "rewrite", "source_file": str(defines), "reason": "defines it"}], + framework_roots=(str(root),), + ) + assert row["source_file"] == str(defines) + assert row["previous_source_file"] == "/gone/wrong.py" + assert row["previous_method"] == "name_grep" + assert row["source_resolution_method"] == "llm_review" + assert row["review_reason"] == "defines it" + assert notes == [f"k001: /gone/wrong.py -> {defines}"] + + def test_an_invented_path_is_refused(self, tree): + """Existence is the floor. Without it a backend rewrites a fiction.""" + root, _ = tree + row = _candidate(source_file="/repo/current.py") + notes = cra.apply_revisions( + [row], + [{"kernel_id": "k001", "action": "rewrite", "source_file": str(root / "ops/invented.py")}], + framework_roots=(str(root),), + ) + assert row["source_file"] == "/repo/current.py" + assert any("rejected unverifiable path" in note for note in notes) + + def test_a_real_path_outside_every_root_is_refused(self, tree, tmp_path): + """Being openable is not enough; it must be framework source.""" + root, _ = tree + outside = tmp_path / "elsewhere.py" + outside.write_text("x = 1\n", encoding="utf-8") + row = _candidate(source_file="/repo/current.py") + notes = cra.apply_revisions( + [row], + [{"kernel_id": "k001", "action": "rewrite", "source_file": str(outside)}], + framework_roots=(str(root),), + ) + assert row["source_file"] == "/repo/current.py" + assert any("rejected unverifiable path" in note for note in notes) + + def test_rewrite_without_a_path_is_ignored(self, tree): + root, _ = tree + row = _candidate(source_file="/repo/current.py") + notes = cra.apply_revisions( + [row], + [{"kernel_id": "k001", "action": "rewrite"}], + framework_roots=(str(root),), + ) + assert row["source_file"] == "/repo/current.py" + assert any("rewrite without a path" in note for note in notes) + + def test_keep_changes_nothing(self, tree): + root, _ = tree + row = _candidate(source_file="/repo/current.py") + before = dict(row) + assert cra.apply_revisions( + [row], [{"kernel_id": "k001", "action": "keep"}], framework_roots=(str(root),) + ) == [] + assert row == before + + @pytest.mark.parametrize("action", ["unresolve", "drop"]) + def test_unresolve_and_drop_clear_the_source(self, tree, action): + """Both mean "do not send a backend here"; an empty source says so.""" + root, _ = tree + row = _candidate(source_file="/repo/wrong.py", source_line=42, source_function="launch") + cra.apply_revisions( + [row], + [{"kernel_id": "k001", "action": action, "reason": "dispatch wrapper"}], + framework_roots=(str(root),), + ) + assert row["source_file"] == "" + assert row["previous_source_file"] == "/repo/wrong.py" + assert row["review_action"] == action + assert "source_line" not in row and "source_function" not in row + + def test_an_unmentioned_candidate_is_left_alone(self, tree): + """Silence is not a verdict; only named rows move.""" + root, _ = tree + rows = [_candidate(), _candidate(kernel_id="k002", source_file="/repo/other.py")] + cra.apply_revisions( + rows, [{"kernel_id": "k001", "action": "unresolve"}], framework_roots=(str(root),) + ) + assert rows[1]["source_file"] == "/repo/other.py" + assert "review_action" not in rows[1] + + def test_unknown_id_and_action_are_reported_not_applied(self, tree): + root, _ = tree + row = _candidate() + notes = cra.apply_revisions( + [row], + [ + {"kernel_id": "k999", "action": "unresolve"}, + {"kernel_id": "k001", "action": "teleport"}, + ], + framework_roots=(str(root),), + ) + assert any("unknown kernel_id" in note for note in notes) + assert any("unknown action" in note for note in notes) + assert "review_action" not in row + + def test_an_authoritative_resolution_is_not_overridable(self, tree): + """The active finder demangles the symbol the binary actually exports. + + Reading the same tree cannot beat knowing that, so a curated resolution + is protected from a session that merely looked at the source. + """ + root, defines = tree + row = _candidate(source_file="/curated/truth.py", source_resolution_method="active_finder") + notes = cra.apply_revisions( + [row], + [{"kernel_id": "k001", "action": "rewrite", "source_file": str(defines)}], + framework_roots=(str(root),), + protected_ids={"k001"}, + ) + assert row["source_file"] == "/curated/truth.py" + assert any("resolved by an authoritative tier" in note for note in notes) + + def test_a_protected_candidate_may_still_be_kept(self, tree): + """Protection blocks changes, not agreement.""" + root, _ = tree + row = _candidate(source_file="/curated/truth.py") + notes = cra.apply_revisions( + [row], + [{"kernel_id": "k001", "action": "keep"}], + framework_roots=(str(root),), + protected_ids={"k001"}, + ) + assert notes == [] + + def test_routability_hints_are_recorded_for_the_caller_to_weigh(self, tree): + """The gate stays deterministic; the session only leaves a hint.""" + root, defines = tree + row = _candidate() + cra.apply_revisions( + [row], + [ + { + "kernel_id": "k001", + "action": "rewrite", + "source_file": str(defines), + "reusable_native_kernel": False, + "skip_reason": "dispatch wrapper", + } + ], + framework_roots=(str(root),), + ) + assert row["review_reusable_hint"] is False + assert row["review_skip_reason"] == "dispatch wrapper" + # Never written directly -- classify_patchability owns this field. + assert "reusable_native_kernel" not in row + + +# --- harnesses the session claims must be openable -------------------------- + + +class TestVerifiedHarnesses: + def test_absent_paths_are_dropped(self, tree): + _, defines = tree + assert cra._verified_harnesses(["/gone/test_pa.py", str(defines)]) == [str(defines)] + + def test_an_explicit_empty_list_is_honoured(self): + """"This kernel has no harness" is an answer worth keeping.""" + assert cra._verified_harnesses([]) == [] + + @pytest.mark.parametrize("proposed", [None, "not-a-list", 42]) + def test_no_proposal_leaves_the_field_alone(self, proposed): + assert cra._verified_harnesses(proposed) is None + + def test_a_verified_list_reaches_the_candidate(self, tree): + root, defines = tree + row = _candidate() + notes = cra.apply_revisions( + [row], + [ + { + "kernel_id": "k001", + "action": "rewrite", + "source_file": str(defines), + "benchmark_files": ["/gone/bench.py", str(defines)], + } + ], + framework_roots=(str(root),), + ) + assert row["review_benchmark_files"] == [str(defines)] + assert any("benchmark_files -> 1 verified path" in note for note in notes) + + +# --- the session may run shell commands, so the tree is checked ------------- + + +class TestSourceFingerprint: + def test_unreadable_paths_are_skipped_not_fatal(self, tree): + _, defines = tree + prints = cra.source_fingerprint([str(defines), "/gone/x.py", ""]) + assert set(prints) == {str(defines)} + + def test_an_untouched_tree_reports_no_drift(self, tree): + _, defines = tree + before = cra.source_fingerprint([str(defines)]) + assert cra.fingerprint_drift(before, cra.source_fingerprint([str(defines)])) == [] + + def test_an_edited_file_is_detected(self, tree): + """A review that changed the code under optimization is discardable. + + The benchmark that follows would otherwise measure an unrecorded edit + and credit it to whatever ran next. + """ + _, defines = tree + before = cra.source_fingerprint([str(defines)]) + defines.write_text("def _gqa_sparse_decode_kernel(): return 1\n", encoding="utf-8") + assert cra.fingerprint_drift(before, cra.source_fingerprint([str(defines)])) == [str(defines)] + + def test_a_deleted_file_is_detected(self, tree): + _, defines = tree + before = cra.source_fingerprint([str(defines)]) + defines.unlink() + assert cra.fingerprint_drift(before, cra.source_fingerprint([str(defines)])) == [str(defines)] + + +# --- the answer is a file, so a half-written one is not mistaken for one ---- + + +class TestLoadRevisions: + def test_a_missing_file_is_reported_as_such(self, tmp_path): + revisions, error = cra.load_revisions(tmp_path / "nope.json") + assert revisions == [] and "not written" in error + + @pytest.mark.parametrize( + ("body", "expected"), + [ + ("{not json", "not valid JSON"), + ("[1, 2]", "not a JSON object"), + ('{"other": []}', "no 'revisions' list"), + ], + ) + def test_an_unusable_body_names_what_was_wrong(self, tmp_path, body, expected): + path = tmp_path / cra.REVISIONS_FILENAME + path.write_text(body, encoding="utf-8") + revisions, error = cra.load_revisions(path) + assert revisions == [] and expected in error + + def test_non_dict_entries_are_dropped(self, tmp_path): + path = tmp_path / cra.REVISIONS_FILENAME + path.write_text(json.dumps({"revisions": [{"kernel_id": "k001"}, "junk"]}), encoding="utf-8") + revisions, error = cra.load_revisions(path) + assert error == "" and revisions == [{"kernel_id": "k001"}] + + +# --- the prompt hands over paths, never contents ---------------------------- + + +class TestBuildReviewPrompt: + def test_it_offers_paths_rather_than_pre_loaded_source(self, tmp_path, tree): + """Pre-loading would bound the review by what was guessed relevant.""" + root, defines = tree + prompt = cra.build_review_prompt( + run_dir=tmp_path, + raw_candidates_path=tmp_path / cra.RAW_CANDIDATES_FILENAME, + revisions_path=tmp_path / cra.REVISIONS_FILENAME, + reference_paths={"tracelens report": "/run/analysis.md"}, + framework_roots=(str(root),), + ) + assert str(tmp_path / cra.RAW_CANDIDATES_FILENAME) in prompt + assert "/run/analysis.md" in prompt + assert str(root) in prompt + # The body of a framework file is never shipped by the prompt itself. + assert defines.read_text(encoding="utf-8") not in prompt + + def test_it_states_the_actions_and_the_measured_field_ban(self, tmp_path): + prompt = cra.build_review_prompt( + run_dir=tmp_path, + raw_candidates_path=tmp_path / "raw.json", + revisions_path=tmp_path / "rev.json", + reference_paths={}, + framework_roots=(), + ) + for token in ("keep", "rewrite", "unresolve", "drop", "gpu_pct", "benchmark_files"): + assert token in prompt + + def test_it_asks_for_dims_and_for_how_they_were_obtained(self, tmp_path): + """An unstated derivation is no more reviewable than the backend's own + guess, which is what the dims are there to replace. + """ + prompt = cra.build_review_prompt( + run_dir=tmp_path, + raw_candidates_path=tmp_path / "raw.json", + revisions_path=tmp_path / "rev.json", + reference_paths={}, + framework_roots=(), + ) + assert "shapes" in prompt + assert cra.REVIEW_BACKFILL_PROVENANCE in prompt + assert cra.REVIEW_DERIVED_PROVENANCE in prompt + assert "State where the dims came from in reason" in prompt + + def test_it_tells_the_session_not_to_echo_the_routability_field(self, tmp_path): + """The table it audits carries the field, so "revisable" reads as + "return it". Saying so cost four located kernels in one run. + """ + prompt = cra.build_review_prompt( + run_dir=tmp_path, + raw_candidates_path=tmp_path / "raw.json", + revisions_path=tmp_path / "rev.json", + reference_paths={}, + framework_roots=(), + ) + assert "Do not copy reusable_native_kernel back" in prompt + assert "a false with no" in prompt and "skip_reason is ignored" in prompt + + def test_it_does_not_send_the_session_after_tracelens_internals(self, tmp_path): + """``analysis.md`` is TraceLens' only supported output; the rest of that + directory is internal and may be deleted. + + Nothing is really given up by staying inside the contract: for every + operator the sidecars describe, ``analysis.md`` carries the same dims and + launcher in its own table, and for a graph-launched operator neither has + anything. + """ + prompt = cra.build_review_prompt( + run_dir=tmp_path, + raw_candidates_path=tmp_path / "raw.json", + revisions_path=tmp_path / "rev.json", + reference_paths={}, + framework_roots=(), + ) + for internal in ("category_data", "priority_data", "perf_report_csvs"): + assert internal not in prompt + assert "analysis.md is TraceLens' only supported output" in prompt + + def test_it_does_not_offer_a_field_the_stamping_pass_recomputes(self, tmp_path): + """Inviting a revision that is then silently overwritten spends the + session's effort on nothing and hides the overwrite from the audit. + """ + prompt = cra.build_review_prompt( + run_dir=tmp_path, + raw_candidates_path=tmp_path / "raw.json", + revisions_path=tmp_path / "rev.json", + reference_paths={}, + framework_roots=(), + ) + assert "recommended_backends" not in prompt + + def test_write_scope_is_stated(self, tmp_path): + prompt = cra.build_review_prompt( + run_dir=tmp_path, + raw_candidates_path=tmp_path / "raw.json", + revisions_path=tmp_path / "rev.json", + reference_paths={}, + framework_roots=(), + ) + assert f"Write nothing outside {tmp_path}" in prompt + + +# --- a mandatory pass that must never take the run down with it ------------- + + +class TestRunCandidateReview: + def _args(self, tmp_path: Path) -> dict: + return { + "run_dir": tmp_path, + "raw_candidates_path": tmp_path / cra.RAW_CANDIDATES_FILENAME, + "reference_paths": {}, + "framework_roots": (), + } + + def test_the_written_file_is_what_marks_success(self, tmp_path): + """Artifact presence decides, as it does for the TraceLens runner. + + A provider can report an error after the answer already landed. + """ + def _runner(*, prompt, run_dir, model, timeout_sec): + (run_dir / cra.REVISIONS_FILENAME).write_text( + json.dumps({"revisions": [{"kernel_id": "k001", "action": "keep"}]}), + encoding="utf-8", + ) + return "transport reset after the write" + + out = cra.run_candidate_review(**self._args(tmp_path), session_runner=_runner) + assert out.ok and out.status == "completed" + assert out.revisions == [{"kernel_id": "k001", "action": "keep"}] + + def test_a_transient_failure_is_retried(self, tmp_path): + """The pass is mandatory; a gateway hiccup should not skip the audit.""" + calls = {"n": 0} + + def _runner(*, prompt, run_dir, model, timeout_sec): + calls["n"] += 1 + if calls["n"] == 1: + return "gateway 502" + (run_dir / cra.REVISIONS_FILENAME).write_text( + json.dumps({"revisions": []}), encoding="utf-8" + ) + return "" + + out = cra.run_candidate_review(**self._args(tmp_path), session_runner=_runner) + assert calls["n"] == 2 and out.ok + + def test_a_definitive_failure_is_reported_not_raised(self, tmp_path): + """Losing the audit costs candidates; raising would cost the run.""" + out = cra.run_candidate_review( + **self._args(tmp_path), + attempts=2, + session_runner=lambda **_kwargs: "gateway down", + ) + assert not out.ok and out.status == "failed" + assert "gateway down" in out.detail + + def test_a_raising_session_is_contained(self, tmp_path): + def _boom(**_kwargs): + raise RuntimeError("https://secret.example/?token=leak") + + out = cra.run_candidate_review( + **self._args(tmp_path), attempts=1, session_runner=_boom + ) + assert not out.ok + assert "leak" not in out.detail and "RuntimeError" in out.detail + + def test_a_stale_answer_cannot_be_mistaken_for_a_fresh_one(self, tmp_path): + """Each attempt clears the file first, so silence never reads as success.""" + (tmp_path / cra.REVISIONS_FILENAME).write_text( + json.dumps({"revisions": [{"kernel_id": "stale", "action": "drop"}]}), + encoding="utf-8", + ) + out = cra.run_candidate_review( + **self._args(tmp_path), attempts=1, session_runner=lambda **_kwargs: "no answer" + ) + assert not out.ok + + +# --- tool scope: the session investigates, it does not patch ---------------- + + +class TestToolScope: + def test_reading_and_searching_are_allowed(self): + for tool in ("Read", "Grep", "Glob"): + assert tool in cra.ALLOWED_TOOLS + + def test_editing_the_framework_tree_is_not(self): + """The tree here is the code under optimization; the agent proposes.""" + assert "Edit" not in cra.ALLOWED_TOOLS + assert "Edit" in cra._DENIED_TOOLS + + def test_sub_agents_and_the_web_are_denied(self): + """Turns and cost stay bounded, and the answers are all local.""" + for tool in ("Task", "WebFetch", "WebSearch"): + assert tool not in cra.ALLOWED_TOOLS + assert tool in cra._DENIED_TOOLS + + +# --- the stage boundary in tracelens_analysis ------------------------------- + +import argparse # noqa: E402 + +import tracelens_analysis as tla # noqa: E402 + + +class TestReviewStageBoundary: + def _args(self) -> argparse.Namespace: + return argparse.Namespace(model_name="m", framework="vllm", source_root=None) + + def test_an_unexpected_fault_costs_the_audit_not_the_run(self, tmp_path, monkeypatch): + """The stage sits at the end of an analysis a benchmark paid for. + + Nothing inside it may propagate: the deterministic table is still + usable, and killing a multi-hour run over an advisory pass trades a + small loss for a total one. + """ + def _boom(*_args, **_kwargs): + raise RuntimeError("unforeseen") + + monkeypatch.setattr(tla, "_run_candidate_review_stage", _boom) + warnings: list[dict] = [] + out = tla.run_candidate_review_stage( + tmp_path, candidates=[], args=self._args(), trace_health_warnings=warnings + ) + assert out == {} + assert warnings[0]["code"] == "candidate_review_failed" + assert warnings[0]["severity"] == "error" + assert warnings[0]["detail"] == "RuntimeError" + + def test_a_clean_stage_passes_its_artifacts_through(self, tmp_path, monkeypatch): + monkeypatch.setattr( + tla, + "_run_candidate_review_stage", + lambda *_a, **_kw: {"kernel_candidates_raw": "/x/raw.json"}, + ) + assert tla.run_candidate_review_stage( + tmp_path, candidates=[], args=self._args() + ) == {"kernel_candidates_raw": "/x/raw.json"} + + def test_the_session_is_only_pointed_at_supported_outputs(self, tmp_path, monkeypatch): + """The reference list is what the session is invited to read. + + TraceLens supports ``analysis.md`` and nothing else in that directory, + so naming a sidecar there both breaks the contract and points the + session at a file that may not exist. + """ + captured: dict = {} + + def _capture(**kwargs): + captured.update(kwargs) + return cra.ReviewOutcome(status="failed", detail="not run") + + monkeypatch.setattr(cra, "run_candidate_review", _capture) + tla._run_candidate_review_stage(tmp_path, candidates=[], args=self._args()) + + offered = " ".join(captured["reference_paths"].values()) + for internal in ("category_data", "priority_data", "perf_report_csvs"): + assert internal not in offered + assert "analysis.md" in offered + + def test_dims_land_even_though_the_path_did_not_move(self, tmp_path, monkeypatch): + """The re-derivation is skipped for rows that did not change, and for a + long time a path was the only thing that could change. + + Operand dims are most often supplied for a kernel the deterministic + tiers already located, so a check keyed on the path alone drops exactly + the proposals that were hardest to get. Two production analyses staged + `review_backfill` dims on `keep` revisions and shipped a table with none. + """ + source = tmp_path / "sparse_attn.py" + source.write_text("def k(): pass\n", encoding="utf-8") + candidate = { + "kernel_id": "k004", + "name": "_gqa_sparse_fwd_kernel", + "source_file": str(source), + "source_type": "python", + "shapes": [], + } + monkeypatch.setattr(tla, "_reusable_roots", lambda: (str(tmp_path).lower(),)) + monkeypatch.setattr( + cra, + "run_candidate_review", + lambda **_kw: cra.ReviewOutcome( + status="completed", + revisions=[ + { + "kernel_id": "k004", + "action": "keep", + "shapes": ["(8192,8,128) bf16", "(8192,1,128) bf16"], + "shape_provenance": cra.REVIEW_BACKFILL_PROVENANCE, + } + ], + ), + ) + + tla._run_candidate_review_stage( + tmp_path, candidates=[candidate], args=self._args() + ) + + assert candidate["shapes"] == ["(8192,8,128) bf16", "(8192,1,128) bf16"] + assert candidate["shape_provenance"] == cra.REVIEW_BACKFILL_PROVENANCE + + +class TestRederiveAfterReview: + def test_a_veto_is_honoured(self, tmp_path, monkeypatch): + """The session may refuse a kernel it knows is not worth a session. + + The deterministic gate has to pass first, or the veto is moot -- this + pins that a kernel the rules accept can still be turned down, and that + the session's reason is the one reported. + """ + root = tmp_path / "vllm" + root.mkdir() + source = root / "k.py" + source.write_text("def k(): pass\n", encoding="utf-8") + monkeypatch.setattr(tla, "_reusable_roots", lambda: (str(root).lower(),)) + + accepted = {"name": "k", "source_file": str(source), "source_type": "python"} + tla._rederive_after_review(accepted) + assert accepted["reusable_native_kernel"] is True, "gate must accept it first" + + vetoed = { + "name": "k", + "source_file": str(source), + "source_type": "python", + "review_reusable_hint": False, + "review_skip_reason": "dispatch wrapper only", + } + tla._rederive_after_review(vetoed) + assert vetoed["reusable_native_kernel"] is False + assert vetoed["skip_reason"] == "dispatch wrapper only" + + def test_a_promotion_is_not(self, tmp_path): + """A hint cannot talk the gate into dispatching what it rejected. + + classify_patchability stays the one gate; letting a permissive hint + through would give it a second, model-written one. + """ + item = {"name": "k", "source_file": "", "review_reusable_hint": True} + tla._rederive_after_review(item) + assert item["reusable_native_kernel"] is False + + def test_a_verified_harness_list_survives_restamping(self, tmp_path): + """Stamping recomputes benchmark_files from the coarse marker table. + + A session that went and looked has the better answer. + """ + source = tmp_path / "attention.py" + source.write_text("def paged_attention(): pass\n", encoding="utf-8") + item = { + "name": "kernel_paged_attention_2d", + "source_file": str(source), + "review_benchmark_files": [str(source)], + } + tla._rederive_after_review(item) + assert item["benchmark_files"] == [str(source)] + + +class TestAdoptReviewedShapes: + def test_dims_are_taken_where_the_trace_recorded_none(self, tmp_path): + """Empty dims are not a neutral state: the backend then picks its own + without any view of the serving configuration. + """ + source = tmp_path / "k.py" + source.write_text("def k(): pass\n", encoding="utf-8") + item = { + "name": "k", + "source_file": str(source), + "shapes": [], + "review_shapes": ["(8192,6144) bf16"], + "review_input_dtypes": ["bf16"], + "review_shape_provenance": "review_backfill", + } + tla._rederive_after_review(item) + assert item["shapes"] == ["(8192,6144) bf16"] + assert item["input_dtypes"] == ["bf16"] + assert item["shape_provenance"] == "review_backfill" + + def test_a_recorded_shape_outranks_a_reviewed_one(self, tmp_path): + """Nothing downstream re-measures the reviewed dims, so a real + measurement is never given up for them. + """ + item = { + "name": "k", + "source_file": "", + "shapes": ["(1,1) fp32"], + "shape_provenance": "torch_trace", + "review_shapes": ["(9,9) bf16"], + "review_shape_provenance": "review_derived", + } + tla._adopt_reviewed_shapes(item) + assert item["shapes"] == ["(1,1) fp32"] + assert item["shape_provenance"] == "torch_trace" + + def test_the_alternate_representations_do_not_outlive_the_dims(self, tmp_path): + """A harness built from a mix of old and new dims still benchmarks + cleanly, which is why the stale halves are dropped rather than kept. + """ + item = { + "name": "k", + "source_file": "", + "shapes": [], + "input_shapes": [["stale"]], + "invocation_cases": [{"operation": "stale"}], + "raw_arg_spec": {"0": "stale"}, + "_input_shapes_synthetic": True, + "review_shapes": ["(8192,6144) bf16"], + } + tla._adopt_reviewed_shapes(item) + for key in ("input_shapes", "invocation_cases", "raw_arg_spec", "_input_shapes_synthetic"): + assert key not in item + + def test_an_unlabelled_adoption_still_records_a_provenance(self, tmp_path): + """The dispatch gate reads this field; leaving it blank reads as + measured, which is the one thing it must not say. + """ + item = {"name": "k", "source_file": "", "shapes": [], "review_shapes": ["(8,8) bf16"]} + tla._adopt_reviewed_shapes(item) + assert item["shape_provenance"] == "review_derived" diff --git a/src/hyperloom/agents/kernel/tests/test_invocation_spec.py b/src/hyperloom/agents/kernel/tests/test_invocation_spec.py index 14c9b4fdf0..00994d09a4 100644 --- a/src/hyperloom/agents/kernel/tests/test_invocation_spec.py +++ b/src/hyperloom/agents/kernel/tests/test_invocation_spec.py @@ -738,3 +738,65 @@ def _rewrite_candidate(rows: list[dict], **extra) -> dict: } candidate.update(extra) return candidate + + +# --- dims the candidate review recovered must reach the driver ------------- + + +def test_spec_reports_missing_inputs_when_the_trace_recorded_no_arguments(): + """A graph replay has no CPU-side parent op, so a correctly resolved kernel + can still arrive with no dims at all -- and the spec has to say so rather + than emit an argument list built from nothing. + """ + spec = invocation_spec.build_invocation_spec( + {"kernel_id": "k001", "name": "kernel_paged_attention_2d", "shapes": []}, + source_file="/repo/attn.py", + ) + # Empty members are compacted away, so an argument list is absent rather + # than present and empty. + assert not spec.get("invocation", {}).get("arguments") + assert "inputs" in spec["missing_fields"] + + +def test_reviewed_dims_reach_the_invocation_spec(): + """The spec is what the backend builds its driver from, so dims recovered by + the candidate review have to arrive here or the correction changes nothing + about what actually gets benchmarked. + + ``shapes`` is the field the review revises; ``build_invocation_spec`` reads + it as the fallback for ``input_shapes``, which is what carries it through. + """ + spec = invocation_spec.build_invocation_spec( + { + "kernel_id": "k001", + "name": "kernel_paged_attention_2d", + "shapes": ["(8192,6144) bf16", "(6144,1536) fp4"], + "input_dtypes": ["bf16", "fp4"], + "shape_provenance": "review_derived", + }, + source_file="/repo/attn.py", + ) + arguments = spec["invocation"]["arguments"] + assert [record["shape"] for record in arguments] == [[8192, 6144], [6144, 1536]] + assert [record["dtype"] for record in arguments] == ["bf16", "fp4"] + assert "inputs" not in spec["missing_fields"] + + +def test_only_the_reviewed_table_can_be_resolved_as_the_candidate_source(tmp_path): + """The pre-review baseline sits in the same directory under its own name. + + Resolving that instead would hand the backend the dims the review corrected, + which is the one outcome the two-artifact split exists to make impossible. + """ + run_dir = tmp_path / "run" + run_dir.mkdir() + (run_dir / "kernel_candidates.raw.json").write_text( + json.dumps({"hot_kernels": [{"kernel_id": "k001", "shapes": []}]}), + encoding="utf-8", + ) + reviewed = {"hot_kernels": [{"kernel_id": "k001", "shapes": ["(8192,6144) bf16"]}]} + (run_dir / "kernel_candidates.json").write_text(json.dumps(reviewed), encoding="utf-8") + + resolved = kernel_optimization.resolve_candidates_path(run_dir) + assert resolved.name == "kernel_candidates.json" + assert kernel_optimization.load_candidates(resolved)[0]["shapes"] == ["(8192,6144) bf16"] diff --git a/src/hyperloom/agents/kernel/tests/test_kernel_optimization_verification.py b/src/hyperloom/agents/kernel/tests/test_kernel_optimization_verification.py index d1d6210a1d..b95b6159b6 100644 --- a/src/hyperloom/agents/kernel/tests/test_kernel_optimization_verification.py +++ b/src/hyperloom/agents/kernel/tests/test_kernel_optimization_verification.py @@ -131,6 +131,38 @@ def test_structured_shape_cases_parse_moe_args(): assert cases["supplementary_shapes"] == [] +def test_captured_shapes_block_claims_measurement_only_for_a_measurement(): + """The block tells the backend not to invent shapes, so it has to be honest + about where these came from: a graph replay records no arguments, and the + dims may have been reconstructed by the candidate review. Whether they were + measured is exactly what a later reader needs when the tuned kernel turns + out not to move end-to-end throughput. + """ + measured = ko._build_captured_shapes_block( + {"shapes": ["(8192,6144) bf16"], "shape_provenance": "torch_trace"} + ) + assert "TraceLens-captured" in measured + assert "do NOT invent" in measured + + for provenance in ("review_backfill", "review_derived"): + reviewed = ko._build_captured_shapes_block( + {"shapes": ["(8192,6144) bf16"], "shape_provenance": provenance} + ) + assert "TraceLens-captured" not in reviewed + assert "reconstructed" in reviewed + assert provenance in reviewed + # The instruction itself does not soften: choosing its own shapes is + # what the backend does when this block is absent. + assert "do NOT invent" in reviewed + + +def test_captured_shapes_block_is_empty_without_dims(): + """No block means the backend picks its own shapes -- the state this whole + path exists to get out of. + """ + assert ko._build_captured_shapes_block({"shapes": []}) == "" + + def test_structured_shape_cases_prefer_input_shapes(): candidate = { "name": "aiter::ck_moe_stage2", @@ -228,6 +260,11 @@ def test_structured_shape_cases_falls_back_to_input_shapes_when_group_rows_empty def test_build_prompt_includes_structured_shape_contract(): + """The runtime-metadata block promotes ``input_shapes`` into a shape contract. + + Rendered through the full prompt: the metadata block is GEAK's, and forge is + handed the same operands through its invocation spec instead. + """ shape = "(15360,8,768) bf16
(128,1536,2048) bf16" candidate = { "name": "aiter::ck_moe_stage2", @@ -238,7 +275,7 @@ def test_build_prompt_includes_structured_shape_contract(): "input_shapes": [{"call_num": 48, "shape": shape}], } - prompt = ko.build_prompt(candidate, _args(), backend="forge") + prompt = ko.build_prompt(candidate, _args()) assert "when `benchmark_shape_cases` is present" in prompt assert '"benchmark_shape_cases"' in prompt @@ -262,7 +299,7 @@ def test_build_prompt_omits_structured_shape_cases_without_program_output(): "shapes": [{"call_num": 48, "shape": "(15360,8,768) bf16"}], } - prompt = ko.build_prompt(candidate, _args(), backend="forge") + prompt = ko.build_prompt(candidate, _args()) metadata_json = prompt.split("```json\n", 1)[1].split("\n```", 1)[0] metadata = json.loads(metadata_json) @@ -271,6 +308,50 @@ def test_build_prompt_omits_structured_shape_cases_without_program_output(): assert "when `benchmark_shape_cases` is present" not in prompt +def test_build_prompt_forge_drops_the_geak_harness_and_keeps_trace_evidence(): + """Forge is handed trace evidence only, not the harness it does not run. + + The dropped sections are not merely redundant there. The deliverable files + land as new untracked paths that forge's workspace guard refuses, costing the + iteration that wrote them, and the A/B recipes stand up a second benchmark + beside the driver its own gate scores. Asserting their absence is what keeps + a later edit from reintroducing either through the shared prompt. + """ + candidate = { + "name": "aiter::ck_moe_stage2", + "source_file": "/tmp/gemm_moe_ck2stages.cu", + "source_type": "hip_cpp", + "kernel_repo": "/tmp/aiter", + "gpu_pct": 24.3, + "device_kernel_name": "ck_moe_stage2_kernel", + "source_resolution_method": "op_to_source", + "input_shapes": [{"call_num": 48, "shape": "(15360,8,768) bf16"}], + } + + forge = ko.build_prompt(candidate, _args(), backend="forge") + full = ko.build_prompt(candidate, _args()) + + # Absent from forge, still present for the backend that needs them. + for token in ( + "optimization_report.md", + "optimized_versions/", + "mini-swe-agent step", + "cpp_extension.load", + "structured context for GEAK", + "IMPORTANT — sandbox rules", + ): + assert token not in forge, token + assert token in full, token + + # Kept: what forge cannot derive from the kernel or its invocation spec. + assert "DEVICE KERNEL FOCUS" in forge + assert "ck_moe_stage2_kernel" in forge + assert "Preserve function name" in forge + assert forge.startswith("# TASK: Optimize the `aiter::ck_moe_stage2` kernel") + # A skipped section must not leave a run of blank lines behind. + assert "\n\n\n" not in forge + + def test_benchmark_available_alone_does_not_pass_correctness(tmp_path): verification = ko.build_verification( _args(micro_speedup=1.3), diff --git a/src/hyperloom/agents/kernel/tests/test_kernel_search_roots.py b/src/hyperloom/agents/kernel/tests/test_kernel_search_roots.py new file mode 100644 index 0000000000..5cf2f6baa9 --- /dev/null +++ b/src/hyperloom/agents/kernel/tests/test_kernel_search_roots.py @@ -0,0 +1,121 @@ +############################################################################### +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT +# +# See LICENSE for license information. +############################################################################### + +"""Guards on the framework roots the grep tier searches for kernel source. + +``KNOWN_SEARCH_ROOTS`` used to be a literal naming one container's layout +(``/sgl-workspace/...`` and a python3.10 venv). On a host that installs the +frameworks anywhere else -- a wheel under ``dist-packages``, a different Python +minor -- every root was absent, so the grep tier searched nothing, returned no +hit for any kernel, and the LLM tiers got an empty shortlist and a validation +gate that rejected every path outside those absent roots. The run still +succeeded: it reported zero routable kernels, which reads exactly like a trace +with nothing worth optimizing, and kernel-opt sat idle with no work to dispatch. + +These tests pin the properties that keep that silent failure from returning: +roots are discovered at runtime, non-existent ones never survive, and a host +with nothing installed says so instead of looking healthy. +""" + +from __future__ import annotations + +import logging +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tools")) + +import tracelens_analysis as tl # noqa: E402 + + +class TestInstalledPackageDir: + def test_locates_a_real_package(self): + located = tl._installed_package_dir("json") + assert located and Path(located).is_dir() + + def test_absent_package_returns_empty(self): + assert tl._installed_package_dir("no_such_package_xyz") == "" + + def test_rejects_non_identifier(self): + """Guards a path fragment being passed where a package name belongs.""" + assert tl._installed_package_dir("/usr/local/lib") == "" + assert tl._installed_package_dir("") == "" + + +class TestDiscoverKernelSearchRoots: + def test_drops_roots_that_do_not_exist(self, monkeypatch, tmp_path): + present = tmp_path / "vllm" + present.mkdir() + monkeypatch.setattr( + tl, + "_resolve_kernel_search_roots", + lambda: (f"{present}/", "/gone/aiter/"), + ) + tl._discover_kernel_search_roots.cache_clear() + assert tl._discover_kernel_search_roots() == (str(present),) + + def test_strips_trailing_separator(self, monkeypatch, tmp_path): + """Callers match these as prefixes without a separator of their own.""" + root = tmp_path / "aiter" + root.mkdir() + monkeypatch.setattr(tl, "_resolve_kernel_search_roots", lambda: (f"{root}/",)) + tl._discover_kernel_search_roots.cache_clear() + assert tl._discover_kernel_search_roots() == (str(root),) + + def test_deduplicates_while_preserving_order(self, monkeypatch, tmp_path): + first = tmp_path / "vllm" + second = tmp_path / "aiter" + first.mkdir() + second.mkdir() + monkeypatch.setattr( + tl, + "_resolve_kernel_search_roots", + lambda: (f"{first}/", f"{second}/", str(first), f"{first}//"), + ) + tl._discover_kernel_search_roots.cache_clear() + assert tl._discover_kernel_search_roots() == (str(first), str(second)) + + def test_falls_back_to_local_discovery_without_the_orchestrator( + self, monkeypatch, tmp_path + ): + """Standalone CLI use must still find the installed frameworks.""" + located = tmp_path / "aiter" + located.mkdir() + monkeypatch.setattr(tl, "_resolve_kernel_search_roots", None) + monkeypatch.setattr(tl, "_KERNEL_SOURCE_PACKAGES", ("aiter",)) + monkeypatch.setattr(tl, "_FALLBACK_SEARCH_ROOTS", ()) + monkeypatch.setattr( + tl, "_installed_package_dir", lambda package: str(located) if package == "aiter" else "" + ) + tl._discover_kernel_search_roots.cache_clear() + assert tl._discover_kernel_search_roots() == (str(located),) + + def test_pinned_layouts_are_a_last_resort_not_a_requirement( + self, monkeypatch, tmp_path + ): + """A pinned root is used only when it exists, never assumed.""" + checkout = tmp_path / "sgl-workspace" / "vllm" + checkout.mkdir(parents=True) + monkeypatch.setattr(tl, "_resolve_kernel_search_roots", None) + monkeypatch.setattr(tl, "_KERNEL_SOURCE_PACKAGES", ()) + monkeypatch.setattr( + tl, "_FALLBACK_SEARCH_ROOTS", (str(checkout), "/sgl-workspace/gone") + ) + tl._discover_kernel_search_roots.cache_clear() + assert tl._discover_kernel_search_roots() == (str(checkout),) + + def test_no_searchable_root_is_reported_loudly(self, monkeypatch, caplog): + """An unsearchable host must not look like a healthy one.""" + monkeypatch.setattr(tl, "_resolve_kernel_search_roots", lambda: ("/gone/vllm/",)) + tl._discover_kernel_search_roots.cache_clear() + with caplog.at_level(logging.WARNING, logger=tl.log.name): + assert tl._discover_kernel_search_roots() == () + assert "no framework source root" in caplog.text + + def teardown_method(self): + """Drop the cached roots so the next test resolves them afresh.""" + tl._discover_kernel_search_roots.cache_clear() diff --git a/src/hyperloom/agents/kernel/tests/test_llm_source_fallback.py b/src/hyperloom/agents/kernel/tests/test_llm_source_fallback.py deleted file mode 100644 index a61f540dcb..0000000000 --- a/src/hyperloom/agents/kernel/tests/test_llm_source_fallback.py +++ /dev/null @@ -1,932 +0,0 @@ -############################################################################### -# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. -# SPDX-License-Identifier: MIT -# -# See LICENSE for license information. -############################################################################### - -"""The LLM tier of source resolution: selection only, with canonical routing. - -This pipeline was broken by an LLM writing a placeholder into a field that was -consumed as a path, so the tier that reintroduces an LLM has to be provably -unable to repeat that: it may only echo back one of the paths it was given, the -answer is checked against the filesystem, and a provider must be selected by -the role override or the canonical credential shape. -""" - -from __future__ import annotations - -import json -import sys -import types -from pathlib import Path - -import pytest - -sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tools")) - -import _llm_source_fallback as lsf # noqa: E402 -from hyperloom.common import llm_config # noqa: E402 - - -@pytest.fixture() -def enabled(): - """Kept as a no-op so the tier's tests read the same after it became default.""" - return None - - -@pytest.fixture() -def provider(monkeypatch): - """Select a provider, which the finalizer settles before it greps.""" - monkeypatch.setenv("HYPERLOOM_LLM_SOURCE_PROVIDER", "openai_compatible") - - -@pytest.fixture() -def files(tmp_path): - real = tmp_path / "kernel.py" - real.write_text("@triton.jit\ndef my_kernel():\n pass\n", encoding="utf-8") - test_file = tmp_path / "test_kernel.py" - test_file.write_text("def test_my_kernel():\n pass\n", encoding="utf-8") - return str(real), str(test_file) - - -def _replies(payload) -> callable: - text = payload if isinstance(payload, str) else json.dumps(payload) - return lambda _prompt, _model, _timeout: text - - -# --- Always on ------------------------------------------------------------------ - - -def test_runs_without_any_opt_in(files): - """The tier is unconditional; no environment setup precedes a call.""" - picked, _conf, _reason = lsf.select_source_via_llm( - "my_kernel", [files[0]], complete=_replies({"source_file": files[0], "confidence": 1.0}) - ) - assert picked == files[0] - - -def test_empty_shortlist_still_short_circuits(files): - """Nothing to choose from means no call, independent of the tier being on.""" - picked, _conf, reason = lsf.select_source_via_llm("my_kernel", [], complete=_replies({})) - assert picked == "" - assert "no candidates" in reason - - -# --- Selection, never generation ---------------------------------------------- - - -def test_accepts_a_candidate_from_the_shortlist(enabled, files): - real, test_file = files - picked, confidence, _reason = lsf.select_source_via_llm( - "my_kernel", - [test_file, real], - complete=_replies({"source_file": real, "confidence": 0.9, "reason": "defines the kernel"}), - ) - assert picked == real - assert confidence == 0.9 - - -def test_rejects_a_path_outside_the_shortlist(enabled, files, tmp_path): - """An invented path is the failure mode this tier must not have.""" - invented = str(tmp_path / "hallucinated.py") - picked, _conf, reason = lsf.select_source_via_llm( - "my_kernel", [files[0]], complete=_replies({"source_file": invented, "confidence": 1.0}) - ) - assert picked == "" - assert "not one of the candidates" in reason - - -def test_rejects_a_shortlisted_path_that_vanished(enabled, tmp_path): - missing = str(tmp_path / "gone.py") - picked, _conf, reason = lsf.select_source_via_llm( - "my_kernel", [missing], complete=_replies({"source_file": missing, "confidence": 1.0}) - ) - assert picked == "" - assert "does not exist" in reason - - -def test_rejects_a_path_outside_the_framework_roots(enabled, files): - picked, _conf, reason = lsf.select_source_via_llm( - "my_kernel", - [files[0]], - framework_roots=("/sgl-workspace/sglang",), - complete=_replies({"source_file": files[0], "confidence": 1.0}), - ) - assert picked == "" - assert "outside every framework root" in reason - - -def test_rejects_a_symlink_that_escapes_the_framework_root(enabled, tmp_path): - """A lexical in-root path may not resolve to a file outside the root.""" - root = tmp_path / "root" - root.mkdir() - outside = tmp_path / "outside.py" - outside.write_text("def kernel():\n pass\n", encoding="utf-8") - link = root / "kernel.py" - link.symlink_to(outside) - picked, _conf, reason = lsf.select_source_via_llm( - "kernel", - [str(link)], - framework_roots=(str(root),), - complete=_replies({"source_file": str(link), "confidence": 1.0}), - ) - assert picked == "" - assert "outside every framework root" in reason - - -def test_accepts_a_symlink_whose_target_remains_inside_the_root(enabled, tmp_path): - """Symlinks stay valid, but what is stored is the target that was checked. - - Root containment is decided on the resolved target, so keeping the link - would record a location whose authorization can be revoked afterwards by - retargeting it. The review tier resolves for the same reason. - """ - root = tmp_path / "root" - root.mkdir() - target = root / "implementation.py" - target.write_text("def kernel():\n pass\n", encoding="utf-8") - link = root / "kernel.py" - link.symlink_to(target) - picked, _conf, _reason = lsf.select_source_via_llm( - "kernel", - [str(link)], - framework_roots=(str(root),), - complete=_replies({"source_file": str(link), "confidence": 1.0}), - ) - assert picked == str(target) - - -def test_no_candidates_short_circuits_without_calling_the_model(enabled): - def _boom(*_args): # pragma: no cover - must not run - raise AssertionError("model called with an empty shortlist") - - picked, _conf, reason = lsf.select_source_via_llm("my_kernel", [], complete=_boom) - assert picked == "" - assert "no candidates" in reason - - -# --- Confidence and malformed replies ----------------------------------------- - - -def test_low_confidence_is_discarded(enabled, files): - picked, confidence, reason = lsf.select_source_via_llm( - "my_kernel", [files[0]], complete=_replies({"source_file": files[0], "confidence": 0.4}) - ) - assert picked == "" - assert confidence == 0.4 - assert "below" in reason - - -@pytest.mark.parametrize( - "confidence", - ["NaN", "Infinity", "-Infinity", '"NaN"', '"Infinity"', '"-Infinity"', "-0.1", "1.1"], -) -def test_non_finite_and_out_of_range_confidence_is_rejected(files, confidence): - """Only finite confidence values in the declared range are accepted.""" - source = files[0] - reply = f'{{"source_file": {json.dumps(source)}, "confidence": {confidence}}}' - picked, parsed_confidence, reason = lsf.select_source_via_llm( - "my_kernel", - [source], - complete=_replies(reply), - ) - assert picked == "" - assert parsed_confidence == 0.0 - assert "finite and in [0, 1]" in reason - - -def test_empty_answer_means_none_of_the_candidates(enabled, files): - picked, _conf, reason = lsf.select_source_via_llm( - "my_kernel", [files[0]], complete=_replies({"source_file": "", "confidence": 0.0}) - ) - assert picked == "" - assert "no candidate" in reason - - -def test_prose_wrapped_json_is_still_parsed(enabled, files): - real = files[0] - reply = f'Sure!\n```json\n{{"source_file": "{real}", "confidence": 0.95}}\n```\n' - picked, _conf, _reason = lsf.select_source_via_llm("my_kernel", [real], complete=_replies(reply)) - assert picked == real - - -def test_non_json_reply_is_rejected(enabled, files): - errors = [] - picked, _conf, reason = lsf.select_source_via_llm( - "my_kernel", - [files[0]], - complete=_replies("I could not determine the file."), - errors=errors, - ) - assert picked == "" - assert "no JSON object" in reason - assert errors == [reason] - - -@pytest.mark.parametrize( - ("reply", "expected_reason"), - [ - (None, "reply is not text"), - ("{invalid}", "unparseable JSON"), - ('{"confidence": "high"}', "confidence is not numeric"), - ], -) -def test_parse_answer_rejects_invalid_payload_shapes(reply, expected_reason): - """Malformed payload variants must fail with their precise safe reason.""" - parsed, source_file, confidence, reason = lsf._parse_answer(reply) - assert parsed is False - assert source_file == "" - assert confidence == 0.0 - assert reason == expected_reason - - -def test_parse_answer_rejects_non_object_json(monkeypatch): - """The defensive parser guard must reject a decoded non-object payload.""" - - class Match: - """Return a fixed JSON array from the regex match surface.""" - - @staticmethod - def group(_index): - """Return the non-object JSON payload.""" - return "[]" - - class Pattern: - """Expose the minimal search interface used by the parser.""" - - @staticmethod - def search(_text): - """Return the fixed match.""" - return Match() - - monkeypatch.setattr(lsf, "_JSON_BLOCK_RE", Pattern()) - parsed, source_file, confidence, reason = lsf._parse_answer("ignored") - assert parsed is False - assert source_file == "" - assert confidence == 0.0 - assert reason == "JSON payload is not an object" - - -def test_model_error_is_swallowed(enabled, files): - def _raise(*_args): - raise RuntimeError("gateway 401") - - errors: list[str] = [] - picked, _conf, reason = lsf.select_source_via_llm( - "my_kernel", - [files[0]], - complete=_raise, - errors=errors, - ) - assert picked == "" - assert "llm call failed" in reason - assert errors == [reason] - - -def test_model_error_is_redacted_from_reason_errors_and_log(files): - """Transport diagnostics retain only a stable type and status code.""" - - class TransportError(RuntimeError): - """Represent a provider failure carrying sensitive response details.""" - - status_code = 401 - - def _raise(*_args): - raise TransportError("https://gateway.example/v1?token=query-secret Authorization: Bearer header-secret") - - errors: list[str] = [] - logs: list[str] = [] - picked, _conf, reason = lsf.select_source_via_llm( - "my_kernel", - [files[0]], - complete=_raise, - errors=errors, - log=logs.append, - ) - recorded = "\n".join([reason, *errors, *logs]) - assert picked == "" - assert "TransportError" in recorded - assert "status_code=401" in recorded - assert "gateway.example" not in recorded - assert "query-secret" not in recorded - assert "header-secret" not in recorded - assert "Authorization" not in recorded - - -def test_exception_label_skips_hostile_and_boolean_codes(): - """Exception metadata inspection must ignore unsafe or ambiguous values.""" - - class HostileMetadataError(RuntimeError): - """Expose unusable metadata before one safe code.""" - - @property - def status_code(self): - """Raise instead of exposing provider response details.""" - raise RuntimeError("secret response body") - - code = True - errno = "E_GATEWAY" - - assert lsf._safe_exception_label(HostileMetadataError()) == ("HostileMetadataError (errno=E_GATEWAY)") - - -# --- Provider routing and audit ----------------------------------------------- - - -def _stub_credential_shape( - monkeypatch: pytest.MonkeyPatch, - *, - anthropic: bool, - openai: bool, -) -> None: - """Make the canonical llm_config predicates report one credential shape.""" - monkeypatch.setattr(llm_config, "is_anthropic_only", lambda: anthropic and not openai) - monkeypatch.setattr(llm_config, "is_openai_only", lambda: openai and not anthropic) - monkeypatch.setattr(llm_config, "has_anthropic_side", lambda: anthropic) - monkeypatch.setattr(llm_config, "has_openai_side", lambda: openai) - - -@pytest.mark.parametrize( - ("configured", "expected"), - [ - ("anthropic", lsf._PROVIDER_CLAUDE), - ("openai", lsf._PROVIDER_OPENAI), - ], -) -def test_explicit_provider_override_wins_credential_shape(monkeypatch, configured, expected): - """The role-specific provider knob wins without consulting credential shape.""" - monkeypatch.setenv(lsf._PROVIDER_ENV, configured) - - def _unexpected_shape_probe(): - raise AssertionError("explicit provider must short-circuit credential inference") - - for predicate in ( - "is_anthropic_only", - "is_openai_only", - "has_anthropic_side", - "has_openai_side", - ): - monkeypatch.setattr(llm_config, predicate, _unexpected_shape_probe) - - assert lsf._resolve_provider() == expected - - -@pytest.mark.parametrize( - ("anthropic", "openai", "expected"), - [ - pytest.param(False, False, None, id="neither"), - pytest.param(False, True, lsf._PROVIDER_OPENAI, id="openai-only"), - pytest.param(True, False, lsf._PROVIDER_CLAUDE, id="anthropic-only"), - pytest.param(True, True, lsf._PROVIDER_OPENAI, id="both"), - ], -) -def test_provider_is_inferred_from_canonical_credential_shape( - monkeypatch, - anthropic, - openai, - expected, -): - """All four shapes route canonically; dual-configured single shots use OpenAI.""" - monkeypatch.delenv(lsf._PROVIDER_ENV, raising=False) - _stub_credential_shape(monkeypatch, anthropic=anthropic, openai=openai) - - if expected is None: - with pytest.raises(RuntimeError, match=lsf._PROVIDER_ENV): - lsf._resolve_provider() - else: - assert lsf._resolve_provider() == expected - - -@pytest.mark.parametrize( - ("anthropic", "openai", "expected_provider"), - [ - pytest.param(False, True, lsf._PROVIDER_OPENAI, id="openai-only"), - pytest.param(True, False, lsf._PROVIDER_CLAUDE, id="anthropic-only"), - pytest.param(True, True, lsf._PROVIDER_OPENAI, id="both"), - ], -) -@pytest.mark.parametrize( - ("preview_value", "expected_preview"), - [ - pytest.param("", False, id="preview-unset"), - pytest.param("1", True, id="preview-opted-in"), - ], -) -def test_provider_inference_does_not_authorize_source_preview( - monkeypatch, - files, - anthropic, - openai, - expected_provider, - preview_value, - expected_preview, -): - """Credential inference changes routing only; source-content egress stays opt-in.""" - monkeypatch.delenv(lsf._PROVIDER_ENV, raising=False) - monkeypatch.delenv(lsf._PREVIEW_ENV, raising=False) - if preview_value: - monkeypatch.setenv(lsf._PREVIEW_ENV, preview_value) - _stub_credential_shape(monkeypatch, anthropic=anthropic, openai=openai) - - audit = lsf.llm_source_audit() - prompt = lsf._build_prompt( - "my_kernel", - [files[0]], - framework_roots=(str(Path(files[0]).parent),), - ) - - assert audit["provider"] == expected_provider - assert audit["source_preview_authorised"] is expected_preview - assert ("@triton.jit" in prompt) is expected_preview - - -def test_network_calls_fail_closed_without_override_or_credentials(monkeypatch): - """A model name alone must never imply a provider or endpoint.""" - monkeypatch.delenv(lsf._PROVIDER_ENV, raising=False) - _stub_credential_shape(monkeypatch, anthropic=False, openai=False) - with pytest.raises(RuntimeError, match=lsf._PROVIDER_ENV): - lsf._resolve_provider() - - -def test_provider_helpers_fail_closed_without_valid_configuration(monkeypatch): - """Unsupported providers must fail and Claude audit hosts must stay generic.""" - with pytest.raises(RuntimeError, match="unsupported"): - lsf._resolve_provider("unknown-provider") - monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False) - monkeypatch.delenv("DEEPSEEK_BASE_URL", raising=False) - assert lsf._endpoint_host(lsf._PROVIDER_CLAUDE) == "provider-default" - - -def test_default_claude_model_handles_role_import_failure(monkeypatch): - """Standalone tools must tolerate an unavailable orchestrator role module.""" - monkeypatch.setitem(sys.modules, "hyperloom.orchestrator.roles.agent_role", None) - assert lsf._default_claude_model() == "" - - -@pytest.mark.parametrize( - ("configured", "expected"), - [ - ("claude", lsf._PROVIDER_CLAUDE), - ("claude_agent_sdk", lsf._PROVIDER_CLAUDE), - ("openai", lsf._PROVIDER_OPENAI), - ("openai-compatible", lsf._PROVIDER_OPENAI), - ], -) -def test_provider_aliases_route_to_native_backends(monkeypatch, configured, expected): - """Explicit aliases normalize to one auditable provider identifier.""" - monkeypatch.setenv(lsf._PROVIDER_ENV, configured) - calls = [] - monkeypatch.setattr( - lsf, - "_complete_claude_sdk", - lambda *_args: calls.append(lsf._PROVIDER_CLAUDE) or "claude", - ) - monkeypatch.setattr( - lsf, - "_complete_openai", - lambda *_args: calls.append(lsf._PROVIDER_OPENAI) or "openai", - ) - assert lsf._complete("prompt", "model", 1.0) in {"claude", "openai"} - assert calls == [expected] - - -def test_openai_provider_adapter_uses_the_sanctioned_client_contract(monkeypatch): - """The OpenAI adapter must get its client from ``llm_config``, not build one. - - Credential resolution is asserted by ``llm_config``'s own tests; what matters - here is that the adapter goes through the contract and sends the request the - fallback expects. - """ - captured = {} - - def _create(**kwargs): - """Capture the completion request and return one SDK-shaped response.""" - captured["request"] = kwargs - message = types.SimpleNamespace(content="provider reply") - return types.SimpleNamespace(choices=[types.SimpleNamespace(message=message)]) - - class Client: - """Minimal OpenAI client surface used by the adapter.""" - - def __init__(self): - """Expose the chat-completions surface the contract calls.""" - completions = types.SimpleNamespace(create=_create) - self.chat = types.SimpleNamespace(completions=completions) - - def _get_client(**kwargs): - captured["client_kwargs"] = kwargs - return Client() - - monkeypatch.setattr("hyperloom.common.llm_config.get_openai_client", _get_client) - - reply = lsf._complete_openai("prompt", "source-model", 7.5) - - assert reply == "provider reply" - assert captured["client_kwargs"] == {} - assert captured["request"] == { - "model": "source-model", - "messages": [ - {"role": "system", "content": lsf._SYSTEM_PROMPT}, - {"role": "user", "content": "prompt"}, - ], - "temperature": 0.0, - "timeout": 7.5, - } - - -def test_message_text_accepts_sdk_text_shapes(): - """Claude SDK string, text, and content-block forms must all be readable.""" - from hyperloom.common.claude_oneshot import message_text - - assert message_text("direct") == ["direct"] - assert message_text(types.SimpleNamespace(text="attribute")) == ["attribute"] - message = types.SimpleNamespace( - content=[ - {"text": "dict"}, - types.SimpleNamespace(text="object"), - {"other": "ignored"}, - ] - ) - assert message_text(message) == ["dict", "object"] - - -def test_claude_provider_rejects_incomplete_sdk(monkeypatch): - """A partial SDK installation must fail before any network request.""" - monkeypatch.setitem(sys.modules, "claude_agent_sdk", types.ModuleType("claude_agent_sdk")) - with pytest.raises(RuntimeError, match="missing query / ClaudeAgentOptions"): - lsf._complete_claude_sdk("prompt", "model", 1.0) - - -def test_claude_provider_uses_a_tool_free_native_sdk_call(monkeypatch): - """Claude routing uses the SDK without granting repository or shell tools.""" - captured = {} - - class Options: - """Capture ClaudeAgentOptions keyword arguments.""" - - def __init__(self, **kwargs): - self.kwargs = kwargs - captured["options"] = self - - async def query(*, prompt, options): - """Yield one SDK-shaped result message.""" - captured["prompt"] = prompt - captured["query_options"] = options - yield types.SimpleNamespace(result='{"source_file": "", "confidence": 0}') - - fake_sdk = types.SimpleNamespace(query=query, ClaudeAgentOptions=Options) - monkeypatch.setitem(sys.modules, "claude_agent_sdk", fake_sdk) - monkeypatch.setattr( - "hyperloom.common.llm_config.claude_sdk_env_options", - lambda **_kwargs: { - "model": "overridden-model", - "tools": ["Read"], - "setting_sources": ["user"], - "skills": ["unsafe-skill"], - "strict_mcp_config": False, - "mcp_servers": {"unsafe": {"command": "sh"}}, - "plugins": [{"type": "local", "path": "/tmp/plugin"}], - "max_turns": 99, - "allowed_tools": ["Bash"], - "disallowed_tools": [], - }, - ) - reply = lsf._complete_claude_sdk("prompt", "claude-model", 1.0) - assert json.loads(reply)["source_file"] == "" - assert captured["prompt"] == "prompt" - assert captured["query_options"] is captured["options"] - options = captured["options"].kwargs - assert options["model"] == "claude-model" - assert options["tools"] == [] - assert options["setting_sources"] == [] - assert options["skills"] == [] - assert options["strict_mcp_config"] is True - assert options["mcp_servers"] == {} - assert options["plugins"] == [] - assert options["max_turns"] == 1 - assert options["allowed_tools"] == [] - assert "Read" in options["disallowed_tools"] - assert "Bash" in options["disallowed_tools"] - - -def test_provider_audit_never_records_url_credentials(monkeypatch): - """Audit metadata contains only the endpoint hostname, never URL secrets.""" - monkeypatch.setenv(lsf._PROVIDER_ENV, "openai-compatible") - monkeypatch.setenv(lsf._MODEL_ENV, "gpt-source") - monkeypatch.setenv( - "OPENAI_BASE_URL", - "https://user:secret@gateway.example/Unified/v1?token=must-not-leave", - ) - audit = lsf.llm_source_audit() - assert audit == { - "provider": lsf._PROVIDER_OPENAI, - "model": "gpt-source", - "endpoint_host": "gateway.example", - "source_preview_authorised": False, - } - assert "secret" not in json.dumps(audit) - assert "token" not in json.dumps(audit) - - -def test_provider_model_settings_do_not_cross(monkeypatch): - """Each native provider reads only its own model fallback.""" - monkeypatch.delenv(lsf._MODEL_ENV, raising=False) - monkeypatch.setenv("OPENAI_MODEL", "gpt-source") - monkeypatch.setenv("CLAUDE_MODEL", "claude-source") - assert lsf._resolve_model("", lsf._PROVIDER_OPENAI) == "gpt-source" - assert lsf._resolve_model("", lsf._PROVIDER_CLAUDE) == "claude-source" - - -def test_prompt_context_and_unreadable_preview_are_explicit(files, tmp_path): - """Prompt context must be retained while preview failures stay non-fatal.""" - prompt = lsf._build_prompt( - "my_kernel", - [files[0]], - context_block="Trace context", - with_preview=False, - ) - assert prompt.startswith("Trace context\n") - assert lsf._preview(str(tmp_path / "missing.py")) == "" - - -def test_selection_reports_unconfigured_provider(monkeypatch, files): - """Missing override and credentials must be reported without a transport.""" - monkeypatch.delenv(lsf._PROVIDER_ENV, raising=False) - _stub_credential_shape(monkeypatch, anthropic=False, openai=False) - errors = [] - logs = [] - - picked, confidence, reason = lsf.select_source_via_llm( - "my_kernel", - [files[0]], - model="source-model", - errors=errors, - log=logs.append, - ) - - assert picked == "" - assert confidence == 0.0 - assert reason.startswith("llm configuration failed: RuntimeError") - assert errors == [reason] - assert logs == ["llm_source_fallback: configuration failed: RuntimeError"] - - -def test_selection_reports_missing_provider_model(monkeypatch, files): - """A configured provider without a model must fail before prompt creation.""" - monkeypatch.setenv(lsf._PROVIDER_ENV, "openai_compatible") - for name in (lsf._MODEL_ENV, "OPENAI_MODEL", "CODEX_MODEL"): - monkeypatch.delenv(name, raising=False) - errors = [] - logs = [] - - picked, confidence, reason = lsf.select_source_via_llm( - "my_kernel", - [files[0]], - errors=errors, - log=logs.append, - ) - - assert picked == "" - assert confidence == 0.0 - assert reason == "no model configured" - assert errors == [reason] - assert logs == [f"llm_source_fallback: no model configured; set ${lsf._MODEL_ENV}"] - - -# --- Wiring into finalization -------------------------------------------------- - - -def test_finalizer_leaves_the_candidate_alone_when_grep_finds_nothing(provider, monkeypatch): - """No shortlist means no call, and the candidate stays unresolved.""" - import tracelens_analysis as tl - - monkeypatch.setattr(tl, "collect_source_candidates_via_grep", lambda *_a, **_k: []) - item = {"name": "some_kernel", "gpu_pct": 40.0, "source_file": ""} - tl._apply_llm_source_fallback(item) - assert item["source_file"] == "" - assert item["source_resolution_reason"] == "llm_fallback_no_shortlist" - - -def test_finalizer_settles_the_provider_before_paying_for_the_shortlist(monkeypatch): - """An unconfigured tier must not run the grep it would only decline to use. - - The shortlist walks every framework root once per keyword, so doing it - first would charge that to every hot kernel on a deployment with neither a - role override nor canonical provider credentials. - """ - import tracelens_analysis as tl - - monkeypatch.delenv("HYPERLOOM_LLM_SOURCE_PROVIDER", raising=False) - _stub_credential_shape(monkeypatch, anthropic=False, openai=False) - grepped: list[str] = [] - monkeypatch.setattr( - tl, - "collect_source_candidates_via_grep", - lambda name, *_a, **_k: grepped.append(name) or [], - ) - item = {"name": "hot_kernel", "gpu_pct": 40.0, "source_file": ""} - tl._apply_llm_source_fallback(item) - assert grepped == [] - assert item["source_resolution_reason"] == "llm_fallback_skipped: no provider configured" - assert item["source_resolution_llm_audit"]["outcome"] == "configuration_error" - - -def test_finalizer_skips_cold_kernels(monkeypatch): - """Below the GPU-share floor the round-trip is not worth its cost.""" - import tracelens_analysis as tl - - called = [] - monkeypatch.setattr(tl, "collect_source_candidates_via_grep", lambda *_a, **_k: called.append(1) or []) - tl._apply_llm_source_fallback({"name": "k", "gpu_pct": 0.5, "source_file": ""}) - assert not called - - -def test_gateway_failure_is_not_recorded_as_a_model_decline(provider, monkeypatch): - """Transport failure and a valid refusal require different operator action.""" - import tracelens_analysis as tl - - monkeypatch.setattr( - tl, - "collect_source_candidates_via_grep", - lambda *_args, **_kwargs: ["/repo/kernel.py"], - ) - monkeypatch.setattr( - lsf, - "llm_source_audit", - lambda **_kwargs: { - "provider": "openai_compatible", - "model": "gpt-source", - "endpoint_host": "gateway.example", - "source_preview_authorised": False, - }, - ) - - def _fail(*_args, errors=None, **_kwargs): - """Return the public failure shape while reporting a transport error.""" - errors.append("llm call failed: gateway 401") - return "", 0.0, "llm call failed: gateway 401" - - monkeypatch.setattr(lsf, "select_source_via_llm", _fail) - item = { - "name": "kernel", - "gpu_pct": 40.0, - "source_file": "", - "source_resolution_reason": "trace_resolver_error: truncated", - } - tl._apply_llm_source_fallback(item) - assert "trace_resolver_error" in item["source_resolution_reason"] - assert "llm_fallback_error" in item["source_resolution_reason"] - assert "llm_fallback_declined" not in item["source_resolution_reason"] - assert item["source_resolution_llm_audit"]["outcome"] == "error" - - -def test_valid_model_refusal_is_still_recorded_as_declined(provider, monkeypatch): - """A parsed no-candidate verdict remains distinct from call failure.""" - import tracelens_analysis as tl - - monkeypatch.setattr( - tl, - "collect_source_candidates_via_grep", - lambda *_args, **_kwargs: ["/repo/kernel.py"], - ) - monkeypatch.setattr(lsf, "llm_source_audit", lambda **_kwargs: {"provider": "test"}) - - def _decline(*_args, errors=None, **_kwargs): - """Return a valid refusal without adding a call error.""" - assert errors == [] - return "", 0.0, "model reported no candidate defines the kernel" - - monkeypatch.setattr(lsf, "select_source_via_llm", _decline) - item = {"name": "kernel", "gpu_pct": 40.0, "source_file": ""} - tl._apply_llm_source_fallback(item) - assert item["source_resolution_reason"].startswith("llm_fallback_declined") - assert item["source_resolution_llm_audit"]["outcome"] == "declined" - - -def test_fallback_provider_audit_is_projected_into_the_artifact(): - """Provider identity remains attached to the decision it produced.""" - import tracelens_analysis as tl - - audit = { - "provider": "claude_agent_sdk", - "model": "claude-source", - "endpoint_host": "provider-default", - "source_preview_authorised": False, - "outcome": "accepted", - } - entry = tl.build_source_resolution_entries( - [ - { - "kernel_id": "k1", - "name": "kernel", - "gpu_pct": 10.0, - "source_file": "/repo/kernel.py", - "source_resolution_method": "llm_fallback", - "source_resolution_llm_audit": audit, - } - ] - )[0] - assert entry["llm_audit"] == audit - assert entry["method"] == "llm_fallback" - - -def test_runtime_api_names_yield_no_shortlist(): - import tracelens_analysis as tl - - assert tl.collect_source_candidates_via_grep("hipGraphLaunch") == [] - - -def test_relaxed_shortlist_is_reachable_after_strict_grep_gives_up( - monkeypatch, - tmp_path, -): - """Short mangled identifiers feed only the bounded LLM shortlist.""" - import tracelens_analysis as tl - - source = tmp_path / "kernel.py" - source.write_text("def gemm():\n pass\n", encoding="utf-8") - monkeypatch.setattr(tl, "KNOWN_SEARCH_ROOTS", [str(tmp_path)]) - tl._GREP_CACHE.clear() - mangled = "_Z2ab4gemm" - assert tl._candidate_keywords(mangled) == [] - assert tl.locate_source_via_grep(mangled) == "" - assert tl.collect_source_candidates_via_grep(mangled) == [str(source)] - - -# --- source egress ------------------------------------------------------------ - - -def test_file_contents_do_not_leave_without_authorisation(monkeypatch, files): - """Shipping file heads to a provider is an operator decision, not a default.""" - real, _ = files - monkeypatch.delenv(lsf._PREVIEW_ENV, raising=False) - prompt = lsf._build_prompt("my_kernel", [real]) - assert "@triton.jit" not in prompt - # The path itself still carries most of the selection signal. - assert real in prompt - - -def test_authorisation_restores_the_preview(monkeypatch, files): - """The tier is not crippled by the default; an operator can opt back in.""" - real, _ = files - monkeypatch.setenv(lsf._PREVIEW_ENV, "1") - assert lsf.source_preview_authorised() is True - assert "@triton.jit" in lsf._build_prompt( - "my_kernel", - [real], - framework_roots=(str(Path(real).parent),), - ) - - -def test_preview_does_not_read_a_symlink_target_outside_the_root(monkeypatch, tmp_path): - """An escaping symlink may be named but its target must never be read.""" - root = tmp_path / "root" - root.mkdir() - outside = tmp_path / "outside.py" - outside.write_text("outside_secret_marker\n", encoding="utf-8") - link = root / "kernel.py" - link.symlink_to(outside) - reads: list[str] = [] - monkeypatch.setenv(lsf._PREVIEW_ENV, "1") - monkeypatch.setattr(lsf, "_preview", lambda path: reads.append(path) or "leaked") - - prompt = lsf._build_prompt( - "kernel", - [str(link)], - framework_roots=(str(root),), - ) - - assert reads == [] - assert "outside_secret_marker" not in prompt - assert "leaked" not in prompt - - -def test_preview_and_storage_reuse_the_same_canonical_target(monkeypatch, tmp_path): - """Retargeting a symlink after preview cannot change the stored path.""" - root = tmp_path / "root" - root.mkdir() - original = root / "implementation.py" - original.write_text("original_kernel_marker\n", encoding="utf-8") - outside = tmp_path / "outside.py" - outside.write_text("outside_kernel_marker\n", encoding="utf-8") - link = root / "kernel.py" - link.symlink_to(original) - monkeypatch.setenv(lsf._PREVIEW_ENV, "1") - - def _retarget_after_preview(prompt, _model, _timeout): - """Retarget the candidate only after its validated preview is built.""" - assert "original_kernel_marker" in prompt - assert "outside_kernel_marker" not in prompt - link.unlink() - link.symlink_to(outside) - return json.dumps({"source_file": str(link), "confidence": 1.0}) - - picked, _confidence, _reason = lsf.select_source_via_llm( - "kernel", - [str(link)], - framework_roots=(str(root),), - complete=_retarget_after_preview, - ) - - assert picked == str(original) diff --git a/src/hyperloom/agents/kernel/tests/test_source_resolution_artifact.py b/src/hyperloom/agents/kernel/tests/test_source_resolution_artifact.py index 2d3a230a07..3cbbff1642 100644 --- a/src/hyperloom/agents/kernel/tests/test_source_resolution_artifact.py +++ b/src/hyperloom/agents/kernel/tests/test_source_resolution_artifact.py @@ -24,7 +24,6 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tools")) import tracelens_analysis as tl # noqa: E402 -from _llm_source_review import build_review_prompt, review_resolution_document # noqa: E402 from hyperloom.common import kernel_source_contract as ksc # noqa: E402 @@ -99,13 +98,16 @@ def test_validate_catches_a_path_that_claims_to_be_rejected(): def test_validate_rejects_non_finite_and_out_of_range_confidence(): """Artifact confidence must remain a finite probability.""" for confidence in (float("nan"), float("inf"), float("-inf"), -0.1, 1.1, "NaN"): - doc = _doc_with( - ksc.make_entry( - kernel_id="k1", - name="n", - gpu_pct=1.0, - confidence=confidence, - ) + doc = ksc.make_document( + [ + ksc.make_entry( + kernel_id="k1", + name="n", + gpu_pct=1.0, + confidence=confidence, + ) + ], + generated_by="test", ) assert any("invalid confidence" in problem for problem in ksc.validate_document(doc)) @@ -151,778 +153,6 @@ def test_written_artifact_satisfies_its_own_contract(tmp_path): assert doc["framework"] == "sglang" -# --- review tier guard rails ------------------------------------------------- - - -def _doc_with(*entries): - return ksc.make_document(list(entries), generated_by="test") - - -def _reply(revisions): - def _complete(_prompt, _model, _timeout): - return json.dumps({"revisions": revisions}) - - return _complete - - -def test_review_may_unresolve_a_confidently_wrong_entry(): - """The failure the deterministic tiers cannot self-detect. - - aten::fill_ has no defining source; the trace tier attributes it to - whichever business file called it, and that path passes every mechanical - check. Only a reviewer looking at the whole entry can reject it. - """ - doc = _doc_with( - ksc.make_entry(kernel_id="k1", name="aten::fill_", gpu_pct=9.0, - source_file="/repo/moe.py", method=ksc.METHOD_TRACE) - ) - out, notes = review_resolution_document( - doc, - framework_roots=("/repo",), - complete=_reply([{"kernel_id": "k1", "action": "unresolve", "reason": "builtin"}]), - model="m", - ) - entry = out["entries"][0] - assert entry["source_file"] == "" - assert entry["method"] == ksc.METHOD_UNRESOLVED - assert entry["previous_source_file"] == "/repo/moe.py" - assert notes - - -def test_review_cannot_invent_a_path(): - """A rewrite must land somewhere verifiable, or the original stands.""" - doc = _doc_with( - ksc.make_entry(kernel_id="k1", name="n", gpu_pct=9.0, - source_file="/repo/a.py", method=ksc.METHOD_TRACE) - ) - out, notes = review_resolution_document( - doc, - framework_roots=("/repo",), - complete=_reply([{"kernel_id": "k1", "action": "rewrite", - "source_file": "/tmp/invented.py", "reason": "guess"}]), - model="m", - ) - entry = out["entries"][0] - assert entry["source_file"] == "/repo/a.py" - assert entry["method"] == ksc.METHOD_TRACE - assert "previous_source_file" not in entry - assert any("rejected unverifiable path" in n for n in notes) - - -def test_review_denies_rewrites_when_the_path_contract_is_unavailable(monkeypatch, tmp_path): - """An unusable guard denies the rewrite instead of waving it through. - - ``path_is_acceptable`` is the only thing standing between a generated path - and ``source_file``. When the contract module cannot be imported the tier - has no way to verify anything, so a rewrite must fail closed -- otherwise - losing the import silently turns the guard off. - """ - from _llm_source_review import _apply_revision - - monkeypatch.setattr("_llm_source_review._KSC", None) - entry = {"kernel_id": "k1", "source_file": "/repo/a.py", "method": "name_grep"} - note = _apply_revision( - entry, - {"kernel_id": "k1", "action": "rewrite", "source_file": "/etc/passwd"}, - (str(tmp_path),), - ) - assert entry["source_file"] == "/repo/a.py" - assert entry["method"] == "name_grep" - assert "previous_source_file" not in entry - assert "path contract unavailable" in note - - -def test_review_stores_the_symlink_target_it_validated(tmp_path): - """Keeping the link would let the location leave the roots after the check.""" - target = tmp_path / "implementation.py" - target.write_text("def kernel(): pass\n", encoding="utf-8") - link = tmp_path / "kernel.py" - link.symlink_to(target) - doc = _doc_with( - ksc.make_entry(kernel_id="k1", name="n", gpu_pct=9.0, - source_file=str(tmp_path / "wrong.py"), method=ksc.METHOD_TRACE) - ) - out, _ = review_resolution_document( - doc, - framework_roots=(str(tmp_path),), - complete=_reply([{"kernel_id": "k1", "action": "rewrite", - "source_file": str(link), "reason": "defines it"}]), - model="m", - ) - assert out["entries"][0]["source_file"] == str(target) - - -def test_review_accepts_a_rewrite_to_a_file_that_exists(tmp_path): - """A rewrite must land on a real file under a known root.""" - real = tmp_path / "right.py" - real.write_text("def kernel(): pass\n", encoding="utf-8") - doc = _doc_with( - ksc.make_entry(kernel_id="k1", name="n", gpu_pct=9.0, - source_file=str(tmp_path / "wrong.py"), method=ksc.METHOD_TRACE) - ) - out, _ = review_resolution_document( - doc, - framework_roots=(str(tmp_path),), - complete=_reply([{"kernel_id": "k1", "action": "rewrite", - "source_file": str(real), "reason": "defines it"}]), - model="m", - ) - entry = out["entries"][0] - assert entry["source_file"] == str(real) - assert entry["method"] == ksc.METHOD_LLM - assert entry["previous_source_file"].endswith("wrong.py") - # Line and function described the old file; carrying them over would lie. - assert entry["source_line"] is None - assert entry["source_function"] == "" - - -def test_line_annotated_rewrite_is_stored_as_an_openable_path(tmp_path): - """Call-site metadata is split from the path before downstream use.""" - real = tmp_path / "right.py" - real.write_text("def kernel(): pass\n", encoding="utf-8") - doc = _doc_with( - ksc.make_entry( - kernel_id="k1", - name="n", - gpu_pct=9.0, - source_file=str(tmp_path / "wrong.py"), - method=ksc.METHOD_TRACE, - ) - ) - out, _ = review_resolution_document( - doc, - framework_roots=(str(tmp_path),), - complete=_reply( - [{ - "kernel_id": "k1", - "action": "rewrite", - "source_file": f"{real}(247): kernel", - }] - ), - model="m", - ) - entry = out["entries"][0] - assert entry["source_file"] == str(real) - assert entry["source_line"] == 247 - assert entry["source_function"] == "kernel" - assert Path(entry["source_file"]).is_file() - - -def test_a_line_suffix_only_difference_is_not_a_rewrite(tmp_path): - """Adding call-site syntax to the same file must not create fake history.""" - real = tmp_path / "same.py" - real.write_text("def kernel(): pass\n", encoding="utf-8") - doc = _doc_with( - ksc.make_entry( - kernel_id="k1", - name="n", - gpu_pct=9.0, - source_file=str(real), - method=ksc.METHOD_TRACE, - ) - ) - out, notes = review_resolution_document( - doc, - framework_roots=(str(tmp_path),), - complete=_reply( - [{ - "kernel_id": "k1", - "action": "rewrite", - "source_file": f"{real}(1): kernel", - }] - ), - model="m", - ) - assert out["entries"][0]["source_file"] == str(real) - assert "previous_source_file" not in out["entries"][0] - assert notes == [] - - -def test_line_annotated_paths_are_still_verifiable(tmp_path): - """TraceLens reports call sites as "path.py(247): fn"; that is a real path. - - Measured on a live session, 29 of 36 resolved entries carried this suffix. - Checking existence without stripping it would reject every one of them. - """ - real = tmp_path / "moe.py" - real.write_text("def kernel(): pass\n", encoding="utf-8") - roots = (str(tmp_path),) - assert ksc.path_is_acceptable(str(real), roots) - assert ksc.path_is_acceptable(f"{real}(247)", roots) - assert ksc.path_is_acceptable(f"{real}(247): kernel", roots) - assert ksc.strip_line_suffix(f"{real}(247): kernel") == str(real) - # Stripping must not resurrect a path that is simply absent. - assert not ksc.path_is_acceptable(f"{tmp_path}/gone.py(1): f", roots) - - -def test_symlink_cannot_escape_a_framework_root(tmp_path): - """Root containment is checked on resolved targets, not lexical paths.""" - root = tmp_path / "root" - root.mkdir() - outside = tmp_path / "outside.py" - outside.write_text("def kernel(): pass\n", encoding="utf-8") - escaping = root / "escaping.py" - escaping.symlink_to(outside) - inside = root / "inside.py" - inside.write_text("def kernel(): pass\n", encoding="utf-8") - internal = root / "internal.py" - internal.symlink_to(inside) - assert not ksc.path_is_acceptable(str(escaping), (str(root),)) - assert ksc.path_is_acceptable(str(internal), (str(root),)) - assert ksc.canonical_source_path(str(escaping), (str(root),)) == "" - assert ksc.canonical_source_path(str(internal), (str(root),)) == str(inside) - - -def test_review_preview_does_not_read_an_escaping_current_symlink(monkeypatch, tmp_path): - """A current entry is previewed only through a validated canonical target.""" - root = tmp_path / "root" - root.mkdir() - outside = tmp_path / "outside.py" - outside.write_text("review_outside_secret\n", encoding="utf-8") - link = root / "kernel.py" - link.symlink_to(outside) - entry = ksc.make_entry( - kernel_id="k1", - name="n", - gpu_pct=9.0, - source_file=str(link), - method=ksc.METHOD_TRACE, - ) - reads: list[str] = [] - monkeypatch.setattr("_llm_source_review._preview", lambda path: reads.append(path) or "leaked") - - prompt = build_review_prompt( - [entry], - with_preview=True, - framework_roots=(str(root),), - ) - - assert reads == [] - assert "review_outside_secret" not in prompt - assert "leaked" not in prompt - - -def test_review_rejects_a_plausible_path_that_does_not_exist(tmp_path): - """Under-a-root is not enough: the file must be there. - - A model can emit a perfectly plausible path inside the framework tree. If - root membership alone qualified it, the backend would be handed a fabricated - file -- the wrong-source failure this pipeline exists to prevent. - """ - invented = tmp_path / "python" / "sglang" / "srt" / "does_not_exist.py" - doc = _doc_with( - ksc.make_entry(kernel_id="k1", name="n", gpu_pct=9.0, - source_file="", method=ksc.METHOD_UNRESOLVED) - ) - out, notes = review_resolution_document( - doc, - framework_roots=(str(tmp_path),), - complete=_reply([{"kernel_id": "k1", "action": "rewrite", - "source_file": str(invented), "reason": "looks right"}]), - model="m", - ) - assert out["entries"][0]["source_file"] == "" - assert any("rejected unverifiable path" in n for n in notes) - - -def test_review_keeps_entries_below_the_gpu_floor_untouched(): - doc = _doc_with( - ksc.make_entry(kernel_id="k1", name="tiny", gpu_pct=0.01, - source_file="/repo/a.py", method=ksc.METHOD_TRACE) - ) - out, notes = review_resolution_document( - doc, framework_roots=("/repo",), complete=_reply([]), model="m" - ) - assert out["entries"][0]["source_file"] == "/repo/a.py" - assert any("GPU share" in n for n in notes) - - -def test_unusable_reply_leaves_the_table_alone(): - doc = _doc_with( - ksc.make_entry(kernel_id="k1", name="n", gpu_pct=9.0, - source_file="/repo/a.py", method=ksc.METHOD_TRACE) - ) - - def _garbage(_p, _m, _t): - return "not json at all" - - out, notes = review_resolution_document( - doc, framework_roots=("/repo",), complete=_garbage, model="m" - ) - assert out["entries"][0]["source_file"] == "/repo/a.py" - assert any("unusable reply" in n for n in notes) - - -def test_call_failure_leaves_the_table_alone(): - doc = _doc_with( - ksc.make_entry(kernel_id="k1", name="n", gpu_pct=9.0, - source_file="/repo/a.py", method=ksc.METHOD_TRACE) - ) - - def _boom(_p, _m, _t): - raise RuntimeError("gateway 401") - - out, notes = review_resolution_document( - doc, framework_roots=("/repo",), complete=_boom, model="m" - ) - assert out["entries"][0]["source_file"] == "/repo/a.py" - assert any("llm call failed" in n for n in notes) - - -def test_review_call_error_is_redacted_from_notes_and_log(): - """Provider failures expose only stable exception metadata.""" - class ProviderError(RuntimeError): - """Represent a provider response carrying sensitive details.""" - - code = "gateway_timeout" - - doc = _doc_with( - ksc.make_entry( - kernel_id="k1", - name="n", - gpu_pct=9.0, - source_file="/repo/a.py", - method=ksc.METHOD_TRACE, - ) - ) - - def _boom(_prompt, _model, _timeout): - raise ProviderError( - "https://gateway.example/v1?token=query-secret " - "Authorization: Bearer header-secret" - ) - - logs: list[str] = [] - out, notes = review_resolution_document( - doc, - framework_roots=("/repo",), - complete=_boom, - model="m", - log=logs.append, - ) - recorded = "\n".join([*notes, *logs]) - assert out["entries"][0]["source_file"] == "/repo/a.py" - assert "ProviderError" in recorded - assert "code=gateway_timeout" in recorded - assert "gateway.example" not in recorded - assert "query-secret" not in recorded - assert "header-secret" not in recorded - assert "Authorization" not in recorded - - -def test_review_discards_staged_changes_when_path_validation_raises(monkeypatch, tmp_path): - """A later validation exception cannot commit an earlier revision.""" - replacement = tmp_path / "replacement.py" - replacement.write_text("def kernel(): pass\n", encoding="utf-8") - doc = _doc_with( - ksc.make_entry( - kernel_id="k1", - name="a", - gpu_pct=9.0, - source_file="/repo/a.py", - method=ksc.METHOD_TRACE, - ), - ksc.make_entry( - kernel_id="k2", - name="b", - gpu_pct=8.0, - source_file="/repo/b.py", - method=ksc.METHOD_TRACE, - ), - ) - original_entries = json.loads(json.dumps(doc["entries"])) - - class PathValidationError(RuntimeError): - """Represent a failing path guard.""" - - original_helper = ksc.canonical_source_path - - def _guard(path, roots): - """Delegate ordinary paths and fail on the second revision.""" - if path == str(replacement): - raise PathValidationError("https://secret.example/?token=path-secret") - return original_helper(path, roots) - - monkeypatch.setattr(ksc, "canonical_source_path", _guard) - out, notes = review_resolution_document( - doc, - framework_roots=(str(tmp_path),), - complete=_reply([ - {"kernel_id": "k1", "action": "unresolve"}, - {"kernel_id": "k2", "action": "rewrite", "source_file": str(replacement)}, - ]), - model="m", - ) - - assert out["entries"] == original_entries - assert out["llm_audit"]["review"]["outcome"] == "validation_error" - assert any("PathValidationError" in note for note in notes) - assert all("path-secret" not in note for note in notes) - - -def test_review_parse_exception_leaves_entries_untouched(monkeypatch): - """A parser exception is advisory and cannot alter the source table.""" - doc = _doc_with( - ksc.make_entry( - kernel_id="k1", - name="n", - gpu_pct=9.0, - source_file="/repo/a.py", - method=ksc.METHOD_TRACE, - ) - ) - original_entries = json.loads(json.dumps(doc["entries"])) - - def _raise(_reply): - """Simulate a parser failure with sensitive response context.""" - raise ValueError("https://secret.example/?token=parse-secret") - - monkeypatch.setattr("_llm_source_review.parse_revisions", _raise) - out, notes = review_resolution_document( - doc, - framework_roots=("/repo",), - complete=lambda *_args: "{}", - model="m", - ) - - assert out["entries"] == original_entries - assert any("ValueError" in note for note in notes) - assert all("parse-secret" not in note for note in notes) - - -def test_revision_for_unknown_kernel_is_reported_not_applied(): - """An id nobody asked about rejects the batch rather than being skipped. - - A reply naming a kernel that was never sent is not a reply about the batch - that was sent, so the entry it does not mention keeps its own resolution - instead of being silently left to a partial review. - """ - doc = _doc_with( - ksc.make_entry(kernel_id="k1", name="n", gpu_pct=9.0, - source_file="/repo/a.py", method=ksc.METHOD_TRACE) - ) - out, notes = review_resolution_document( - doc, - framework_roots=("/repo",), - complete=_reply([{"kernel_id": "ghost", "action": "unresolve"}]), - model="m", - ) - assert out["entries"][0]["source_file"] == "/repo/a.py" - assert any("unknown kernel_id" in n for n in notes) - - -def test_one_unreadable_gpu_pct_does_not_disable_the_whole_review(tmp_path): - """A single bad row must not cost every other row its review. - - The caller wraps this tier in a blanket handler, so an exception raised - while ranking would surface as "review skipped" for the entire run. An - unreadable share ranks as zero, which drops that row below the floor and - leaves it untouched -- the rest of the table is still reviewed. - """ - real = tmp_path / "right.py" - real.write_text("def kernel(): pass\n", encoding="utf-8") - hot = ksc.make_entry(kernel_id="k1", name="hot", gpu_pct=9.0, - source_file=str(tmp_path / "wrong.py"), method=ksc.METHOD_TRACE) - broken = ksc.make_entry(kernel_id="k2", name="broken", gpu_pct=1.0, - source_file="/repo/b.py", method=ksc.METHOD_TRACE) - broken["gpu_pct"] = "12%" - out, _ = review_resolution_document( - _doc_with(hot, broken), - framework_roots=(str(tmp_path), "/repo"), - complete=_reply([ - {"kernel_id": "k1", "action": "rewrite", "source_file": str(real)}, - ]), - model="m", - ) - assert out["entries"][0]["source_file"] == str(real) - assert out["entries"][1]["source_file"] == "/repo/b.py" - - -def test_missing_review_id_rejects_the_entire_batch(): - """A truncated reply is not equivalent to an explicit keep decision.""" - doc = _doc_with( - ksc.make_entry( - kernel_id="k1", name="a", gpu_pct=9.0, - source_file="/repo/a.py", method=ksc.METHOD_TRACE, - ), - ksc.make_entry( - kernel_id="k2", name="b", gpu_pct=8.0, - source_file="/repo/b.py", method=ksc.METHOD_TRACE, - ), - ) - out, notes = review_resolution_document( - doc, - complete=_reply([{"kernel_id": "k1", "action": "unresolve"}]), - model="m", - ) - assert [entry["source_file"] for entry in out["entries"]] == [ - "/repo/a.py", - "/repo/b.py", - ] - assert any("missing=" in note for note in notes) - assert out["llm_audit"]["review"]["outcome"] == "protocol_error" - - -def test_duplicate_review_id_rejects_the_entire_batch(): - """Repeated revisions cannot overwrite their own audit history.""" - doc = _doc_with( - ksc.make_entry( - kernel_id="k1", name="a", gpu_pct=9.0, - source_file="/repo/a.py", method=ksc.METHOD_TRACE, - ) - ) - revisions = [ - {"kernel_id": "k1", "action": "unresolve"}, - {"kernel_id": "k1", "action": "unresolve"}, - ] - out, notes = review_resolution_document( - doc, - complete=_reply(revisions), - model="m", - ) - assert out["entries"][0]["source_file"] == "/repo/a.py" - assert any("duplicate=" in note for note in notes) - - -def test_revision_for_an_unsent_entry_rejects_the_entire_batch(): - """An entry below the review floor cannot be modified by an extra ID.""" - doc = _doc_with( - ksc.make_entry( - kernel_id="hot", name="a", gpu_pct=9.0, - source_file="/repo/a.py", method=ksc.METHOD_TRACE, - ), - ksc.make_entry( - kernel_id="cold", name="b", gpu_pct=0.1, - source_file="/repo/b.py", method=ksc.METHOD_TRACE, - ), - ) - revisions = [ - {"kernel_id": "hot", "action": "keep"}, - {"kernel_id": "cold", "action": "unresolve"}, - ] - out, notes = review_resolution_document( - doc, - complete=_reply(revisions), - model="m", - ) - assert out["entries"][1]["source_file"] == "/repo/b.py" - assert any("extra/unknown kernel_id" in note for note in notes) - - -# --- revisions must reach the pipeline, not just the artifact --------------- - - -def test_revision_is_folded_back_onto_the_candidate(): - """A review that only edits the artifact is inert. - - Dispatch reads kernel_candidates.json; the artifact is an audit view. Unless - the revision is written back, the review tier changes nothing that runs. - """ - candidates = [ - {"kernel_id": "k1", "name": "foo_kernel", "gpu_pct": 9.0, - "source_file": "/sgl-workspace/aiter/csrc/wrong.cpp", "source_type": "hip_cpp"} - ] - entries = [ - ksc.make_entry( - kernel_id="k1", name="foo_kernel", gpu_pct=9.0, - source_file="/sgl-workspace/aiter/ops/right.py", method=ksc.METHOD_LLM, - ) - ] - entries[0]["previous_source_file"] = "/sgl-workspace/aiter/csrc/wrong.cpp" - entries[0]["previous_method"] = ksc.METHOD_TRACE - - assert tl.apply_resolution_entries_to_candidates(entries, candidates) == 1 - got = candidates[0] - assert got["source_file"] == "/sgl-workspace/aiter/ops/right.py" - assert got["source_path"] == "/sgl-workspace/aiter/ops/right.py" - assert got["source_resolution_previous_file"].endswith("wrong.cpp") - assert got["source_resolution_previous_method"] == ksc.METHOD_TRACE - # source_type drives reusable_native_kernel, so it must be recomputed. - assert got["source_type"] == "python" - assert "reusable_native_kernel" in got - assert "skip_reason" in got - - -def test_unreviewed_entries_leave_candidates_untouched(): - """Only entries carrying previous_source_file were revised.""" - candidates = [{"kernel_id": "k1", "name": "n", "gpu_pct": 9.0, - "source_file": "/repo/a.py", "source_type": "python"}] - entries = [ksc.make_entry(kernel_id="k1", name="n", gpu_pct=9.0, - source_file="/repo/a.py", method=ksc.METHOD_TRACE)] - assert tl.apply_resolution_entries_to_candidates(entries, candidates) == 0 - assert candidates[0]["source_file"] == "/repo/a.py" - assert "source_resolution_previous_file" not in candidates[0] - - -def test_unresolve_clears_the_candidate_source(): - """Dropping to unresolved must also clear it downstream, not only here.""" - candidates = [{"kernel_id": "k1", "name": "aten::fill_", "gpu_pct": 9.0, - "source_file": "/repo/moe.py", "source_type": "python"}] - entries = [ksc.make_entry(kernel_id="k1", name="aten::fill_", gpu_pct=9.0, - source_file="", method=ksc.METHOD_UNRESOLVED)] - entries[0]["previous_source_file"] = "/repo/moe.py" - entries[0]["previous_method"] = ksc.METHOD_TRACE - - assert tl.apply_resolution_entries_to_candidates(entries, candidates) == 1 - got = candidates[0] - assert got["source_file"] == "" - assert got["reusable_native_kernel"] is False - - -def _aiter_candidate(**over): - """A candidate carrying the curated metadata an op_to_source hit stamps.""" - item = { - "kernel_id": "k1", - "name": "fused_moe", - "gpu_pct": 9.0, - "source_file": "/repo/aiter/impl.cu", - "source_type": "hip_cpp", - "kernel_sources": ["/repo/aiter/impl.cu"], - "kernel_kind": "aiter_ck", - "source_framework": "aiter", - "prebuilt_binary": "/repo/aiter/impl.co", - "runtime_backend": "aiter", - "launcher_source_file": "/repo/sglang/launch.py", - "source_promoted_from_launcher": True, - "tracelens_launcher_path": "/repo/sglang/launch.py(10): launch", - "kernel_path": "/repo/sglang/launch.py(10): launch", - "vendor_dispatch_wrapper": True, - "runtime_generated_kernel": True, - "source_resolution_confidence": 0.91, - "op_to_source_kind": "dispatch", - "op_to_source_patchable": True, - } - item.update(over) - return item - - -def _rewrite_to(path): - entry = ksc.make_entry( - kernel_id="k1", name="fused_moe", gpu_pct=9.0, - source_file=path, method=ksc.METHOD_LLM, - ) - entry["previous_source_file"] = "/repo/aiter/impl.cu" - entry["previous_method"] = ksc.METHOD_TRACE - return entry - - -def test_a_rewrite_clears_metadata_describing_the_old_source(): - """Otherwise the candidate describes two sources and readers disagree. - - forge_submit._resolve_framework consults source_framework before it looks - at source_file, so a stale value routes a vLLM rewrite as aiter. - """ - item = _aiter_candidate() - assert tl.apply_resolution_entries_to_candidates([_rewrite_to("/repo/vllm/new.py")], [item]) == 1 - assert item["source_file"] == "/repo/vllm/new.py" - for stale in ( - "kernel_sources", - "kernel_kind", - "source_framework", - "prebuilt_binary", - "runtime_backend", - "launcher_source_file", - "source_promoted_from_launcher", - "tracelens_launcher_path", - "kernel_path", - "vendor_dispatch_wrapper", - "source_resolution_confidence", - "op_to_source_kind", - "op_to_source_patchable", - ): - assert stale not in item, stale - assert item["runtime_generated_kernel"] is False - - -def test_a_stale_aiter_asm_kind_cannot_skip_a_rewritten_kernel(): - """classify_patchability reads kernel_kind; keeping it skips the new source.""" - item = _aiter_candidate(kernel_kind="aiter_asm") - tl.apply_resolution_entries_to_candidates([_rewrite_to("/repo/vllm/new.py")], [item]) - assert "aiter_asm" not in str(item.get("skip_reason") or "") - - -def test_a_curated_resolution_is_not_overridable(): - """op_to_source.json names the real compute core; a file head cannot outrank it.""" - item = _aiter_candidate( - op_to_source_status="resolved", - source_resolution_method=ksc.METHOD_CURATED, - ) - entry = _rewrite_to("/repo/vllm/new.py") - assert tl.apply_resolution_entries_to_candidates([entry], [item]) == 0 - assert item["source_file"] == "/repo/aiter/impl.cu" - assert item["kernel_kind"] == "aiter_ck" - # The artifact must not advertise a revision that was not applied. - assert entry["review_rejected"] == "curated_resolution_not_overridable" - assert entry["source_file"] == "/repo/aiter/impl.cu" - - -@pytest.mark.parametrize("status", ["non_rewritable", "no_kernel"]) -def test_a_curated_negative_verdict_is_not_overridable(status): - """A curated terminal miss is authoritative even when it keeps a launcher.""" - item = _aiter_candidate( - op_to_source_status=status, - source_resolution_method=ksc.METHOD_CURATED, - ) - entry = _rewrite_to("/repo/vllm/new.py") - - assert tl.apply_resolution_entries_to_candidates([entry], [item]) == 0 - assert item["source_file"] == "/repo/aiter/impl.cu" - assert entry["review_rejected"] == "curated_resolution_not_overridable" - - -def test_audit_history_survives_artifact_rebuild(tmp_path, monkeypatch): - """Applied revisions remain reversible after candidates are re-projected.""" - old = tmp_path / "old.py" - new = tmp_path / "new.py" - old.write_text("def old(): pass\n", encoding="utf-8") - new.write_text("def new(): pass\n", encoding="utf-8") - candidates = [{ - "kernel_id": "k1", - "name": "kernel", - "gpu_pct": 9.0, - "source_file": str(old), - "source_type": "python", - "source_resolution_method": ksc.METHOD_TRACE, - }] - - def _review(doc, *, log_path): - """Inject one reviewed rewrite without making a network call.""" - assert log_path is None - entry = doc["entries"][0] - entry["previous_source_file"] = entry["source_file"] - entry["previous_method"] = entry["method"] - entry["source_file"] = str(new) - entry["method"] = ksc.METHOD_LLM - - monkeypatch.setattr(tl, "_review_source_resolution", _review) - out_path = tmp_path / ksc.SOURCE_RESOLUTION_FILENAME - assert tl.write_source_resolution_artifact(candidates, out_path) == out_path - entry = json.loads(out_path.read_text(encoding="utf-8"))["entries"][0] - assert entry["source_file"] == str(new) - assert entry["previous_source_file"] == str(old) - assert entry["previous_method"] == ksc.METHOD_TRACE - assert candidates[0]["source_resolution_previous_file"] == str(old) - assert candidates[0]["source_resolution_previous_method"] == ksc.METHOD_TRACE - - -def test_review_runs_without_any_opt_in(tmp_path): - """The tier is unconditional; nothing in the environment gates it.""" - real = tmp_path / "right.py" - real.write_text("def kernel(): pass\n", encoding="utf-8") - doc = _doc_with( - ksc.make_entry(kernel_id="k1", name="n", gpu_pct=9.0, - source_file=str(tmp_path / "wrong.py"), method=ksc.METHOD_TRACE) - ) - out, _ = review_resolution_document( - doc, - framework_roots=(str(tmp_path),), - complete=_reply([{"kernel_id": "k1", "action": "rewrite", - "source_file": str(real), "reason": "defines it"}]), - model="m", - ) - assert out["entries"][0]["source_file"] == str(real) - - # --- degrade, don't abort, against an older installed contract module ------- # # tracelens_analysis.py runs as a standalone subprocess and imports the diff --git a/src/hyperloom/agents/kernel/tests/test_source_resolution_guards.py b/src/hyperloom/agents/kernel/tests/test_source_resolution_guards.py index 1cac9c89f2..ffcd66541e 100644 --- a/src/hyperloom/agents/kernel/tests/test_source_resolution_guards.py +++ b/src/hyperloom/agents/kernel/tests/test_source_resolution_guards.py @@ -209,7 +209,7 @@ def test_successful_grep_replaces_a_prior_rejection_method(monkeypatch): ) item = {"name": "k", "source_file": "Not found", "duration_us": 1.0} - got = tl._finalize_candidates([item], allow_model_tiers=False)[0] + got = tl._finalize_candidates([item])[0] assert got["source_file"] == "/repo/pkg/kernel.py" assert got["source_resolution_method"] == "name_grep" @@ -468,7 +468,6 @@ def test_trace_launcher_caller_does_not_override_grep_definition( got = tl._finalize_candidates( [_wiring_candidate()], trace_files=[_wiring_trace(tmp_path, f"{launcher}(42): launch")], - allow_model_tiers=False, )[0] assert got["source_file"] == definition @@ -479,28 +478,26 @@ def test_trace_launcher_caller_does_not_override_grep_definition( assert "trace launcher differs from grep source" in got["source_resolution_reason"] -def test_unconfirmed_trace_launcher_continues_to_model_fallback( +def test_unconfirmed_trace_launcher_leaves_the_source_empty( monkeypatch, tmp_path, ): - """An unconfirmed launcher must leave source empty for the next tier.""" - fallback_calls = [] - - def _record_fallback(item): - """Record that finalization reached the model fallback tier.""" - fallback_calls.append(item["name"]) + """An unconfirmed launcher is evidence, never an attribution. + Finalization stops here rather than guessing from the symbol; the + whole-table review that follows can weigh the blank against the launcher + frame and the rest of the table. + """ launcher = "/repo/model/launcher.py" monkeypatch.setattr(tl, "locate_source_via_grep", lambda _name: "") - monkeypatch.setattr(tl, "_apply_llm_source_fallback", _record_fallback) got = tl._finalize_candidates( [_wiring_candidate()], trace_files=[_wiring_trace(tmp_path, f"{launcher}(42): launch")], )[0] - assert fallback_calls == [got["name"]] assert got["source_file"] == "" + assert got["reusable_native_kernel"] is False assert got["trace_launcher_file"] == launcher assert got.get("source_resolution_method") != "trace_python_stack" assert "trace launcher unconfirmed by name grep" in got["source_resolution_reason"] @@ -522,19 +519,14 @@ def test_wiring_without_trace_files_falls_back_quietly(tmp_path): assert "trace_resolver_error" not in str(got.get("source_resolution_reason", "")) -def test_deterministic_finalization_never_calls_model_tiers( - monkeypatch, - tmp_path, -): - """The deterministic route's no-LLM contract is enforced below the CLI.""" - monkeypatch.setattr(tl, "locate_source_via_grep", lambda _name: "") +def test_finalization_never_calls_a_model(monkeypatch, tmp_path): + """Source resolution is wholly deterministic; review is a later stage. - def _model_called(*_args, **_kwargs): - """Fail if either model tier is reached.""" - raise AssertionError("deterministic route called a model tier") - - monkeypatch.setattr(tl, "_apply_llm_source_fallback", _model_called) - monkeypatch.setattr(tl, "_review_source_resolution", _model_called) + Nothing under ``_finalize_candidates`` may reach a provider, so the + deterministic route gets its no-LLM guarantee from the code rather than + from a flag it has to remember to pass. + """ + monkeypatch.setattr(tl, "locate_source_via_grep", lambda _name: "") artifact = tmp_path / "kernel_source_resolution.json" item = { "name": "zz_no_source_kernel", @@ -542,13 +534,9 @@ def _model_called(*_args, **_kwargs): "duration_us": 100.0, "gpu_pct": 10.0, } - got = tl._finalize_candidates( - [item], - source_resolution_out=artifact, - allow_model_tiers=False, - )[0] + got = tl._finalize_candidates([item], source_resolution_out=artifact)[0] assert got["source_file"] == "" - assert "deterministic route" in got["source_resolution_reason"] + assert got["skip_reason"] == "source file not resolved" assert artifact.is_file() doc = json.loads(artifact.read_text(encoding="utf-8")) assert "llm_audit" not in doc diff --git a/src/hyperloom/agents/kernel/tests/test_tracelens_csv.py b/src/hyperloom/agents/kernel/tests/test_tracelens_csv.py index 31e43e8565..0bfc701c14 100644 --- a/src/hyperloom/agents/kernel/tests/test_tracelens_csv.py +++ b/src/hyperloom/agents/kernel/tests/test_tracelens_csv.py @@ -507,13 +507,37 @@ def test_unknown_source_root_is_not_reusable_native(): assert tla.recommend_backends(candidate) == [] -def test_known_rmsnorm_harness_is_registered_without_repo_root(): +def test_known_rmsnorm_harness_is_registered_without_repo_root(monkeypatch, tmp_path): + """A curated harness is found from the kernel name alone, with no repo root. + + The hint is checkout-relative, so it is resolved against the search roots + rather than a pinned ``/sgl-workspace`` path, and only a file that is really + there is reported: a harness list naming paths nobody can open reads + downstream as a runnable harness. + """ + harness = tmp_path / "aiter" / "op_tests" / "test_rmsnorm2d.py" + harness.parent.mkdir(parents=True) + harness.write_text("def test_rmsnorm2d(): pass\n", encoding="utf-8") + monkeypatch.setattr(tla, "KNOWN_SEARCH_ROOTS", (str(tmp_path / "aiter"),)) + tla._harness_search_bases.cache_clear() + files = tla.find_benchmark_files( "_ZN5aiter24add_rmsnorm_quant_kernelIDF16bDF16bLi256EEEv", "", "/sgl-workspace/aiter/csrc/kernels/rmsnorm_quant_kernels.cu", ) - assert any("rmsnorm" in path.lower() for path in files) + tla._harness_search_bases.cache_clear() + + assert files == [str(harness)] + + +def test_absent_curated_harness_is_not_reported(monkeypatch, tmp_path): + """A hint that resolves nowhere yields nothing, not an unopenable path.""" + monkeypatch.setattr(tla, "KNOWN_SEARCH_ROOTS", (str(tmp_path / "aiter"),)) + tla._harness_search_bases.cache_clear() + files = tla.find_benchmark_files("kernel_paged_attention_2d", "", "/pkg/attention.py") + tla._harness_search_bases.cache_clear() + assert files == [] def test_125_finalize_adds_kernel_category_for_attention(): diff --git a/src/hyperloom/agents/kernel/tests/test_vendor_operator_playbook_mori.py b/src/hyperloom/agents/kernel/tests/test_vendor_operator_playbook_mori.py index 7f7a9b8940..eb9eeff6f4 100644 --- a/src/hyperloom/agents/kernel/tests/test_vendor_operator_playbook_mori.py +++ b/src/hyperloom/agents/kernel/tests/test_vendor_operator_playbook_mori.py @@ -313,6 +313,29 @@ def test_finalize_candidates_fills_source_file_for_real_vendor_binary_shape(): assert tla.looks_like_source_path(source_file) +def test_playbook_anchor_overrides_a_same_word_grep_collision(): + """A registry match outranks whatever the grep tier guessed. + + These operators reduce to the keywords "dispatch" and "combine", which + collide with unrelated vendor files (``mxfp4_moe_aux_dispatch.h``, + ``fmha_fwd_d64_bf16_combine.cu``) once the search roots actually resolve. + The registry is a curated statement that the operator is tuned through a + task bundle, so handing a backend the colliding path would rewrite the + wrong file. + """ + collision = "/usr/local/lib/python3.12/dist-packages/aiter_meta/csrc/x_dispatch.h" + candidates = [ + _mori_dispatch_candidate(source_file=collision), + _mori_combine_candidate(source_file=collision), + ] + out = tla._finalize_candidates(candidates, total_dur=1000.0) + + for item in out: + assert item["patch_strategy"] == "vendor_playbook" + assert item["source_file"] != collision + assert str(item["source_file"]).endswith("mori_ep_config.py") + + # --- 3 & 4. forge_submit.submit() vendor-playbook route + one-session dedup -- diff --git a/src/hyperloom/agents/kernel/tools/_candidate_review_agent.py b/src/hyperloom/agents/kernel/tools/_candidate_review_agent.py new file mode 100644 index 0000000000..24f40cb3c6 --- /dev/null +++ b/src/hyperloom/agents/kernel/tools/_candidate_review_agent.py @@ -0,0 +1,845 @@ +############################################################################### +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT +# +# See LICENSE for license information. +############################################################################### + +"""Agent review of the deterministically produced kernel-candidate table. + +The deterministic tiers fail in a way no "fill in the blanks" pass can catch: +they do not come up empty, they come up *confidently wrong*. A launcher frame +proves who launched a kernel, not who defines it; a keyword grep on ``dispatch`` +lands on whichever vendor header mentions the word. Both produce a real, +existing, root-resident path that passes every mechanical check. + +Reviewing that needs the context the tiers do not have -- what the model is, +how it is being served, what the trace actually recorded -- and enough of the +framework tree to confirm a file really defines the kernel it is credited with. +Rather than pre-loading any of that into a prompt, this hands the agent the +*paths* and lets it read what it needs. + +Three properties keep the added freedom bounded: + +* **Proposals only.** The agent may revise where a kernel lives, whether it is + worth dispatching, and the operand dims to tune it against. It may not touch + the trace's own measurements -- GPU share, duration, launch count -- nor the + keys the row is identified by. Those carry the impact ranking, the closing + gain figure and the attempt ledger. :data:`IMMUTABLE_FIELDS` is enforced here, + not requested in prose. + + Operand dims are the exception, and deliberately so: a graph replay has no + CPU-side parent op, so the profiler records no arguments for exactly the + kernels that dominate a captured model, and the field arrives empty. Refusing + the review's answer there does not preserve a measurement -- it hands the + choice to a tuning backend that cannot see the serving configuration. Review + dims therefore carry their own provenance so a later reader can still tell a + recovered shape from a computed one. +* **Nothing is taken on faith.** A revised path must exist under a known + framework root. This is not a correctness check; it stops an invented path + from being written. +* **Nothing is destroyed.** Every revision records ``previous_source_file`` and + ``previous_method``, and the pre-review table is kept beside the reviewed one, + so a bad review is auditable and reversible. + +The session may run shell commands (demangling a mangled vendor symbol is +exactly the job), so the framework tree is fingerprinted before and after. A +review that modified the code under optimization is discarded rather than +applied: the benchmark that follows would otherwise measure an unrecorded edit. +""" + +from __future__ import annotations + +import asyncio +import json +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable, Sequence + +try: + from hyperloom.common import kernel_source_contract as _KSC +except ImportError: # pragma: no cover - standalone invocation + _KSC = None # type: ignore[assignment] + +try: + from hyperloom.common.kernel_shape_contract import ( + REVIEW_BACKFILL_PROVENANCE, + REVIEW_DERIVED_PROVENANCE, + REVIEW_SHAPE_PROVENANCE, + ) +except ImportError: # pragma: no cover - standalone invocation + REVIEW_BACKFILL_PROVENANCE = "review_backfill" + REVIEW_DERIVED_PROVENANCE = "review_derived" + REVIEW_SHAPE_PROVENANCE = frozenset( + {REVIEW_BACKFILL_PROVENANCE, REVIEW_DERIVED_PROVENANCE} + ) + +#: Written by the agent; its presence is what marks the session successful. +REVISIONS_FILENAME = "kernel_candidates_revisions.json" + +#: The pre-review table, kept so a bad review can be told from a bad parse. +RAW_CANDIDATES_FILENAME = "kernel_candidates.raw.json" + +#: Rejected outright. Two different reasons, both fatal to accept: +#: +#: ``gpu_pct`` / ``duration_us`` / ``call_count`` are event durations read +#: straight off the trace -- there is no parsing ambiguity to correct. They also +#: feed the dispatch floor and the final gain accounting, so a plausible edit +#: here lets the review talk a kernel past the gate it is not supposed to open, +#: and makes the closing report unfalsifiable. +#: +#: ``kernel_id`` / ``name`` / ``device_kernel_name`` are join keys, not facts. +#: ``kernel_id`` is half the attempt-ledger identity; ``name`` is what every +#: shape and metric lookup keys on. Revising either detaches the row from its +#: own history and from the CSVs the review is meant to consult. +IMMUTABLE_FIELDS: frozenset[str] = frozenset( + { + "call_count", + "device_kernel_name", + "duration_us", + "gpu_pct", + "kernel_id", + "name", + } +) + +#: Recomputed from the operand dims, so a revision naming one is dropped with a +#: note rather than silently overwritten downstream. The session supplies +#: ``shapes`` and ``input_dtypes``; everything the harness builder needs is +#: derived from those by :func:`_rederive_after_review`, which keeps the three +#: representations from drifting apart. +DERIVED_SHAPE_FIELDS: frozenset[str] = frozenset( + { + "input_shapes", + "invocation_cases", + "raw_arg_spec", + } +) + +_ACTION_KEEP = "keep" +_ACTION_REWRITE = "rewrite" +_ACTION_UNRESOLVE = "unresolve" +_ACTION_DROP = "drop" +_ACTIONS = frozenset({_ACTION_KEEP, _ACTION_REWRITE, _ACTION_UNRESOLVE, _ACTION_DROP}) + +#: Read and search freely; run shell commands; write only the revision file. +#: ``Edit`` is withheld deliberately -- the agent proposes, it does not patch, +#: and the framework tree here is the code under optimization. +ALLOWED_TOOLS: tuple[str, ...] = ("Read", "Grep", "Glob", "Bash", "Write") + +_DENIED_TOOLS: tuple[str, ...] = ( + "Edit", + "NotebookEdit", + "Task", + "TaskOutput", + "TaskStop", + "WebFetch", + "WebSearch", + "AskUserQuestion", + "EnterPlanMode", + "ExitPlanMode", + "SlashCommand", +) + +_MAX_TURNS = 120 +_DEFAULT_TIMEOUT_SEC = 900.0 +_DEFAULT_ATTEMPTS = 2 + +_SYSTEM_PROMPT = ( + "You audit an automated mapping from GPU kernel symbols to the source that " + "defines them, and decide which kernels are worth handing to a kernel " + "optimizer. Investigate with the tools available: read the candidate table, " + "grep the framework tree, demangle symbols, consult the model config and " + "serving arguments. Verify before you revise -- a file that merely calls a " + "kernel is not the file that defines it. Never modify anything outside the " + "output directory you are given; the framework tree is the code under " + "optimization and is checked for tampering. Report findings only by writing " + "the revisions file you are asked for." +) + + +@dataclass +class ReviewOutcome: + """What one review session produced. + + Attributes: + status: ``completed``, ``skipped`` or a failure label recorded in the + audit and surfaced as a trace-health warning. + revisions: The parsed revision records (empty unless ``completed``). + notes: One human-readable line per applied or rejected revision. + detail: Failure detail, or ``""`` on success. + revisions_path: Where the agent wrote its answer, when it did. + """ + + status: str = "skipped" + revisions: list[dict[str, Any]] = field(default_factory=list) + notes: list[str] = field(default_factory=list) + detail: str = "" + revisions_path: Path | None = None + + @property + def ok(self) -> bool: + """Whether the session produced a usable revision set.""" + return self.status == "completed" + + +def _safe_exception_label(exc: BaseException) -> str: + """Return a stable exception label without leaking message contents.""" + label = type(exc).__name__ + for attribute in ("status_code", "code", "errno"): + value = getattr(exc, attribute, None) + if isinstance(value, bool): + continue + if isinstance(value, int): + return f"{label} ({attribute}={value})" + return label + + +# --------------------------------------------------------------------------- +# Framework-tree tamper check +# --------------------------------------------------------------------------- + + +def source_fingerprint(paths: Sequence[str]) -> dict[str, list[Any]]: + """Record size and mtime for each readable path. + + Scoped to the files the candidates actually name rather than the whole + framework tree: those are the ones a source-resolution session has reason + to open, and hashing gigabytes of installed packages to guard a few dozen + files would cost more than the session it protects. + + Args: + paths (Sequence[str]): Candidate source paths, absolute. + + Returns: + dict[str, list[Any]]: ``{path: [size, mtime_ns]}`` for readable paths. + """ + out: dict[str, list[Any]] = {} + for raw in paths: + path = str(raw or "").strip() + if not path or path in out: + continue + try: + stat = os.stat(path) + except OSError: + continue + out[path] = [stat.st_size, stat.st_mtime_ns] + return out + + +def fingerprint_drift(before: dict[str, list[Any]], after: dict[str, list[Any]]) -> list[str]: + """Return the paths whose size or mtime changed between two fingerprints.""" + drifted: list[str] = [] + for path, signature in before.items(): + if after.get(path) != signature: + drifted.append(path) + return sorted(drifted) + + +# --------------------------------------------------------------------------- +# Prompt +# --------------------------------------------------------------------------- + + +def build_review_prompt( + *, + run_dir: Path, + raw_candidates_path: Path, + revisions_path: Path, + reference_paths: dict[str, str], + framework_roots: Sequence[str], + context_block: str = "", +) -> str: + """Render the review request as a set of paths to investigate. + + Deliberately carries no file contents. Pre-loading the framework tree would + bound the review by whatever was guessed to be relevant, whereas the agent + can follow the evidence -- and only ships what it actually opened. + """ + lines = [ + "Audit the kernel-candidate table produced by the deterministic " + "analysis stage and correct it where the evidence disagrees.", + "", + f"Candidate table to audit: {raw_candidates_path}", + "", + "Reference material (read what you need):", + ] + for label, path in reference_paths.items(): + if path: + lines.append(f" {label}: {path}") + if framework_roots: + lines.append(" framework source roots:") + lines.extend(f" {root}" for root in framework_roots) + if context_block: + lines += ["", context_block] + lines += [ + "", + "For every entry in hot_kernels decide one of:", + " keep the current source_file plausibly defines this kernel", + " rewrite the current path is wrong and you verified a better one", + " unresolve this kernel has no single defining source, or the path is", + " wrong and you could not determine the right one", + " drop this entry is not a kernel worth optimizing at all", + "", + "Rules:", + " - Verify a rewrite by opening the file and confirming it defines the", + " kernel. A file that only calls or dispatches to it does not count.", + " - Prefer unresolve over a guess. A wrong path costs an entire", + " optimization attempt; an empty one just falls through.", + " - You may revise source_file, reusable_native_kernel, skip_reason,", + " benchmark_files, shapes and input_dtypes. You may not revise what", + " the trace measured (gpu_pct, duration_us, call_count) or the", + " identity it is keyed by (kernel_id, name, device_kernel_name);", + " those are ignored if present. input_shapes, invocation_cases and", + " raw_arg_spec are recomputed from shapes, so do not send them.", + " - benchmark_files comes from a curated table keyed by coarse name", + " markers, so it often names a harness for the wrong member of a", + " kernel family. Replace it with harnesses you located and can open,", + " or with an empty list when this kernel has none. Paths that do not", + " exist are dropped.", + " - Entries you do not mention are left exactly as they are.", + " - Do not copy reusable_native_kernel back from the table you were", + " given. Every unresolved row carries false there, and returning that", + " value alongside a corrected path refuses the kernel you just found.", + " Send the field only to veto a kernel the rules would otherwise", + " accept, and always with a skip_reason saying why; a false with no", + " skip_reason is ignored. Routability is recomputed from the path you", + " give, so a rewrite needs nothing else from you.", + "", + "On shapes -- read this before proposing any:", + " A graph replay has no CPU-side parent op, so the profiler records no", + " arguments for a graph-launched kernel and shapes arrives empty. That", + " is not harmless: with no shapes the tuning backend picks its own, and", + " it cannot see the serving configuration. A prefill kernel serving an", + " 8192-token input has been tuned at sequence length 512 this way, which", + " measured a large speedup that vanished end to end.", + " So an empty shapes is worth filling. Two ways, in this order:", + f" {REVIEW_BACKFILL_PROVENANCE} another row of analysis.md already records", + " the dims. A composite operator that kept its module", + " attribution lists the device kernels it launches in", + " its Kernel Name cell, and carries the arguments this", + " row is missing. Confirm the row really launches this", + " kernel before taking its dims -- a neighbouring", + " instantiation of the same kernel family is a", + " different problem size, not this one.", + f" {REVIEW_DERIVED_PROVENANCE} you computed the dims from the model config,", + " the serving arguments and the kernel signature.", + " analysis.md is TraceLens' only supported output. Do not go looking for", + " its intermediate files; they are internal and may not be there.", + " Set shape_provenance to whichever applies; any other value is read as", + f" {REVIEW_DERIVED_PROVENANCE}. Give one shape string per operand, dims first and", + ' dtype second, e.g. "(8192,6144) bf16".', + " State where the dims came from in reason -- which operator's row, or", + " which config fields and what arithmetic -- so a reader can check it in", + " seconds. An unstated derivation is not reviewable, and a shape nothing", + " can check is what the tuning backend already produces on its own.", + "", + f"Write your answer to {revisions_path} as JSON:", + ' {"revisions": [{"kernel_id": "k001", "action": "rewrite",', + ' "source_file": "/abs/path.py",', + ' "reusable_native_kernel": true,', + ' "skip_reason": "",', + ' "shapes": ["(8192,6144) bf16", "(6144,1536) fp4"],', + ' "input_dtypes": ["bf16", "fp4"],', + f' "shape_provenance": "{REVIEW_DERIVED_PROVENANCE}",', + ' "reason": "one sentence citing what you checked"}]}', + "", + "Use action keep with shapes when the recorded path is already right and", + "only the dims are missing.", + "", + f"Write nothing outside {run_dir}. Do not modify framework source.", + ] + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Session drivers +# --------------------------------------------------------------------------- + + +def _resolve_backend() -> str: + """Return ``codex`` or ``claude`` from the configured credentials.""" + from hyperloom.common import llm_config # noqa: PLC0415 + + if llm_config.is_anthropic_only(): + return "claude" + if llm_config.has_openai_side(): + return "codex" + return "claude" + + +async def _run_claude_session( + prompt: str, + *, + run_dir: Path, + model: str, + timeout_sec: float, + log: Callable[[str], None] | None, +) -> str: + """Drive one tool-enabled Claude Agent SDK session; return any SDK error.""" + import claude_agent_sdk as sdk # type: ignore[import-not-found] # noqa: PLC0415 + + from hyperloom.common.llm_config import claude_sdk_env_options # noqa: PLC0415 + + kwargs: dict[str, Any] = dict(claude_sdk_env_options(model=model)) + kwargs.update( + { + "model": model, + "system_prompt": _SYSTEM_PROMPT, + "max_turns": _MAX_TURNS, + "allowed_tools": list(ALLOWED_TOOLS), + "disallowed_tools": list(_DENIED_TOOLS), + "cwd": str(run_dir), + } + ) + try: + options = sdk.ClaudeAgentOptions(**kwargs) + except TypeError: + kwargs.pop("cwd", None) + options = sdk.ClaudeAgentOptions(**kwargs) + + async def _drive() -> None: + async for message in sdk.query(prompt=prompt, options=options): + if log is not None: + for text in _message_text(message): + if text.strip(): + log(f"[review-agent] {text.strip()[:400]}") + + try: + await asyncio.wait_for(_drive(), timeout=max(60.0, timeout_sec)) + except Exception as exc: # noqa: BLE001 - artifact presence decides success + return _safe_exception_label(exc) + return "" + + +def _message_text(message: Any) -> list[str]: + """Best-effort text extraction that never breaks the session loop.""" + try: + from hyperloom.common.claude_oneshot import message_text # noqa: PLC0415 + + return list(message_text(message)) + except Exception: # noqa: BLE001 - logging aid only + return [] + + +async def _run_codex_session( + prompt: str, + *, + run_dir: Path, + model: str, + timeout_sec: float, +) -> str: + """Drive one Codex Agent SDK turn scoped to ``run_dir``; return any error.""" + from hyperloom.common.codex_session import ( # noqa: PLC0415 + CodexSessionError, + run_codex_turn, + ) + + try: + result = await run_codex_turn( + prompt=prompt, + developer_instructions=_SYSTEM_PROMPT, + cwd=run_dir, + model=model, + timeout_sec=max(60.0, timeout_sec), + writable_roots=(run_dir,), + ) + except CodexSessionError as exc: + return _safe_exception_label(exc) + return str(getattr(result, "error", "") or "") + + +def _resolve_model(backend: str) -> str: + """Resolve the session model from the configured environment.""" + explicit = str(os.environ.get("HYPERLOOM_LLM_SOURCE_MODEL") or "").strip() + if explicit: + return explicit + if backend == "codex": + return str(os.environ.get("CODEX_MODEL") or "").strip() or "gpt-5-codex" + return str(os.environ.get("CLAUDE_MODEL") or "").strip() or "claude-opus-5" + + +# --------------------------------------------------------------------------- +# Revision loading and application +# --------------------------------------------------------------------------- + + +def load_revisions(revisions_path: Path) -> tuple[list[dict[str, Any]], str]: + """Read the revision file the agent wrote. + + Returns: + tuple[list[dict[str, Any]], str]: ``(revisions, error)``; ``error`` is + empty when the file parsed into a revision list. + """ + try: + payload = json.loads(Path(revisions_path).read_text(encoding="utf-8")) + except OSError: + return [], "revisions file was not written" + except (TypeError, ValueError): + return [], "revisions file is not valid JSON" + if not isinstance(payload, dict): + return [], "revisions file is not a JSON object" + revisions = payload.get("revisions") + if not isinstance(revisions, list): + return [], "revisions file has no 'revisions' list" + return [r for r in revisions if isinstance(r, dict)], "" + + +def _verified_harnesses(proposed: Any) -> list[str] | None: + """Keep the proposed harness paths that exist, or ``None`` when unset. + + The curated harness table is keyed by coarse name markers, so it offers a + plausible file for a whole family of kernels rather than the one that + exercises this kernel. The session can look, which is worth more than the + marker match -- but only files it can name and that are really there + survive, since a non-empty list reads downstream as a runnable harness. + """ + if not isinstance(proposed, list): + return None + return [ + str(entry) + for entry in proposed + if isinstance(entry, str) and entry.strip() and os.path.isfile(entry.strip()) + ] + + +def _acceptable_path(picked: str, roots: Sequence[str]) -> str: + """Return the canonical form of ``picked``, or ``""`` when unverifiable.""" + if _KSC is None: + return "" + bare = _KSC.strip_line_suffix(picked) + return _KSC.canonical_source_path(bare, tuple(roots)) or "" + + +def _proposed_strings(value: Any) -> list[str] | None: + """Keep the non-empty strings in a proposed list, or ``None`` when unset.""" + if not isinstance(value, list): + return None + return [entry.strip() for entry in value if isinstance(entry, str) and entry.strip()] + + +def _proposed_shape_provenance(value: Any) -> str: + """Narrow a claimed shape provenance to one the review is allowed to assert. + + A session that located a recorded shape and one that computed a shape from + the model config are both useful, but only the first is a measurement. The + claim is therefore restricted to the two review values, and anything else -- + including a session naming ``torch_trace`` -- degrades to the derived label + rather than being taken at its word. Laundering a derivation as a + measurement would strip the one signal that tells a later reader whether a + disappointing end-to-end result is worth blaming on the shape. + """ + claimed = str(value or "").strip().lower() + if claimed in REVIEW_SHAPE_PROVENANCE: + return claimed + return REVIEW_DERIVED_PROVENANCE + + +def _record_shape_proposal( + entry: dict[str, Any], + revision: dict[str, Any], + *, + kernel_id: str, + notes: list[str], +) -> None: + """Stage proposed operand dims for the deterministic re-derivation pass. + + Held under ``review_*`` keys rather than written straight onto the row, so + the stamping pass stays the only thing that decides what the harness builder + finally sees -- the same split the routability hint already uses. + """ + shapes = _proposed_strings(revision.get("shapes")) + if shapes is None: + return + if not shapes: + notes.append(f"{kernel_id}: empty shapes proposal ignored") + return + entry["review_shapes"] = shapes + entry["review_shape_provenance"] = _proposed_shape_provenance( + revision.get("shape_provenance") + ) + dtypes = _proposed_strings(revision.get("input_dtypes")) + if dtypes: + entry["review_input_dtypes"] = dtypes + notes.append( + f"{kernel_id}: shapes -> {len(shapes)} operand(s) " + f"({entry['review_shape_provenance']})" + ) + + +def _record_judgement_proposals( + entry: dict[str, Any], + revision: dict[str, Any], + *, + kernel_id: str, + notes: list[str], +) -> None: + """Stage the routability hint and the verified harness list. + + A veto is only taken with a stated reason. The table handed to the session + already carries ``reusable_native_kernel``, and an unresolved row carries + ``false``; a session correcting that row's path has been observed returning + the field unchanged while its own prose argued the file is editable. Nothing + distinguishes that echo from an intended refusal except the reason the + prompt asks for alongside it, and refusing on the echo threw away every + kernel the review had just located. + + The asymmetry is deliberate and matches the one below it: a permissive hint + is ignored because ``classify_patchability`` still has to agree, so it costs + nothing to drop. A restrictive one has no second gate behind it. + """ + proposed_skip = revision.get("skip_reason") + skip_text = proposed_skip.strip() if isinstance(proposed_skip, str) else "" + if isinstance(proposed_skip, str): + entry["review_skip_reason"] = skip_text + proposed_reusable = revision.get("reusable_native_kernel") + if isinstance(proposed_reusable, bool): + if proposed_reusable or skip_text: + entry["review_reusable_hint"] = proposed_reusable + else: + notes.append( + f"{kernel_id}: veto ignored, no skip_reason given " + "(a refusal has to say why)" + ) + harnesses = _verified_harnesses(revision.get("benchmark_files")) + if harnesses is not None: + entry["review_benchmark_files"] = harnesses + notes.append(f"{kernel_id}: benchmark_files -> {len(harnesses)} verified path(s)") + + +def apply_revisions( + candidates: list[dict[str, Any]], + revisions: Sequence[dict[str, Any]], + *, + framework_roots: Sequence[str], + protected_ids: frozenset[str] | set[str] = frozenset(), +) -> list[str]: + """Apply the agent's proposals to ``candidates`` in place. + + Only the judgement fields move. Derived state (``source_type``, + ``kernel_repo``, backends, category, routability) is deliberately left for + the caller to recompute through the deterministic stamping pass, so + :func:`classify_patchability` stays the single gate rather than gaining a + second, model-written one. + + Args: + candidates: The finalized candidate rows, mutated in place. + revisions: Revision records parsed from the agent's answer. + framework_roots: Roots a revised path must resolve under. + protected_ids: Candidates resolved by an authoritative tier. The active + finder demangles the device symbol and pins the source in the + installed tree; reading the same tree cannot beat knowing which + symbol the binary actually exports, so those are left alone. + + Returns: + list[str]: One note per applied or rejected revision. + """ + by_id = { + str(c.get("kernel_id") or ""): c for c in candidates if isinstance(c, dict) + } + notes: list[str] = [] + for revision in revisions: + kernel_id = str(revision.get("kernel_id") or "").strip() + entry = by_id.get(kernel_id) + if entry is None: + notes.append(f"{kernel_id or '(no id)'}: unknown kernel_id, ignored") + continue + action = str(revision.get("action") or "").strip().lower() + if kernel_id in protected_ids and action != _ACTION_KEEP: + notes.append(f"{kernel_id}: {action} refused, resolved by an authoritative tier") + continue + if action not in _ACTIONS: + notes.append(f"{kernel_id}: unknown action {action!r}, ignored") + continue + touched = sorted(IMMUTABLE_FIELDS.intersection(revision) - {"kernel_id"}) + if touched: + notes.append(f"{kernel_id}: ignored measured field(s) {', '.join(touched)}") + derived = sorted(DERIVED_SHAPE_FIELDS.intersection(revision)) + if derived: + notes.append( + f"{kernel_id}: ignored derived field(s) {', '.join(derived)}; " + "propose shapes instead" + ) + reason = str(revision.get("reason") or "").strip() + + # Operand dims are worth having whether or not the path moved, and the + # rows that most need them are the ones the deterministic tiers already + # located: under graph capture a replay records no arguments, so a + # correctly resolved kernel can still arrive with no shape at all. + if action == _ACTION_KEEP: + _record_shape_proposal(entry, revision, kernel_id=kernel_id, notes=notes) + _record_judgement_proposals(entry, revision, kernel_id=kernel_id, notes=notes) + continue + + previous_file = str(entry.get("source_file") or "") + previous_method = str(entry.get("source_resolution_method") or "") + + if action in (_ACTION_UNRESOLVE, _ACTION_DROP): + entry["previous_source_file"] = previous_file + entry["previous_method"] = previous_method + entry["source_file"] = "" + entry.pop("source_line", None) + entry.pop("source_function", None) + entry["source_resolution_method"] = "llm_review" + entry["review_action"] = action + entry["review_reason"] = reason or "no defining source" + notes.append(f"{kernel_id}: {action} (was {previous_file or '(none)'})") + continue + + picked = str(revision.get("source_file") or "").strip() + if not picked: + notes.append(f"{kernel_id}: rewrite without a path, ignored") + continue + canonical = _acceptable_path(picked, framework_roots) + if not canonical: + notes.append(f"{kernel_id}: rejected unverifiable path {picked!r}") + continue + previous_bare = _KSC.strip_line_suffix(previous_file) if _KSC else previous_file + # A rewrite that lands on the path already recorded is not a correction, + # but the rest of the same revision still is: falling through keeps a + # confirmed location from costing the shapes proposed alongside it. + if canonical != previous_bare: + entry["previous_source_file"] = previous_file + entry["previous_method"] = previous_method + entry["source_file"] = canonical + entry.pop("source_line", None) + entry.pop("source_function", None) + entry["source_resolution_method"] = "llm_review" + entry["review_action"] = action + entry["review_reason"] = reason or "no reason given" + notes.append(f"{kernel_id}: {previous_file or '(none)'} -> {canonical}") + + _record_shape_proposal(entry, revision, kernel_id=kernel_id, notes=notes) + _record_judgement_proposals(entry, revision, kernel_id=kernel_id, notes=notes) + return notes + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + + +def run_candidate_review( + *, + run_dir: Path, + raw_candidates_path: Path, + reference_paths: dict[str, str], + framework_roots: Sequence[str], + context_block: str = "", + timeout_sec: float = _DEFAULT_TIMEOUT_SEC, + attempts: int = _DEFAULT_ATTEMPTS, + log: Callable[[str], None] | None = None, + session_runner: Callable[..., str] | None = None, +) -> ReviewOutcome: + """Run the review session and return its parsed revisions. + + Retries a failed session once by default: the pass is mandatory on the agent + route, and a gateway hiccup should not be the reason a run's candidate table + goes unaudited. A definitive failure is reported rather than raised -- the + deterministic table is still usable, and killing a multi-hour optimization + over an advisory pass would trade a small loss for a total one. + + Args: + run_dir: Session directory; the only place the agent may write. + raw_candidates_path: The pre-review candidate table to audit. + reference_paths: Labelled artifact paths offered to the agent. + framework_roots: Roots a revised path must resolve under. + context_block: Rendered model/serving context, or ``""``. + timeout_sec: Wall-clock bound per attempt. + attempts: Total attempts, including the first. + log: Optional diagnostics callback. + session_runner: Injection point for the session call (tests). + + Returns: + ReviewOutcome: The session result; never raises. + """ + + def _say(message: str) -> None: + if log is not None: + log(f"candidate_review: {message}") + + revisions_path = Path(run_dir) / REVISIONS_FILENAME + prompt = build_review_prompt( + run_dir=Path(run_dir), + raw_candidates_path=Path(raw_candidates_path), + revisions_path=revisions_path, + reference_paths=reference_paths, + framework_roots=framework_roots, + context_block=context_block, + ) + + try: + backend = _resolve_backend() + model = _resolve_model(backend) + except Exception as exc: # noqa: BLE001 - configuration is reported, not raised + detail = _safe_exception_label(exc) + _say(f"configuration failed: {detail}") + return ReviewOutcome(status="configuration_error", detail=detail) + + last_detail = "" + for attempt in range(1, max(1, int(attempts)) + 1): + revisions_path.unlink(missing_ok=True) + _say(f"attempt {attempt}/{attempts} via {backend} ({model})") + try: + if session_runner is not None: + error = session_runner( + prompt=prompt, + run_dir=Path(run_dir), + model=model, + timeout_sec=timeout_sec, + ) + elif backend == "codex": + error = asyncio.run( + _run_codex_session( + prompt, + run_dir=Path(run_dir), + model=model, + timeout_sec=timeout_sec, + ) + ) + else: + error = asyncio.run( + _run_claude_session( + prompt, + run_dir=Path(run_dir), + model=model, + timeout_sec=timeout_sec, + log=log, + ) + ) + except Exception as exc: # noqa: BLE001 - advisory pass, never fatal + error = _safe_exception_label(exc) + + revisions, parse_error = load_revisions(revisions_path) + if not parse_error: + # The SDK can report an error after the answer landed; the artifact + # is what decides, exactly as the TraceLens skill runner does. + return ReviewOutcome( + status="completed", + revisions=revisions, + revisions_path=revisions_path, + ) + last_detail = error or parse_error + _say(f"attempt {attempt} unusable: {last_detail}") + + return ReviewOutcome(status="failed", detail=last_detail or "no revisions produced") + + +__all__ = [ + "ALLOWED_TOOLS", + "DERIVED_SHAPE_FIELDS", + "IMMUTABLE_FIELDS", + "RAW_CANDIDATES_FILENAME", + "REVIEW_BACKFILL_PROVENANCE", + "REVIEW_DERIVED_PROVENANCE", + "REVIEW_SHAPE_PROVENANCE", + "REVISIONS_FILENAME", + "ReviewOutcome", + "apply_revisions", + "build_review_prompt", + "fingerprint_drift", + "load_revisions", + "run_candidate_review", + "source_fingerprint", +] diff --git a/src/hyperloom/agents/kernel/tools/_llm_source_fallback.py b/src/hyperloom/agents/kernel/tools/_llm_source_fallback.py deleted file mode 100644 index c0d3495986..0000000000 --- a/src/hyperloom/agents/kernel/tools/_llm_source_fallback.py +++ /dev/null @@ -1,543 +0,0 @@ -############################################################################### -# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. -# SPDX-License-Identifier: MIT -# -# See LICENSE for license information. -############################################################################### - -"""Last-resort semantic pick of a kernel's source file, from a fixed shortlist. - -This is the final tier of source resolution, behind the curated dictionary, the -trace-derived launcher, and the name grep. It exists for models or frameworks -where all three come up empty -- not for any case seen so far. - -Two properties keep it from becoming another source of nondeterminism, which -matters because an unresolved-launcher sentinel produced by an LLM is what broke -this pipeline in the first place: - -* **Selection, never generation.** The model receives a shortlist gathered by a - relaxed grep and may only return one of those exact strings. A path it invents - is rejected outright. -* **Last in line.** It sees a candidate only after the curated dictionary, the - trace-derived launcher and the grep have all come up empty, and only when the - kernel is worth at least 5% of GPU time. - -Every accepted answer is stamped ``source_resolution_method="llm_fallback"`` so -it can be audited apart from deterministic resolutions. -""" - -from __future__ import annotations - -import asyncio -import json -import math -import os -import re -from typing import Any, Callable -from urllib.parse import urlsplit - -try: - from hyperloom.common import kernel_source_contract as _KSC -except ImportError: # pragma: no cover - standalone invocation - _KSC = None # type: ignore[assignment] - -# An answer below this confidence is discarded: a coin-flip pick would send a -# backend at the wrong file and burn a whole optimization attempt. -_MIN_CONFIDENCE = 0.7 - -# Head of each shortlisted file shown to the model; enough to tell a kernel -# definition from a test or a dispatch shim. -_PREVIEW_LINES = 40 -_PREVIEW_CHARS = 2000 - -# Shipping those heads sends repository source to the model provider. That is an -# operator's data-egress decision, not a default this tier may assume, so the -# preview is off unless explicitly authorised. Without it the tiers still see -# the candidate paths, which carry most of the signal. -_PREVIEW_ENV = "HYPERLOOM_LLM_SOURCE_PREVIEW" -_PROVIDER_ENV = "HYPERLOOM_LLM_SOURCE_PROVIDER" -_MODEL_ENV = "HYPERLOOM_LLM_SOURCE_MODEL" - -_PROVIDER_CLAUDE = "claude_agent_sdk" -_PROVIDER_OPENAI = "openai_compatible" -_PROVIDER_ALIASES = { - "anthropic": _PROVIDER_CLAUDE, - "claude": _PROVIDER_CLAUDE, - "claude-agent-sdk": _PROVIDER_CLAUDE, - "claude_agent_sdk": _PROVIDER_CLAUDE, - "openai": _PROVIDER_OPENAI, - "openai-compatible": _PROVIDER_OPENAI, - "openai_compatible": _PROVIDER_OPENAI, -} - -# Claude Code must behave as a plain completion client in this tier. Explicitly -# denying every built-in tool prevents a source-resolution request from reading -# anything beyond the prompt assembled under the egress policy above. -_CLAUDE_DISALLOWED_TOOLS = ( - "Bash", - "BashOutput", - "KillShell", - "Read", - "Write", - "Edit", - "NotebookEdit", - "Glob", - "Grep", - "Agent", - "Task", - "TaskOutput", - "TaskStop", - "WebFetch", - "WebSearch", - "TodoWrite", - "AskUserQuestion", - "EnterPlanMode", - "ExitPlanMode", - "Skill", - "SlashCommand", -) - -_DEFAULT_TIMEOUT_SEC = 60.0 - -_JSON_BLOCK_RE = re.compile(r"\{.*\}", re.DOTALL) - -_SYSTEM_PROMPT = ( - "You identify which source file implements a GPU kernel. " - "You are given the kernel symbol and a shortlist of candidate files. " - "Choose the file that DEFINES the kernel body. Reject files that merely " - "call it, test it, or dispatch to it, and reject a CPU implementation when " - "the kernel runs on GPU. " - 'Answer with JSON only: {"source_file": "", ' - '"confidence": <0..1>, "reason": ""}. ' - 'If none of the candidates defines the kernel, return "source_file": "".' -) - - -def _default_claude_model() -> str: - """The project-wide default Claude model, or "" when it cannot be read. - - Deliberately without a hardcoded fallback id: a pinned default here would - go stale the moment the project moves (it already outlived claude-opus-4-8), - and a tier quietly running an older model than the rest of the pipeline is - worse than one that declines to run. An empty result is reported as a - resolution failure by the caller rather than guessing. - """ - try: - from hyperloom.orchestrator.roles.agent_role import ( # noqa: PLC0415 - DEFAULT_CLAUDE_MODEL, - ) - except Exception: # noqa: BLE001 - tools also run outside the package - return "" - return str(DEFAULT_CLAUDE_MODEL or "") - - -def _resolve_provider(provider: str = "") -> str: - """Resolve the explicit provider or infer it from canonical credential shape. - - The role-specific override wins. Otherwise Anthropic-only deployments use - the native Claude path, while any configured OpenAI side uses the OpenAI - path. The latter preserves the project's OpenAI default for dual-configured - single-shot roles. - """ - raw = str(provider or os.environ.get(_PROVIDER_ENV) or "").strip().lower() - supported = "claude_agent_sdk, openai_compatible" - if raw: - resolved = _PROVIDER_ALIASES.get(raw) - if resolved: - return resolved - raise RuntimeError(f"unsupported {_PROVIDER_ENV}={raw!r}; choose one of: {supported}") - - from hyperloom.common import llm_config # noqa: PLC0415 - keep standalone import-light - - if llm_config.is_anthropic_only(): - return _PROVIDER_CLAUDE - if llm_config.has_openai_side(): - return _PROVIDER_OPENAI - raise RuntimeError( - f"{_PROVIDER_ENV} is not set and no provider credentials are configured; choose one of: {supported}" - ) - - -def _resolve_model(model: str = "", provider: str = "") -> str: - """Resolve a model without borrowing another provider's model setting.""" - explicit = str(model or os.environ.get(_MODEL_ENV) or "").strip() - if explicit: - return explicit - if provider == _PROVIDER_OPENAI: - return str(os.environ.get("OPENAI_MODEL") or os.environ.get("CODEX_MODEL") or "").strip() - return str(os.environ.get("CLAUDE_MODEL") or _default_claude_model()).strip() - - -def _endpoint_host(provider: str) -> str: - """Return a credential-free endpoint identifier for artifact auditing.""" - if provider == _PROVIDER_OPENAI: - raw = str(os.environ.get("OPENAI_BASE_URL") or "").strip() - default = "api.openai.com" - else: - raw = str(os.environ.get("ANTHROPIC_BASE_URL") or os.environ.get("DEEPSEEK_BASE_URL") or "").strip() - default = "provider-default" - if not raw: - return default - return urlsplit(raw).hostname or "configured" - - -def llm_source_provider_configured() -> bool: - """Whether source resolution has an explicit or inferred provider. - - Lets a caller skip the work it would only do in order to build a request -- - the shortlist grep walks every framework root -- when neither the role - override nor a canonical credential side can select a provider. - """ - try: - _resolve_provider() - except RuntimeError: - return False - return True - - -def llm_source_audit(*, provider: str = "", model: str = "") -> dict[str, Any]: - """Return non-secret provider metadata suitable for an audit artifact.""" - try: - selected_provider = _resolve_provider(provider) - except RuntimeError: - return { - "provider": "unconfigured", - "model": str(model or os.environ.get(_MODEL_ENV) or "").strip(), - "endpoint_host": "", - "source_preview_authorised": source_preview_authorised(), - } - return { - "provider": selected_provider, - "model": _resolve_model(model, selected_provider), - "endpoint_host": _endpoint_host(selected_provider), - "source_preview_authorised": source_preview_authorised(), - } - - -def source_preview_authorised() -> bool: - """Whether the operator opted into sending file heads to the model.""" - return str(os.environ.get(_PREVIEW_ENV, "")).strip().lower() in ("1", "true", "yes", "on") - - -def _preview(path: str) -> str: - """First lines of ``path``, for telling an implementation from a shim.""" - try: - with open(path, encoding="utf-8", errors="replace") as fh: - head = "".join(next(fh, "") for _ in range(_PREVIEW_LINES)) - except OSError: - return "" - return head[:_PREVIEW_CHARS] - - -def _canonical_candidates( - candidates: list[str], - framework_roots: tuple[str, ...], - *, - require_roots: bool = True, -) -> dict[str, str]: - """Map candidates to canonical targets after filesystem validation.""" - canonical: dict[str, str] = {} - for path in candidates: - if require_roots: - target = _KSC.canonical_source_path(path, framework_roots) if _KSC else "" - else: - bare = _KSC.strip_line_suffix(path) if _KSC else path - target = os.path.realpath(bare) if bare and os.path.isfile(bare) else "" - if target: - canonical[path] = target - return canonical - - -def _build_prompt_with_targets( - kernel_name: str, - candidates: list[str], - context_block: str = "", - with_preview: bool | None = None, - *, - framework_roots: tuple[str, ...] = (), -) -> tuple[str, dict[str, str]]: - """Render the prompt and return its prevalidated canonical targets.""" - if with_preview is None: - with_preview = source_preview_authorised() - canonical_paths = _canonical_candidates( - candidates, - framework_roots, - require_roots=bool(framework_roots), - ) - preview_paths = canonical_paths if framework_roots else {} - parts = [] - if context_block: - parts.append(context_block + "\n") - parts += [f"Kernel symbol: {kernel_name}", "", "Candidates:"] - for index, path in enumerate(candidates, 1): - canonical = preview_paths.get(path, "") - if with_preview and canonical: - parts.append(f"\n[{index}] {path}\n```\n{_preview(canonical)}\n```") - else: - parts.append(f"\n[{index}] {path}") - return "\n".join(parts), canonical_paths - - -def _build_prompt( - kernel_name: str, - candidates: list[str], - context_block: str = "", - with_preview: bool | None = None, - *, - framework_roots: tuple[str, ...] = (), -) -> str: - """Render the shortlist, with validated previews when authorised.""" - prompt, _ = _build_prompt_with_targets( - kernel_name, - candidates, - context_block, - with_preview, - framework_roots=framework_roots, - ) - return prompt - - -def _parse_answer(text: str) -> tuple[bool, str, float, str]: - """Extract ``(parsed, source_file, confidence, reason)`` from a model reply. - - ``parsed`` separates "the reply was unreadable" from "the model answered that - no candidate fits". Collapsing the two would report a malformed reply as a - considered verdict and send triage the wrong way. - """ - if not isinstance(text, str): - return False, "", 0.0, "reply is not text" - match = _JSON_BLOCK_RE.search(text or "") - if not match: - return False, "", 0.0, "no JSON object in reply" - try: - payload = json.loads(match.group(0)) - except (TypeError, ValueError): - return False, "", 0.0, "unparseable JSON" - if not isinstance(payload, dict): - return False, "", 0.0, "JSON payload is not an object" - try: - confidence = float(payload.get("confidence") or 0.0) - except (TypeError, ValueError): - return False, "", 0.0, "confidence is not numeric" - if not math.isfinite(confidence) or not 0.0 <= confidence <= 1.0: - return False, "", 0.0, "confidence must be finite and in [0, 1]" - return ( - True, - str(payload.get("source_file") or "").strip(), - confidence, - str(payload.get("reason") or "").strip(), - ) - - -def _validate( - picked: str, - candidates: list[str], - canonical_paths: dict[str, str], - framework_roots: tuple[str, ...], -) -> tuple[str, str]: - """Return the prevalidated canonical pick and any rejection reason.""" - if not picked: - return "", "model reported no candidate defines the kernel" - if picked not in candidates: - # The whole point of a shortlist is that the answer comes from it. - return "", f"path is not one of the candidates: {picked!r}" - canonical = canonical_paths.get(picked, "") - if canonical: - return canonical, "" - if framework_roots: - return "", f"path does not exist or is outside every framework root: {picked!r}" - return "", f"path does not exist: {picked!r}" - - -def _safe_exception_label(exc: BaseException) -> str: - """Return a stable exception type and optional non-secret error code.""" - label = type(exc).__name__ - for attribute in ("status_code", "code", "errno"): - try: - value = getattr(exc, attribute, None) - except Exception: # noqa: BLE001 - hostile exception properties stay private - continue - if isinstance(value, bool): - continue - if isinstance(value, int): - return f"{label} ({attribute}={value})" - if isinstance(value, str) and re.fullmatch(r"[A-Za-z0-9_.-]{1,64}", value): - return f"{label} ({attribute}={value})" - return label - - -def _complete_openai(prompt: str, model: str, timeout_sec: float) -> str: - """Run one completion against the configured OpenAI-compatible endpoint.""" - from hyperloom.common import llm_config # noqa: PLC0415 - optional dependency - - return llm_config.chat_completion( - llm_config.get_openai_client(), - model=model, - messages=[ - {"role": "system", "content": _SYSTEM_PROMPT}, - {"role": "user", "content": prompt}, - ], - temperature=0.0, - timeout=timeout_sec, - ).text - - -def _complete_claude_sdk(prompt: str, model: str, timeout_sec: float) -> str: - """Run one tool-free completion through the native Claude Agent SDK.""" - try: - import claude_agent_sdk as sdk # type: ignore[import-not-found] # noqa: PLC0415 - except ImportError as exc: - raise RuntimeError("claude_agent_sdk is not installed") from exc - if not (hasattr(sdk, "query") and hasattr(sdk, "ClaudeAgentOptions")): - raise RuntimeError("claude_agent_sdk missing query / ClaudeAgentOptions") - - from hyperloom.common.claude_oneshot import message_text # noqa: PLC0415 - from hyperloom.common.llm_config import claude_sdk_env_options # noqa: PLC0415 - - kwargs: dict[str, Any] = dict(claude_sdk_env_options(model=model)) - kwargs.update( - { - "model": model, - "system_prompt": _SYSTEM_PROMPT, - "tools": [], - "setting_sources": [], - "skills": [], - "strict_mcp_config": True, - "mcp_servers": {}, - "plugins": [], - "max_turns": 1, - "allowed_tools": [], - "disallowed_tools": list(_CLAUDE_DISALLOWED_TOOLS), - } - ) - options = sdk.ClaudeAgentOptions(**kwargs) - - async def _drive() -> str: - """Collect the final result, falling back to streamed text blocks.""" - final = "" - chunks: list[str] = [] - async for message in sdk.query(prompt=prompt, options=options): - result = getattr(message, "result", None) - if isinstance(result, str) and result.strip(): - final = result - continue - chunks.extend(message_text(message)) - return final.strip() or "".join(chunks).strip() - - try: - asyncio.get_running_loop() - except RuntimeError: - return asyncio.run(asyncio.wait_for(_drive(), timeout=max(0.1, float(timeout_sec)))) - raise RuntimeError("claude_agent_sdk completion cannot run inside an active event loop") - - -def _complete(prompt: str, model: str, timeout_sec: float) -> str: - """Route one completion through the selected native provider.""" - provider = _resolve_provider() - if provider == _PROVIDER_CLAUDE: - return _complete_claude_sdk(prompt, model, timeout_sec) - return _complete_openai(prompt, model, timeout_sec) - - -def select_source_via_llm( - kernel_name: str, - candidates: list[str], - *, - framework_roots: tuple[str, ...] = (), - model: str = "", - timeout_sec: float = _DEFAULT_TIMEOUT_SEC, - context_block: str = "", - log: Callable[[str], None] | None = None, - complete: Callable[[str, str, float], str] | None = None, - errors: list[str] | None = None, -) -> tuple[str, float, str]: - """Pick the file that defines ``kernel_name`` from ``candidates``. - - Args: - kernel_name: The kernel symbol being resolved. - candidates: Shortlist of on-disk paths from the relaxed grep. An empty - list short-circuits: with nothing to choose from there is nothing to - ask, and inventing a path is not allowed. - framework_roots: Accepted path roots; a pick outside them is rejected. - model: Chat model; defaults to ``$HYPERLOOM_LLM_SOURCE_MODEL``, then - the selected provider's model setting. - timeout_sec: Per-call ceiling. There is no retry -- a failure here is - advisory and the candidate simply stays unresolved. - log: Optional ``callable(str)`` for diagnostics. - complete: Injection point for the completion call (tests). - errors: Optional output list for configuration, transport and parsing - failures. A valid model refusal does not append to it. - - Returns: - ``(source_file, confidence, reason)``; ``source_file`` is ``""`` on any - failure, including a low-confidence answer. - """ - - def _say(message: str) -> None: - if callable(log): - log(f"llm_source_fallback: {message}") - - if not kernel_name or not candidates: - return "", 0.0, "no candidates to choose from" - - shortlist = [str(c) for c in candidates if str(c).strip()] - caller = complete or _complete - try: - provider = _resolve_provider() if complete is None else "" - chosen_model = _resolve_model(model, provider) - except RuntimeError as exc: - detail = _safe_exception_label(exc) - _say(f"configuration failed: {detail}") - reason = f"llm configuration failed: {detail}" - if errors is not None: - errors.append(reason) - return "", 0.0, reason - if not chosen_model: - _say(f"no model configured; set ${_MODEL_ENV}") - reason = "no model configured" - if errors is not None: - errors.append(reason) - return "", 0.0, reason - try: - prompt, canonical_paths = _build_prompt_with_targets( - kernel_name, - shortlist, - context_block, - framework_roots=framework_roots, - ) - reply = caller( - prompt, - chosen_model, - timeout_sec, - ) - except Exception as exc: # noqa: BLE001 - advisory tier, never fatal - detail = _safe_exception_label(exc) - _say(f"call failed: {detail}") - reason = f"llm call failed: {detail}" - if errors is not None: - errors.append(reason) - return "", 0.0, reason - - parsed, picked, confidence, reason = _parse_answer(reply) - if not parsed: - _say(f"rejected: {reason}") - if errors is not None: - errors.append(reason) - return "", 0.0, reason - canonical, why = _validate(picked, shortlist, canonical_paths, framework_roots) - if not canonical: - _say(f"rejected: {why}") - return "", confidence, why - if confidence < _MIN_CONFIDENCE: - _say(f"rejected: confidence {confidence:.2f} < {_MIN_CONFIDENCE}") - return "", confidence, f"confidence {confidence:.2f} below {_MIN_CONFIDENCE}" - - _say(f"accepted {canonical} (confidence={confidence:.2f})") - return canonical, confidence, reason - - -__all__ = [ - "llm_source_audit", - "llm_source_provider_configured", - "select_source_via_llm", -] diff --git a/src/hyperloom/agents/kernel/tools/_llm_source_review.py b/src/hyperloom/agents/kernel/tools/_llm_source_review.py deleted file mode 100644 index cfed6e3a9e..0000000000 --- a/src/hyperloom/agents/kernel/tools/_llm_source_review.py +++ /dev/null @@ -1,398 +0,0 @@ -############################################################################### -# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. -# SPDX-License-Identifier: MIT -# -# See LICENSE for license information. -############################################################################### - -"""Model review of an already-produced source-resolution artifact. - -The deterministic tiers fail in a way that a fallback cannot catch: they do not -come up empty, they come up *confidently wrong*. Measured over historical -sessions, only 59% of verifiable resolutions actually mention the kernel they -claim to define, and ``aten::fill_`` alone has been resolved to four unrelated -business files -- every one of them a real, existing, root-resident source file -that passes every mechanical check. A tier gated on "source_file is empty" -never sees any of those. - -So this tier reviews *every* entry rather than only the blanks, and may rewrite -a location the deterministic tiers already filled in. - -Two properties keep the added freedom from becoming a new failure mode: - -* **Nothing is taken on faith.** A rewritten path must exist on disk or sit - under a known framework root (:func:`path_is_acceptable`). This does not - check correctness -- it stops an invented path from being written. -* **Nothing is destroyed.** Every revision keeps ``previous_source_file`` and - ``previous_method``, so a bad review is auditable and reversible. -""" - -from __future__ import annotations - -import copy -import json -import re -from typing import Any, Callable - -from _llm_source_context import launcher_stack -from _llm_source_fallback import ( - _complete, - _preview, - _resolve_model, - _resolve_provider, - _safe_exception_label, - llm_source_audit, - source_preview_authorised, -) - -try: - from hyperloom.common import kernel_source_contract as _KSC -except ImportError: # pragma: no cover - standalone invocation - _KSC = None # type: ignore[assignment] - -#: Reviewing a long tail of sub-percent kernels costs tokens and changes -#: nothing anyone will act on. -_DEFAULT_MIN_GPU_PCT = 1.0 - -#: Cap on entries per request, so one call stays within a sane context. -_DEFAULT_MAX_ENTRIES = 40 - -_DEFAULT_TIMEOUT_SEC = 180.0 - -_JSON_BLOCK_RE = re.compile(r"\{[\s\S]*\}") - -_ACTION_KEEP = "keep" -_ACTION_REWRITE = "rewrite" -_ACTION_UNRESOLVE = "unresolve" -_ACTIONS = frozenset({_ACTION_KEEP, _ACTION_REWRITE, _ACTION_UNRESOLVE}) - -_PROMPT_HEADER = """\ -You are auditing an automated mapping from GPU kernel symbols to the source -file that defines each kernel. The mapping was produced by heuristics that are -known to fail in a specific way: when a kernel has no source of its own (a -PyTorch built-in such as aten::fill_, or a bare launch API), the heuristic -often attributes it to whichever business file happened to call it. Such an -entry looks perfectly plausible -- the path exists and holds real code -- but -the file does not define that kernel. - -For each entry decide one of: - keep the file plausibly defines this kernel - rewrite the file is wrong and you know a better path - unresolve this kernel has no single defining source file, or the current - path is wrong and you do not know the right one - -Prefer "unresolve" over a guess. A wrong path costs an entire optimization -attempt; an empty one just falls through. - -Reply with JSON only: -{"revisions": [{"kernel_id": "...", "action": "keep|rewrite|unresolve", - "source_file": "...", "reason": "..."}]} -"source_file" is required for "rewrite" and ignored otherwise. Include every -entry you were given. -""" - - -def _gpu_pct(entry: dict[str, Any]) -> float: - """GPU share of one entry, treating an unreadable value as zero. - - Ranking and the floor both read this. Letting a malformed value raise would - escape into the caller's blanket handler and switch the whole tier off for - the run, so one bad row would cost every other row its review. - """ - try: - return float(entry.get("gpu_pct") or 0.0) - except (TypeError, ValueError): - return 0.0 - - -def _entry_block( - entry: dict[str, Any], - with_preview: bool, - framework_roots: tuple[str, ...], -) -> str: - """Render one entry, with a head of its current file when readable.""" - src = str(entry.get("source_file") or "") - lines = [ - f"kernel_id: {entry.get('kernel_id')}", - f"symbol: {entry.get('name')}", - f"gpu_pct: {entry.get('gpu_pct')}", - f"current_source_file: {src or '(unresolved)'}", - f"resolved_by: {entry.get('method')}", - ] - stack = launcher_stack(entry) - if stack: - # The call site is often the only thing separating a kernel from the - # business file that merely calls it. - lines.append("launcher_stack:\n" + "\n".join(f" {f}" for f in stack)) - bare = _KSC.strip_line_suffix(src) if _KSC else src - canonical = ( - _KSC.canonical_source_path(bare, framework_roots) - if with_preview and bare and _KSC - else "" - ) - if canonical: - lines.append(f"file_head:\n```\n{_preview(canonical)}\n```") - return "\n".join(lines) - - -def build_review_prompt( - entries: list[dict[str, Any]], - *, - with_preview: bool | None = None, - context_block: str = "", - framework_roots: tuple[str, ...] = (), -) -> str: - """Render the review request for ``entries``. - - ``with_preview`` defaults to the operator's egress decision (see - :func:`source_preview_authorised`) rather than to True, so repository source - is never shipped to the provider by omission. - """ - if with_preview is None: - with_preview = source_preview_authorised() - head = _PROMPT_HEADER - if context_block: - head = f"{head}\n{context_block}\n" - blocks = [_entry_block(e, with_preview, framework_roots) for e in entries] - return head + "\nEntries:\n\n" + "\n\n---\n\n".join(blocks) - - -def parse_revisions(text: str) -> tuple[bool, list[dict[str, Any]], str]: - """Extract ``(parsed, revisions, error)`` from a model reply. - - ``parsed`` distinguishes an unreadable reply from a reply that revised - nothing; conflating them would report a broken call as a clean audit. - """ - if not isinstance(text, str): - return False, [], "reply is not text" - match = _JSON_BLOCK_RE.search(text or "") - if not match: - return False, [], "no JSON object in reply" - try: - payload = json.loads(match.group(0)) - except (TypeError, ValueError): - return False, [], "unparseable JSON" - if not isinstance(payload, dict): - return False, [], "JSON payload is not an object" - revisions = payload.get("revisions") - if not isinstance(revisions, list): - return False, [], "payload has no 'revisions' list" - out = [r for r in revisions if isinstance(r, dict)] - return True, out, "" - - -def _apply_revision( - entry: dict[str, Any], - revision: dict[str, Any], - roots: tuple[str, ...], -) -> str: - """Apply one revision in place; return a note describing what happened.""" - action = str(revision.get("action") or "").strip().lower() - reason = str(revision.get("reason") or "").strip() - if action not in _ACTIONS: - return f"{entry.get('kernel_id')}: ignored unknown action {action!r}" - if action == _ACTION_KEEP: - return "" - previous_file = str(entry.get("source_file") or "") - previous_method = str(entry.get("method") or "") - - if action == _ACTION_UNRESOLVE: - if not previous_file: - return "" - entry["previous_source_file"] = previous_file - entry["previous_method"] = previous_method - entry["source_file"] = "" - entry["source_line"] = None - entry["source_function"] = "" - entry["method"] = _KSC.METHOD_UNRESOLVED if _KSC else "unresolved" - entry["reason"] = f"llm_review unresolved: {reason or 'no defining source'}" - return f"{entry.get('kernel_id')}: unresolved (was {previous_file})" - - picked_raw = str(revision.get("source_file") or "").strip() - if not picked_raw: - return f"{entry.get('kernel_id')}: rewrite without a path, ignored" - if _KSC is None: - # The mechanical floor lives in the contract module. Without it a - # rewrite cannot be verified at all, and writing an unverifiable path is - # the failure this tier exists to prevent -- so an unusable guard denies - # the rewrite instead of waving it through. - return f"{entry.get('kernel_id')}: rejected rewrite, path contract unavailable" - picked, source_line, source_function = _KSC.split_line_suffix(picked_raw) - previous_bare = _KSC.strip_line_suffix(previous_file) - if picked == previous_bare: - return "" - # The mechanical floor: a rewrite may not conjure a location. This is not a - # correctness check, only a guard against invention. - canonical = _KSC.canonical_source_path(picked, roots) - if not canonical: - return f"{entry.get('kernel_id')}: rejected unverifiable path {picked_raw!r}" - if canonical == previous_bare: - return "" - entry["previous_source_file"] = previous_file - entry["previous_method"] = previous_method - entry["source_file"] = canonical - entry["source_line"] = source_line - entry["source_function"] = source_function - entry["method"] = _KSC.METHOD_LLM - entry["reason"] = f"llm_review rewrote: {reason or 'no reason given'}" - return f"{entry.get('kernel_id')}: {previous_file or '(none)'} -> {canonical}" - - -def review_resolution_document( - doc: dict[str, Any], - *, - framework_roots: tuple[str, ...] = (), - min_gpu_pct: float = _DEFAULT_MIN_GPU_PCT, - max_entries: int = _DEFAULT_MAX_ENTRIES, - model: str | None = None, - timeout_sec: float = _DEFAULT_TIMEOUT_SEC, - context_block: str = "", - log: Callable[[str], None] | None = None, - complete: Callable[[str, str, float], str] | None = None, -) -> tuple[dict[str, Any], list[str]]: - """Review and possibly revise the resolution table in ``doc``. - - Args: - doc: A source-resolution document; revised in place and returned. - framework_roots: Roots a rewritten path may live under when it does not - exist on this host. - min_gpu_pct: Entries below this GPU share are not worth a review. - max_entries: Hardest-hitting N entries reviewed in one call. - model: Chat model; defaults to ``$HYPERLOOM_LLM_SOURCE_MODEL``, then - the selected provider's model setting. - timeout_sec: Per-call ceiling; there is no retry. - log: Optional ``callable(str)`` for diagnostics. - complete: Injection point for the completion call (tests). - - Returns: - ``(doc, notes)``; ``notes`` records every applied and rejected revision. - """ - - def _say(message: str) -> None: - if callable(log): - log(f"llm_source_review: {message}") - - entries = doc.get("entries") if isinstance(doc, dict) else None - if not isinstance(entries, list) or not entries: - return doc, ["no entries to review"] - - reviewable_items = [ - (index, entry) - for index, entry in enumerate(entries) - if isinstance(entry, dict) and _gpu_pct(entry) >= min_gpu_pct - ] - reviewable_items.sort(key=lambda item: -_gpu_pct(item[1])) - reviewable_items = reviewable_items[: max(1, int(max_entries))] - reviewable = [entry for _, entry in reviewable_items] - if not reviewable: - return doc, [f"no entry at or above {min_gpu_pct}% GPU share"] - - sent_ids = [str(entry.get("kernel_id") or "") for entry in reviewable] - duplicate_sent = sorted( - kernel_id for kernel_id in set(sent_ids) if sent_ids.count(kernel_id) > 1 - ) - if not all(sent_ids) or duplicate_sent: - note = ( - "entries sent for review have missing or duplicate kernel_id values" - + (f": {duplicate_sent}" if duplicate_sent else "") - ) - return doc, [note] - try: - provider = _resolve_provider() if complete is None else "" - chosen_model = _resolve_model(model or "", provider) - except RuntimeError as exc: - audit = llm_source_audit(model=model or "") - audit["outcome"] = "configuration_error" - doc.setdefault("llm_audit", {})["review"] = audit - detail = _safe_exception_label(exc) - _say(f"configuration failed: {detail}") - return doc, [f"llm configuration failed: {detail}"] - if not chosen_model: - audit = llm_source_audit(model=model or "") - audit["outcome"] = "configuration_error" - doc.setdefault("llm_audit", {})["review"] = audit - _say("no source-resolution model configured") - return doc, ["no model configured"] - - audit = llm_source_audit(model=chosen_model) - audit["outcome"] = "requested" - doc.setdefault("llm_audit", {})["review"] = audit - caller = complete or _complete - try: - reply = caller( - build_review_prompt( - reviewable, - context_block=context_block, - framework_roots=framework_roots, - ), - chosen_model, - timeout_sec, - ) - except Exception as exc: # noqa: BLE001 - advisory tier, never fatal - audit["outcome"] = "call_error" - detail = _safe_exception_label(exc) - _say(f"call failed: {detail}") - return doc, [f"llm call failed: {detail}"] - - try: - parsed, revisions, error = parse_revisions(reply) - except Exception as exc: # noqa: BLE001 - malformed replies remain advisory - detail = _safe_exception_label(exc) - audit["outcome"] = "unusable_reply" - _say(f"unusable reply: {detail}") - return doc, [f"unusable reply: {detail}"] - if not parsed: - audit["outcome"] = "unusable_reply" - _say(f"unusable reply: {error}") - return doc, [f"unusable reply: {error}"] - - received_ids = [str(revision.get("kernel_id") or "") for revision in revisions] - duplicate_ids = sorted( - kernel_id - for kernel_id in set(received_ids) - if received_ids.count(kernel_id) > 1 - ) - missing_ids = sorted(set(sent_ids) - set(received_ids)) - extra_ids = sorted(set(received_ids) - set(sent_ids)) - if duplicate_ids or missing_ids or extra_ids: - details = [] - if missing_ids: - details.append(f"missing={missing_ids[:8]}") - if duplicate_ids: - details.append(f"duplicate={duplicate_ids[:8]}") - if extra_ids: - details.append(f"extra/unknown kernel_id={extra_ids[:8]}") - note = "revision set does not match entries sent: " + "; ".join(details) - audit["outcome"] = "protocol_error" - doc["reviewed_by"] = "llm_source_review" - doc["review_notes"] = [note] - _say(note) - return doc, [note] - - notes: list[str] = [] - try: - staged_entries = copy.deepcopy(entries) - staged_by_id = { - kernel_id: staged_entries[index] - for kernel_id, (index, _) in zip(sent_ids, reviewable_items) - } - for revision in revisions: - # The batch check above already rejected any id that was not sent, - # so every revision still here names a staged entry. - entry = staged_by_id[str(revision.get("kernel_id") or "")] - note = _apply_revision(entry, revision, framework_roots) - if note: - notes.append(note) - except Exception as exc: # noqa: BLE001 - discard the entire staged batch - detail = _safe_exception_label(exc) - note = f"revision validation failed: {detail}" - audit["outcome"] = "validation_error" - _say(note) - return doc, [note] - - doc["entries"] = staged_entries - doc["reviewed_by"] = "llm_source_review" - doc["review_notes"] = notes - audit["outcome"] = "completed" - _say(f"reviewed {len(reviewable)} entr(ies), {len(notes)} change(s)") - return doc, notes diff --git a/src/hyperloom/agents/kernel/tools/kernel_optimization.py b/src/hyperloom/agents/kernel/tools/kernel_optimization.py index 914465cbf8..e178cf43a5 100755 --- a/src/hyperloom/agents/kernel/tools/kernel_optimization.py +++ b/src/hyperloom/agents/kernel/tools/kernel_optimization.py @@ -41,6 +41,16 @@ sys.path.pop(0) +try: + from hyperloom.common.kernel_shape_contract import ( + REVIEW_SHAPE_PROVENANCE as _REVIEW_SHAPE_PROVENANCE, + ) +except ImportError: # pragma: no cover - standalone invocation + # Runs as a subprocess against the installed hyperloom, which may predate + # the constant. An unrecognised provenance then reads as measured, which is + # the pre-existing wording rather than a new claim. + _REVIEW_SHAPE_PROVENANCE = frozenset({"review_backfill", "review_derived"}) + def update_status( status_path: Path, @@ -831,16 +841,19 @@ def _structured_benchmark_shape_cases(candidate: dict[str, Any]) -> dict[str, An def _build_captured_shapes_block(candidate: dict[str, Any]) -> str: """Fallback shapes block when no TraceLens ``task_group`` is attached. - Surfaces the candidate's TraceLens-captured argument shapes so GEAK binds - its harness to the exact shapes the kernel saw during serving. Returns ``""`` - when no captured shapes exist. + Pins the harness to the candidate's argument shapes so the backend does not + pick its own. The heading states where the dims came from: a graph replay + records no arguments, so the hottest kernels of a captured model can only be + given dims the candidate review recovered or computed. Telling the backend + those were measured would misrepresent the one thing worth knowing when the + tuned kernel later fails to move end-to-end throughput. Args: candidate: The kernel candidate dict, possibly carrying captured shapes. Returns: - The captured-shapes prompt block, or ``""`` when no shapes exist. + The shapes prompt block, or ``""`` when the candidate has no shapes. """ shapes = candidate.get("shapes") or candidate.get("kernel_shapes") rendered = _format_shapes_for_case(shapes) @@ -848,13 +861,29 @@ def _build_captured_shapes_block(candidate: dict[str, Any]) -> str: return "" bound = str(candidate.get("bound_type") or candidate.get("bound") or "").strip() bound_line = f" (bound: {bound})" if bound else "" + provenance = str(candidate.get("shape_provenance") or "").strip().lower() + if provenance in _REVIEW_SHAPE_PROVENANCE: + heading = "Benchmark shapes (reconstructed for this serving run)" + origin = ( + "The profiler recorded no arguments for this kernel because it runs\n" + "inside a captured graph; these dims were reconstructed from the run's\n" + f"own artifacts and serving configuration (provenance: {provenance}).\n" + "They are the workload this kernel actually serves, so optimizing\n" + "against them is what produces an end-to-end gain -- whereas shapes you\n" + "choose yourself cannot account for the serving configuration:\n" + ) + else: + heading = "Benchmark shapes (TraceLens-captured from the serving run)" + origin = ( + "They are what the kernel saw during sglang/vLLM serving, so optimizing\n" + "against them is what produces an end-to-end gain on the workload:\n" + ) return ( - "\n## Benchmark shapes (TraceLens-captured from the serving run)\n\n" + f"\n## {heading}\n\n" "Build your harness shape sweep / `get_inputs()` from EXACTLY these\n" - f"captured argument shapes{bound_line} -- do NOT invent shapes. They are what\n" - "the kernel saw during sglang/vLLM serving, so optimizing against them is\n" - "what produces an end-to-end gain on the workload:\n" - f"- args: {rendered}\n" + f"argument shapes{bound_line} -- do NOT invent shapes.\n" + + origin + + f"- args: {rendered}\n" "Correctness golden: the ORIGINAL kernel's output on these shapes " "(baseline / `fn=` injection); do not hand-derive a reference from scratch.\n" + _build_kernel_contract_block(candidate) @@ -1207,8 +1236,8 @@ def _build_hypothesis_block(candidate: dict[str, Any]) -> str: "for the corresponding P-item. Treat them as starting points —", "verify each against the source / a quick micro-benchmark before", "committing to a direction. If your measurements contradict any", - "hypothesis, follow the data and document the discrepancy in", - "`optimization_report.md`.", + "hypothesis, follow the data and record the discrepancy in your", + "final summary.", "", ] for entry in all_prose: @@ -1256,7 +1285,7 @@ def _build_hypothesis_block(candidate: dict[str, Any]) -> str: "verify the reasoning against the source / a quick micro-benchmark", "before committing to the recommended direction. If your", "measurements contradict the hypothesis, follow the data and", - "document the discrepancy in `optimization_report.md`.", + "record the discrepancy in your final summary.", "", ] if identification: @@ -1511,22 +1540,26 @@ def build_prompt( *, backend: str | None = None, ) -> str: - """Render the full optimization prompt handed to a rewrite backend. + """Render the optimization prompt handed to a rewrite backend. - Assembles the hardware/budget/source-attribution preamble, the source - listing, semantically-ordered benchmark references, and the TraceLens - benchmark-cases / priority / hypothesis blocks into one prompt string. + Two shapes, because the backends own different amounts of the run. GEAK is + handed the whole harness: budget protocol, sandbox rules, the deliverable + contract Hyperloom later parses, and the A/B recipes it needs because it + brings no benchmark of its own. Forge brings its own driver, clock, gate and + artifact export, so it is handed only what it cannot derive -- the trace + evidence and the source-attribution guards. See the ``backend == "forge"`` + branch for why the omitted sections are not merely redundant there. Args: candidate (dict[str, Any]): Kernel candidate dict supplying source, benchmarks, shapes, and TraceLens context. args (argparse.Namespace): Parsed CLI args (source override, GPU count, kernel id, etc.). - backend (str | None): Target backend name, used to tailor backend- - specific prompt sections; None for the generic prompt. + backend (str | None): Target backend name. ``"forge"`` selects the slim + prompt; every other value (including None) renders the full one. Returns: - str: The fully rendered prompt text for the rewrite backend. + str: The rendered prompt text for the rewrite backend. """ source_file = args.source_file or candidate.get("source_file", "") source_block = "" @@ -1746,8 +1779,11 @@ def build_prompt( is_multinode_run = int(os.environ.get("INFERENCE_OPTIMIZER_NODES", "0") or 0) >= 2 except ValueError: is_multinode_run = False + # Held separately from ``safety``: the GPU-less sandbox applies to whichever + # backend runs, while the rest of ``safety`` describes the GEAK harness only. + multinode_block = "" if is_multinode_run: - safety += ( + multinode_block = ( "\nMULTI-NODE SANDBOX (no local GPU): every compile + benchmark\n" "step MUST be dispatched to a GPU-bearing RayJob pod. Do NOT\n" "call `hipcc`, `torch.cuda.*`, or `torch.utils.cpp_extension.load`\n" @@ -1769,6 +1805,7 @@ def build_prompt( "Treat `kernel-bench` as your only measurement gate; everything\n" "else (code edits, correctness reasoning) still happens locally.\n" ) + safety += multinode_block if not is_multigpu: safety += "- Use the provided benchmark/test files above for correctness/perf measurement.\n" elif num_gpus >= 2: @@ -1783,6 +1820,40 @@ def build_prompt( "rank's slice of the algorithm (e.g. local reduce + memcpy) so you can " "still measure compute/IO improvements.\n" ) + if (backend or "").strip().lower() == "forge": + # Forge already owns everything the sections below describe, and two of + # them fight it. The deliverables (``optimization_report.md``, a copy + # under ``optimized_versions/``) arrive as new untracked paths that its + # workspace guard refuses, which costs the iteration that wrote them -- + # while Hyperloom writes both itself from forge's published manifest and + # git diff, so nothing reads the agent's copies. The A/B recipes + # (standalone hipcc program, ``cpp_extension.load``) tell the agent to + # stand up a second benchmark beside the driver its own in-session gate + # scores, and a number from the wrong benchmark is worse than no number. + # The budget protocol narrates a mini-swe-agent cost meter forge has no + # header for, and the runtime-metadata block is labelled for GEAK's + # parser -- forge reads the invocation spec instead, which carries the + # same operands in more detail. + # + # What is left is what forge cannot derive: the trace evidence, and the + # two source-attribution guards that keep a rewrite off a @compile_ops + # wrapper and on the right ``__global__``. + forge_sections = [ + f"# TASK: Optimize the `{kernel_name}` kernel", + f"kernel_name: {kernel_name}\nkernel_url: {source_file}", + promotion_block, + device_symbol_block, + hypothesis_block, + extra_context_block, + benchmark_cases_block, + priority_block, + "Preserve function name, signature, decorators, and numerical behavior.", + repo_block, + multinode_block, + ] + # Most of these blocks render empty for any given kernel; joining them + # unfiltered leaves runs of blank lines where a section was skipped. + return "\n\n".join(part.strip() for part in forge_sections if part.strip()) tracelens_context_block = "" # Fall back to the full analysis.md only when no hypothesis_block was rendered. if not hypothesis_block.strip(): diff --git a/src/hyperloom/agents/kernel/tools/tracelens_analysis.py b/src/hyperloom/agents/kernel/tools/tracelens_analysis.py index 3a2e8969f8..b0c6c7aedf 100755 --- a/src/hyperloom/agents/kernel/tools/tracelens_analysis.py +++ b/src/hyperloom/agents/kernel/tools/tracelens_analysis.py @@ -43,6 +43,13 @@ except ImportError: _resolve_flydsl_source_roots = None +try: + from hyperloom.orchestrator.framework.paths import ( + resolve_kernel_search_roots as _resolve_kernel_search_roots, + ) +except ImportError: + _resolve_kernel_search_roots = None + try: from apply_kernel_patch import known_target_roots as _known_target_roots except ImportError: @@ -137,6 +144,15 @@ except ImportError: # pragma: no cover - standalone invocation _KSC = None # type: ignore[assignment] +try: + from hyperloom.common.kernel_shape_contract import ( + REVIEW_DERIVED_PROVENANCE as _REVIEW_DERIVED_PROVENANCE, + ) +except ImportError: # pragma: no cover - standalone invocation + # This script also runs against an installed hyperloom that may predate the + # constant; the literal keeps the review's dims labelled either way. + _REVIEW_DERIVED_PROVENANCE = "review_derived" + log = logging.getLogger(__name__) # Duplicated from kernel_source_contract.SOURCE_RESOLUTION_FILENAME: the @@ -1720,15 +1736,102 @@ def is_vendor_dispatch_wrapper(name: str, source_file: str) -> bool: return any(sig in text for sig in _VENDOR_DISPATCH_SIGS) -KNOWN_SEARCH_ROOTS = ( +#: Packages whose trees hold rewritable kernel source. Located at runtime so a +#: wheel install, an editable checkout and a serving image all resolve, rather +#: than only the one layout a literal happens to name. +_KERNEL_SOURCE_PACKAGES: tuple[str, ...] = ( + "aiter", + "aiter_meta", + "sglang", + "sgl_kernel", + "vllm", +) + +#: Last-resort checkout layouts for a host where nothing above is importable. +#: Kept small on purpose: a pinned path cannot follow a package across +#: container images or Python versions, and a list of them going stale in +#: silence is what emptied this tier and stalled kernel-opt entirely. +_FALLBACK_SEARCH_ROOTS: tuple[str, ...] = ( "/sgl-workspace/aiter", "/sgl-workspace/sglang/sgl-kernel", "/sgl-workspace/sglang/python/sglang", "/sgl-workspace/vllm", - "/opt/venv/lib/python3.10/site-packages/sglang", - "/opt/venv/lib/python3.10/site-packages/aiter", - "/opt/venv/lib/python3.10/site-packages/vllm", ) + + +def _installed_package_dir(package: str) -> str: + """Locate a package's directory without importing it. + + Args: + package (str): Importable package name. + + Returns: + str: The package directory, or ``""`` when it is not on this + interpreter's path. + """ + if not package or not package.isidentifier(): + return "" + try: + spec = importlib.util.find_spec(package) + except (AttributeError, ImportError, ValueError): + return "" + if spec is None: + return "" + for location in list(getattr(spec, "submodule_search_locations", None) or []): + candidate = str(location).rstrip("/") + if candidate: + return candidate + origin = str(getattr(spec, "origin", "") or "") + return os.path.dirname(origin) if origin else "" + + +@lru_cache(maxsize=1) +def _discover_kernel_search_roots() -> tuple[str, ...]: + """Resolve the framework trees to grep for kernel source, at runtime. + + Prefers the orchestrator's centralised resolver so this tool agrees with + PolicyGate and patch application on where framework source lives. When that + package is not importable (standalone CLI use) it falls back to locating + each known package itself, then to the pinned checkout layouts. + + Non-existent roots are dropped: grepping them returns nothing and is + indistinguishable from a kernel that genuinely has no source here. + + Returns: + tuple[str, ...]: Existing roots without trailing separators, + de-duplicated in discovery order. + """ + discovered: list[str] = [] + if _resolve_kernel_search_roots is not None: + discovered.extend(_resolve_kernel_search_roots()) + else: + discovered.extend( + location + for location in ( + _installed_package_dir(package) for package in _KERNEL_SOURCE_PACKAGES + ) + if location + ) + discovered.extend(_FALLBACK_SEARCH_ROOTS) + roots: list[str] = [] + seen: set[str] = set() + for root in discovered: + normalized = str(root or "").rstrip("/") + if not normalized or normalized in seen or not os.path.isdir(normalized): + continue + seen.add(normalized) + roots.append(normalized) + if not roots: + log.warning( + "no framework source root exists on this host (looked for %s); " + "kernel source resolution will find nothing and every hot kernel " + "will be reported as non-routable", + ", ".join(_KERNEL_SOURCE_PACKAGES), + ) + return tuple(roots) + + +KNOWN_SEARCH_ROOTS = _discover_kernel_search_roots() # Extensions a grep hit may be admitted under. Deliberately narrow, and kept in # lockstep with source_type_for(): a suffix admitted here but unclassified there # lands as source_type="unknown", which classify_patchability rejects. Worse, it @@ -2560,112 +2663,183 @@ def find_repo_root(source_file: str) -> str: _BENCHMARK_DIRS = ("op_tests", "tests", "benchmarks", "benchmark", "test", "perf") +#: Curated harness lookups, keyed by marker substrings in a kernel's name or +#: source path. Paths are *checkout-relative* on purpose: the same harness sits +#: at ``/sgl-workspace/aiter/op_tests/...`` in a serving image and is absent from +#: a wheel install, so a pinned absolute path is either right on one host or a +#: fabrication on every other. _KNOWN_HARNESS_HINTS: tuple[tuple[tuple[str, ...], tuple[str, ...]], ...] = ( # --- Normalization --- ( - ("rmsnorm_quant", "add_rmsnorm_quant", "rmsnorm", "add_rmsnorm"), ( - "/sgl-workspace/aiter/op_tests/test_rmsnorm2dFusedAddQuant.py", - "/sgl-workspace/aiter/op_tests/test_rmsnorm2d.py", - "/sgl-workspace/aiter/op_tests/op_benchmarks/triton/bench_rmsnorm.py", - "/sgl-workspace/sglang/sgl-kernel/benchmark/bench_rmsnorm.py", - "/sgl-workspace/aiter/op_tests/triton_tests/normalization/test_rmsnorm.py", - "/sgl-workspace/aiter/op_tests/triton_tests/normalization/test_fused_add_rmsnorm_pad.py", + "rmsnorm_quant", + "add_rmsnorm_quant", + "rmsnorm", + "add_rmsnorm", + ), + ( + "aiter/op_tests/test_rmsnorm2dFusedAddQuant.py", + "aiter/op_tests/test_rmsnorm2d.py", + "aiter/op_tests/op_benchmarks/triton/bench_rmsnorm.py", + "sglang/sgl-kernel/benchmark/bench_rmsnorm.py", + "aiter/op_tests/triton_tests/normalization/test_rmsnorm.py", + "aiter/op_tests/triton_tests/normalization/test_fused_add_rmsnorm_pad.py", ), ), # --- Activation --- ( - ("activation", "act_and_mul", "silu"), ( - "/sgl-workspace/aiter/op_tests/test_activation.py", - "/sgl-workspace/aiter/op_tests/op_benchmarks/triton/bench_ff_a16w16_fused.py", - "/sgl-workspace/sglang/sgl-kernel/tests/test_activation.py", - "/sgl-workspace/sglang/sgl-kernel/benchmark/bench_activation.py", - "/sgl-workspace/sglang/python/sglang/jit_kernel/tests/test_activation.py", - "/sgl-workspace/sglang/python/sglang/jit_kernel/benchmark/bench_activation.py", + "activation", + "act_and_mul", + "silu", + ), + ( + "aiter/op_tests/test_activation.py", + "aiter/op_tests/op_benchmarks/triton/bench_ff_a16w16_fused.py", + "sglang/sgl-kernel/tests/test_activation.py", + "sglang/sgl-kernel/benchmark/bench_activation.py", + "sglang/python/sglang/jit_kernel/tests/test_activation.py", + "sglang/python/sglang/jit_kernel/benchmark/bench_activation.py", ), ), # --- Attention --- ( - ("paged_attention", "fmha", "attention"), ( - "/sgl-workspace/aiter/op_tests/test_pa.py", - "/sgl-workspace/aiter/op_tests/op_benchmarks/triton/bench_pa_decode.py", - "/sgl-workspace/aiter/op_tests/op_benchmarks/triton/bench_pa_prefill.py", + "paged_attention", + "fmha", + "attention", + ), + ( + "aiter/op_tests/test_pa.py", + "aiter/op_tests/op_benchmarks/triton/bench_pa_decode.py", + "aiter/op_tests/op_benchmarks/triton/bench_pa_prefill.py", ), ), # --- MLA decode --- ( - ("mla_decode", "pseudo_mla", "mla_persistent"), ( - "/sgl-workspace/aiter/op_tests/test_mla.py", - "/sgl-workspace/aiter/op_tests/test_mla_persistent.py", - "/sgl-workspace/aiter/op_tests/op_benchmarks/triton/bench_mla_decode.py", + "mla_decode", + "pseudo_mla", + "mla_persistent", + ), + ( + "aiter/op_tests/test_mla.py", + "aiter/op_tests/test_mla_persistent.py", + "aiter/op_tests/op_benchmarks/triton/bench_mla_decode.py", ), ), # --- MoE CK two-stage --- ( - ("ck_moe_stage", "moe_2stage", "moe_stage1", "moe_stage2"), ( - "/sgl-workspace/aiter/op_tests/test_moe_2stage.py", - "/sgl-workspace/aiter/op_tests/op_benchmarks/triton/bench_moe.py", + "ck_moe_stage", + "moe_2stage", + "moe_stage1", + "moe_stage2", + ), + ( + "aiter/op_tests/test_moe_2stage.py", + "aiter/op_tests/op_benchmarks/triton/bench_moe.py", ), ), # --- MoE FP8 blockscale (ASM) --- ( - ("fmoe_fp8_blockscale", "moe_blockscale"), ( - "/sgl-workspace/aiter/op_tests/test_moe_blockscale.py", - "/sgl-workspace/aiter/op_tests/triton_tests/moe/test_moe_gemm_a8w8_blockscale.py", + "fmoe_fp8_blockscale", + "moe_blockscale", + ), + ( + "aiter/op_tests/test_moe_blockscale.py", + "aiter/op_tests/triton_tests/moe/test_moe_gemm_a8w8_blockscale.py", ), ), # --- GEMM A8W8 blockscale --- ( - ("gemm_a8w8_blockscale",), ( - "/sgl-workspace/aiter/op_tests/test_gemm_a8w8_blockscale.py", - "/sgl-workspace/aiter/op_tests/op_benchmarks/triton/bench_gemm_a8w8_blockscale.py", + "gemm_a8w8_blockscale", + ), + ( + "aiter/op_tests/test_gemm_a8w8_blockscale.py", + "aiter/op_tests/op_benchmarks/triton/bench_gemm_a8w8_blockscale.py", ), ), # --- Quantization --- ( - ("dynamic_per_token_scaled_quant", "per_token_quant"), ( - "/sgl-workspace/aiter/op_tests/test_quant.py", - "/sgl-workspace/aiter/op_tests/triton_tests/quant/test_quant.py", + "dynamic_per_token_scaled_quant", + "per_token_quant", + ), + ( + "aiter/op_tests/test_quant.py", + "aiter/op_tests/triton_tests/quant/test_quant.py", ), ), # --- Batch-invariant addmm (Triton) --- ( - ("batch_invariant", "addmm"), - ("/sgl-workspace/sglang/test/registered/unit/batch_invariant_ops/test_batch_invariant_ops.py",), + ( + "batch_invariant", + "addmm", + ), + ( + "sglang/test/registered/unit/batch_invariant_ops/test_batch_invariant_ops.py", + ), ), ) -def _known_harness_files(name: str, source_file: str, *, require_exists: bool = True) -> list[Path]: +@lru_cache(maxsize=1) +def _harness_search_bases() -> tuple[str, ...]: + """Directories a checkout-relative harness path may be joined onto. + + A hint reads ``aiter/op_tests/...``, so the base is whatever holds the + ``aiter`` checkout. Each resolved search root contributes both itself and + its parent, because a root is the package directory on a wheel install + (``.../dist-packages/aiter``) and the checkout itself in a serving image + (``/sgl-workspace/aiter``); one join is the right one and the other simply + does not exist. + + Returns: + tuple[str, ...]: Existing base directories, de-duplicated. + """ + bases: list[str] = [] + seen: set[str] = set() + for root in KNOWN_SEARCH_ROOTS: + trimmed = root.rstrip("/") + for base in (os.path.dirname(trimmed), trimmed): + if base and base not in seen and os.path.isdir(base): + seen.add(base) + bases.append(base) + return tuple(bases) + + +def _known_harness_files(name: str, source_file: str) -> list[Path]: """Return curated benchmark/test harnesses matching a kernel. - Looks up :data:`_KNOWN_HARNESS_HINTS` by marker substrings found in the - kernel name / source path. By default it returns only hinted harnesses that - exist on disk; callers without a repo root can request the curated hint list - itself so tests and downstream prompts remain stable in minimal containers - where ``/sgl-workspace`` is absent. + Resolves each checkout-relative hint against the bases that exist here and + keeps only files actually present. A list naming paths that cannot be + opened is worse than an empty one, because every reader downstream -- the + dispatch prompt included -- treats a non-empty list as a harness it can run. Args: name (str): Kernel symbol/name. source_file (str): Resolved source-file path (may be empty). - require_exists (bool): When True, only paths present on disk are - returned. When False, matching curated hints are returned as-is. Returns: - list[Path]: Curated harness files, possibly empty. + list[Path]: Existing curated harness files, possibly empty. """ blob = f"{name} {source_file}".lower() out: list[Path] = [] - for markers, paths in _KNOWN_HARNESS_HINTS: - if any(marker in blob for marker in markers): - out.extend(Path(p) for p in paths if (not require_exists or Path(p).exists())) + seen: set[str] = set() + bases = _harness_search_bases() + for markers, relatives in _KNOWN_HARNESS_HINTS: + if not any(marker in blob for marker in markers): + continue + for relative in relatives: + for base in bases: + candidate = os.path.join(base, relative) + if candidate not in seen and os.path.isfile(candidate): + seen.add(candidate) + out.append(Path(candidate)) + break return out @@ -2699,10 +2873,9 @@ def find_benchmark_files(name: str, repo_root: str, source_file: str = "") -> li Returns: Up to ten matching harness paths, with multi-GPU tests demoted. """ + known = _known_harness_files(name, source_file) if not repo_root: - known = _known_harness_files(name, source_file, require_exists=False) return [str(p) for p in known[:10]] - known = _known_harness_files(name, source_file) keywords = _candidate_keywords(name) # Add the source stem (and no-underscore variant) for repos that name tests differently. if source_file: @@ -3108,8 +3281,6 @@ def is_multigpu_kernel(name: str, source_file: str) -> bool: def analyze_trace_files( trace_files: list[Path], top_k: int, - *, - allow_model_tiers: bool = True, ) -> list[dict[str, Any]]: """Aggregate GPU kernels across raw trace files into top-K candidates. @@ -3120,9 +3291,6 @@ def analyze_trace_files( Args: trace_files (list[Path]): Trace files (optionally gzipped) to scan. top_k (int): Number of hottest kernels to keep. - allow_model_tiers (bool): Whether source resolution may call a model. - The deterministic route sets this to false, so the "no model calls" - promise holds on this path too. Returns: list[dict[str, Any]]: Finalized hot-kernel candidate dicts. @@ -3177,7 +3345,6 @@ def analyze_trace_files( top, total_dur=total_dur, trace_files=trace_files, - allow_model_tiers=allow_model_tiers, ) @@ -4252,8 +4419,14 @@ def _stamp_candidate_metadata(item: dict[str, Any], op_cat_map: dict[str, str] | # playbook candidate has no rewritable device source, so point that # field at the task bundle's anchor file instead of leaving it # empty (which would otherwise fall through as "missing_native_source"). - if not str(item.get("source_file") or "").strip(): - item["source_file"] = resolve_kernel_anchor_path(playbook) + # + # The anchor also overrides whatever the grep tier guessed. A registry + # match is a curated statement that this operator is tuned through a + # task bundle, whereas the guess can be a same-word collision -- + # ``mori::EpDispatchCombineOp::dispatch`` reduces to the keyword + # "dispatch" and lands on an unrelated vendor header. Handing that path + # to a backend would rewrite the wrong file. + item["source_file"] = resolve_kernel_anchor_path(playbook) item["benchmark_files"] = find_benchmark_files( item["name"], item.get("kernel_repo", ""), item.get("source_file", "") ) @@ -4269,10 +4442,6 @@ def _stamp_candidate_metadata(item: dict[str, Any], op_cat_map: dict[str, str] | item.setdefault("source_path", item.get("source_file", "")) -#: A candidate below this GPU share is not worth an LLM round-trip. -_LLM_FALLBACK_MIN_GPU_PCT = 5.0 - - #: Populated once per run from the CLI args so both model tiers see the same #: serving configuration. Empty when the analysis runs without that context. _RUNTIME_CONTEXT: dict[str, Any] = {} @@ -4358,93 +4527,6 @@ def _append_resolution_reason(item: dict[str, Any], reason: str) -> None: item["source_resolution_reason"] = f"{current}; {reason}" -def _apply_llm_source_fallback(item: dict[str, Any]) -> None: - """Last-resort LLM pick for a candidate every deterministic tier missed. - - No-op unless the operator enabled the tier and the candidate is hot enough to - justify the call. Mutates ``item`` in place only on an accepted answer. - """ - name = str(item.get("name") or "") - try: - from _llm_source_fallback import ( # noqa: PLC0415 - llm_source_audit, - llm_source_provider_configured, - select_source_via_llm, - ) - - try: - gpu_pct = float(item.get("gpu_pct") or 0.0) - except (TypeError, ValueError): - gpu_pct = 0.0 - if gpu_pct < _LLM_FALLBACK_MIN_GPU_PCT: - _append_resolution_reason( - item, - f"llm_fallback_skipped: gpu_pct {gpu_pct:.2f} < {_LLM_FALLBACK_MIN_GPU_PCT}", - ) - return - # Settle the provider before gathering the shortlist. That grep walks - # every framework root once per keyword, and on a deployment that never - # configured a provider the tier would pay for it on every hot kernel - # only to decline the call. - if not llm_source_provider_configured(): - audit = llm_source_audit() - audit["outcome"] = "configuration_error" - item["source_resolution_llm_audit"] = audit - _append_resolution_reason(item, "llm_fallback_skipped: no provider configured") - return - shortlist = collect_source_candidates_via_grep(name) - if not shortlist: - _append_resolution_reason(item, "llm_fallback_no_shortlist") - log.info("LLM source fallback: no grep shortlist for %r", name) - return - audit = llm_source_audit() - audit["outcome"] = "requested" - item["source_resolution_llm_audit"] = audit - call_errors: list[str] = [] - picked, confidence, reason = select_source_via_llm( - name, - shortlist, - framework_roots=tuple(KNOWN_SEARCH_ROOTS), - context_block=_source_context_block(), - log=_forward_to_log, - errors=call_errors, - ) - if picked: - audit["outcome"] = "accepted" - item["source_file"] = picked - item["source_resolution_method"] = "llm_fallback" - item["source_resolution_confidence"] = confidence - item["source_resolution_reason"] = reason - return - if call_errors: - audit["outcome"] = "error" - failure = f"llm_fallback_error: {call_errors[0]}" - _append_resolution_reason(item, failure) - log.warning("LLM source fallback failed for %r (%s)", name, failure) - return - # Answered but not accepted: invented path, low confidence, or refusal. - audit["outcome"] = "declined" - _append_resolution_reason( - item, f"llm_fallback_declined: {reason or 'no candidate accepted'}" - ) - log.info( - "LLM source fallback declined for %r over %d shortlist entr(ies): %s", - name, - len(shortlist), - reason or "no candidate accepted", - ) - except Exception as exc: # noqa: BLE001 - advisory tier, never breaks finalization - # Import errors, gateway 401s and timeouts all land here; without a - # trail they look identical to the tier being switched off. - reason = f"llm_fallback_error: {type(exc).__name__}: {exc}" - audit = item.get("source_resolution_llm_audit") - if isinstance(audit, dict): - audit["outcome"] = "error" - log.warning("LLM source fallback failed for %r (%s)", name, reason) - _append_resolution_reason(item, reason) - return - - @functools.lru_cache(maxsize=64) def _package_parent_dir(package: str) -> str: """Directory holding ``package``'s own directory, resolved at runtime. @@ -4602,7 +4684,6 @@ def _finalize_candidates( trace_files: list[Path] | None = None, log_path: Path | str | None = None, source_resolution_out: Path | str | None = None, - allow_model_tiers: bool = True, model_name: str = "", ) -> list[dict[str, Any]]: """Apply shared post-processing to parsed candidate rows. @@ -4624,8 +4705,6 @@ def _finalize_candidates( trace_files: Optional raw trace files. When given, Python launcher frames are retained as evidence and accepted as source only when name grep independently resolves the same file. - allow_model_tiers: Whether fallback and artifact review may call an LLM. - The deterministic CLI route sets this to false. model_name: Runtime model identity recorded in the resolution artifact. Returns: @@ -4777,14 +4856,10 @@ def _finalize_candidates( item, f"trace launcher unconfirmed by name grep: {trace_source}", ) - if not item.get("source_file"): - if allow_model_tiers: - _apply_llm_source_fallback(item) - else: - _append_resolution_reason( - item, - "llm_fallback_skipped: deterministic route", - ) + # An unresolved candidate stops here. The agent review pass runs + # once over the finished table rather than per kernel, so it can + # weigh a blank against the rest of the evidence instead of + # guessing from a symbol alone. # Promote a tiny pybind shim TU to the real device code. item["kernel_repo"] = find_repo_root(item.get("source_file", "")) item["source_file"] = upgrade_pybind_shim_source( @@ -4823,8 +4898,6 @@ def _finalize_candidates( framework=framework or "", model_name=model_name, log_path=log_path, - op_cat_map=op_cat_map, - allow_review=allow_model_tiers, ) return top @@ -4986,101 +5059,6 @@ def _is_curated_resolution(item: dict[str, Any]) -> bool: } -def apply_resolution_entries_to_candidates( - entries: list[dict[str, Any]], - candidates: list[dict[str, Any]], - op_cat_map: dict[str, str] | None = None, -) -> int: - """Fold reviewed entries back onto the candidates, and re-classify. - - Without this the review tier is inert: it revises the audit artifact while - every downstream stage keeps reading ``kernel_candidates.json``. Re-running - ``_stamp_candidate_metadata`` matters as much as copying the path -- a - rewrite changes ``source_type``, which decides ``reusable_native_kernel`` - and therefore whether the kernel is dispatched at all. - - Two invariants keep a rewritten candidate internally consistent: - - * A curated resolution is never overwritten (see - :func:`_is_curated_resolution`). - * Otherwise every field derived from the previous path is cleared before the - new one is stamped, so no downstream reader can pick up metadata that - describes the source the candidate no longer points at. - - Returns: - The number of candidates actually changed. - """ - by_id = {str(c.get("kernel_id")): c for c in candidates if isinstance(c, dict)} - changed = 0 - for entry in entries: - if not isinstance(entry, dict) or "previous_source_file" not in entry: - continue - item = by_id.get(str(entry.get("kernel_id") or "")) - if item is None: - continue - new_source = str(entry.get("source_file") or "") - if new_source == str(item.get("source_file") or ""): - continue - if _is_curated_resolution(item): - entry["review_rejected"] = "curated_resolution_not_overridable" - entry["source_file"] = str(item.get("source_file") or "") - entry["source_line"] = item.get("source_line") - entry["source_function"] = str(item.get("source_function") or "") - entry["method"] = _candidate_resolution_method(item) - entry["reason"] = str(item.get("source_resolution_reason") or "") - continue - for key in _SOURCE_DERIVED_METADATA: - item.pop(key, None) - item["source_file"] = new_source - item["source_path"] = new_source - item["source_line"] = entry.get("source_line") - item["source_function"] = str(entry.get("source_function") or "") - item["source_resolution_method"] = str(entry.get("method") or "") - item["source_resolution_reason"] = str(entry.get("reason") or "") - item["source_resolution_previous_file"] = str(entry.get("previous_source_file") or "") - item["source_resolution_previous_method"] = str(entry.get("previous_method") or "") - item["kernel_repo"] = find_repo_root(new_source) if new_source else "" - item["source_type"] = source_type_for(item.get("name", ""), new_source) - if item["source_type"] != "vendor_binary" and is_vendor_dispatch_wrapper( - item.get("name", ""), new_source - ): - item["source_type"] = "vendor_binary" - item["vendor_dispatch_wrapper"] = True - item["runtime_generated_kernel"] = is_runtime_generated_kernel( - item.get("name", ""), new_source - ) - _stamp_candidate_metadata(item, op_cat_map) - changed += 1 - return changed - - -def _review_source_resolution(doc: dict[str, Any], *, log_path: Path | str | None) -> None: - """Let the opt-in review tier revise the table before it is written. - - Runs on the whole table rather than only the blanks: the deterministic - tiers' failure mode is a confidently wrong path, which a blanks-only tier - can never reach. Revisions carry their own guard rails (see the module) and - are applied in place; any failure leaves the deterministic result standing. - """ - try: - from _llm_source_review import review_resolution_document # noqa: PLC0415 - - _, notes = review_resolution_document( - doc, - framework_roots=tuple(KNOWN_SEARCH_ROOTS), - context_block=_source_context_block(), - log=_forward_to_log, - ) - summary = f"source-resolution review: {len(notes)} change(s)" - log.info("%s", summary) - if log_path: - append_log(log_path, summary) - for note in notes: - append_log(log_path, f" review: {note}") - except Exception as exc: # noqa: BLE001 - advisory tier, never fatal - log.warning("source-resolution review failed (%r); keeping deterministic result", exc) - - def write_source_resolution_artifact( candidates: list[dict[str, Any]], out_path: Path | str, @@ -5088,8 +5066,6 @@ def write_source_resolution_artifact( framework: str = "", model_name: str = "", log_path: Path | str | None = None, - op_cat_map: dict[str, str] | None = None, - allow_review: bool = True, ) -> Path | None: """Write the source-resolution artifact next to the candidate report. @@ -5106,37 +5082,6 @@ def write_source_resolution_artifact( model_name=model_name, framework=framework, ) - if allow_review: - _review_source_resolution(doc, log_path=log_path) - # Fold any revision back onto the candidates: they, not this file, are - # what the dispatch stage reads. - reviewed_entries = [ - entry for entry in (doc.get("entries") or []) if isinstance(entry, dict) - ] - applied = apply_resolution_entries_to_candidates( - reviewed_entries, candidates, op_cat_map - ) - if applied: - review_metadata = { - str(entry.get("kernel_id") or ""): { - key: entry[key] - for key in ( - "previous_source_file", - "previous_method", - "review_rejected", - ) - if key in entry - } - for entry in reviewed_entries - } - rebuilt = build_source_resolution_entries(candidates) - for entry in rebuilt: - entry.update(review_metadata.get(str(entry.get("kernel_id") or ""), {})) - doc["entries"] = rebuilt - note = f"source-resolution review applied to {applied} candidate(s)" - log.info("%s", note) - if log_path: - append_log(log_path, note) problems = _KSC.validate_document(doc) if problems: log.warning( @@ -6947,6 +6892,320 @@ def build_audit_summary( } +def run_candidate_review_stage( + run_dir: Path, + *, + candidates: list[dict[str, Any]], + args: argparse.Namespace, + log_path: Path | str | None = None, + trace_health_warnings: list[dict[str, Any]] | None = None, +) -> dict[str, str]: + """Run the review stage, converting any unexpected fault into a warning. + + The stage is advisory by construction, and it sits at the end of an + analysis that a multi-hour benchmark paid for. An unforeseen fault in it + must cost the audit, not the run, so nothing escapes this boundary. + """ + try: + return _run_candidate_review_stage( + run_dir, + candidates=candidates, + args=args, + log_path=log_path, + trace_health_warnings=trace_health_warnings, + ) + except Exception as exc: # noqa: BLE001 - never let the audit fail the run + log.warning("candidate review stage failed (%r); keeping the deterministic table", exc) + if trace_health_warnings is not None: + trace_health_warnings.append( + { + "code": "candidate_review_failed", + "severity": "error", + "status": "internal_error", + "detail": type(exc).__name__, + "message": ( + "The candidate review stage raised " + f"{type(exc).__name__}; kernel_candidates.json is the " + "unreviewed deterministic result." + ), + } + ) + return {} + + +#: Everything the review can stage without moving ``source_file``. The +#: re-derivation is skipped for rows that did not change, and a path is only one +#: of the things that can: operand dims are most often supplied for a kernel the +#: deterministic tiers already located, so keying the check on the path alone +#: drops exactly the proposals that were hardest to obtain. +_REVIEW_STAGED_PROPOSALS = ( + "review_shapes", + "review_input_dtypes", + "review_reusable_hint", + "review_benchmark_files", +) + +#: Rebuilt from ``shapes``, so they must not outlive the dims they described. +_REVIEW_STALE_SHAPE_FIELDS = ( + "input_shapes", + "invocation_cases", + "raw_arg_spec", + "shape_donor_operation", + "_input_shapes_synthetic", +) + + +def _adopt_reviewed_shapes(item: dict[str, Any]) -> None: + """Take the operand dims the review supplied, if it supplied any. + + Only fires where the deterministic stage came up empty. A recorded shape + outranks a reviewed one even when the review is confident, because the + reviewed dims can be arithmetic over the serving configuration and nothing + downstream re-measures them; the integration benchmark hours later is the + first thing that would notice they were wrong. + + The alternate representations are dropped rather than translated. They + describe the previous dims, and a harness built from a mix of the two would + be wrong in a way that still benchmarks cleanly. + """ + proposed = item.get("review_shapes") + if not isinstance(proposed, list) or not proposed: + return + if item.get("shapes"): + return + item["shapes"] = list(proposed) + item["shape_provenance"] = str( + item.get("review_shape_provenance") or _REVIEW_DERIVED_PROVENANCE + ) + reviewed_dtypes = item.get("review_input_dtypes") + if isinstance(reviewed_dtypes, list) and reviewed_dtypes: + item["input_dtypes"] = list(reviewed_dtypes) + for key in _REVIEW_STALE_SHAPE_FIELDS: + item.pop(key, None) + + +def _rederive_after_review(item: dict[str, Any], op_cat_map: dict[str, str] | None = None) -> None: + """Recompute everything that follows from ``source_file`` after a revision. + + The review returns a location, not a verdict. Re-running the deterministic + stamping keeps :func:`classify_patchability` the only gate that decides + routability, so the vendor-binary, dispatch-wrapper and runtime-generated + rejections still apply to a path the model supplied. + """ + new_source = str(item.get("source_file") or "") + for key in _SOURCE_DERIVED_METADATA: + item.pop(key, None) + item["source_path"] = new_source + item["kernel_repo"] = find_repo_root(new_source) if new_source else "" + item["source_type"] = source_type_for(item.get("name", ""), new_source) + if item["source_type"] != "vendor_binary" and is_vendor_dispatch_wrapper( + item.get("name", ""), new_source + ): + item["source_type"] = "vendor_binary" + item["vendor_dispatch_wrapper"] = True + item["runtime_generated_kernel"] = is_runtime_generated_kernel(item.get("name", ""), new_source) + _adopt_reviewed_shapes(item) + _stamp_candidate_metadata(item, op_cat_map) + # Stamping recomputes benchmark_files from the curated marker table, which + # is coarser than a session that went and looked. Its verified answer wins. + reviewed_harnesses = item.get("review_benchmark_files") + if isinstance(reviewed_harnesses, list): + item["benchmark_files"] = list(reviewed_harnesses) + # A restrictive hint is honoured, a permissive one is not. The reviewer can + # veto a kernel it knows is not worth a tuning session, but it cannot talk + # the gate into dispatching something the deterministic rules rejected. + if item.get("review_reusable_hint") is False and item.get("reusable_native_kernel"): + item["reusable_native_kernel"] = False + item["skip_reason"] = ( + str(item.get("review_skip_reason") or "").strip() + or f"review: {item.get('review_reason') or 'not worth a tuning session'}" + ) + + +def _run_candidate_review_stage( + run_dir: Path, + *, + candidates: list[dict[str, Any]], + args: argparse.Namespace, + log_path: Path | str | None = None, + trace_health_warnings: list[dict[str, Any]] | None = None, +) -> dict[str, str]: + """Audit the deterministic candidate table with one agent session. + + The deterministic tiers resolve a kernel from its symbol alone; they cannot + tell a file that defines a kernel from one that merely launches it, and they + have no view of the model or how it is being served. This hands that table + to an agent together with the paths of everything the run already produced, + and folds back the revisions it can verify. + + Mandatory on the agent route, but never fatal: a definitive failure records + an ``error``-severity trace-health warning and leaves the deterministic + table standing. Losing the audit costs some candidates; failing the run + would cost the hours of benchmarking that produced the trace. + + Args: + run_dir: The per-run output directory. + candidates: The finalized candidate rows, revised in place. + args: Parsed CLI args (model name, framework, source root). + log_path: Optional log file for diagnostics. + trace_health_warnings: Warning sink surfaced to the Coordinator. + + Returns: + dict[str, str]: Artifact paths produced by this stage. + """ + artifacts: dict[str, str] = {} + warnings = trace_health_warnings if trace_health_warnings is not None else [] + + def _note(message: str) -> None: + log.info("%s", message) + if log_path: + append_log(log_path, message) + + try: + from _candidate_review_agent import ( # noqa: PLC0415 + RAW_CANDIDATES_FILENAME, + REVISIONS_FILENAME, + apply_revisions, + fingerprint_drift, + run_candidate_review, + source_fingerprint, + ) + except ImportError as exc: # pragma: no cover - packaging fault + warnings.append( + { + "code": "candidate_review_unavailable", + "severity": "error", + "message": ( + "The candidate review agent could not be imported " + f"({type(exc).__name__}); the candidate table is the " + "unreviewed deterministic result." + ), + } + ) + return artifacts + + tracelens_dir = run_dir / "tracelens" + raw_path = run_dir / RAW_CANDIDATES_FILENAME + routable = [c for c in candidates if isinstance(c, dict) and c.get("reusable_native_kernel") is True] + atomic_write_json( + raw_path, + { + "model_name": args.model_name, + "framework": args.framework, + "source": "tracelens_analysis:deterministic", + "hot_kernels": candidates, + "routable_kernels": routable, + }, + ) + artifacts["kernel_candidates_raw"] = str(raw_path) + + source_paths = [str(c.get("source_file") or "") for c in candidates if isinstance(c, dict)] + before = source_fingerprint(source_paths) + + # Only ``analysis.md`` is a supported TraceLens output; everything else in + # that directory is internal and may be removed without notice. The rest of + # the list is Hyperloom's own or the model's, so it is ours to offer. + # + # Little is lost by not pointing at the sidecars: for every operator they + # describe, ``analysis.md`` carries the same operand dims and launcher in + # its own table, and for a graph-launched operator neither has anything -- + # the replay has no CPU-side parent op, so nothing recorded the arguments. + reference_paths = { + "source resolution audit": str(run_dir / _SOURCE_RESOLUTION_NAME), + "tracelens report": str(tracelens_dir / "analysis.md"), + "trace input manifest": str(run_dir / "trace_input_manifest.json"), + "model directory": str(_RUNTIME_CONTEXT.get("model_path") or ""), + } + outcome = run_candidate_review( + run_dir=run_dir, + raw_candidates_path=raw_path, + reference_paths={k: v for k, v in reference_paths.items() if v}, + framework_roots=KNOWN_SEARCH_ROOTS, + context_block=_source_context_block(), + log=_forward_to_log, + ) + + if not outcome.ok: + warnings.append( + { + "code": "candidate_review_failed", + "severity": "error", + "status": outcome.status, + "detail": outcome.detail, + "message": ( + f"The mandatory candidate review did not complete " + f"({outcome.status}: {outcome.detail}). kernel_candidates.json " + "is the unreviewed deterministic result; a wrongly resolved " + "kernel will not have been caught." + ), + } + ) + _note(f"candidate review failed ({outcome.status}): {outcome.detail}") + return artifacts + + drifted = fingerprint_drift(before, source_fingerprint(source_paths)) + if drifted: + warnings.append( + { + "code": "candidate_review_touched_source", + "severity": "error", + "paths": drifted[:16], + "message": ( + "The candidate review session modified framework source " + f"({len(drifted)} file(s)); its revisions were discarded. " + "Any benchmark run after this point would measure an " + "unrecorded edit." + ), + } + ) + _note(f"candidate review discarded: it modified {len(drifted)} source file(s)") + return artifacts + + op_cat_map = load_op_category_map(tracelens_dir / "perf_report_csvs") + before_state = { + str(c.get("kernel_id") or ""): str(c.get("source_file") or "") for c in candidates + } + protected_ids = { + str(c.get("kernel_id") or "") + for c in candidates + if isinstance(c, dict) and _is_curated_resolution(c) + } + notes = apply_revisions( + candidates, + outcome.revisions, + framework_roots=KNOWN_SEARCH_ROOTS, + protected_ids=protected_ids, + ) + changed = 0 + for item in candidates: + if not isinstance(item, dict): + continue + kernel_id = str(item.get("kernel_id") or "") + if str(item.get("source_file") or "") == before_state.get(kernel_id) and not any( + item.get(key) is not None for key in _REVIEW_STAGED_PROPOSALS + ): + continue + _rederive_after_review(item, op_cat_map) + changed += 1 + + revisions_path = run_dir / REVISIONS_FILENAME + atomic_write_json( + revisions_path, + { + "status": outcome.status, + "revisions": outcome.revisions, + "applied_notes": notes, + "candidates_changed": changed, + "raw_candidates": str(raw_path), + }, + ) + artifacts["kernel_candidates_revisions"] = str(revisions_path) + _note(f"candidate review applied {changed} change(s) over {len(outcome.revisions)} revision(s)") + for line in notes: + _note(f" review: {line}") + return artifacts + + def write_reports( run_dir: Path, *, @@ -7486,6 +7745,25 @@ def main() -> int: orchestrator_error = "" # Structured trace-health findings surfaced to the Coordinator. trace_health_warnings: list[dict[str, Any]] = [] + # Without a single searchable root every kernel resolves to "" and the whole + # run reports zero routable candidates -- a host misconfiguration that reads + # exactly like a trace with nothing worth optimizing. Say so up front. + if not KNOWN_SEARCH_ROOTS: + trace_health_warnings.append( + { + "code": "no_framework_source_root", + "severity": "error", + "packages": list(_KERNEL_SOURCE_PACKAGES), + "message": ( + "No framework source root exists on this host (looked for " + f"{', '.join(_KERNEL_SOURCE_PACKAGES)}). Source resolution " + "cannot grep anything, so every hot kernel will be reported " + "as non-routable and kernel-opt will have nothing to " + "dispatch. Install the framework in this interpreter's " + "environment or point $FRAMEWORK_REPO_PATH at its checkout." + ), + } + ) try: update_status( @@ -8080,7 +8358,6 @@ def _collect(prefix: str) -> list[Path]: trace_files=trace_files, log_path=log_path, source_resolution_out=(run_dir / _SOURCE_RESOLUTION_NAME), - allow_model_tiers=False, model_name=args.model_name, ) append_log( @@ -8347,11 +8624,7 @@ def _collect(prefix: str) -> list[Path]: log_path, "dry-run: parsing raw trace for hot kernels (production code path raises here — see #203)", ) - candidates = analyze_trace_files( - trace_files, - args.top_k, - allow_model_tiers=not use_deterministic, - ) + candidates = analyze_trace_files(trace_files, args.top_k) else: raise RuntimeError( "No hot-kernel candidates produced by any TraceLens " @@ -8373,12 +8646,21 @@ def _collect(prefix: str) -> list[Path]: framework=args.framework or "", model_name=args.model_name or "", log_path=log_path, - allow_review=False, ) if source_resolution_path.is_file(): artifacts["kernel_source_resolution"] = str( source_resolution_path ) + if not use_deterministic: + artifacts.update( + run_candidate_review_stage( + run_dir, + candidates=candidates, + args=args, + log_path=log_path, + trace_health_warnings=trace_health_warnings, + ) + ) artifacts.update( write_reports( run_dir, diff --git a/src/hyperloom/common/kernel_shape_contract.py b/src/hyperloom/common/kernel_shape_contract.py index 31a275b208..d605abfff5 100644 --- a/src/hyperloom/common/kernel_shape_contract.py +++ b/src/hyperloom/common/kernel_shape_contract.py @@ -5,9 +5,25 @@ from __future__ import annotations -# Operand-dim provenances trusted by the kernel-opt dispatch gate. -DISPATCHABLE_SHAPE_PROVENANCE = frozenset({"torch_trace", "capture_backfill", "tuning_csv"}) +# Operand dims read out of the trace. +MEASURED_SHAPE_PROVENANCE = frozenset({"torch_trace", "capture_backfill", "tuning_csv"}) + +# Operand dims the candidate-review session supplied. Split by how it got them, +# because the two are worth different confidence when a tuned kernel later fails +# to move end-to-end throughput: a backfill is a recorded shape the deterministic +# lookup merely failed to join to this row, whereas a derivation is arithmetic +# over the model config and serving arguments and can be wrong in ways nothing +# detects until the integration benchmark runs. +REVIEW_BACKFILL_PROVENANCE = "review_backfill" +REVIEW_DERIVED_PROVENANCE = "review_derived" +REVIEW_SHAPE_PROVENANCE = frozenset({REVIEW_BACKFILL_PROVENANCE, REVIEW_DERIVED_PROVENANCE}) + +# Provenances the kernel-opt dispatch gate accepts. Review-supplied dims are +# admitted deliberately: under CUDA graph capture a replay has no cpu_op parent, +# so the trace records no arguments at all for the hottest kernels, and refusing +# the review's answer does not fall back to a measured shape -- it falls back to +# the tuning backend inventing one with no view of the serving configuration. +DISPATCHABLE_SHAPE_PROVENANCE = MEASURED_SHAPE_PROVENANCE | REVIEW_SHAPE_PROVENANCE # Alias used by the kernel-opt predispatch validator. ALLOWED_SHAPE_PROVENANCE = DISPATCHABLE_SHAPE_PROVENANCE - diff --git a/src/hyperloom/common/tests/test_kernel_shape_contract.py b/src/hyperloom/common/tests/test_kernel_shape_contract.py index 2c9f144cd7..726731d8ce 100644 --- a/src/hyperloom/common/tests/test_kernel_shape_contract.py +++ b/src/hyperloom/common/tests/test_kernel_shape_contract.py @@ -6,6 +6,8 @@ from hyperloom.common.kernel_shape_contract import ( ALLOWED_SHAPE_PROVENANCE, DISPATCHABLE_SHAPE_PROVENANCE, + MEASURED_SHAPE_PROVENANCE, + REVIEW_SHAPE_PROVENANCE, ) @@ -19,3 +21,20 @@ def test_capture_backfill_is_dispatchable(): def test_geometry_provenance_not_dispatchable(): assert "launch_grid" not in DISPATCHABLE_SHAPE_PROVENANCE + + +def test_reviewed_dims_are_dispatchable(): + """Under graph capture the trace records no arguments for the hottest + kernels, so refusing the review's dims does not fall back to a measured + shape -- it falls back to the tuning backend inventing one. + """ + assert REVIEW_SHAPE_PROVENANCE + assert REVIEW_SHAPE_PROVENANCE <= DISPATCHABLE_SHAPE_PROVENANCE + + +def test_measured_and_reviewed_stay_distinguishable(): + """Collapsing the two would remove the only signal that says whether a + disappointing end-to-end result is worth blaming on the shape. + """ + assert not MEASURED_SHAPE_PROVENANCE & REVIEW_SHAPE_PROVENANCE + assert DISPATCHABLE_SHAPE_PROVENANCE == MEASURED_SHAPE_PROVENANCE | REVIEW_SHAPE_PROVENANCE diff --git a/src/hyperloom/inference_optimizer/cli/__init__.py b/src/hyperloom/inference_optimizer/cli/__init__.py index 78010bcdc1..d125dafe87 100644 --- a/src/hyperloom/inference_optimizer/cli/__init__.py +++ b/src/hyperloom/inference_optimizer/cli/__init__.py @@ -2222,15 +2222,6 @@ async def _run_optimize(args: argparse.Namespace) -> int: print("Recipe sediment : ENABLED (KEEP/REVERT provenance written to persistent recipe)") else: print("Recipe sediment : DISABLED (--no-recipe-sediment)") - from hyperloom.orchestrator.kernel.request_handlers import set_allow_empty_kernel_shape - - allow_empty_kernel_shape = bool(getattr(args, "allow_empty_kernel_shape", False)) - set_allow_empty_kernel_shape(allow_empty_kernel_shape) - if allow_empty_kernel_shape: - print("Kernel shape : empty-shape dispatch ALLOWED (--allow-empty-kernel-shape)") - else: - print("Kernel shape : non-empty trace shape REQUIRED for kernel-opt dispatch") - # Resolve critic backend + runtime root before _build_backends; abort rc=2 if --critic-agent runtime unreachable. critic_choice = _resolve_critic_choice(args) if critic_choice == "mock" and args.critic_protocol != "auto": diff --git a/src/hyperloom/inference_optimizer/cli/parser.py b/src/hyperloom/inference_optimizer/cli/parser.py index 6de63ab826..9816f65313 100644 --- a/src/hyperloom/inference_optimizer/cli/parser.py +++ b/src/hyperloom/inference_optimizer/cli/parser.py @@ -1109,17 +1109,6 @@ def _build_parser() -> argparse.ArgumentParser: # Integration toggles. Roofline refresh is unconditional (fires at PRELUDE # and every 10% cumulative_gain_validated crossing). - opt.add_argument( - "--allow-empty-kernel-shape", - dest="allow_empty_kernel_shape", - action="store_true", - default=False, - help="Escape hatch (default off): allow kernel optimization to " - "dispatch a candidate with no trace-anchored shape. Normally " - "a shapeless candidate is rejected with a structured error so " - "the run returns to ``trace_analyze`` instead of burning a " - "kernel-optimization budget on an unanchored kernel.", - ) opt.add_argument( "--enable-roofline", dest="enable_roofline", diff --git a/src/hyperloom/inference_optimizer/tests/test_coordinator_gemm_promote_units.py b/src/hyperloom/inference_optimizer/tests/test_coordinator_gemm_promote_units.py index eab1443462..9423f3c29d 100644 --- a/src/hyperloom/inference_optimizer/tests/test_coordinator_gemm_promote_units.py +++ b/src/hyperloom/inference_optimizer/tests/test_coordinator_gemm_promote_units.py @@ -2455,7 +2455,7 @@ async def _append_and_seq(*_args, **_kwargs): coord.phase_kernel._geak_enabled = lambda: False coord._gemm_tuning_required_before_kernel_opt = lambda: True coord.phase_machine._record_phase_entry_evidence = lambda **_kwargs: None - coord.phase_kernel._should_continue_kernel_after_gemm = lambda: False + coord.phase_kernel._kernel_opt_work_remains = lambda: False async def _noop(*_args, **_kwargs): return None @@ -2536,7 +2536,7 @@ async def _noop(*_args, **_kwargs): coord.phase_machine._kernel_enabled = lambda: True coord.phase_kernel._geak_enabled = lambda: False coord.phase_machine._record_phase_entry_evidence = lambda **_kwargs: None - coord.phase_kernel._should_continue_kernel_after_gemm = lambda: False + coord.phase_kernel._kernel_opt_work_remains = lambda: False coord.phase_kernel._maybe_reprofile_for_kernel = _noop monkeypatch.setattr(krh_mod, "_resolve_gemm_tuning_backend", lambda _p: "forge") diff --git a/src/hyperloom/inference_optimizer/tests/test_coordinator_sync_helpers_coverage_unit.py b/src/hyperloom/inference_optimizer/tests/test_coordinator_sync_helpers_coverage_unit.py index 09b9214f3d..34f8f6a3ce 100644 --- a/src/hyperloom/inference_optimizer/tests/test_coordinator_sync_helpers_coverage_unit.py +++ b/src/hyperloom/inference_optimizer/tests/test_coordinator_sync_helpers_coverage_unit.py @@ -520,9 +520,9 @@ def test_gemm_tuning_required_before_kernel_opt(coord: Coordinator, monkeypatch) assert coord._gemm_tuning_required_before_kernel_opt() is False -def test_should_continue_kernel_after_gemm(coord: Coordinator) -> None: +def test_kernel_opt_work_remains(coord: Coordinator) -> None: coord.shared_state.continue_kernel_after_gemm = False - assert coord._should_continue_kernel_after_gemm() is False + assert coord._kernel_opt_work_remains() is False # -- canonical id helpers -------------------------------------------------- diff --git a/src/hyperloom/inference_optimizer/tests/test_decision_framework.py b/src/hyperloom/inference_optimizer/tests/test_decision_framework.py index a95d0e112f..1c505265a2 100644 --- a/src/hyperloom/inference_optimizer/tests/test_decision_framework.py +++ b/src/hyperloom/inference_optimizer/tests/test_decision_framework.py @@ -415,7 +415,8 @@ async def test_run_optimization_handler_forwards_verification_evidence( "kernel_id": "k006", "name": "aiter_native_kernel", "source_file": "/sgl-workspace/aiter/csrc/kernels/rmsnorm_quant_kernels.cu", - "reusable_native_kernel": true + "reusable_native_kernel": true, + "gpu_pct": 12.5 } ] } diff --git a/src/hyperloom/inference_optimizer/tests/test_framework_paths_units.py b/src/hyperloom/inference_optimizer/tests/test_framework_paths_units.py index 02cf0ce5ea..35632d4d03 100644 --- a/src/hyperloom/inference_optimizer/tests/test_framework_paths_units.py +++ b/src/hyperloom/inference_optimizer/tests/test_framework_paths_units.py @@ -169,6 +169,59 @@ def test_source_file_allowlist_excludes_flydsl(self, monkeypatch): assert not any("flydsl" in root.lower() for root in allowlist) +class TestResolveKernelSearchRoots: + def test_drops_roots_that_do_not_exist(self, monkeypatch, tmp_path): + """A pinned root that no longer exists must not reach the caller. + + Grepping an absent directory yields no hits, which is indistinguishable + from a kernel whose source is genuinely absent -- the exact failure that + silently emptied kernel-opt's candidate list. + """ + present = tmp_path / "vllm" + present.mkdir() + monkeypatch.setattr( + fp, + "_discover_installed_framework_roots", + lambda: (f"{present}/", "/gone/aiter/"), + ) + monkeypatch.setattr(fp, "_DEFAULT_SOURCE_ROOTS", ()) + monkeypatch.setattr(fp, "resolve_flydsl_source_roots", lambda: ()) + assert fp.resolve_kernel_search_roots() == (f"{present}/",) + + def test_excludes_bare_site_packages_parents(self, monkeypatch, tmp_path): + """Only package dirs, never the whole site-packages tree. + + The allowlist reports the parent so an editability check can contain any + installed file; grepping it would scan every wheel on the host. + """ + parent = tmp_path / "dist-packages" + (parent / "vllm").mkdir(parents=True) + monkeypatch.setattr(fp, "_discover_installed_package_roots", lambda: (f"{parent}/",)) + monkeypatch.setattr( + fp, "_discover_installed_framework_roots", lambda: (f"{parent}/vllm/",) + ) + roots = fp.resolve_kernel_search_roots() + assert f"{parent}/vllm/" in roots + assert f"{parent}/" not in roots + + def test_empty_when_nothing_is_installed(self, monkeypatch): + """No searchable root is reported as such, not as a silent success.""" + monkeypatch.setattr(fp, "_discover_installed_framework_roots", lambda: ()) + monkeypatch.setattr(fp, "_discover_scriptable_repo_roots", lambda: ()) + monkeypatch.setattr(fp, "_discover_explicit_framework_root", lambda: ()) + monkeypatch.setattr(fp, "_DEFAULT_SOURCE_ROOTS", ("/gone/vllm/",)) + monkeypatch.setattr(fp, "resolve_flydsl_source_roots", lambda: ("/gone/flydsl/",)) + assert fp.resolve_kernel_search_roots() == () + + def test_includes_explicit_framework_checkout(self, monkeypatch, tmp_path): + """An editable checkout is invisible to importlib; the env var finds it.""" + checkout = tmp_path / "my-vllm" + checkout.mkdir() + monkeypatch.setattr(fp, "_discover_installed_framework_roots", lambda: ()) + monkeypatch.setenv(fp.GENERIC_FRAMEWORK_ROOT_ENV, str(checkout)) + assert f"{checkout}/" in fp.resolve_kernel_search_roots() + + class TestFlydslExtraSourceDirs: def test_lists_only_roots_that_exist(self, monkeypatch, tmp_path): monkeypatch.setenv("FLYDSL_ROOT", str(tmp_path / "missing")) diff --git a/src/hyperloom/inference_optimizer/tests/test_kernel_request_handlers_helpers.py b/src/hyperloom/inference_optimizer/tests/test_kernel_request_handlers_helpers.py index d18d71a71f..4ee82f5ead 100644 --- a/src/hyperloom/inference_optimizer/tests/test_kernel_request_handlers_helpers.py +++ b/src/hyperloom/inference_optimizer/tests/test_kernel_request_handlers_helpers.py @@ -293,3 +293,79 @@ def test_in_flight_kernel_ids_scans_running(tmp_path: Path) -> None: # malformed -> skipped (status_dir / "ko-4.json").write_text("{bad", encoding="utf-8") assert krh._in_flight_kernel_ids(tmp_path) == {"k7", "k8"} + + +# -- unattempted_skip_reason / gate-rejected dispatch ---------------------- +def test_unattempted_skip_reason_covers_the_bookkeeping_reasons() -> None: + """Only reasons meaning "no backend ran" count as unattempted.""" + assert krh.unattempted_skip_reason("below_min_gpu_pct=5.0") + assert krh.unattempted_skip_reason("group_exhausted") + assert krh.unattempted_skip_reason("opfanout_merged_into=k002") + assert not krh.unattempted_skip_reason("") + assert not krh.unattempted_skip_reason("non_reusable_kernel") + + +def test_batch_candidates_reports_its_skip_reasons(tmp_path: Path) -> None: + """The filter's reasons reach the caller, not just the log.""" + artifact = tmp_path / "kernel_candidates.json" + artifact.write_text( + json.dumps( + { + "hot_kernels": [ + { + "kernel_id": "k001", + "name": "cold_kernel", + "gpu_pct": 1.0, + "reusable_native_kernel": True, + "source_file": "/pkg/k.py", + } + ] + } + ), + encoding="utf-8", + ) + skipped: dict[str, str] = {} + selected = krh._batch_kernel_candidates( + {"candidates_path": str(artifact)}, + skipped_out=skipped, + ) + assert selected == [] + assert skipped["k001"].startswith("below_min_gpu_pct") + + +def test_gate_rejected_named_kernel_is_skipped_not_failed(tmp_path: Path) -> None: + """A threshold is not an optimization failure. + + Recording one spends the source's retry quota on a decision no backend made, + and the report then explains a technical failure that never happened. + """ + import asyncio + + artifact = tmp_path / "kernel_candidates.json" + artifact.write_text( + json.dumps( + { + "hot_kernels": [ + { + "kernel_id": "k001", + "name": "cold_kernel", + "gpu_pct": 1.0, + "reusable_native_kernel": True, + "source_file": "/pkg/k.py", + } + ] + } + ), + encoding="utf-8", + ) + session = tmp_path / "session" + session.mkdir() + out = asyncio.run( + krh.run_optimization_handler( + {"candidates_path": str(artifact), "kernel_id": "k001"}, + session_dir=session, + ) + ) + assert out["status"] == "skipped" + assert out["reason"].startswith("below_min_gpu_pct") + assert out["kernel_id"] == "k001" diff --git a/src/hyperloom/inference_optimizer/tests/test_kernel_shape_validation.py b/src/hyperloom/inference_optimizer/tests/test_kernel_shape_validation.py index c1c3635734..cbbe8ecb54 100644 --- a/src/hyperloom/inference_optimizer/tests/test_kernel_shape_validation.py +++ b/src/hyperloom/inference_optimizer/tests/test_kernel_shape_validation.py @@ -1,10 +1,12 @@ # SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. # SPDX-License-Identifier: MIT -"""Kernel-opt dispatch shape / path validation. +"""Kernel-opt dispatch shape-provenance / path validation. -A candidate may only dispatch with a non-empty trace-anchored shape and -existing source/workspace paths; ``dry_run`` and an escape flag bypass the gate. +A shapeless candidate dispatches: the backend's driver preparation recovers the +operand dims the trace never recorded. Provenance is validated only for a shape +that is present, and source/workspace paths must exist; ``dry_run`` bypasses all +of it. """ from __future__ import annotations @@ -27,12 +29,40 @@ def _candidate(**over): return base -def test_empty_shape_is_rejected(tmp_path: Path): +def test_empty_shape_dispatches(tmp_path: Path): + # A graph-launched kernel records no cpu_op parent, so the trace carries no + # argument dims for it. Refusing the dispatch does not produce a measured + # shape, it only makes the hottest kernels of a captured model permanently + # unoptimizable. payload = {"kernel_id": "k001", "candidate": _candidate(shapes=[])} - res = krh._validate_kernel_shape_and_paths(payload, session_dir=tmp_path) - assert res is not None - assert res["error_class"] == "empty_kernel_shape" - assert res["status"] == "failed" + assert ( + krh._validate_kernel_shape_and_paths( + payload, + session_dir=tmp_path, + ) + is None + ) + + +@pytest.mark.parametrize("provenance", ["unresolved", "launch_grid", "tile_name"]) +def test_empty_shape_dispatches_whatever_provenance_says( + tmp_path: Path, + provenance: str, +): + # On a shapeless row the marker names why the dims are absent, so it must + # not be read as an untrusted operand dim -- that would close the removed + # empty-shape gate from the provenance side. + payload = { + "kernel_id": "k001", + "candidate": _candidate(shapes=[], shape_provenance=provenance), + } + assert ( + krh._validate_kernel_shape_and_paths( + payload, + session_dir=tmp_path, + ) + is None + ) def test_non_empty_shape_passes(tmp_path: Path): @@ -111,46 +141,6 @@ def test_existing_source_path_passes(tmp_path: Path): ) -def test_escape_flag_allows_empty_shape(tmp_path: Path): - payload = { - "kernel_id": "k001", - "allow_empty_kernel_shape": True, - "candidate": _candidate(shapes=[]), - } - assert ( - krh._validate_kernel_shape_and_paths( - payload, - session_dir=tmp_path, - ) - is None - ) - - -def test_escape_env_no_longer_allows_empty_shape(tmp_path: Path, monkeypatch): - monkeypatch.setenv("HYPERLOOM_ALLOW_EMPTY_KERNEL_SHAPE", "1") - payload = {"kernel_id": "k001", "candidate": _candidate(shapes=[])} - krh.set_allow_empty_kernel_shape(False) - result = krh._validate_kernel_shape_and_paths( - payload, - session_dir=tmp_path, - ) - assert result is not None - assert result["error_class"] == "empty_kernel_shape" - - -def test_escape_process_flag_allows_empty_shape(tmp_path: Path): - krh.set_allow_empty_kernel_shape(True) - payload = {"kernel_id": "k001", "candidate": _candidate(shapes=[])} - assert ( - krh._validate_kernel_shape_and_paths( - payload, - session_dir=tmp_path, - ) - is None - ) - krh.set_allow_empty_kernel_shape(False) - - def test_dry_run_bypasses_validation(tmp_path: Path): payload = { "kernel_id": "k001", @@ -168,7 +158,10 @@ def test_dry_run_bypasses_validation(tmp_path: Path): @pytest.mark.asyncio -async def test_run_optimization_single_rejects_empty_shape(tmp_path: Path): +async def test_run_optimization_single_passes_the_shape_gate(tmp_path: Path): + # The shapeless candidate reaches the path check instead of being turned + # away for its missing dims, so the failure reported is the one the payload + # actually has. payload = { "kernel_id": "k001", "source_file": "/sgl-workspace/sglang/kernels/x.py", @@ -183,7 +176,7 @@ async def test_run_optimization_single_rejects_empty_shape(tmp_path: Path): } res = await krh._run_optimization_single(payload, session_dir=tmp_path) assert res["status"] == "failed" - assert res["error_class"] == "empty_kernel_shape" + assert res["error_class"] == "missing_source_path" def test_finalize_candidates_stamps_trace_provenance(): diff --git a/src/hyperloom/inference_optimizer/tests/test_prelude_roofline.py b/src/hyperloom/inference_optimizer/tests/test_prelude_roofline.py index 9135d79377..03300d4537 100644 --- a/src/hyperloom/inference_optimizer/tests/test_prelude_roofline.py +++ b/src/hyperloom/inference_optimizer/tests/test_prelude_roofline.py @@ -370,6 +370,180 @@ async def _skip_reprofile() -> None: assert fusion_calls == 1 +@pytest.mark.asyncio +async def test_on_enter_kernel_skips_gemm_but_still_dispatches_kernel_opt( + coord: Coordinator, monkeypatch +): + """The phase dispatches its own kernel_opt on both entry routes. + + The dispatch sat on the GEMM route alone, so skipping GEMM tuning removed + the phase's source-level kernel work too -- two unrelated settings, with + nothing in the log connecting them. A run then held eight routable + candidates, cleared the dispatch floor, and reached SWEEP having optimized + nothing, because the only remaining path was an orchestration request that + was never made. + """ + monkeypatch.setenv("INFERENCE_OPTIMIZER_SKIP_GEMM_TUNING", "1") + monkeypatch.setattr(coord.phase_machine, "_kernel_enabled", lambda: True) + monkeypatch.setattr(coord.phase_kernel, "_geak_enabled", lambda: False) + monkeypatch.setattr(coord.phase_kernel, "_fusion_required_before_kernel_opt", lambda: False) + assert coord._gemm_tuning_required_before_kernel_opt() is False + + dispatched = 0 + + async def _skip_reprofile() -> None: + return None + + async def _dispatch() -> None: + nonlocal dispatched + dispatched += 1 + + monkeypatch.setattr(coord.phase_kernel, "_maybe_reprofile_for_kernel", _skip_reprofile) + monkeypatch.setattr(coord.phase_kernel, "_kernel_opt_work_remains", lambda: True) + monkeypatch.setattr(coord.phase_kernel, "_run_kernel_opt_entry_batch", _dispatch) + + await coord._on_enter_kernel(from_phase="EXPLORE") + + assert dispatched == 1 + + +@pytest.mark.asyncio +async def test_kernel_entry_does_not_dispatch_without_untried_candidates( + coord: Coordinator, monkeypatch +): + """Nothing routable left is the one reason to hand the phase back.""" + monkeypatch.setenv("INFERENCE_OPTIMIZER_SKIP_GEMM_TUNING", "1") + monkeypatch.setattr(coord.phase_machine, "_kernel_enabled", lambda: True) + monkeypatch.setattr(coord.phase_kernel, "_geak_enabled", lambda: False) + monkeypatch.setattr(coord.phase_kernel, "_fusion_required_before_kernel_opt", lambda: False) + + dispatched = 0 + + async def _skip_reprofile() -> None: + return None + + async def _dispatch() -> None: + nonlocal dispatched + dispatched += 1 + + monkeypatch.setattr(coord.phase_kernel, "_maybe_reprofile_for_kernel", _skip_reprofile) + monkeypatch.setattr(coord.phase_kernel, "_kernel_opt_work_remains", lambda: False) + monkeypatch.setattr(coord.phase_kernel, "_run_kernel_opt_entry_batch", _dispatch) + + await coord._on_enter_kernel(from_phase="EXPLORE") + + assert dispatched == 0 + + +def test_a_trace_recorded_with_task_params_is_not_stale(coord: Coordinator): + """The two writers of ``last_profile_workload`` disagree by construction. + + The roofline path records through ``record_profile_workload(task_params)`` + and fills ``server_args`` / ``extra_envs``; the kernel-entry path records + through ``profile_workload_context()`` and leaves them empty. Comparing the + whole dict therefore reported a change on every first KERNEL entry -- a full + re-profile plus a second TraceLens pass, with the serving configuration + provably unchanged -- and then stopped, because the re-profile it forced had + rewritten the record in the other writer's shape. + """ + state = coord.shared_state + state.current_best = { + "extra_server_args": "--block-size 128 --enable-expert-parallel", + "extra_envs": {"VLLM_ROCM_USE_AITER": "1"}, + } + state.last_profile_status = "succeeded" + state.last_profile_workload = state.profile_workload_context( + { + "base_extra_args": "--block-size 128 --enable-expert-parallel", + "base_extra_envs": {"VLLM_ROCM_USE_AITER": "1"}, + } + ) + # The record and a freshly built context differ, exactly as in production. + assert state.last_profile_workload != state.profile_workload_context() + + assert coord.phase_kernel._profile_workload_changed() is False + + +def test_a_trace_of_a_different_workload_is_still_stale(coord: Coordinator): + """Only the parameterization is forgiven; the workload itself still counts.""" + state = coord.shared_state + state.last_profile_status = "succeeded" + state.last_profile_workload = state.profile_workload_context() + assert coord.phase_kernel._profile_workload_changed() is False + + state.isl = int(state.isl or 0) + 4096 + + assert coord.phase_kernel._profile_workload_changed() is True + + +def _recorded_under(state, *, server_args: str, envs: dict) -> None: + """Record a profile the way the roofline path does: with the task params.""" + state.current_best = {"extra_server_args": server_args, "extra_envs": dict(envs)} + state.last_profile_status = "succeeded" + state.last_profile_workload = state.profile_workload_context( + {"base_extra_args": server_args, "base_extra_envs": dict(envs)} + ) + + +def _reprofiles(coord: Coordinator) -> bool: + """Whether the two staleness checks together call for a re-profile.""" + phase = coord.phase_kernel + signature = phase._current_profile_config_signature() + return phase._profile_config_changed(signature) or phase._profile_workload_changed() + + +_BASE_ARGS = "--block-size 128 --enable-expert-parallel" +_BASE_ENVS = {"VLLM_ROCM_USE_AITER": "1"} + + +@pytest.mark.parametrize( + ("label", "mutate", "expected"), + [ + ("nothing moved", lambda s: None, False), + ( + "explore added a server arg", + lambda s: s.current_best.__setitem__( + "extra_server_args", _BASE_ARGS + " --max-num-batched-tokens 16384" + ), + True, + ), + ( + "explore added an env", + lambda s: s.current_best.__setitem__( + "extra_envs", {**_BASE_ENVS, "VLLM_ROCM_USE_AITER_MOE": "1"} + ), + True, + ), + ( + "an env changed value", + lambda s: s.current_best.__setitem__( + "extra_envs", {**_BASE_ENVS, "VLLM_ROCM_USE_AITER": "0"} + ), + True, + ), + ("the workload changed", lambda s: setattr(s, "isl", int(s.isl or 0) + 4096), True), + ], +) +def test_serving_config_changes_still_force_a_reprofile( + coord: Coordinator, label, mutate, expected +): + """Forgiving the parameterization must not forgive a real config change. + + A configuration EXPLORE found and integrated changes which kernels run, so a + trace taken before it is genuinely stale. Those changes reach + ``_profile_config_changed``, which reads them from ``current_best`` on both + sides; only the recording-shape mismatch was taken out of + ``_profile_workload_changed``. This pins the boundary between the two. + """ + state = coord.shared_state + _recorded_under(state, server_args=_BASE_ARGS, envs=_BASE_ENVS) + assert _reprofiles(coord) is False, "the recorded trace starts fresh" + + mutate(state) + + assert _reprofiles(coord) is expected, label + + @pytest.mark.asyncio async def test_kernel_entry_reprofile_skips_when_unchanged(coord: Coordinator): """Projected tput matching the last measured trace (cur == measured) skips the reprofile.""" diff --git a/src/hyperloom/inference_optimizer/tests/test_profile_and_kernel_handlers.py b/src/hyperloom/inference_optimizer/tests/test_profile_and_kernel_handlers.py index 93f906bab0..70b674cf9e 100644 --- a/src/hyperloom/inference_optimizer/tests/test_profile_and_kernel_handlers.py +++ b/src/hyperloom/inference_optimizer/tests/test_profile_and_kernel_handlers.py @@ -4515,27 +4515,73 @@ def test_batch_candidates_default_min_gpu_pct_matches_sharedstate_gate( session_dir, _candidates_factory, ): - """``_batch_kernel_candidates`` default (10.0) must match ``SharedState.untried_hot_reusable_kernels``'s gate so a sub-threshold kernel can't sneak in via task_group fallback. - - k006/k008 straddle the shared ``_DEFAULT_HOT_KERNEL_MIN_GPU_PCT`` (10.0) by a - hair, so the two gates drifting apart in either direction fails this test. + """The dispatch batch and the phase-advance gate must apply the same GPU floor. + + ``_batch_kernel_candidates`` decides what a dispatch actually runs; + ``SharedState.untried_hot_reusable_kernels`` decides whether KERNEL still + owes work and what the report calls unattempted. Drift either way is a live + defect: a lower batch floor advances the phase past kernels it would still + dispatch, a lower state floor holds the phase open on kernels no dispatch + will ever pick up. So assert the two selections are equal, not merely that + each contains what this test expected. + + The straddling pair is derived from the shipped default instead of written + in. Hardcoding it is what let this test go on asserting a 10% boundary at a + 5% default, and pinning the default's own value is already + ``test_untried_hot_kernels_returns_only_reusable_above_threshold``'s job. + + Every candidate gets its own source file. When two share one, the batch + side's op-fanout dedup merges the weaker away and the exclusion passes for a + reason the floor had no part in -- which is how the 10% assertion above kept + passing at a 5% default. """ - cpath = _candidates_factory( - [ - {"kernel_id": "k001", "gpu_pct": 38.0, "reusable_native_kernel": True, "source_file": "/p/moe_op.py"}, - {"kernel_id": "k006", "gpu_pct": 9.87, "reusable_native_kernel": True, "source_file": "/p/rmsnorm.py"}, - {"kernel_id": "k008", "gpu_pct": 10.13, "reusable_native_kernel": True, "source_file": "/p/rmsnorm.py"}, - ] + from hyperloom.orchestrator.state.kernel_decision_settings import ( + _DEFAULT_HOT_KERNEL_MIN_GPU_PCT as floor, ) - # Default 10.0 filters out k006 (9.87) but keeps k001 (38) and k008 (10.13). - out = krh._batch_kernel_candidates( - {"candidates_path": cpath}, - session_dir=session_dir, + from hyperloom.orchestrator.state.shared_state import SharedState + + margin = 0.13 + assert floor - margin > 0, f"floor {floor} too small to straddle by {margin}" + hot_kernels = [ + { + "kernel_id": "k001", + "gpu_pct": round(floor + 30.0, 2), + "reusable_native_kernel": True, + "source_file": "/p/moe_op.py", + }, + { + "kernel_id": "k006", + "gpu_pct": round(floor - margin, 2), + "reusable_native_kernel": True, + "source_file": "/p/rmsnorm.py", + }, + { + "kernel_id": "k008", + "gpu_pct": round(floor + margin, 2), + "reusable_native_kernel": True, + "source_file": "/p/silu_and_mul.py", + }, + ] + cpath = _candidates_factory(hot_kernels) + + state = SharedState.load_or_init(session_dir) + state.last_trace_analyze = {"hot_kernels": hot_kernels, "task_groups": []} + state.save(session_dir) + + skipped: dict[str, str] = {} + batch_ids = sorted( + c.get("kernel_id") + for c in krh._batch_kernel_candidates( + {"candidates_path": cpath}, + session_dir=session_dir, + skipped_out=skipped, + ) ) - out_ids = sorted(c.get("kernel_id") for c in out) - assert "k006" not in out_ids, out_ids - assert "k001" in out_ids - assert "k008" in out_ids + + assert batch_ids == ["k001", "k008"], (batch_ids, skipped) + # The floor, not the dedup, has to be what dropped the sub-threshold row. + assert "below_min_gpu_pct" in skipped.get("k006", ""), skipped + assert sorted(state.untried_hot_reusable_kernels()) == batch_ids def test_in_flight_kernel_ids_returns_running_only(session_dir): diff --git a/src/hyperloom/inference_optimizer/tests/test_shared_state_kernel_opt.py b/src/hyperloom/inference_optimizer/tests/test_shared_state_kernel_opt.py index 920dc098cf..59beeb49c8 100644 --- a/src/hyperloom/inference_optimizer/tests/test_shared_state_kernel_opt.py +++ b/src/hyperloom/inference_optimizer/tests/test_shared_state_kernel_opt.py @@ -517,16 +517,16 @@ def test_untried_hot_kernels_returns_only_reusable_above_threshold(state: Shared {"kernel_id": "k002", "gpu_pct": 37.0, "reusable_native_kernel": True, "source_file": "/p/moe_op.py"}, { "kernel_id": "k003", - "gpu_pct": 9.5, + "gpu_pct": 4.5, "reusable_native_kernel": True, "source_file": "/p/rmsnorm.py", - }, # just BELOW the 10% threshold + }, # just BELOW the 5% threshold { "kernel_id": "k004", - "gpu_pct": 10.5, + "gpu_pct": 5.5, "reusable_native_kernel": True, "source_file": "/p/rmsnorm.py", - }, # just ABOVE the 10% threshold + }, # just ABOVE the 5% threshold { "kernel_id": "k006", "gpu_pct": 15.0, @@ -537,7 +537,7 @@ def test_untried_hot_kernels_returns_only_reusable_above_threshold(state: Shared ) untried = state.untried_hot_reusable_kernels() assert set(untried) == {"k001", "k002", "k004"} - assert "k003" not in untried # below _DEFAULT_HOT_KERNEL_MIN_GPU_PCT (10.0) + assert "k003" not in untried # below _DEFAULT_HOT_KERNEL_MIN_GPU_PCT (5.0) assert "k006" not in untried # non-reusable @@ -545,9 +545,11 @@ def test_untried_hot_kernels_reproduces_log1_session_164910Z(state: SharedState) """Replay of a real trace: 2 of its reusable hot kernels report untried. The gpu_pct values below are verbatim from the recorded session and must NOT - be tuned to the gate. Under the 10% ``_DEFAULT_HOT_KERNEL_MIN_GPU_PCT`` only - k001 (23.7) and k002 (37.3) clear it; k004 (9.7), k005 (2.8) and k003 (1.3) - are below it and k006/k007 are non-reusable. + be tuned to the gate. Under the 5% ``_DEFAULT_HOT_KERNEL_MIN_GPU_PCT`` k001 + (23.7), k002 (37.3) and k004 (9.7) clear it; k005 (2.8) and k003 (1.3) are + below it and k006/k007 are non-reusable. k004 is the case the old 10% gate + dropped: a real hot kernel that no wrapper-free operator in this trace could + have reached. """ _set_trace( state, @@ -587,8 +589,8 @@ def test_untried_hot_kernels_reproduces_log1_session_164910Z(state: SharedState) ], ) untried = state.untried_hot_reusable_kernels() - assert set(untried) == {"k001", "k002"} - assert "k004" not in untried # 9.7% sits below the 10% gate + assert set(untried) == {"k001", "k002", "k004"} + assert "k003" not in untried # 1.3% sits below the 5% gate assert "k006" not in untried # non-reusable despite 15.6% assert untried[0] == "k002" # strongest-first diff --git a/src/hyperloom/inference_optimizer/tests/test_task_group_dispatch_accounting.py b/src/hyperloom/inference_optimizer/tests/test_task_group_dispatch_accounting.py index 2fd324a98c..17c6c13122 100644 --- a/src/hyperloom/inference_optimizer/tests/test_task_group_dispatch_accounting.py +++ b/src/hyperloom/inference_optimizer/tests/test_task_group_dispatch_accounting.py @@ -856,7 +856,7 @@ async def fake_single(payload, *, session_dir, timeout_override_sec=None): monkeypatch.setattr( krh, "_batch_kernel_candidates", - lambda _payload, session_dir: [candidate], + lambda _payload, session_dir, skipped_out=None: [candidate], ) monkeypatch.setattr(krh, "_run_optimization_single", fake_single) diff --git a/src/hyperloom/orchestrator/framework/paths.py b/src/hyperloom/orchestrator/framework/paths.py index fcd24b3916..afce53947b 100644 --- a/src/hyperloom/orchestrator/framework/paths.py +++ b/src/hyperloom/orchestrator/framework/paths.py @@ -527,6 +527,34 @@ def resolve_patch_target_roots() -> tuple[str, ...]: ) +def resolve_kernel_search_roots() -> tuple[str, ...]: + """Roots to grep when locating the source that defines a GPU kernel. + + Deliberately narrower than :func:`resolve_source_file_allowlist`, which also + reports the bare site/dist-packages parents. Those are correct for a "may + this file be edited" containment test and wrong for a recursive grep: they + pull in every installed package (torch included), which costs seconds per + keyword and matches unrelated code. + + Only roots that exist on this host are returned. An empty result therefore + means "there is nothing here to search", which a caller must surface as a + misconfiguration -- grepping absent directories yields no hits and is + indistinguishable from a kernel whose source genuinely is not present. + + Returns: + tuple[str, ...]: Existing framework package dirs, editable checkouts and + FlyDSL roots, de-duplicated in discovery order. + """ + merged = _merge_roots( + _discover_installed_framework_roots(), + _discover_scriptable_repo_roots(), + _discover_explicit_framework_root(), + _DEFAULT_SOURCE_ROOTS, + resolve_flydsl_source_roots(), + ) + return tuple(root for root in merged if Path(root.rstrip("/")).is_dir()) + + def probe_framework_source_roots_for_env() -> str: """Colon-separated roots for ``INFERENCE_OPTIMIZER_FRAMEWORK_SOURCE_ROOTS``. @@ -631,6 +659,7 @@ def source_file_candidates(value: str) -> tuple[str, ...]: __all__ = [ "probe_framework_source_roots_for_env", + "resolve_kernel_search_roots", "resolve_patch_target_roots", "resolve_rocm_hip_source_roots", "resolve_session_framework_root", diff --git a/src/hyperloom/orchestrator/kernel/request_handlers.py b/src/hyperloom/orchestrator/kernel/request_handlers.py index ba25e71de8..f8c3ddff40 100644 --- a/src/hyperloom/orchestrator/kernel/request_handlers.py +++ b/src/hyperloom/orchestrator/kernel/request_handlers.py @@ -293,7 +293,13 @@ def _reusable_source_roots() -> tuple[str, ...]: _DEFAULT_KERNEL_PHASE_BACKEND_ORDER = ("geak",) # Soft cap on concurrent kernel-backend coroutines (pin with KERNEL_OPT_MAX_PARALLEL). _DEFAULT_KERNEL_BATCH_PARALLEL = 8 -_DEFAULT_BACKEND_BUDGET_MINUTES = 60.0 +# forge-loop holds back a finalize reserve of half this window, so the figure +# here buys only half as much search as it reads. At 60 a campaign completed one +# iteration -- planning alone took 16 of its 30 usable minutes -- and terminated +# on budget_exhausted with nothing kept, which reads as "the kernel cannot be +# optimized" rather than "the kernel was tried once". 90 leaves ~45 usable +# minutes, enough for a second iteration to act on what the first measured. +_DEFAULT_BACKEND_BUDGET_MINUTES = 90.0 # Minimum wall-clock a fallback backend needs; below this the ladder stops. _KERNEL_LADDER_MIN_BACKEND_SEC = 180 # Outer subprocess cap for the whole GEMM-tuning run (all shapes/tuners); sized @@ -900,35 +906,23 @@ def _validate_reusable_native_kernel(payload: dict) -> HandlerResult | None: return None -_ALLOW_EMPTY_KERNEL_SHAPE = False - - -def set_allow_empty_kernel_shape(value: bool) -> None: - """Set the process-local empty-shape escape hatch used by CLI runs.""" - global _ALLOW_EMPTY_KERNEL_SHAPE - _ALLOW_EMPTY_KERNEL_SHAPE = bool(value) - - -def _allow_empty_kernel_shape(payload: dict) -> bool: - """Escape hatch (default off) via ``payload['allow_empty_kernel_shape']`` or the CLI process flag. - - Args: - payload: Request payload that may carry ``allow_empty_kernel_shape``. - - Returns: - ``True`` when empty kernel shapes are explicitly permitted. - """ - if bool(payload.get("allow_empty_kernel_shape")): - return True - return _ALLOW_EMPTY_KERNEL_SHAPE - - def _validate_kernel_shape_and_paths( payload: dict, *, session_dir: Path, ) -> HandlerResult | None: - """Reject a kernel-opt dispatch with no trace-anchored shape or a missing source/workspace path. + """Reject a kernel-opt dispatch with untrusted shape provenance or a missing source/workspace path. + + A shapeless candidate is NOT rejected. A graph-launched kernel records no + CPU-side parent, so the profiler strips its argument dims; refusing to + dispatch those means the hottest kernels of a captured model can never be + optimized. The invocation spec reports the absent operands under its + ``missing`` list and the backend's driver preparation recovers them from the + kernel source, the tests it names and the deployment context the spec + carries. Provenance is still checked, but only for a shape that is actually + there: on a shapeless row the marker names why the dims are absent + (``unresolved``), and reading that as an untrusted operand dim would close + the same door from the other side. Args: payload: Kernel-opt dispatch payload to validate. @@ -945,24 +939,9 @@ def _validate_kernel_shape_and_paths( kernel_id = str(payload.get("kernel_id") or "") name = str(candidate.get("name") or payload.get("kernel_name") or kernel_id) - shapes = candidate.get("shapes") - if not isinstance(shapes, list): - shapes = [] + has_shapes = bool(candidate.get("shapes")) provenance = str(candidate.get("shape_provenance") or payload.get("shape_provenance") or "").strip() - if not shapes and not _allow_empty_kernel_shape(payload): - return { - "status": "failed", - "error_class": "empty_kernel_shape", - "error": ( - "selected kernel candidate has no trace-anchored shape; " - "re-run trace_analyze to capture shapes before optimizing " - "(or pass --allow-empty-kernel-shape to override)" - ), - "kernel_id": kernel_id, - "kernel_name": name, - "shape_provenance": provenance, - } - if provenance and provenance not in _ALLOWED_SHAPE_PROVENANCE: + if has_shapes and provenance and provenance not in _ALLOWED_SHAPE_PROVENANCE: return { "status": "failed", "error_class": "untrusted_shape_provenance", @@ -5815,7 +5794,12 @@ async def run_optimization_handler( return data_guard if payload.get("_single_kernel"): return await _run_optimization_single(payload, session_dir=session_dir) - candidates = _batch_kernel_candidates(payload, session_dir=session_dir) + dispatch_skips: dict[str, str] = {} + candidates = _batch_kernel_candidates( + payload, + session_dir=session_dir, + skipped_out=dispatch_skips, + ) if len(candidates) <= 1: single_payload = dict(payload) kernel_id_pinned = False @@ -5836,12 +5820,46 @@ async def run_optimization_handler( ) else: # No routable candidate: canonicalize an aliased id against the full set. + all_candidates = _all_kernel_candidates(payload) canon = _resolve_candidate_id( single_payload.get("kernel_id"), - _all_kernel_candidates(payload), + all_candidates, ) if canon: single_payload["kernel_id"] = canon + # The filter dropped this kernel for a reason it already knows. + # When that reason means "never dispatched", say so instead of + # falling through to the validation guards: a failure recorded + # here spends the source's retry quota on a decision no backend + # made, and the report then explains a technical failure that + # never happened. + skip_reason = dispatch_skips.get(canon, "") + if unattempted_skip_reason(skip_reason): + return { + "status": "skipped", + "reason": skip_reason, + "kernel_id": canon, + "kernels_considered": len(all_candidates), + "message": ( + f"kernel {canon} was not dispatched: {skip_reason}" + ), + } + # Otherwise the guards below decide, and they need the candidate + # to report against. Without it the attempt ledger files this + # kernel under an empty source and splits its identity from the + # one a later dispatch would use. + named = next( + ( + row + for row in all_candidates + if isinstance(row, dict) and str(row.get("kernel_id") or "") == canon + ), + None, + ) + if named is not None: + single_payload.setdefault("candidate", named) + if named.get("source_file"): + single_payload.setdefault("source_file", named["source_file"]) elif not _names_specific_kernel(single_payload): # Empty eligible queue and no specific target (e.g. the post-GEMM # auto pass): finish cleanly as "skipped", not a failure. @@ -6253,10 +6271,28 @@ def _kernel_dispatch_attempt_cap(entry: dict[str, Any], *, max_failures: int) -> return _DEFAULT_KERNEL_OPT_DISPATCH_ATTEMPTS +#: Skip reasons that mean "never dispatched", not "the optimizer tried and +#: failed". A caller that named such a kernel must report it as skipped: +#: recording an attempt for it spends the source's retry quota on a decision no +#: backend ever made, and reads in the report as a technical failure when the +#: cause was a threshold or a sibling already holding the task. +_UNATTEMPTED_SKIP_PREFIXES: tuple[str, ...] = ( + "below_min_gpu_pct", + "group_exhausted", + "opfanout_merged_into", +) + + +def unattempted_skip_reason(reason: str) -> bool: + """Whether ``reason`` means the kernel was never handed to a backend.""" + return str(reason or "").startswith(_UNATTEMPTED_SKIP_PREFIXES) + + def _batch_kernel_candidates( payload: dict, *, session_dir: Path | None = None, + skipped_out: dict[str, str] | None = None, ) -> list[dict[str, Any]]: """Select the reusable native kernels to dispatch for a batch run. @@ -6592,6 +6628,8 @@ def _is_live( len(selected), skipped, ) + if skipped_out is not None: + skipped_out.update(skipped) return selected @@ -6976,15 +7014,31 @@ async def _run_optimization_single( kernel_id = payload.get("kernel_id") if not kernel_id: return {"status": "failed", "error": "missing 'kernel_id' in payload"} + # A guard result is recorded as an attempt, and the attempt ledger keys on + # kernel_id plus source_file. Carrying the source through means a rejection + # and a later real dispatch share one identity instead of splitting into an + # empty-source row nothing can reconcile. + guard_source = str( + payload.get("source_file") + or (payload.get("candidate") or {}).get("source_file") + or "" + ) + + def _with_source(guard: HandlerResult) -> HandlerResult: + """Stamp the resolved source onto a guard result that omitted it.""" + if guard_source and not guard.get("source_file"): + guard = {**guard, "source_file": guard_source} + return guard + guard = _validate_reusable_native_kernel(payload) if guard is not None: - return guard + return _with_source(guard) shape_guard = _validate_kernel_shape_and_paths( payload, session_dir=session_dir, ) if shape_guard is not None: - return shape_guard + return _with_source(shape_guard) root_err = _kernel_agent_root_error() if root_err: return {"status": "failed", "error_class": "kernel_agent_root_missing", "error": root_err} diff --git a/src/hyperloom/orchestrator/loop/coordinator.py b/src/hyperloom/orchestrator/loop/coordinator.py index fff0cb5723..19a6ddd586 100644 --- a/src/hyperloom/orchestrator/loop/coordinator.py +++ b/src/hyperloom/orchestrator/loop/coordinator.py @@ -1005,8 +1005,8 @@ def router(self) -> IntentRouter: "_replace_latest_gemm_tuning_attempt": "phase_kernel", "_gemm_e2e_candidates": "phase_kernel", "_validate_gemm_tuning_e2e": "phase_kernel", - "_should_continue_kernel_after_gemm": "phase_kernel", - "_run_kernel_opt_after_gemm": "phase_kernel", + "_kernel_opt_work_remains": "phase_kernel", + "_run_kernel_opt_entry_batch": "phase_kernel", "_current_tput_from_validated_gain": "phase_kernel", "_last_measured_roofline_tput": "phase_kernel", "_needs_roofline_for_watermark": "phase_kernel", diff --git a/src/hyperloom/orchestrator/phases/kernel.py b/src/hyperloom/orchestrator/phases/kernel.py index 091a3c97f5..f9a46f6c53 100644 --- a/src/hyperloom/orchestrator/phases/kernel.py +++ b/src/hyperloom/orchestrator/phases/kernel.py @@ -289,21 +289,37 @@ def _profile_config_changed(self, signature: str) -> bool: return previous != signature def _profile_workload_changed(self) -> bool: - """Whether the latest trace predates the active serving workload.""" + """Whether the latest trace predates the active serving workload. + + Compares only the fields that identify the profiled workload. The rest + of the context records how the profile task was parameterized, and the + two writers disagree there by construction: the roofline path records + through ``record_profile_workload(task_params)`` and fills them, while + the kernel-entry path records through ``profile_workload_context()`` and + leaves them empty. A whole-dict comparison therefore reported a change + on every first entry -- costing a full re-profile and a second TraceLens + pass, roughly fifty minutes, with the serving configuration provably + unchanged -- and then stopped reporting one, because the re-profile it + forced had rewritten the record in the other writer's shape. + + The serving configuration is not compared here; that is + :meth:`_profile_config_changed`, which reads it from ``current_best`` on + both sides and is symmetric for the same reason this now is. + """ status = str( getattr(self.shared_state, "last_profile_status", "") or "" ).strip().lower() if status and status != "succeeded": return True recorded = getattr(self.shared_state, "last_profile_workload", None) - expected = self.shared_state.profile_workload_context() if not isinstance(recorded, dict) or not recorded: return bool( getattr(self.shared_state, "last_profile_trace", "") or getattr(self.shared_state, "last_trace_analyze", None) or getattr(self.shared_state, "roofline_snapshots", None) ) - return recorded != expected + identity = self.shared_state.profile_workload_identity + return identity(recorded) != identity(self.shared_state.profile_workload_context()) async def _maybe_reprofile_for_kernel(self) -> None: """Reprofile inline when projected tput diverges from the last measured trace, so GEAK targets the live bottleneck.""" @@ -455,9 +471,7 @@ async def _on_enter_kernel(self, *, from_phase: str) -> None: await self._run_geak_kernel_phase(from_phase=from_phase) return if not self._gemm_tuning_required_before_kernel_opt(): - await self._maybe_reprofile_for_kernel() - await self._maybe_run_forge_fusion_before_kernel_opt() - await self._maybe_run_collective_before_kernel_opt() + await self._finish_kernel_entry() return # Refresh the snapshot before GEMM tuning targets the bottleneck. @@ -529,12 +543,8 @@ async def _on_enter_kernel(self, *, from_phase: str) -> None: "tuned_file": result.get("tuned_file"), }, ) - # Capture explore + GEMM-tuning gains before inline GEAK. - await self._maybe_reprofile_for_kernel() - await self._maybe_run_forge_fusion_before_kernel_opt() - await self._maybe_run_collective_before_kernel_opt() - if self._should_continue_kernel_after_gemm(): - await self._run_kernel_opt_after_gemm() + # Capture explore + GEMM-tuning gains before the entry batch. + await self._finish_kernel_entry() async def _run_bf16_dense_gemm_fallback( self, @@ -3454,8 +3464,30 @@ async def _validate_gemm_tuning_e2e(self, result: dict[str, Any]) -> None: result["micro_decision"] = "candidate_no_e2e_gain" self._replace_latest_gemm_tuning_attempt(result) - def _should_continue_kernel_after_gemm(self) -> bool: - """Decide whether to run source-level kernel_opt right after GEMM tuning. + async def _finish_kernel_entry(self) -> None: + """Close out KERNEL entry on either route: re-profile, run the + independently gated stages, then dispatch whatever kernel_opt work the + candidate table already justifies. + + The dispatch used to sit on the GEMM route alone, so skipping GEMM + tuning silently removed the phase's own kernel_opt as well. The two + settings are unrelated -- one tunes GEMM shape tables, the other + rewrites source-level kernels -- and nothing in the log connected them, + so a run could hold eight routable candidates, clear the dispatch floor, + and still reach SWEEP having optimized nothing, waiting on an + orchestration request that never came. + + What the dispatch needs is untried routable candidates. That is what it + asks for, on both routes. + """ + await self._maybe_reprofile_for_kernel() + await self._maybe_run_forge_fusion_before_kernel_opt() + await self._maybe_run_collective_before_kernel_opt() + if self._kernel_opt_work_remains(): + await self._run_kernel_opt_entry_batch() + + def _kernel_opt_work_remains(self) -> bool: + """Whether KERNEL entry should dispatch source-level kernel_opt itself. Returns: bool: ``True`` when the ``continue_kernel_after_gemm`` flag is set @@ -3465,15 +3497,21 @@ def _should_continue_kernel_after_gemm(self) -> bool: return False return bool(self.shared_state.untried_hot_reusable_kernels()) - async def _run_kernel_opt_after_gemm(self) -> None: - """Run the source-level kernel optimization batch after GEMM tuning.""" + async def _run_kernel_opt_entry_batch(self) -> None: + """Dispatch the source-level kernel optimization batch at KERNEL entry. + + No ``kernel_id`` is named, so the handler's own filter decides the set: + every candidate that clears the dispatch floor and has retries left goes + in one batch. Naming one here would put the phase back in the business + of picking, which is the part that stalls when nobody picks. + """ cached = self.shared_state.last_trace_analyze or {} candidates_path = str(cached.get("candidates_path") or "") if not candidates_path: - log.info("KERNEL entry: skip kernel_opt after GEMM; no candidates_path") + log.info("KERNEL entry: skip kernel_opt; no candidates_path") return log.info( - "KERNEL entry: continuing to source-level kernel_opt after GEMM tuning", + "KERNEL entry: dispatching the source-level kernel_opt batch", ) try: from ..kernel.request_handlers import run_optimization_handler @@ -3503,7 +3541,7 @@ async def _run_kernel_opt_after_gemm(self) -> None: "kind": "run_optimization_done", "status": result.get("status", "ok") if isinstance(result, dict) else "failed", "result": result, - "source": "kernel_entry_auto_after_gemm", + "source": "kernel_entry_auto", }, priority=1, ) diff --git a/src/hyperloom/orchestrator/prompts/prompt_builder.py b/src/hyperloom/orchestrator/prompts/prompt_builder.py index f0b9258004..90c572c83e 100644 --- a/src/hyperloom/orchestrator/prompts/prompt_builder.py +++ b/src/hyperloom/orchestrator/prompts/prompt_builder.py @@ -778,8 +778,12 @@ def _idea_generation_lines() -> list[str]: request{target_agent: 'kernel_agent', kind: 'run_optimization', params: {kernel_id: , source_file: , - candidates_path: , - budget_minutes: 60}} + candidates_path: }} + + Budget policy: DO NOT add a `budget_minutes` field. The Coordinator owns + the per-optimization wall clock and applies the same value it uses for its + own dispatch; naming one here pins the backend to this template's number + instead, which is how a raised operator budget got silently discarded. Backend policy: DO NOT add a `backends` field. Current GEAK owns the KERNEL phase by default. Forge per-kernel mode is available only when the @@ -806,7 +810,7 @@ def _idea_generation_lines() -> list[str]: **Multi-KEEP queue:** `pending_keep_kernels` (sorted strongest-first) lists queued KEEPs; integrate `[0]` each tick. Do NOT propose `report` while it is non-empty, nor while `untried_hot_reusable_kernels` - (reusable hot kernels with zero attempts and `gpu_pct >= 10%`, the + (reusable hot kernels with zero attempts and `gpu_pct >= 5%`, the default that `HYPERLOOM_KERNEL_OPT_MIN_GPU_PCT` overrides) remain — drain them with `run_optimization{candidates_path: }` (the batch handler fans out automatically). diff --git a/src/hyperloom/orchestrator/state/kernel_decision_settings.py b/src/hyperloom/orchestrator/state/kernel_decision_settings.py index 4aee429600..03aaab14c5 100644 --- a/src/hyperloom/orchestrator/state/kernel_decision_settings.py +++ b/src/hyperloom/orchestrator/state/kernel_decision_settings.py @@ -45,7 +45,12 @@ def resolve_kernel_opt_max_failures() -> int: # Holds KERNEL phase-advance open (kernel_work_pending), filters the dispatch # batch queue, and drives the advisory 'untried hot kernels' report annotation. # It does NOT block ``report``. -_DEFAULT_HOT_KERNEL_MIN_GPU_PCT = 10.0 +# Tuned for a model whose hot kernels are spread rather than concentrated. A +# 60-layer sparse-MoE decoder splits its work across so many operators that +# nothing but a graph-launch wrapper reaches double digits, so a 10% floor +# admitted no real kernel at all and left the batch dispatcher idle while the +# orchestrator picked candidates one at a time. +_DEFAULT_HOT_KERNEL_MIN_GPU_PCT = 5.0 def resolve_hot_kernel_min_gpu_pct() -> float: diff --git a/src/hyperloom/orchestrator/state/shared_state.py b/src/hyperloom/orchestrator/state/shared_state.py index bad544e0e0..a28613c5f0 100644 --- a/src/hyperloom/orchestrator/state/shared_state.py +++ b/src/hyperloom/orchestrator/state/shared_state.py @@ -53,7 +53,7 @@ import os import shlex import time -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Mapping from contextlib import contextmanager from dataclasses import asdict, dataclass, field from datetime import datetime, timezone @@ -1085,6 +1085,42 @@ class SharedState(_RenderMixin, _ExploreStateMixin): # breakdown instrumentation. Plain class attr => not serialized. _session_dir = None + #: Fields of :meth:`profile_workload_context` that say *what was profiled*. + #: The rest -- ``server_args``, ``extra_envs``, ``remove_args``, + #: ``unset_envs``, ``args_mode`` -- say how the profile task was + #: parameterized, and are only populated when the recorder had those params + #: to hand. Two call sites record the same trace differently for that reason + #: alone, so comparing them makes a perfectly fresh trace read as stale. + #: ``serving_config`` is excluded here too: it has its own comparison, which + #: comes from ``current_best`` on both sides and is therefore symmetric. + PROFILE_WORKLOAD_IDENTITY_KEYS: tuple[str, ...] = ( + "framework", + "precision", + "model_path", + "tp", + "conc", + "isl", + "osl", + "max_model_len", + ) + + @classmethod + def profile_workload_identity(cls, context: Any) -> dict[str, Any]: + """Project a workload context down to what identifies the profiled run. + + Args: + context (Any): A :meth:`profile_workload_context` result, or + anything else (treated as carrying no identity). + + Returns: + dict[str, Any]: The identity fields, missing ones included as + ``None`` so a recorded context and a freshly built one compare + equal when they describe the same workload. + """ + if not isinstance(context, Mapping): + return {} + return {key: context.get(key) for key in cls.PROFILE_WORKLOAD_IDENTITY_KEYS} + def profile_workload_context( self, overrides: dict[str, Any] | None = None,