feat(working-memory): reserve FLUX.2 denoise working memory from measured lines

invokeai/backend/model_manager/load/load_base.py

@@ -261,10 +261,8 @@ def model_on_device(
                 f"Working memory for {self._cache_record.key}: {working_memory.estimator} reserves "
                 f"{working_mem_bytes / 2**20:.0f} MB ({working_memory.source})"
             )
-            scope = probe.current_scope()
-            if scope is not None and self.compute_device.type != "cpu":
-                calibration = scope
-                working_mem_bytes = probe.calibration_reservation(working_mem_bytes)
+            calibration, working_mem_bytes = probe.lock_reservation(working_mem_bytes, self.compute_device)
+            if calibration is not None:
                 message += f"; the calibration lock reserves {working_mem_bytes / 2**20:.0f} MB"
             InvokeAILogger.get_logger(__name__).debug(message)
         self._lock_on_device(working_mem_bytes, offload_others=calibration is not None)

invokeai/backend/util/vae_working_memory.py

@@ -120,9 +120,10 @@ def _vae_mid_block_score_matrix_bytes(
 # for it: a line in the output's pixel count per operation and the dtype it runs in, `k_<operation>_<dtype>`
 # bytes per pixel plus `b_<operation>_<dtype>` bytes the process pays whatever the size, shipped per convolution
 # backend and allocator (`platform_for`). The values ship in `backend/working_memory/data/vae.json` with the points
-# they are fitted on (`backend/working_memory/fit.py`); `scripts/calibrate_flux2_working_memory.py --only vae`
-# measures a build against them. Every spec here declares that line (`fit.convolution_line`) and opts in to user
-# calibration (`user_calibrated`), so a calibration run's decode windows refit these lines per device.
+# they are fitted on (`backend/working_memory/fit.py`): calibration runs in a running server measure them on a card,
+# and `scripts/working_memory/build_shipped.py` turns their exports into shipped points. Every spec here declares that
+# line (`fit.convolution_line`) and opts in to user calibration (`user_calibrated`), so a calibration run's decode
+# windows refit these lines per device.
 #
 # None of the three scales out of the line. The demand does not follow the element size: on MIOpen fp16 needs about a
 # third of what bf16 does, and on both backends fp32 needs less per value byte than bf16. The allocator matters as

invokeai/backend/working_memory/bounds.py

@@ -44,12 +44,20 @@
         "vae.sd-miopen,native,windows-encode-fp16-1280x1280-scores-RX9060XT-demand_peak": 1.78,
         "vae.flux-miopen,native,windows-decode-bf16-512x512-scores-RX9060XT-demand_peak": 1.33,
         "vae.flux2-miopen,native,windows-decode-bf16-512x512-scores-RX9060XT-demand_peak": 1.33,
+        "denoise.flux2-cudnn,native-denoise-bf16-h3072-8100+77tok-RTX4090-demand_peak": 1.42,
+        "denoise.flux2-miopen,native,windows-denoise-bf16-h3072-8100+77tok-RX9060XT-demand_peak": 1.40,
     }
 )
 """In-fit points, by `point_id`, whose envelope sits more than 1.3x above them, with the ratio they are held to
-instead. All are MIOpen with the native allocator on Windows, where the demand is concave in the size: first-run driver
-spikes at small sizes and a second run reserving up to 1.6x the first at large ones, so a line that covers every point
-cannot stay within 1.3x of all of them. Windows ROCm defaults to expandable segments; this is the override's price."""
+instead. All are the native allocator, never expandable segments or `cudaMallocAsync`:
+
+- The VAE points are MIOpen on Windows, where the demand is concave in the size: first-run driver spikes at small sizes
+  and a second run reserving up to 1.6x the first at large ones, so a line that covers every point cannot stay within
+  1.3x of all of them. Windows ROCm defaults to expandable segments; this is the override's price.
+- The FLUX.2 denoise points are the Klein 4B width at 1440px with a 77-token caption, 8177 tokens, on both cards: the
+  native allocator served them from 1676 MiB (RTX 4090) and 1879 MiB (RX 9060 XT), while the 512-token caption at
+  8612 tokens took 2382 and 2577 MiB. The line has to cover the latter. Under `cudaMallocAsync` and expandable
+  segments the two lie on one line."""
 
 
 def point_id(point: MeasuredPoint, spec: "EstimatorSpec") -> str:

invokeai/backend/working_memory/fit.py

@@ -3,9 +3,9 @@
 An estimator's fitted terms price a line per operation and dtype: `k_<suffix>` bytes per unit of the line's size plus
 `b_<suffix>` bytes, the base a process pays whatever the size. Its `EstimatorSpec.line` says which line an operation's
 features price on, their size on it, and whether a measurement of them can join its fit (`Line`). The AutoencoderKL
-estimators price their convolutions on a line in the output's pixels (`convolution_line`); the Z-Image denoise prices
-its forward on one in attended tokens. Each carried pair is the upper envelope of the fresh-process points it is
-resolved for (see `upper_envelope`), plus `SHIPPED_MARGIN`, raised where the line's other points need it
+estimators price their convolutions on a line in the output's pixels (`convolution_line`); the Z-Image and FLUX.2 denoise
+price their forwards on ones in attended tokens. Each carried pair is the upper envelope of the fresh-process points it
+is resolved for (see `upper_envelope`), plus `SHIPPED_MARGIN`, raised where the line's other points need it
 (`raised_pair`). A point joins the fit of the entry its estimator resolves to on its platform -- the lender's, where the
 estimator borrows -- when it is `in_fit` (not measured out of regime) and its line admits it.
 
@@ -43,7 +43,7 @@
 from typing import TYPE_CHECKING, Final, NamedTuple, Optional
 
 from invokeai.backend.working_memory.coefficients import CoefficientEntry, MeasuredPoint, WorkingMemoryData, find_entry
-from invokeai.backend.working_memory.estimate import FeatureValue
+from invokeai.backend.working_memory.estimate import FeatureValue, WorkingMemoryTerm
 
 if TYPE_CHECKING:
     from invokeai.backend.working_memory.spec import EstimatorSpec
@@ -88,6 +88,25 @@ class Line(NamedTuple):
 features that name no line."""
 
 
+def off_line_reasons(
+    terms: Callable[[Mapping[str, FeatureValue], Mapping[str, float]], Iterable[WorkingMemoryTerm]],
+    features: Mapping[str, FeatureValue],
+    coefficients: Iterable[str],
+) -> list[str]:
+    """What puts an operation off the line its fitted terms price: an attention mask, which narrows the kernels SDPA can
+    pick (``has_attention_mask``), and each analytic term of `terms` that is not zero, as ``<name>=<bytes>``. Empty for
+    an operation whose demand is the line's alone. The terms are priced with every coefficient at zero, so only what
+    the operation adds on top of the line remains; a term added to an estimator takes the operations it prices out of
+    the fit by itself. The estimators label their points with the reasons in this order."""
+    unpriced = dict.fromkeys(coefficients, 0.0)
+    reasons = ["attention_mask"] if features.get("has_attention_mask") else []
+    return reasons + [
+        f"{term.name}={int(term.bytes)}"
+        for term in terms(features, unpriced)
+        if term.coefficient is None and term.bytes
+    ]
+
+
 def convolution_line(features: Mapping[str, FeatureValue]) -> Optional[Line]:
     """The AutoencoderKL convolution line ``<operation>_<dtype>``, in output pixels: every VAE estimator's.
 

invokeai/backend/working_memory/measurement.py

@@ -21,11 +21,12 @@
 Two kinds of window take these definitions. The in-server probe's (`probe`) follow an operation through co-locks, nested
 windows and releases, in segments, and have reasons of their own. A *fresh window* (`measure`) is one operation alone:
 the calibration scripts measure their points that way, each in a fresh process with the server's runtime
-(`invokeai.app.run_app.apply_server_runtime`).
+(`invokeai.app.run_app.apply_server_runtime`). Such a point is a cold and a warm window of the same operation
+(`measure_cold_and_warm`), and its bytes are the larger of the two (`measured_bytes_and_method`).
 """
 
 import time
-from collections.abc import Callable, Iterable
+from collections.abc import Callable, Iterable, Sequence
 from dataclasses import dataclass
 from typing import Final, Literal, Optional
 
@@ -181,3 +182,42 @@ def measure(device: torch.device, operation: Callable[[], object]) -> FreshWindo
         return fresh_window(first, final, oom=oom, spill=spill)
     finally:
         device_counters.close()
+
+
+FreshPointMethod = Literal["demand_peak", "reserved_peak"]
+"""The `MeasuredPoint.method` of a fresh-process point."""
+
+
+def measure_cold_and_warm(device: torch.device, operation: Callable[[], object]) -> tuple[FreshWindow, ...]:
+    """A fresh-process point's windows: `operation` measured twice (`measure`), cold and then warm.
+
+    The first run after a load is a real peak in a server (kernel modules, library workspaces), and so is a later one,
+    which can demand more: a warm Z-Image forward under `cudaMallocAsync` read 800 MiB against its cold run's 720. A
+    cold run that ran out of memory ends the point, whose warm run would measure what the failure left behind.
+    """
+    cold = measure(device, operation)
+    if "oom" in cold.reasons:
+        return (cold,)
+    return (cold, measure(device, operation))
+
+
+def point_reasons(windows: Sequence[FreshWindow]) -> list[ReadingReason]:
+    """Why a fresh-process point is invalid: its windows' reasons, each once, in order. Empty for a valid point."""
+    return list(dict.fromkeys(reason for window in windows for reason in window.reasons))
+
+
+def measured_bytes_and_method(windows: Sequence[FreshWindow]) -> tuple[int, FreshPointMethod]:
+    """A valid fresh-process point's bytes and `MeasuredPoint.method`, from its windows (`measure_cold_and_warm`).
+
+    The largest of each window's figure: its demand where it read the driver, otherwise the growth of torch's counters
+    alone. A window without the driver's figure does not discard another's, which on ROCm maps a few hundred MiB above
+    torch's reserve on a first forward. The method names what the largest was read from: `demand_peak` for a demand,
+    `reserved_peak` for torch's counters.
+    """
+    figures: list[tuple[int, FreshPointMethod]] = []
+    for window in windows:
+        if window.driver is not None and window.demand is not None:
+            figures.append((window.demand, "demand_peak"))
+        else:
+            figures += [(delta, "reserved_peak") for delta in (window.reserved, window.allocated) if delta is not None]
+    return max(figures, key=lambda figure: figure[0])

invokeai/backend/working_memory/probe.py

@@ -61,6 +61,7 @@
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Final, Literal, Optional
 
+import torch
 from pydantic import BaseModel, ConfigDict
 
 from invokeai.app.services.session_processor.session_processor_common import CanceledException
@@ -595,6 +596,17 @@ def current_scope() -> Optional[ProbeScope]:
     return _current_scope.get()
 
 
+def lock_reservation(reservation_bytes: int, device: torch.device) -> tuple[Optional[ProbeScope], int]:
+    """What a lock on `device` keeps free for an estimate of `reservation_bytes`, and the calibration scope it locks
+    under: inside a scope open on this thread and on a GPU, that scope and the calibration reservation
+    (`calibration_reservation`); otherwise None and the estimate's bytes. The model loader applies it to a measured
+    lock, and a node to every lock beside one that must not use up the window's headroom."""
+    scope = current_scope()
+    if scope is None or device.type == "cpu":
+        return None, reservation_bytes
+    return scope, calibration_reservation(reservation_bytes)
+
+
 class _EmptyCacheObserver:
     """Turns `TorchDevice.empty_cache` into segment boundaries of the open windows."""
 

invokeai/backend/working_memory/spec.py

@@ -19,6 +19,7 @@
     "invokeai.backend.moss_image.working_memory",
     "invokeai.backend.util.vae_working_memory",
     "invokeai.backend.z_image.denoise_working_memory",
+    "invokeai.backend.flux2.denoise_working_memory",
 )
 """Every module that registers an `EstimatorSpec`. The registry holds only the specs of modules imported so far, so a
 caller that needs all of them imports these (`import_estimator_modules`); `tests/backend/working_memory/` fails on a
@@ -46,7 +47,7 @@ class EstimatorSpec:
     line: LineFunction = field(compare=False, repr=False)
     """The line an operation's features price on, their size on it, and whether a measurement of them fits it: the
     AutoencoderKL convolution line in output pixels for every VAE estimator (`fit.convolution_line`), one in attended
-    tokens for the Z-Image denoise. Declared by every estimator, so a new one cannot inherit another's line, and with it
+    tokens for each denoise estimator. Declared by every estimator, so a new one cannot inherit another's line, and with it
     which of its measurements join a fit."""
     user_calibrated: bool = False
     """Whether a user's calibration moves this estimator: its windows fit user lines, its estimates read them, the

invokeai/backend/z_image/denoise_working_memory.py

@@ -61,7 +61,6 @@
 """
 
 from collections.abc import Mapping
-from types import MappingProxyType
 from typing import Final, Optional, cast
 
 import torch
@@ -71,7 +70,7 @@
 from invokeai.backend.util.attention import sdpa_score_matrix_bytes
 from invokeai.backend.util.fp8 import peak_fp8_cast_transient_bytes
 from invokeai.backend.working_memory.estimate import FeatureValue, WorkingMemoryEstimate, WorkingMemoryTerm
-from invokeai.backend.working_memory.fit import Line
+from invokeai.backend.working_memory.fit import Line, off_line_reasons
 from invokeai.backend.working_memory.platform import platform_for
 from invokeai.backend.working_memory.spec import EstimatorSpec, register
 
@@ -86,9 +85,6 @@
 """The dtypes `TorchDevice.choose_bfloat16_safe_dtype` runs the transformer in."""
 _ELEMENT_SIZES: Final[Mapping[str, int]] = {"bf16": 2, "fp16": 2, "fp32": 4}
 _COEFFICIENTS: Final = frozenset(f"{part}_denoise_{dtype}" for part in ("k", "b") for dtype in _ELEMENT_SIZES)
-_UNPRICED: Final[Mapping[str, float]] = MappingProxyType(dict.fromkeys(_COEFFICIENTS, 0.0))
-"""Coefficient values that price the fitted terms at nothing: what remains of a forward's terms is what it adds on top
-of its line."""
 
 
 def padded_tokens(tokens: int) -> int:
@@ -117,9 +113,7 @@ def _line(features: Mapping[str, FeatureValue]) -> Optional[Line]:
     dtype, tokens = features.get("dtype"), features.get("tokens")
     if not isinstance(dtype, str) or dtype not in _ELEMENT_SIZES or not isinstance(tokens, int):
         return None
-    on_top = [term for term in _terms(features, _UNPRICED) if term.coefficient is None and term.bytes]
-    off_line = ["attention_mask"] if features.get("has_attention_mask") else []
-    off_line += [f"{term.name}={int(term.bytes)}" for term in on_top]
+    off_line = off_line_reasons(_terms, features, _COEFFICIENTS)
     shape = f"{features.get('image_tokens')}+{features.get('text_tokens')}tok"
     return Line(f"denoise_{dtype}", tokens, not off_line, "-".join(("denoise", dtype, shape, *off_line)))