nvalchemi.distributed.resolve_cap#

nvalchemi.distributed.resolve_cap(state, key, real, *, initial_factor, grow_factor=1.30, stride=16, extra=0, strict_gt=True)[source]#

Grow-only fixed-shape capacity for key, >= real + extra.

A cap is sized on first sight with initial_factor headroom, regrows with grow_factor only when the real count would overflow, and is always rounded up to a multiple of stride so small MD-step fluctuation lands in the same bucket — keeping the compiled graph from recompiling. The cap only ever grows, so a hot path reuses one compiled graph.

Parameters:

state (dict[str, int]) – Mutable dict holding the persistent caps across forwards (the caller owns its lifetime).
key (str) – Which capacity (e.g. "atoms" / "edges" / "max_send").
real (int) – The real count needed this step (before padding).
initial_factor (float) – Headroom multiplier applied the first time key is sized, set to cover the equilibrated peak from the first compile (e.g. edges climb ~25% through equilibration, atoms barely move).
grow_factor (float) – Headroom multiplier applied on a later overflow.
stride (int) – Bucket size; the cap is rounded up to a multiple of this (16 for kernel-friendly shapes; coarser counts that swing more use a larger one).
extra (int) – Slots reserved beyond real (e.g. UMA reserves 2 for the dead-edge anchor pair).
strict_gt (bool) – Overflow test: real + extra > cap when True (the default; edges / send), or >= cap when False (atoms, which need a strictly-larger cap because the dead row sits at cap - 1).

Returns:

The (possibly grown) capacity for key, recorded in state.

Return type:

int