Make symbol-scoped Triton launch configurations tunable

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

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 [],

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:

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))
 
 

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

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