th#1043: joint precision fit for shared-component multi-lane loads

src/gen_worker/executor.py

@@ -505,6 +505,43 @@ def _model_failure_vocab(exc: BaseException) -> str:
     return "load_failed"
 
 
+def _shared_lanes_need_fp8(
+    slot_sizes: Dict[str, Dict[str, int]],
+    shared_components: typing.Iterable[str],
+    free_vram_bytes: int,
+    *,
+    margin_bytes: int = 2 * (1 << 30),
+) -> bool:
+    """Joint VRAM-fit decision for a shared-component multi-lane record
+    (th#1043): loading each lane's precision reactively, one at a time,
+    lets the FIRST lane to load consume all free VRAM at native precision
+    — the shared-component invariant then refuses the offload placement
+    the STARVED lane needs, hard-failing a fit that was achievable all
+    along. Decide precision for the WHOLE group against its combined
+    footprint (shared components counted once) before any lane loads.
+
+    True when native precision doesn't fit the group but fp8-storage
+    (denoiser weights ~halved) does — every lane in the group should force
+    fp8 storage. False when native precision already fits (no forcing
+    needed) or fp8 storage still wouldn't fit (per-lane reactive sizing
+    makes its own honest call instead).
+    """
+    shared = list(shared_components)
+    if len(slot_sizes) < 2 or not shared:
+        return False
+    first = next(iter(slot_sizes.values()))
+    shared_bytes = sum(first.get(c, 0) for c in shared)
+    exclusive_bytes = sum(
+        b for sizes in slot_sizes.values()
+        for comp, b in sizes.items() if comp not in shared
+    )
+    needed = shared_bytes + exclusive_bytes
+    if needed + margin_bytes <= free_vram_bytes:
+        return False
+    fp8_needed = shared_bytes + 0.5 * exclusive_bytes
+    return fp8_needed + margin_bytes <= free_vram_bytes
+
+
 def _is_corrupt_load_error(exc: BaseException) -> bool:
     """Errors a truncated/corrupt snapshot produces at weights-load time
     (gw#408). Broad on purpose: the digest re-verify gate downstream
@@ -5201,6 +5238,34 @@ def _component_share_plan(
         )
         return plan
 
+    def _shared_group_force_fp8(
+        self, spec: EndpointSpec, share_plan: Optional[Dict[str, Dict[str, Any]]],
+    ) -> set:
+        """Slots that must force fp8 denoiser storage to fit their shared-
+        component group jointly resident (th#1043) — empty when the group
+        fits at native precision, or fits nobody's precision at all."""
+        plan = share_plan or {}
+        slots = [s for s, m in plan.items() if m]
+        if len(slots) < 2:
+            return set()
+        shared_components = sorted({c for m in plan.values() for c in m})
+        slot_sizes: Dict[str, Dict[str, int]] = {}
+        for slot in slots:
+            binding = spec.models.get(slot)
+            if binding is None:
+                return set()
+            slot_sizes[slot] = self.store.component_sizes(wire_ref(binding))
+        free = self.store.residency.free_vram_bytes()
+        if not _shared_lanes_need_fp8(slot_sizes, shared_components, free):
+            return set()
+        logger.info(
+            "th#1043: shared-lane group %s for %s doesn't fit resident at "
+            "native precision (free=%.1fGiB) — forcing fp8 storage on every "
+            "lane before any of them loads",
+            slots, spec.name, free / (1 << 30),
+        )
+        return set(slots)
+
     @staticmethod
     def _model_index_components(path: str) -> set:
         """Component names the snapshot's model_index.json declares — the
@@ -5245,6 +5310,7 @@ async def _injection_kwargs(
         result = _InjectionResult(kwargs=kwargs, loaded=loaded)
         compile_artifact = compile_selection.path if compile_selection else None
         share_plan = self._component_share_plan(spec, paths, hints)
+        force_fp8_slots = self._shared_group_force_fp8(spec, share_plan)
         if server is not None:
             for pname, ann in hints.items():
                 if ann is ServerHandle:
@@ -5334,6 +5400,8 @@ async def _injection_kwargs(
                         slot=slot, ref=ref, mode=mode, components=injected,
                         declared_vram_gb=float(
                             getattr(spec.resources, "vram_gb", 0) or 0),
+                        force_storage_dtype=(
+                            "fp8" if slot in force_fp8_slots else ""),
                     )
                 except Exception as exc:
                     # Corruption-shaped load failure (gw#408): digest-verify
@@ -5358,6 +5426,8 @@ async def _injection_kwargs(
                         slot=slot, ref=ref, mode=mode, components=injected,
                         declared_vram_gb=float(
                             getattr(spec.resources, "vram_gb", 0) or 0),
+                        force_storage_dtype=(
+                            "fp8" if slot in force_fp8_slots else ""),
                     )
                 pipe = sl.obj
                 # Reconcile the load outcomes into ServePlan/FnDegraded via

src/gen_worker/models/provision.py

@@ -84,6 +84,7 @@ def load_slot(
     components: Optional[Dict[str, Any]] = None,
     device: str = "",
     declared_vram_gb: float = 0.0,
+    force_storage_dtype: str = "",
 ) -> SlotLoad:
     """Typed slot injection: the slot receives exactly what its ``setup``
     annotation says — a ``str``/``Path`` local path, or a constructed
@@ -96,7 +97,11 @@ def load_slot(
     degraded floors — the executor's knowledge; the CLI passes ``auto``).
     ``device="cpu"`` (CLI ``--device cpu``) skips placement entirely.
     ``components`` are preloaded shared modules (gw#479) forwarded to
-    ``from_pretrained``.
+    ``from_pretrained``. ``force_storage_dtype`` overrides the binding's own
+    storage_dtype (th#1043): a joint multi-lane fit decision made BEFORE any
+    lane in a shared-component group loads, so the first lane to load never
+    greedily consumes free VRAM at native precision and starves a sibling
+    lane into an offload placement the shared-component invariant refuses.
     """
     if annotation is None or annotation is str:
         return SlotLoad(obj=path)
@@ -110,7 +115,7 @@ def load_slot(
     from .memory import place_pipeline
 
     dtype = str(getattr(binding, "dtype", "") or "")
-    storage_dtype = str(getattr(binding, "storage_dtype", "") or "")
+    storage_dtype = force_storage_dtype or str(getattr(binding, "storage_dtype", "") or "")
     out = SlotLoad(obj=None, is_pipeline=True, ran=(dtype or "bf16"))
 
     # th#737: a cast directive on a denoiser-less diffusers tree is a

tests/test_shared_lane_precision_th1043.py

@@ -0,0 +1,54 @@
+"""th#1043: joint precision fit for shared-component multi-lane records.
+
+Real qwen-image evidence: two lanes (t2i/edit) share a ~15.5GB text
+encoder+VAE, each carries an exclusive ~40GB bf16 transformer. Loading each
+lane's precision reactively — one at a time, against whatever free VRAM
+happens to be left at that moment — let the FIRST lane consume all headroom
+at native precision, starving the second into an offload placement the
+shared-component invariant then refuses outright (RetryableError,
+"shared-component lanes require resident placement"). The fix decides
+precision for the WHOLE group up front, against its COMBINED footprint.
+"""
+
+from __future__ import annotations
+
+from gen_worker.executor import _shared_lanes_need_fp8
+
+_GiB = 1024 ** 3
+
+# Real th#1043 shapes: shared text_encoder+vae ~15.5GB, each transformer
+# ~40GB bf16 (no fp8 checkpoint exists yet for this family).
+_QWEN_SIZES = {
+    "t2i": {"text_encoder": 15 * _GiB, "vae": int(0.5 * _GiB), "transformer": 40 * _GiB},
+    "edit": {"text_encoder": 15 * _GiB, "vae": int(0.5 * _GiB), "transformer": 40 * _GiB},
+}
+_SHARED = ["text_encoder", "vae"]
+
+
+def test_starved_group_forces_fp8_when_it_fits():
+    # A100-80GB: ~79GiB usable. Native (15.5 + 40 + 40 = 95.5GiB) doesn't
+    # fit; fp8'd transformers (15.5 + 20 + 20 = 55.5GiB) do.
+    free = 79 * _GiB
+    assert _shared_lanes_need_fp8(_QWEN_SIZES, _SHARED, free) is True
+
+
+def test_group_that_fits_natively_is_left_alone():
+    # A card with enough headroom for both lanes at native precision needs
+    # no forcing.
+    free = 200 * _GiB
+    assert _shared_lanes_need_fp8(_QWEN_SIZES, _SHARED, free) is False
+
+
+def test_group_that_cannot_fit_even_at_fp8_is_left_to_the_ladder():
+    # Too small even halved: forcing fp8 here would just relocate the
+    # failure, not fix it — leave it to the existing per-lane ladder.
+    free = 30 * _GiB
+    assert _shared_lanes_need_fp8(_QWEN_SIZES, _SHARED, free) is False
+
+
+def test_single_lane_group_never_forces():
+    assert _shared_lanes_need_fp8({"t2i": _QWEN_SIZES["t2i"]}, _SHARED, 1 * _GiB) is False
+
+
+def test_no_shared_components_never_forces():
+    assert _shared_lanes_need_fp8(_QWEN_SIZES, [], 1 * _GiB) is False