Resource contract for collective and gated models

A collective regime and a value-dependent transition each buy something a singleton model does not have, and each costs something a singleton model does not pay. This page names the workloads those costs are measured on, says which axis each one is allowed to grow along, and states what a regression is.

It is a contract, not a report: the numbers live in the ASV history, and what is written down here is the shape each cost is allowed to have. A change that moves a level is reviewed against the history; a change that moves an order is a defect whatever the level.

The workloads¶

All of them are in benchmarks/asv/bench_collective_household.py. All but the last two run over the marriage market of lcm_examples.collective_household: two singles who marry under mutual consent, a household with a participation constraint on each partner, and a dissolution edge keyed by the continuing household. ReferenceChainSolve is not that model — it builds a synthetic chain of collective links with no gated edge, no consent and no dissolution, so that reference depth is the only thing varying.

CollectiveHouseholdConstruct and CollectiveHouseholdSolve publish their timing and host-memory numbers through setup_cache-backed track_* metrics, so one build and one cold call are shared by every metric of that class. The two parameterized workloads cannot share a setup across parameters without the device-memory subprocess accepting parameters, so they time track_execution_time in process, as the median of warm calls after setup’s cold call, and keep ASV-native peakmem_execution for host memory.

The budgets¶

Each budget is an order, because that is what a benchmark suite can defend across machines and backends. A level is defended by the ASV history on one machine.

• Model construction is O(regimes × phases × declarations) and involves no device work. It is allowed to grow with the number of declarations a model makes and with nothing else. Construction that grew with a grid size would mean a grid was materialized during the scan.

• First-solve assembly is O(regimes × periods) programs plus O(edges × periods) folds. It contains no gate evaluator: those execute on the simulation side, while this workload only solves. Forward gate compilation is outside the contract’s measured surface.

• Warm solve is O(periods × regimes × cells), the same order as a singleton model of the same total grid size. A collective regime multiplies the cell count by its stakeholder count; a gated edge adds one fold over the target’s grid per period. It may not grow with the number of subjects, which appears nowhere in the solve.

• Simulation is O(periods × (regimes + edges) × subjects) and therefore linear in the cohort. The two cohort sizes exist to make a super-linear term visible; a slope above one between them is the regression this workload is for. They are two orders of magnitude apart deliberately — an intermediate point cost run time without separating anything the endpoints do not.

• Memory, host and device, is O(largest single V array + working set) and does not accumulate across periods. Backward induction frees each period’s intermediates, so a peak that grew with the number of periods would mean it stopped.

• The shape cache is bounded by the model, not by the run. A gate evaluator’s population call is keyed on (callable, cohort size) and every other program on (callable, dedup key), all of which are properties of the model and its declared batch size. Repeated solve() / simulate() calls at one cohort size may not add entries; a cache that grew per call would recompile per call.

What the first-solve number is, and is not¶

track_compilation_time is not a compilation time. The timed region is a whole model.solve() call, so the number is everything that first solve does: tracing, lowering, either compiling each program or loading it from the persistent cache, and then running the backward induction itself. Which of compile-or-load happened is invisible in the number, and so is the execution term.

Measured on this workload — benchmarks/asv/bench_collective_household.py at commit 1feb6b0a, on an NVIDIA GeForce RTX 3070 (8 GiB) under the CUDA backend at float32 (jax_enable_x64 left at its default), with JAX_COMPILATION_CACHE_DIR pointed at a fresh directory asserted empty before the run and holding 60 entries after:

So the cache covers a little under half of a cold first solve, and execution is about 2% of it — small here, but it is a term that grows with the grid while the assembly terms grow with the number of programs, so the split is not fixed. The saving from a cache hit also grows with the model, because compilation grows with program size while tracing and lowering grow with program count.

Three things follow, all about how to read a change in this line:

• A movement is a movement in first-solve cost, not in compiler work, unless the cache state was the same on both sides. Pin it — point JAX_COMPILATION_CACHE_DIR at a fresh directory — for any comparison meant to be about compilation.

• A number from one machine is not comparable to one from another whose cache holds a different set of programs. The ASV history is per machine for this reason.

• On a model whose grids are large enough for execution to dominate, this line stops being an assembly measurement at all. Read it next to track_execution_time, which is the same solve with the programs already in memory.

Where a pointwise reoptimization mode would sit¶

off_grid="pointwise" reads the operands at the landing point and gates them there, in both phases, using the kernels the model already has. A future reoptimization mode — recomputing the target’s own optimum at the realized point — is a different kernel and is kept statically separate from the default one: it must be selectable per edge, must not appear in a model that did not ask for it, and carries its own entries in this table before it is offered. Its cost is O(subjects × target action grid) per edge per period, which is a different order from the default’s O(subjects), so the two may not share a budget line.

Running them¶

These workloads run under the ordinary ASV tasks; Benchmarking has the machine registration, the run and compare invocations, and the publish workflows.

One caveat is specific to this suite. The GPU peak-memory companion needs a CUDA environment and a device that publishes memory statistics, and there is no skip path: nothing here raises ASV’s NotImplementedError or detects a device, so on a machine without one it fails rather than abstains. Run it only where a GPU is present; the host memory rows are the portable ones. Because its subprocess driver takes no arguments, it names the single parameterization it measures rather than sweeping them.