From fbc78f4f10e0725bde2e8e86fff238190aef16b6 Mon Sep 17 00:00:00 2001 From: Vincent Ouyang Date: Tue, 20 Oct 2026 16:19:41 -0700 Subject: [PATCH 1/3] Let tasks declare configuration-only editable symbols A symbol-scoped task can now list editable symbols in configuration_symbols. Protected code may call them at launch time, so the harness guard requires each to stay an undecorated function that passes the existing configuration-factory body check. Without the declaration the guard does not restrict an editable plain function's body, so a launch-time helper could run arbitrary host code. The field must be a subset of symbols and cannot name an entrypoint. Undeclared editable functions keep the existing policy. The validator context and the agent prompt report the declared symbols. --- docs/how-to/add-task.md | 10 ++- src/harness_guard.py | 11 ++++ src/task_prompt.py | 3 + src/task_spec.py | 20 ++++-- tests/test_guard_context_clarity.py | 3 +- tests/test_harness_guard_v2.py | 94 +++++++++++++++++++++++++++++ 6 files changed, 134 insertions(+), 7 deletions(-) diff --git a/docs/how-to/add-task.md b/docs/how-to/add-task.md index f85356502..9f8d2b091 100644 --- a/docs/how-to/add-task.md +++ b/docs/how-to/add-task.md @@ -158,7 +158,11 @@ candidate: A mapping's `scope` is `file`, `symbols`, or `tree`. `symbols` requires a nonempty symbol list; `allow_new_helpers` defaults to `false` and applies only to -implementation helpers under a symbol-scoped boundary. `tree` permits an +implementation helpers under a symbol-scoped boundary. `configuration_symbols` +optionally names a subset of `symbols`, never an entrypoint, that only builds +configuration data, for example a `_launch_config` helper returning tile sizes, +`num_warps` and `num_stages`. Protected code may call such a helper and must +validate what it returns before use. `tree` permits an implementation subtree, not its tests or build/evaluation policy. Broad tree access must not override a protected harness path. Generated build artifacts are not editable source declarations. @@ -180,7 +184,9 @@ protected function definitions. Existing editable non-entrypoint configuration helpers must pass a restricted body check: configuration data/local assignments, returns, and branches using protected predicates or supported configuration operations. Initial aliases, kernel entrypoints, and newly introduced factories -do not gain permission to run at definition time. The guard reports its +do not gain permission to run at definition time. Symbols listed in +`configuration_symbols` must stay undecorated functions that pass the same body +check, because protected code calls them at launch time. The guard reports its effective policy in the validator context. It is a structural integrity check, not isolation of arbitrary Python code; reference/test manipulation remains forbidden and subject to semantic review. diff --git a/src/harness_guard.py b/src/harness_guard.py index da5b1dee0..6fa011c6a 100644 --- a/src/harness_guard.py +++ b/src/harness_guard.py @@ -379,6 +379,16 @@ def _v2_symbol_digest(path: Path, edit: EditScope, initial_names: frozenset[str] raise RuntimeError(f"Protected test/harness policy rejected {path.name}: {exc}") from exc continue kept.append(node) + # Protected host code calls declared configuration symbols at launch time, + # so each must remain a configuration factory, never arbitrary host code. + for name in edit.configuration_symbols: + node = functions.get(name) + if (node is None or node.decorator_list + or not _config_factory_body_allowed(node, bindings, protected_functions, + set(initial_names) | definition_names)): + raise RuntimeError(f"Protected test/harness policy rejected {path.name}: configuration " + f"symbol {name!r} must be an undecorated function that only builds " + "configuration data") tree.body = kept return hashlib.sha256(ast.dump(tree, include_attributes=False).encode("utf-8")).hexdigest() @@ -690,6 +700,7 @@ def _describe_v2_snapshot(snapshot: WorkspaceSnapshot) -> dict[str, object]: else "sha256_bytes"), "editable_symbols": list(edit.symbols) if edit else [], "allow_new_helpers": edit.allow_new_helpers if edit else False, + "configuration_symbols": list(edit.configuration_symbols) if edit else [], "initial_top_level_names": sorted(snapshot.initial_symbols.get(path, ())), "definition_policy": "compiler_decorators_no_test_hooks_v1" if edit else "byte_protected", "allowed_compiler_decorators": sorted(_COMPILER_DECORATORS) if edit else [], diff --git a/src/task_prompt.py b/src/task_prompt.py index 6ec68a5ec..d8161206f 100644 --- a/src/task_prompt.py +++ b/src/task_prompt.py @@ -34,6 +34,9 @@ def build_task_prompt(spec: TaskSpec, workspace: Path, *, target_gpu: str, if edit.scope == "symbols": detail = "only symbols " + ", ".join(edit.symbols) detail += "; new implementation helpers allowed" if edit.allow_new_helpers else "; no new helpers" + if edit.configuration_symbols: + detail += ("; " + ", ".join(edit.configuration_symbols) + + " may only build configuration data (validated by the protected caller)") elif edit.scope == "tree": detail = "implementation subtree; protected task/harness files remain protected" else: diff --git a/src/task_spec.py b/src/task_spec.py index 96842b52d..a0609b8cd 100644 --- a/src/task_spec.py +++ b/src/task_spec.py @@ -124,6 +124,9 @@ class EditScope: scope: str = "file" symbols: tuple[str, ...] = () allow_new_helpers: bool = False + # Editable symbols that may only build configuration data. Protected code + # may call them at launch time, so the harness guard restricts their bodies. + configuration_symbols: tuple[str, ...] = () def contains(self, path: str) -> bool: return path == self.path or ( @@ -181,7 +184,8 @@ def _candidate(value: Any) -> CandidateSpec: edits = [] for raw in obj["editable"]: edit = {"path": raw} if isinstance(raw, str) else _mapping( - raw, "candidate.editable[]", {"path", "scope", "symbols", "allow_new_helpers"}) + raw, "candidate.editable[]", {"path", "scope", "symbols", "allow_new_helpers", + "configuration_symbols"}) path = relative_path(edit.get("path"), "candidate.editable[].path") scope = _choice(edit.get("scope", "file"), ("file", "symbols", "tree"), "scope") symbols = edit.get("symbols", []) @@ -190,12 +194,16 @@ def _candidate(value: Any) -> CandidateSpec: helpers = edit.get("allow_new_helpers", False) if type(helpers) is not bool: raise TaskConfigError("allow_new_helpers must be a boolean") + configuration = edit.get("configuration_symbols", []) + if (not isinstance(configuration, list) or len(configuration) != len(set(configuration)) + or any(not isinstance(s, str) or s not in symbols for s in configuration)): + raise TaskConfigError("configuration_symbols must be unique names listed in symbols") if scope == "symbols": if not symbols or len(symbols) != len(set(symbols)): raise TaskConfigError("Symbol scope requires nonempty unique symbols") - elif "symbols" in edit or "allow_new_helpers" in edit: - raise TaskConfigError("symbols/allow_new_helpers only apply to symbol scope") - item = EditScope(path, scope, tuple(symbols), helpers) + elif {"symbols", "allow_new_helpers", "configuration_symbols"} & set(edit): + raise TaskConfigError("symbols/allow_new_helpers/configuration_symbols only apply to symbol scope") + item = EditScope(path, scope, tuple(symbols), helpers, tuple(configuration)) if any(e.contains(path) or item.contains(e.path) for e in edits): raise TaskConfigError(f"Overlapping editable declarations: {path}") edits.append(item) @@ -219,6 +227,10 @@ def _candidate(value: Any) -> CandidateSpec: if parsed in parsed_entries: raise TaskConfigError("Duplicate entrypoint") parsed_entries.append(parsed) + for edit in edits: + if any(entry.file == edit.path and entry.symbol in edit.configuration_symbols + for entry in parsed_entries): + raise TaskConfigError("A configuration symbol cannot also be an entrypoint") return CandidateSpec(language, state, initial, tuple(edits), tuple(parsed_entries)) diff --git a/tests/test_guard_context_clarity.py b/tests/test_guard_context_clarity.py index 58d5799d1..14d764c64 100644 --- a/tests/test_guard_context_clarity.py +++ b/tests/test_guard_context_clarity.py @@ -89,7 +89,8 @@ def test_raw_file_metadata_corresponds_to_byte_protection(tmp_path): assert policy == { "digest_mode": "sha256_bytes", "digest": hashlib.sha256(runner.read_bytes()).hexdigest(), - "editable_symbols": [], "allow_new_helpers": False, "initial_top_level_names": [], + "editable_symbols": [], "allow_new_helpers": False, "configuration_symbols": [], + "initial_top_level_names": [], "definition_policy": "byte_protected", "allowed_compiler_decorators": [], } runner.write_text(runner.read_text() + "# even a byte-only change is protected\n") diff --git a/tests/test_harness_guard_v2.py b/tests/test_harness_guard_v2.py index baf20e944..992c4612a 100644 --- a/tests/test_harness_guard_v2.py +++ b/tests/test_harness_guard_v2.py @@ -140,3 +140,97 @@ def test_new_import_shadow_files_follow_described_guard_policy(tmp_path, name): assert extra.exists() verify_workspace_harness(before) assert not extra.exists() + + + +LAUNCH_SOURCE = ("import math\n" + "def _launch_config(n):\n return {'BLOCK': 64, 'num_warps': 4}\n" + "def launch(n):\n cfg = _launch_config(n)\n return cfg['BLOCK'] * n\n") +LAUNCH_RETURN = " return {'BLOCK': 64, 'num_warps': 4}\n" + + +def launch_workspace(tmp_path, *, declared=True): + root = tmp_path / "workspace" + (root / "source").mkdir(parents=True) + source = root / "source/kernel.py" + source.write_text(LAUNCH_SOURCE) + edit = {"path": "source/kernel.py", "scope": "symbols", "symbols": ["_launch_config"], "allow_new_helpers": True} + if declared: + edit["configuration_symbols"] = ["_launch_config"] + config = {"schema_version": 2, "candidate": {"language": "triton", "editable": [edit]}, + "evaluation": {"runner": ["python3", "evaluate.py"]}} + (root / "config.yaml").write_text(yaml.safe_dump(config)) + (root / "evaluate.py").write_text("# protected custom evaluation runner\n") + return root, source + + +def test_declared_configuration_symbol_may_be_retuned(tmp_path): + root, source = launch_workspace(tmp_path) + before = snapshot_workspace_harness(root) + assert describe_workspace_harness(root, snapshot=before)["protected_path_policies"]["source/kernel.py"][ + "configuration_symbols"] == ["_launch_config"] + tuned = (" if n >= 4096:\n return {'BLOCK': 256, 'num_warps': 8}\n" + " block = 128\n return {'BLOCK': block, 'num_warps': 4}\n") + source.write_text(source.read_text().replace(LAUNCH_RETURN, tuned)) + verify_workspace_harness(before) + + +HOST_CODE = [ + " math.sqrt(n)\n", + " import os\n", + " for _ in range(2):\n pass\n", + " launch.cache = 1\n", + " globals()['TOLERANCE'] = 1\n", + " with open('x') as f:\n pass\n", +] + + +@pytest.mark.parametrize("host", HOST_CODE) +def test_declared_configuration_symbol_cannot_run_host_code(tmp_path, host): + root, source = launch_workspace(tmp_path) + before = snapshot_workspace_harness(root) + source.write_text(source.read_text().replace(LAUNCH_RETURN, host + LAUNCH_RETURN)) + with pytest.raises(RuntimeError, match="only builds configuration data"): + verify_workspace_harness(before) + + +@pytest.mark.parametrize("change", [ + lambda text: text.replace("def _launch_config(n):", "@staticmethod\ndef _launch_config(n):"), + lambda text: text.replace("def _launch_config(n):\n" + LAUNCH_RETURN, ""), +]) +def test_declared_configuration_symbol_must_remain_a_plain_function(tmp_path, change): + root, source = launch_workspace(tmp_path) + before = snapshot_workspace_harness(root) + source.write_text(change(source.read_text())) + with pytest.raises(RuntimeError, match="source/kernel.py|kernel.py"): + verify_workspace_harness(before) + + +def test_undeclared_editable_functions_keep_the_existing_definition_policy(tmp_path): + # Python implementation tasks let protected code call editable host code. + root, source = launch_workspace(tmp_path, declared=False) + before = snapshot_workspace_harness(root) + source.write_text(source.read_text().replace(LAUNCH_RETURN, HOST_CODE[0] + LAUNCH_RETURN)) + verify_workspace_harness(before) + + +@pytest.mark.parametrize("edit,message", [ + ({"symbols": ["compute"], "configuration_symbols": ["other"]}, "listed in symbols"), + ({"symbols": ["compute"], "configuration_symbols": "compute"}, "listed in symbols"), + ({"symbols": ["compute"], "configuration_symbols": ["compute"], "entry": True}, "entrypoint"), + ({"scope": "file", "configuration_symbols": []}, "symbol scope"), +]) +def test_configuration_symbols_schema(tmp_path, edit, message): + from src.task_spec import TaskConfigError, load_task_spec + entry = edit.pop("entry", False) + raw = {"path": "source/kernel.py", "scope": "symbols", **edit} + candidate = {"language": "triton", "initial_state": "implemented", "initial_language": "triton", + "editable": [raw]} + if entry: + candidate["entrypoints"] = [{"file": "source/kernel.py", "kind": "function", "symbol": "compute"}] + config = {"schema_version": 2, "description": "d", "candidate": candidate, + "evaluation": {"runner": ["python3", "evaluate.py"]}} + path = tmp_path / "config.yaml" + path.write_text(yaml.safe_dump(config)) + with pytest.raises(TaskConfigError, match=message): + load_task_spec(path, task_id="t") From b275555057201ff611a473d1fd1f91dc11484c02 Mon Sep 17 00:00:00 2001 From: Vincent Ouyang Date: Tue, 20 Oct 2026 16:19:42 -0700 Subject: [PATCH 2/3] Make the triton_bmm launch configuration tunable Move the wrapper's per-dtype tile sizes, num_warps and num_stages into the editable _launch_config configuration symbol. The protected wrapper accepts exactly those integer keys within bounds and still derives the grid from the validated tiles. The initial values are unchanged, so baseline launches and timing stay the same. Record exact inverses of the source and config edits so the older migration tests keep comparing the previously reviewed bytes, and extend the boundary and domain-control tests to the new editable helper. --- tasks/triton2triton/vllm/triton_bmm/README.md | 19 +- .../triton2triton/vllm/triton_bmm/config.yaml | 3 + .../vllm/triton_bmm/source/triton_bmm.py | 64 ++++--- tests/reviewed_contract_history.py | 179 ++++++++++++++++++ tests/test_bmm_domain_controls.py | 26 ++- tests/test_triton_module_edit_boundaries.py | 8 +- tests/test_triton_task_migration_v2.py | 20 +- 7 files changed, 272 insertions(+), 47 deletions(-) diff --git a/tasks/triton2triton/vllm/triton_bmm/README.md b/tasks/triton2triton/vllm/triton_bmm/README.md index 111dd1529..b2ca5b57a 100644 --- a/tasks/triton2triton/vllm/triton_bmm/README.md +++ b/tasks/triton2triton/vllm/triton_bmm/README.md @@ -49,17 +49,24 @@ imports, initialization, host dispatch and pointer preparation remain protected, along with the existing numerical checks and benchmark protocol. The structural guard enforces this boundary; it does not isolate arbitrary Python execution. -The protected wrapper also fixes the launch configuration: the grid and the tile -sizes, `num_warps` and `num_stages` it passes at launch. Optimization notes above -that mention tuning those values refer to the original whole-file contract; under -this boundary, optimize the kernel bodies and implementation helpers instead. +### Tunable launch configuration + +`_launch_config(dtype_name, batch, m, n, k)` is also editable. It returns the +`bmm_kernel` tile sizes (`BLOCK_SIZE_M`, `BLOCK_SIZE_N`, `BLOCK_SIZE_K`), +`num_warps` and `num_stages`, and may branch on the dtype name or problem size. +Its body may only build this configuration data; the harness guard rejects other +host code in it. The protected wrapper requires exactly these five integer keys, +tile sizes that are powers of two in [16, 256], `num_warps` in {1, 2, 4, 8, 16} +and `num_stages` in [1, 4], and derives the grid from the returned tiles. The +initial values are the original ones: 128/128/32 tiles, 8 warps and 2 stages for +16-bit dtypes, and 64/64/32 tiles, 4 warps and 2 stages for float32. ### Correctness-only domain controls The protected `bmm_triton` wrapper defines the public domain. It accepts two 3-D tensors with matching batch and inner dimensions and one shared dtype: -float16, bfloat16 or float32, the only keys of its launch-configuration table. +float16, bfloat16 or float32, the only dtypes it accepts. It forwards all six operand strides to the kernel without copying, so transposed, offset, padded, element-strided and batch-broadcast (stride-0) views are valid inputs. It always allocates a new contiguous `(B, M, N)` output @@ -70,7 +77,7 @@ Eleven unscored `control-domain-*` rows (`case_index` 20000-20010, checks `[correctness]`) come from the protected `_domain_controls.py` table; the adapter verifies that table and these rows against `workloads.json`. They cover bfloat16 and float32 as well as float16, and B, M, N and K values that are not -multiples of the wrapper tiles (128/128/32 for 16-bit dtypes, 64/64/32 for +multiples of the initial tiles (128/128/32 for 16-bit dtypes, 64/64/32 for float32), including `M = K = 1` and other small odd shapes. Operand layouts include transposed A, B or both, a batch-inner `(M, B, K).transpose(0, 1)` view, offset views of padded storage, element stride 2, batch stride 2 and a diff --git a/tasks/triton2triton/vllm/triton_bmm/config.yaml b/tasks/triton2triton/vllm/triton_bmm/config.yaml index 7a65f0e41..2558238e1 100644 --- a/tasks/triton2triton/vllm/triton_bmm/config.yaml +++ b/tasks/triton2triton/vllm/triton_bmm/config.yaml @@ -8,7 +8,10 @@ candidate: scope: symbols symbols: - bmm_kernel + - _launch_config allow_new_helpers: true + configuration_symbols: + - _launch_config entrypoints: - file: source/triton_bmm.py kind: function diff --git a/tasks/triton2triton/vllm/triton_bmm/source/triton_bmm.py b/tasks/triton2triton/vllm/triton_bmm/source/triton_bmm.py index 12083ba5f..6791bcb55 100644 --- a/tasks/triton2triton/vllm/triton_bmm/source/triton_bmm.py +++ b/tasks/triton2triton/vllm/triton_bmm/source/triton_bmm.py @@ -101,6 +101,44 @@ def bmm_kernel( tl.store(c_ptrs, c, mask=c_mask) +def _launch_config(dtype_name, batch, m, n, k): + """Return the bmm_kernel tile sizes, num_warps and num_stages. + + Tune these values, optionally per dtype or problem size. The protected + wrapper validates them and derives the grid from the tile sizes. + """ + if dtype_name == "float32": + return {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, + "num_stages": 2, "num_warps": 4} + return {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, + "num_stages": 2, "num_warps": 8} + + +_DTYPE_NAMES = {torch.bfloat16: "bfloat16", torch.float16: "float16", torch.float32: "float32"} +_TILE_KEYS = ("BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K") + + +def _checked_launch_config(cfg): + """Accept only tile sizes, num_warps and num_stages within supported bounds.""" + expected = set(_TILE_KEYS) | {"num_warps", "num_stages"} + if not isinstance(cfg, dict) or set(cfg) != expected: + raise ValueError(f"Launch configuration must set exactly {sorted(expected)}") + checked = {} + for key in sorted(expected): + value = cfg[key] + if type(value) is not int: + raise ValueError(f"Launch setting {key} must be an int") + checked[key] = value + for key in _TILE_KEYS: + if not (16 <= checked[key] <= 256 and checked[key] & (checked[key] - 1) == 0): + raise ValueError(f"{key} must be a power of two in [16, 256]") + if checked["num_warps"] not in (1, 2, 4, 8, 16): + raise ValueError("num_warps must be 1, 2, 4, 8 or 16") + if not 1 <= checked["num_stages"] <= 4: + raise ValueError("num_stages must be in [1, 4]") + return checked + + def bmm_triton(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: """Wrapper for batched matrix multiplication kernel.""" assert a.ndim == 3 and b.ndim == 3, "Inputs must be 3D" @@ -114,31 +152,7 @@ def bmm_triton(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: c = torch.empty((B_size, M, N), device=a.device, dtype=dtype) - configs = { - torch.bfloat16: { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 32, - "num_stages": 2, - "num_warps": 8, - }, - torch.float16: { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 32, - "num_stages": 2, - "num_warps": 8, - }, - torch.float32: { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 32, - "num_stages": 2, - "num_warps": 4, - }, - } - - cfg = configs[dtype] + cfg = _checked_launch_config(_launch_config(_DTYPE_NAMES[dtype], B_size, M, N, K)) grid = ( B_size, triton.cdiv(M, cfg["BLOCK_SIZE_M"]) * triton.cdiv(N, cfg["BLOCK_SIZE_N"]), diff --git a/tests/reviewed_contract_history.py b/tests/reviewed_contract_history.py index 315e54a99..7df686ab6 100644 --- a/tests/reviewed_contract_history.py +++ b/tests/reviewed_contract_history.py @@ -1186,6 +1186,185 @@ def historical_attention_sampling_file(relative, text): return _exact_reviewed_file_inverse(text, record) if record is not None else text +# Tunable launch configurations: each record moves a wrapper's tuning values into +# its editable `_launch_config` helper and adds the protected validator. Full +# current/before bytes bound every inverse; the boundary inverse applies after it. +TRITON_LAUNCH_CONFIG_FILES = {'tasks/triton2triton/vllm/triton_bmm/config.yaml': {'after_sha256': 'd67c2742d3e92815508850fb6ff378fcfb60960f0d8075e98ea0585a64617d0f', + 'before_sha256': '6e04236bbc34f35a60c1dcaaa47f218ac2e42d75a4cf85778240feaa8088462e', + 'reverse_edits': [(10, + 11, + [' - _launch_config\n'], + []), + (12, + 14, + [' configuration_symbols:\n', + ' - _launch_config\n'], + [])]}, + 'tasks/triton2triton/vllm/triton_bmm/source/triton_bmm.py': {'after_sha256': '6a072c2f033459d042d0662fb5ed58b43a2d8f8c1c65f218d0be2bfa5f740c06', + 'before_sha256': 'cb7ef5165d6e7adab326435a534adf667eeeda98d497bf5f86727b15c8ef7feb', + 'reverse_edits': [(103, + 141, + ['def ' + '_launch_config(dtype_name, ' + 'batch, m, n, k):\n', + ' """Return the ' + 'bmm_kernel tile sizes, ' + 'num_warps and ' + 'num_stages.\n', + '\n', + ' Tune these values, ' + 'optionally per dtype or ' + 'problem size. The ' + 'protected\n', + ' wrapper validates ' + 'them and derives the grid ' + 'from the tile sizes.\n', + ' """\n', + ' if dtype_name == ' + '"float32":\n', + ' return ' + '{"BLOCK_SIZE_M": 64, ' + '"BLOCK_SIZE_N": 64, ' + '"BLOCK_SIZE_K": 32,\n', + ' ' + '"num_stages": 2, ' + '"num_warps": 4}\n', + ' return ' + '{"BLOCK_SIZE_M": 128, ' + '"BLOCK_SIZE_N": 128, ' + '"BLOCK_SIZE_K": 32,\n', + ' "num_stages": ' + '2, "num_warps": 8}\n', + '\n', + '\n', + '_DTYPE_NAMES = ' + '{torch.bfloat16: ' + '"bfloat16", ' + 'torch.float16: "float16", ' + 'torch.float32: ' + '"float32"}\n', + '_TILE_KEYS = ' + '("BLOCK_SIZE_M", ' + '"BLOCK_SIZE_N", ' + '"BLOCK_SIZE_K")\n', + '\n', + '\n', + 'def ' + '_checked_launch_config(cfg):\n', + ' """Accept only tile ' + 'sizes, num_warps and ' + 'num_stages within ' + 'supported bounds."""\n', + ' expected = ' + 'set(_TILE_KEYS) | ' + '{"num_warps", ' + '"num_stages"}\n', + ' if not ' + 'isinstance(cfg, dict) or ' + 'set(cfg) != expected:\n', + ' raise ' + 'ValueError(f"Launch ' + 'configuration must set ' + 'exactly ' + '{sorted(expected)}")\n', + ' checked = {}\n', + ' for key in ' + 'sorted(expected):\n', + ' value = ' + 'cfg[key]\n', + ' if type(value) is ' + 'not int:\n', + ' raise ' + 'ValueError(f"Launch ' + 'setting {key} must be an ' + 'int")\n', + ' checked[key] = ' + 'value\n', + ' for key in ' + '_TILE_KEYS:\n', + ' if not (16 <= ' + 'checked[key] <= 256 and ' + 'checked[key] & ' + '(checked[key] - 1) == ' + '0):\n', + ' raise ' + 'ValueError(f"{key} must ' + 'be a power of two in [16, ' + '256]")\n', + ' if ' + 'checked["num_warps"] not ' + 'in (1, 2, 4, 8, 16):\n', + ' raise ' + 'ValueError("num_warps ' + 'must be 1, 2, 4, 8 or ' + '16")\n', + ' if not 1 <= ' + 'checked["num_stages"] <= ' + '4:\n', + ' raise ' + 'ValueError("num_stages ' + 'must be in [1, 4]")\n', + ' return checked\n', + '\n', + '\n'], + []), + (154, + 155, + [' cfg = ' + '_checked_launch_config(_launch_config(_DTYPE_NAMES[dtype], ' + 'B_size, M, N, K))\n'], + [' configs = {\n', + ' torch.bfloat16: ' + '{\n', + ' ' + '"BLOCK_SIZE_M": 128,\n', + ' ' + '"BLOCK_SIZE_N": 128,\n', + ' ' + '"BLOCK_SIZE_K": 32,\n', + ' "num_stages": ' + '2,\n', + ' "num_warps": ' + '8,\n', + ' },\n', + ' torch.float16: ' + '{\n', + ' ' + '"BLOCK_SIZE_M": 128,\n', + ' ' + '"BLOCK_SIZE_N": 128,\n', + ' ' + '"BLOCK_SIZE_K": 32,\n', + ' "num_stages": ' + '2,\n', + ' "num_warps": ' + '8,\n', + ' },\n', + ' torch.float32: ' + '{\n', + ' ' + '"BLOCK_SIZE_M": 64,\n', + ' ' + '"BLOCK_SIZE_N": 64,\n', + ' ' + '"BLOCK_SIZE_K": 32,\n', + ' "num_stages": ' + '2,\n', + ' "num_warps": ' + '4,\n', + ' },\n', + ' }\n', + '\n', + ' cfg = ' + 'configs[dtype]\n'])]}} + + +def historical_triton_launch_config_file(relative, text): + record = TRITON_LAUNCH_CONFIG_FILES.get(relative) + return _exact_reviewed_file_inverse(text, record) if record is not None else text + + def historical_triton_device_boundary_file(relative, text): + text = historical_triton_launch_config_file(relative, text) record = TRITON_DEVICE_BOUNDARY_FILES.get(relative) return _exact_reviewed_file_inverse(text, record) if record is not None else text diff --git a/tests/test_bmm_domain_controls.py b/tests/test_bmm_domain_controls.py index dfe390fdd..3e6a81ed7 100644 --- a/tests/test_bmm_domain_controls.py +++ b/tests/test_bmm_domain_controls.py @@ -78,7 +78,12 @@ def scored_shapes(): def bmm_module(*, fault=None, seen=None): """Execute the actual protected wrapper around a CPU model of bmm_kernel.""" tree = ast.parse((TASK / 'source/triton_bmm.py').read_text()) - wrapper = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == 'bmm_triton') + # The wrapper with its launch-configuration helpers and tables; only the + # device kernel is modeled and the imports are supplied below. + wrapper = [n for n in tree.body if not isinstance(n, (ast.Import, ast.ImportFrom)) + and not (isinstance(n, ast.FunctionDef) and n.name == 'bmm_kernel')] + assert {n.name for n in wrapper if isinstance(n, ast.FunctionDef)} == { + '_launch_config', '_checked_launch_config', 'bmm_triton'} scored = scored_shapes() class Kernel: @@ -110,7 +115,7 @@ def launch(a, b, c, batch, m, n, k, sab, sam, sak, sbb, sbk, sbn, scb, scm, scn, namespace = {'torch': torch, 'triton': SimpleNamespace(cdiv=lambda x, y: (x + y - 1) // y), 'bmm_kernel': Kernel()} - exec(compile(ast.Module(body=[wrapper], type_ignores=[]), '', 'exec'), namespace) + exec(compile(ast.Module(body=wrapper, type_ignores=[]), '', 'exec'), namespace) return SimpleNamespace(bmm_triton=namespace['bmm_triton']) @@ -373,8 +378,10 @@ def test_scored_rows_kernel_checks_and_timing_unchanged_from_base(env): runner = (TASK / 'scripts/task_runner.py').read_text() assert runner.count(DISPATCH) == 1 assert runner.replace(DISPATCH, '', 1) == base_text('scripts/task_runner.py') + from reviewed_contract_history import historical_triton_launch_config_file for relative in ('source/triton_bmm.py', '_arena_checks.py'): - assert (TASK / relative).read_text() == base_text(relative) + current = historical_triton_launch_config_file(f'tasks/{TASK_ID}/{relative}', (TASK / relative).read_text()) + assert current == base_text(relative) old_adapter = base_text('_arena_eval.py').splitlines(True) new_adapter = (TASK / '_arena_eval.py').read_text().splitlines(True) opcodes = difflib.SequenceMatcher(a=old_adapter, b=new_adapter, autojunk=False).get_opcodes() @@ -384,14 +391,15 @@ def test_scored_rows_kernel_checks_and_timing_unchanged_from_base(env): assert (harness.WARMUP_ITERATIONS, harness.BENCHMARK_ITERATIONS) == (10, 100) -@pytest.mark.parametrize('change', ['kernel', 'helper', 'controls', 'manifest', 'wrapper', 'runner']) +@pytest.mark.parametrize('change', ['kernel', 'helper', 'launch_config', 'launch_host_code', + 'controls', 'manifest', 'wrapper', 'runner']) def test_guard_protects_control_module_and_keeps_kernel_editable(tmp_path, change): work = tmp_path / 'task' shutil.copytree(TASK, work, ignore=shutil.ignore_patterns('__pycache__', 'build')) spec = load_task_spec(work / 'config.yaml', task_id=TASK_ID) edit, = spec.candidate.editable assert edit.path == 'source/triton_bmm.py' and edit.scope == 'symbols' - assert edit.symbols == ('bmm_kernel',) and edit.allow_new_helpers + assert edit.symbols == ('bmm_kernel', '_launch_config') and edit.allow_new_helpers assert not edit.contains('_domain_controls.py') guard = snapshot_workspace_harness(work, task_spec=spec) assert '_domain_controls.py' in guard.digests @@ -403,6 +411,12 @@ def test_guard_protects_control_module_and_keeps_kernel_editable(tmp_path, chang node.body.append(ast.Pass()) elif change == 'helper': tree.body.extend(ast.parse('@triton.jit\ndef tile_helper(x):\n return x\n').body) + elif change in ('launch_config', 'launch_host_code'): + node = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == '_launch_config') + tuned = ('if m >= 512:\n return {"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64,' + ' "num_stages": 3, "num_warps": 8}\n') + host = 'torch.cuda.synchronize()\n' + node.body[1:1] = ast.parse(tuned if change == 'launch_config' else host).body elif change == 'controls': (work / '_domain_controls.py').write_text('CASES = []\n') elif change == 'manifest': @@ -413,7 +427,7 @@ def test_guard_protects_control_module_and_keeps_kernel_editable(tmp_path, chang runner = work / 'scripts/task_runner.py' runner.write_text(runner.read_text().replace(DISPATCH, '', 1)) source.write_text(ast.unparse(ast.fix_missing_locations(tree))) - if change in ('kernel', 'helper'): + if change in ('kernel', 'helper', 'launch_config'): verify_workspace_harness(guard, discard_added=False) else: with pytest.raises(RuntimeError): diff --git a/tests/test_triton_module_edit_boundaries.py b/tests/test_triton_module_edit_boundaries.py index 04106aade..e1fa19d41 100644 --- a/tests/test_triton_module_edit_boundaries.py +++ b/tests/test_triton_module_edit_boundaries.py @@ -308,8 +308,12 @@ def test_identical_actual_gate_payload_old_guard_accepts_new_guard_rejects(tmp_p @pytest.mark.parametrize('row',TASKS,ids=lambda r:r['task'].split('/')[-1]) def test_all_original_device_functions_are_declared_and_host_definitions_protected(row): root=ROOT/'tasks'/row['task'];spec=load_task_spec(root/'config.yaml',task_id=row['task']);edit,=spec.candidate.editable - assert edit.scope=='symbols'and list(edit.symbols)==row['symbols']and edit.allow_new_helpers is True - assert all(e.symbol in row['symbols']for e in spec.candidate.entrypoints) source=ast.parse((root/row['source']).read_text()) + # A tunable launch configuration is the only editable host definition: an + # undecorated `_launch_config` declared after the device functions. + launch=[n for n in source.body if isinstance(n,ast.FunctionDef)and n.name=='_launch_config'] + assert all(not n.decorator_list for n in launch) + assert edit.scope=='symbols'and list(edit.symbols)==row['symbols']+[n.name for n in launch]and edit.allow_new_helpers is True + assert all(e.symbol in row['symbols']for e in spec.candidate.entrypoints) actual=[n.name for n in source.body if isinstance(n,ast.FunctionDef)and any('triton.jit'in ast.unparse(d)for d in n.decorator_list)] assert actual==row['symbols'] diff --git a/tests/test_triton_task_migration_v2.py b/tests/test_triton_task_migration_v2.py index b96fa9264..62495981f 100644 --- a/tests/test_triton_task_migration_v2.py +++ b/tests/test_triton_task_migration_v2.py @@ -288,10 +288,14 @@ def test_vllm_v2_preserves_all_original_cases_checks_sources_and_helpers(path): assert manifest['migration']['original_harness_sha256'] == hashlib.sha256(before.encode()).hexdigest() for edit in spec.candidate.editable: source = task/edit.path + # Undo only reviewed launch-configuration moves first; every check + # below then compares the previously reviewed source bytes. + from reviewed_contract_history import historical_triton_launch_config_file + source_text = historical_triton_launch_config_file(source.relative_to(ROOT).as_posix(), source.read_text()) original = subprocess.check_output(['git','show',f'{BASE}:{source.relative_to(ROOT).as_posix()}'],cwd=ROOT) if task.name == 'triton_decode_attn_grouped_stage1': from reviewed_contract_history import historical_attention_sampling_file - restored = historical_attention_sampling_file(source.relative_to(ROOT).as_posix(), source.read_text()) + restored = historical_attention_sampling_file(source.relative_to(ROOT).as_posix(), source_text) assert restored.encode() == original continue if task.name == 'triton_topk_log_softmax': @@ -338,7 +342,7 @@ def test_vllm_v2_preserves_all_original_cases_checks_sources_and_helpers(path): # 32-assignment tile, including top-k 33/65 tails. The dedicated # boundary tests reject truncation. The wrapper adds the tile count # and normalizes strided inputs; retain every other source byte. - current = source.read_text() + current = source_text before_tree, after_tree = ast.parse(original), ast.parse(current) before_kernel = next(n for n in before_tree.body if isinstance(n, ast.FunctionDef) and n.name == 'pack_bitmatrix') @@ -385,7 +389,7 @@ def test_vllm_v2_preserves_all_original_cases_checks_sources_and_helpers(path): if task.name == 'triton_ssd_bmm': # Only remove the causal early return that skipped legitimate # full-matrix output tiles. All remaining source stays byte-exact. - current = source.read_bytes() + current = source_text.encode() assert hashlib.sha256(current).hexdigest() == ( '421143d84f536ca07fff4e70837e9ff7518751e2b9509294d4f6d865bb7a67c3') old = (b' if IS_CAUSAL:\n' @@ -396,7 +400,7 @@ def test_vllm_v2_preserves_all_original_cases_checks_sources_and_helpers(path): if task.name == 'triton_fla_l2norm': # Conditional v8 pin: the new wide-row kernel and its sole wrapper # dispatch are the complete reviewed delta; GPU proof is pending. - current = source.read_bytes() + current = source_text.encode() assert hashlib.sha256(current).hexdigest() == ( 'c6b03db7c1dda337dd2f0c60261070f9b4d11a31b436bc2a73a5675f1aa166a4') old_tree, new_tree = ast.parse(original), ast.parse(current) @@ -430,7 +434,7 @@ def test_vllm_v2_preserves_all_original_cases_checks_sources_and_helpers(path): if task.name == 'triton_kda_gate': # Preserve the entire historical JIT and public wrapper AST except # the three reviewed view normalizations before flat pointer reads. - current = source.read_bytes() + current = source_text.encode() assert hashlib.sha256(current).hexdigest() == ( '227cb972264fa01bb0029edd4eaf096e18aab9ff84a29f25cc0e4b66ec1e8ec1') old_tree, new_tree = ast.parse(original), ast.parse(current) @@ -456,7 +460,7 @@ def test_vllm_v2_preserves_all_original_cases_checks_sources_and_helpers(path): # remain identical. Only four wrapper-side view normalizations # admit valid noncontiguous inputs; scored contiguous calls keep # their original pointers and device work. - current = source.read_bytes() + current = source_text.encode() assert hashlib.sha256(current).hexdigest() == ( 'e92eab87cb20d42421b3c67c050bb14a9bef2fa3eb7df92c9318439b42f4f93b') old_tree, new_tree = ast.parse(original), ast.parse(current) @@ -477,7 +481,7 @@ def test_vllm_v2_preserves_all_original_cases_checks_sources_and_helpers(path): # Keep the kernel, wrapper signature/decorators, and all other # source AST at their original identities; the focused stride # controls reject an implementation that ignores either operand. - current = source.read_bytes() + current = source_text.encode() assert hashlib.sha256(current).hexdigest() == ( 'c65a20b3878c2a520cf37645ae4d18d95cff4d0202c4399bb263833485fe4913') old_tree, new_tree = ast.parse(original), ast.parse(current) @@ -502,7 +506,7 @@ def wrapper(tree): new_wrapper.body = [node for node in new_wrapper.body if node not in copies] assert ast.dump(new_tree, include_attributes=False) == ast.dump(old_tree, include_attributes=False) continue - assert source.read_bytes() == original + assert source_text.encode() == original assert 'Evaluation contract' in (task/'README.md').read_text() From 0cd82d6e6241ef1a5ea34da5ebe29c3baf59c1d8 Mon Sep 17 00:00:00 2001 From: Vincent Ouyang Date: Tue, 27 Oct 2026 10:04:19 -0700 Subject: [PATCH 3/3] Keep launch configuration helpers out of decorator factories Only protected launch callers may supply metadata to declared configuration helpers. Prevent editable compiler decorators from invoking those helpers with arbitrary host objects, and cover the BMM definition-time regression. --- docs/how-to/add-task.md | 5 +++- src/harness_guard.py | 5 +++- tests/test_configuration_symbol_policy.py | 36 +++++++++++++++++++++++ 3 files changed, 44 insertions(+), 2 deletions(-) diff --git a/docs/how-to/add-task.md b/docs/how-to/add-task.md index a88013299..024cfb7dd 100644 --- a/docs/how-to/add-task.md +++ b/docs/how-to/add-task.md @@ -193,7 +193,10 @@ configuration data, literal collections, arithmetic and branches. Unshadowed function calls, arbitrary attributes, callbacks and comprehensions are rejected; defaults must be literal data and annotations use simple builtin types or strings. The protected caller must supply plain scalar/container metadata and -validate the returned configuration before launch. Every editable Python file, +validate the returned configuration before launch. Declared launch helpers are +not definition-time factories and cannot be called from editable compiler +decorators, which could supply arbitrary objects instead of metadata. +Every editable Python file, whatever its scope, also may not add runtime patching: calls to `globals()`, `vars()`, `setattr`/`delattr`, `exec`/`eval`/`__import__` (including qualified `builtins` calls and imported aliases), `importlib` or frame diff --git a/src/harness_guard.py b/src/harness_guard.py index 053095b69..9abbaf51d 100644 --- a/src/harness_guard.py +++ b/src/harness_guard.py @@ -485,7 +485,10 @@ def _v2_symbol_digest(path: Path, edit: EditScope, initial_names: frozenset[str] # editable target or newly introduced helper is not a protected factory. protected_functions = (set(functions) & initial_names) - set(edit.symbols) configuration_factories = set(protected_functions) - for name in (set(functions) & initial_names & set(edit.symbols)) - entrypoint_symbols: + # Declared launch helpers receive plain metadata only from protected callers. + # Do not let editable decorators invoke them with arbitrary host objects. + for name in ((set(functions) & initial_names & set(edit.symbols)) + - entrypoint_symbols - set(edit.configuration_symbols)): if _config_factory_body_allowed(functions[name], bindings, protected_functions, set(initial_names) | definition_names): configuration_factories.add(name) diff --git a/tests/test_configuration_symbol_policy.py b/tests/test_configuration_symbol_policy.py index f6419dbac..b1da9c34a 100644 --- a/tests/test_configuration_symbol_policy.py +++ b/tests/test_configuration_symbol_policy.py @@ -4,9 +4,14 @@ editable candidate or requiring Triton/GPU execution. """ +import ast +from pathlib import Path +import shutil + import pytest from src.harness_guard import snapshot_workspace_harness, verify_workspace_harness +from src.task_spec import load_task_spec from test_harness_guard_v2 import LAUNCH_RETURN, launch_workspace @@ -144,3 +149,34 @@ def test_branch_local_data_must_be_definitely_assigned(tmp_path): )) with pytest.raises(RuntimeError, match="configuration symbol"): verify_workspace_harness(before, discard_added=False) + + +def test_bmm_launch_configuration_cannot_be_used_as_a_definition_time_factory(tmp_path): + """An editable decorator must not supply helper objects to the scalar caller.""" + original = Path(__file__).resolve().parents[1] / 'tasks/triton2triton/vllm/triton_bmm' + root = tmp_path / 'task' + shutil.copytree(original, root, ignore=shutil.ignore_patterns('__pycache__', 'build')) + spec = load_task_spec(root / 'config.yaml', task_id='triton2triton/vllm/triton_bmm') + before = snapshot_workspace_harness(root, task_spec=spec) + source = root / 'source/triton_bmm.py' + tree = ast.parse(source.read_text()) + config = next(node for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == '_launch_config') + tree.body.remove(config) + # This is valid pure-data code when the protected caller supplies its normal + # scalar metadata. An editable definition-time caller could pass a class. + config.body.insert(1, ast.parse("if dtype_name == 'probe':\n probe = m[0]\n").body[0]) + kernel = next(node for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == 'bmm_kernel') + kernel.decorator_list = [ast.parse( + "triton.jit(do_not_specialize=[" + "_launch_config('probe', 1, ConfigProbe, 16, 16)['num_stages']])", + mode='eval', + ).body] + index = tree.body.index(kernel) + tree.body[index:index] = [ast.parse(CONFIG_PROBE).body[0], config] + source.write_text(ast.unparse(ast.fix_missing_locations(tree))) + # Reject it structurally, without importing candidate code or triggering its + # __class_getitem__ during decorator evaluation. + with pytest.raises(RuntimeError, match='Unapproved definition-time factory'): + verify_workspace_harness(before, discard_added=False)