diff --git a/docs/how-to/add-task.md b/docs/how-to/add-task.md index 579e00d49..024cfb7dd 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,19 @@ 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. Every editable Python file, +do not gain permission to run at definition time. Symbols listed in +`configuration_symbols` use a stricter launch-time policy: they must stay +undecorated functions whose values come only from metadata parameters, local +configuration data, literal collections, arithmetic and branches. Unshadowed +`min`, `max`, `abs`, `bool`, `int`, `float`, `len`, `triton.cdiv` and +`triton.next_power_of_2` calls are allowed. Global objects, helper/protected +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. 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 e2df60afe..9abbaf51d 100644 --- a/src/harness_guard.py +++ b/src/harness_guard.py @@ -245,6 +245,123 @@ def _config_factory_body_allowed(node: ast.FunctionDef, bindings: dict[str, str] return True +def _launch_configuration_allowed(node: ast.FunctionDef, bindings: dict[str, str], + bound_names: set[str]) -> bool: + """Keep declared launch configuration functions on plain metadata and data. + + The protected caller supplies scalar/container metadata, never tensors or + user-defined objects. Unlike legacy definition-time factories, these helpers + cannot read global objects or call protected functions: even an apparently + passive subscript can invoke a helper's __class_getitem__ or __getitem__. + """ + if node.decorator_list: + return False + args = node.args + parameters = [*args.posonlyargs, *args.args, *args.kwonlyargs, + *([args.vararg] if args.vararg else []), + *([args.kwarg] if args.kwarg else [])] + locals_ = {arg.arg for arg in parameters} + locals_.update(item.id for item in ast.walk(node) + if isinstance(item, ast.Name) and isinstance(item.ctx, ast.Store)) + imports = {key: value for key, value in bindings.items() if key not in locals_} + builtins = {"min", "max", "abs", "bool", "int", "float", "len"} - ( + bound_names | set(bindings) | locals_) + + def literal(value: ast.AST | None) -> bool: + if isinstance(value, ast.Constant): + return type(value.value) in (str, bytes, int, float, bool, type(None)) + if isinstance(value, (ast.Tuple, ast.List, ast.Set)): + return all(literal(item) for item in value.elts) + if isinstance(value, ast.Dict): + return all(key is not None and literal(key) and literal(item) + for key, item in zip(value.keys, value.values)) + return (isinstance(value, ast.UnaryOp) and isinstance(value.op, (ast.UAdd, ast.USub)) + and literal(value.operand)) + + def annotation(value: ast.AST | None) -> bool: + if value is None or (isinstance(value, ast.Constant) + and (value.value is None or isinstance(value.value, str))): + return True + return (isinstance(value, ast.Name) + and value.id in {"bool", "int", "float", "str", "bytes", "dict", "list", "tuple", "set"} + and value.id not in bound_names | set(bindings)) + + if (not all(literal(value) for value in [*args.defaults, *args.kw_defaults] if value is not None) + or not all(annotation(value) for value in [node.returns, *[arg.annotation for arg in parameters]])): + return False + + def expression(value: ast.AST | None, available: set[str]) -> bool: + if isinstance(value, ast.Constant): + return literal(value) + if isinstance(value, ast.Name): + return value.id in available + if isinstance(value, (ast.Tuple, ast.List, ast.Set)): + return all(expression(item, available) for item in value.elts) + if isinstance(value, ast.Dict): + return all(key is not None and expression(key, available) and expression(item, available) + for key, item in zip(value.keys, value.values)) + if isinstance(value, ast.BinOp): + return (isinstance(value.op, (ast.Add, ast.Sub, ast.Mult, ast.Div, ast.FloorDiv, + ast.Mod, ast.Pow, ast.LShift, ast.RShift, + ast.BitOr, ast.BitAnd, ast.BitXor)) + and expression(value.left, available) and expression(value.right, available)) + if isinstance(value, ast.UnaryOp): + return (isinstance(value.op, (ast.UAdd, ast.USub, ast.Not, ast.Invert)) + and expression(value.operand, available)) + if isinstance(value, ast.BoolOp): + return all(expression(item, available) for item in value.values) + if isinstance(value, ast.Compare): + return expression(value.left, available) and all( + expression(item, available) for item in value.comparators) + if isinstance(value, ast.IfExp): + return all(expression(item, available) for item in (value.test, value.body, value.orelse)) + if isinstance(value, ast.Subscript): + return expression(value.value, available) and expression(value.slice, available) + if isinstance(value, ast.Slice): + return all(item is None or expression(item, available) + for item in (value.lower, value.upper, value.step)) + if isinstance(value, ast.Call): + safe = ((isinstance(value.func, ast.Name) and value.func.id in builtins) + or _qualified_import(value.func, imports) in _LAUNCH_LAMBDA_CALLS) + return (safe and all(expression(item, available) for item in value.args) + and all(item.arg is not None and expression(item.value, available) + for item in value.keywords)) + return False + + def block(statements: list[ast.stmt], available: set[str]) -> tuple[set[str], bool] | None: + available = set(available) + falls_through = True + for statement in statements: + if isinstance(statement, ast.Return): + if not expression(statement.value, available): + return None + falls_through = False + elif isinstance(statement, (ast.Assign, ast.AnnAssign)): + targets = statement.targets if isinstance(statement, ast.Assign) else [statement.target] + if (not all(isinstance(target, ast.Name) for target in targets) + or not expression(statement.value, available) + or (isinstance(statement, ast.AnnAssign) and not annotation(statement.annotation))): + return None + available.update(target.id for target in targets) + elif isinstance(statement, ast.If): + if not expression(statement.test, available): + return None + yes, no = block(statement.body, available), block(statement.orelse, available) + if yes is None or no is None: + return None + continuing = [names for names, continues in (yes, no) if continues] + if continuing: + available = set.intersection(*continuing) + else: + falls_through = False + elif not (isinstance(statement, ast.Pass) + or (isinstance(statement, ast.Expr) and isinstance(statement.value, ast.Constant))): + return None + return available, falls_through + + return block(node.body, {arg.arg for arg in parameters}) is not None + + def _validate_editable_definition(node: ast.AST, bindings: dict[str, str], *, initial_names: frozenset[str], definition_names: set[str], configuration_factories: set[str], @@ -368,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) @@ -385,6 +505,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 _launch_configuration_allowed(node, bindings, + 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() @@ -1036,6 +1166,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..fe8154cc1 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,17 @@ 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 any(not isinstance(s, str) or s not in symbols for s in configuration) + or len(configuration) != len(set(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 +228,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/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_configuration_symbol_policy.py b/tests/test_configuration_symbol_policy.py new file mode 100644 index 000000000..b1da9c34a --- /dev/null +++ b/tests/test_configuration_symbol_policy.py @@ -0,0 +1,182 @@ +"""Configuration-only helpers consume scalar metadata, never helper objects. + +These structural checks use the public workspace guard without importing an +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 + + +CONFIG_PROBE = ( + "class ConfigProbe:\n" + " def __class_getitem__(cls, key):\n" + " print('host helper executed')\n" + " return key\n" +) + + +def guarded_config(tmp_path, *, imports=""): + root, source = launch_workspace(tmp_path) + if imports: + source.write_text(imports + source.read_text()) + return source, snapshot_workspace_harness(root) + + +@pytest.mark.parametrize("body", [ + # All returned leaves originate in caller-provided scalar metadata or literals. + (" if n >= 4096:\n" + " block = 256\n" + " warps = 8\n" + " else:\n" + " block = 64\n" + " warps = 4\n" + " return {'BLOCK': block, 'num_warps': warps}\n"), + (" choices = {'small': [32, 64], 'large': (128, 256)}\n" + " block = choices['large'][1] if n > 4096 else choices['small'][0]\n" + " return {'BLOCK': block, 'num_warps': 4}\n"), + (" block: int = ((n + 31) // 32) * 32\n" + " block = max(16, min(256, block))\n" + " enabled = bool(n > 0 and n <= 4096)\n" + " warps = int(float(4)) if enabled else abs(-8)\n" + " count = len((block, warps))\n" + " return {'BLOCK': block, 'num_warps': warps + count - 2}\n"), + (" if n > 4096:\n" + " return {'BLOCK': 256, 'num_warps': 8}\n" + " else:\n" + " block = 64\n" + " return {'BLOCK': block, 'num_warps': 4}\n"), + (" if n > 4096:\n" + " block = 256\n" + " else:\n" + " return {'BLOCK': 64, 'num_warps': 4}\n" + " return {'BLOCK': block, 'num_warps': 8}\n"), +]) +def test_configuration_data_arithmetic_branches_and_local_collections(tmp_path, body): + source, before = guarded_config(tmp_path) + source.write_text(source.read_text().replace(LAUNCH_RETURN, body)) + verify_workspace_harness(before, discard_added=False) + + +@pytest.mark.parametrize("imports,expression", [ + ("import triton\n", "triton.next_power_of_2(triton.cdiv(n, 32))"), + ("from triton import cdiv as ceil_div\n", "ceil_div(n, 32)"), +]) +def test_unshadowed_pure_metadata_calls_are_allowed(tmp_path, imports, expression): + source, before = guarded_config(tmp_path, imports=imports) + body = f" block = {expression}\n return {{'BLOCK': block, 'num_warps': 4}}\n" + source.write_text(source.read_text().replace(LAUNCH_RETURN, body)) + verify_workspace_harness(before, discard_added=False) + + +def test_scalar_signature_annotations_and_literal_defaults_are_allowed(tmp_path): + source, before = guarded_config(tmp_path) + source.write_text(source.read_text().replace( + "def _launch_config(n):", + "def _launch_config(n: int, mode: 'str' = 'default') -> dict:", + )) + verify_workspace_harness(before, discard_added=False) + + +@pytest.mark.parametrize("body", [ + " block = ConfigProbe[0]\n" + LAUNCH_RETURN, + " alias = ConfigProbe\n block = alias[0]\n" + LAUNCH_RETURN, + " return ConfigProbe\n", + " block = launch(n)\n" + LAUNCH_RETURN, + " block: ConfigProbe[0] = 64\n" + LAUNCH_RETURN, + " block = [item for item in (64, 128)][0]\n" + LAUNCH_RETURN, + " block = (lambda: 64)\n" + LAUNCH_RETURN, + " block = f'{ConfigProbe}'\n" + LAUNCH_RETURN, + " block = min(n, 64, key=ConfigProbe)\n" + LAUNCH_RETURN, + " block = min(*(n, 64))\n" + LAUNCH_RETURN, + " block = min(**{'iterable': (n, 64)})\n" + LAUNCH_RETURN, +]) +def test_configuration_symbols_cannot_reach_host_objects_or_callbacks(tmp_path, body): + source, before = guarded_config(tmp_path) + # The helper is itself legal implementation code. Reading it from the + # configuration function must still fail, before a subscript can dispatch it. + source.write_text(source.read_text().replace(LAUNCH_RETURN, body) + CONFIG_PROBE) + with pytest.raises(RuntimeError, match="configuration symbol"): + verify_workspace_harness(before, discard_added=False) + + +@pytest.mark.parametrize("signature", [ + "def _launch_config(n, extra=ConfigProbe[0]):", + "def _launch_config(n: ConfigProbe[0]):", + "def _launch_config(n) -> ConfigProbe[0]:", +]) +def test_configuration_signature_cannot_execute_implicit_helper_code(tmp_path, signature): + source, before = guarded_config(tmp_path) + changed = source.read_text().replace( + "def _launch_config(n):", CONFIG_PROBE + signature, + ) + source.write_text(changed) + with pytest.raises(RuntimeError, match="configuration symbol"): + verify_workspace_harness(before, discard_added=False) + + +@pytest.mark.parametrize("signature,body", [ + ("def _launch_config(n, min=None):", " block = min(n, 64)\n"), + ("def _launch_config(n):", " min = n\n block = min(n, 64)\n"), + ("def _launch_config(n, triton=None):", " block = triton.cdiv(n, 32)\n"), + ("def _launch_config(n):", " triton = n\n block = triton.cdiv(n, 32)\n"), + ("def _launch_config(n, ceil_div=None):", " block = ceil_div(n, 32)\n"), + ("def _launch_config(n):", " ceil_div = n\n block = ceil_div(n, 32)\n"), +]) +def test_metadata_call_names_cannot_be_shadowed(tmp_path, signature, body): + source, before = guarded_config( + tmp_path, imports="import triton\nfrom triton import cdiv as ceil_div\n", + ) + source.write_text(source.read_text().replace("def _launch_config(n):", signature) + .replace(LAUNCH_RETURN, body + LAUNCH_RETURN)) + with pytest.raises(RuntimeError, match="configuration symbol"): + verify_workspace_harness(before, discard_added=False) + + +def test_branch_local_data_must_be_definitely_assigned(tmp_path): + source, before = guarded_config(tmp_path) + source.write_text(source.read_text().replace( + LAUNCH_RETURN, + " if n > 4096:\n block = 256\n" + " return {'BLOCK': block, 'num_warps': 4}\n", + )) + 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) 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 200963abc..887f77153 100644 --- a/tests/test_harness_guard_v2.py +++ b/tests/test_harness_guard_v2.py @@ -414,3 +414,99 @@ 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": [{}]}, "listed in symbols"), + ({"symbols": ["compute"], "configuration_symbols": [[]]}, "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") diff --git a/tests/test_triton_module_edit_boundaries.py b/tests/test_triton_module_edit_boundaries.py index 4e314d22f..85ab43d62 100644 --- a/tests/test_triton_module_edit_boundaries.py +++ b/tests/test_triton_module_edit_boundaries.py @@ -310,8 +310,12 @@ def test_identical_actual_gate_payload_is_rejected_with_either_editable_scope(tm @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()