Honor memory_mode over legacy cpu_offload and prefetch dense-quant transformer shards

plan_diffusion_memory only applies the legacy cpu_offload override when no
memory_mode was supplied, matching the documented API contract that
memory_mode overrides cpu_offload when set; an explicit fast request now
stays resident even if the old flag is also enabled.

The transformer-quant dense path fetches the base repo's transformer/
shards inside the locked finalize phase, where unload and cancellation
cannot preempt the multi-GB download. The load worker now widens the
preemptible prefetch to include those shards when that path can actually
run: quant requested and supported for the device, scheme resolvable, and
no pre-quantized checkpoint shortcutting the dense build.

studio/backend/core/inference/diffusion.py

@@ -219,6 +219,34 @@ class DiffusionBackend:
 
         return hf_hub_download(repo_id, gguf_filename, token = hf_token)
 
+    def _dense_quant_prefetch_needed(self, fam: DiffusionFamily, kwargs: dict) -> bool:
+        """True when ``load_pipeline`` may take the dense transformer-quant path, so
+        the prefetch should also pull the base repo's ``transformer/`` shards.
+
+        Those shards are excluded from the prefetch by default (the GGUF supplies
+        the transformer), but ``_load_dense_quant_pipeline`` fetches them with
+        ``from_pretrained(subfolder = "transformer")`` under the load lock during
+        "finalizing", after the previous pipeline was already evicted, where
+        unload/cancellation cannot preempt the download. Mirrors the dense-path
+        gates in ``load_pipeline``: quant requested and supported for this device,
+        and no pre-quantized checkpoint that would shortcut the dense build."""
+        mode = normalize_transformer_quant(kwargs.get("transformer_quant"))
+        if mode is None:
+            return False
+        try:
+            target = self._resolve_device_target(fam)
+            if not dense_transformer_supported(target):
+                return False
+            scheme = select_transformer_quant_scheme(target, mode)
+            if scheme is None:
+                return False
+            source = resolve_prequant_source(
+                fam, scheme, path_override = kwargs.get("transformer_prequant_path")
+            )
+            return source is None
+        except Exception:  # noqa: BLE001 — widening the prefetch is best-effort only
+            return False
+
     def _prefetch_files(
         self,
         repo_id: str,
@@ -364,7 +392,16 @@ class DiffusionBackend:
             )
             kwargs["base_repo"] = base
             expected, base_files = self._estimate_download_bytes(
-                kwargs["repo_id"], kwargs.get("gguf_filename"), base, kwargs.get("hf_token")
+                kwargs["repo_id"],
+                kwargs.get("gguf_filename"),
+                base,
+                kwargs.get("hf_token"),
+                # The dense transformer-quant path downloads the base repo's
+                # transformer/ shards via from_pretrained(subfolder="transformer")
+                # INSIDE the locked finalize phase, where unload/cancellation cannot
+                # preempt the multi-GB pull. When that path can actually run, pull the
+                # shards here in the preemptible prefetch instead.
+                include_transformer = self._dense_quant_prefetch_needed(fam, kwargs),
             )
             with self._lock:
                 # Stamp progress only if this load is still current; a superseding
@@ -433,7 +470,12 @@ class DiffusionBackend:
 
     @staticmethod
     def _estimate_download_bytes(
-        repo_id: str, gguf_filename: Optional[str], base_repo: str, hf_token: Optional[str]
+        repo_id: str,
+        gguf_filename: Optional[str],
+        base_repo: str,
+        hf_token: Optional[str],
+        *,
+        include_transformer: bool = False,
     ) -> tuple[int, list[str]]:
         """Total download size for the progress bar, plus the base-repo files to
         fetch (the prefetch reuses this list, so the base is listed only once)."""
@@ -452,7 +494,7 @@ class DiffusionBackend:
                 total += sum(s.size or 0 for s in info.siblings if s.rfilename == gguf_filename)
             base_info = api.model_info(base_repo, files_metadata = True, token = hf_token)
             for s in base_info.siblings:
-                if _base_file_downloaded(s.rfilename):
+                if _base_file_downloaded(s.rfilename, include_transformer = include_transformer):
                     base_files.append(s.rfilename)
                     total += s.size or 0
         except Exception as exc:  # noqa: BLE001 — estimate is best-effort
@@ -1156,16 +1198,18 @@ def _hf_base_model(repo_id: str, hf_token: Optional[str]) -> Optional[str]:
     return base if isinstance(base, str) and base.strip() else None
 
 
-def _base_file_downloaded(rfilename: str) -> bool:
+def _base_file_downloaded(rfilename: str, *, include_transformer: bool = False) -> bool:
     """True for base-repo files ``from_pretrained`` actually fetches.
 
     The transformer is supplied by the GGUF, and repo docs (``assets/``, the
     top-level README/PDF/images) are never downloaded — counting them would peg
     the progress estimate above what lands on disk, so the bar would sit short of
     100% for the whole pipeline-load phase instead of advancing to "finalizing".
-    """
+    ``include_transformer`` admits the ``transformer/`` shards for loads where the
+    dense transformer-quant path will fetch them anyway (see
+    ``_dense_quant_prefetch_needed``)."""
     if rfilename.startswith("transformer/"):
-        return False
+        return include_transformer
     if "/" not in rfilename:  # top-level: only the pipeline manifest is fetched
         return rfilename == "model_index.json"
     return not rfilename.startswith("assets/")

studio/backend/core/inference/diffusion_memory.py

@@ -367,7 +367,17 @@ def plan_diffusion_memory(
         policy = OFFLOAD_MODEL
         reasons.append("companions exceed budget; whole-module offload of every component")
 
-    if explicit_offload and policy == OFFLOAD_NONE and can_offload and not device_memory.is_unified:
+    # The legacy cpu_offload flag only applies when NO memory_mode was supplied:
+    # the API documents memory_mode as overriding cpu_offload when set, so an
+    # explicit `fast` request must stay resident even if the caller also left the
+    # old flag enabled.
+    if (
+        explicit_offload
+        and normalize_memory_mode(requested_mode) is None
+        and policy == OFFLOAD_NONE
+        and can_offload
+        and not device_memory.is_unified
+    ):
         policy = OFFLOAD_MODEL
         reasons.append("explicit cpu_offload overrides resident placement")
 

studio/backend/tests/test_diffusion_backend.py

@@ -1193,6 +1193,52 @@ def test_transformer_quant_unsupported_scheme_skips_dense_download(
     assert _FakeTransformer.last["path"]  # GGUF from_single_file used
 
 
+def test_base_file_downloaded_include_transformer_flag():
+    # Default: transformer/ shards are the GGUF's job, so they are excluded from
+    # the prefetch list; the dense transformer-quant path opts them back in.
+    from core.inference.diffusion import _base_file_downloaded
+
+    assert _base_file_downloaded("transformer/diffusion_pytorch_model-00001.safetensors") is False
+    assert (
+        _base_file_downloaded(
+            "transformer/diffusion_pytorch_model-00001.safetensors", include_transformer = True
+        )
+        is True
+    )
+    # The flag must not admit anything else that is normally excluded.
+    assert _base_file_downloaded("assets/teaser.png", include_transformer = True) is False
+    assert _base_file_downloaded("README.md", include_transformer = True) is False
+
+
+def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch):
+    # The transformer/ prefetch only widens when the dense quant path can really
+    # run: quant requested + device supported + scheme resolvable + no prequant
+    # checkpoint shortcutting the dense build.
+    from core.inference import diffusion as dmod
+
+    backend = DiffusionBackend()
+    _force_cuda_target(backend, monkeypatch)
+    fam = detect_family("unsloth/Z-Image-Turbo-GGUF")
+    monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
+    monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8")
+    monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None)
+
+    assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is True
+    # No quant requested -> never widen.
+    assert backend._dense_quant_prefetch_needed(fam, {}) is False
+    # A resolvable pre-quantized checkpoint shortcuts the dense download.
+    monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: object())
+    assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False
+    # Unsupported scheme bails before the dense path (and so must the prefetch).
+    monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None)
+    monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: None)
+    assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False
+    # Device without dense support (e.g. non-CUDA) never widens.
+    monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8")
+    monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: False)
+    assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False
+
+
 def test_companion_cache_bytes_local_dir_excludes_transformer(tmp_path):
     # A LOCAL diffusers base: sum the on-disk VAE / text-encoder weights so auto memory
     # planning sees the resident companions, but exclude transformer/ (the GGUF supplies

studio/backend/tests/test_diffusion_memory.py

@@ -270,6 +270,21 @@ def test_explicit_cpu_offload_overrides_resident_auto_choice():
     assert any("explicit cpu_offload" in r for r in plan.reasons)
 
 
+def test_explicit_memory_mode_wins_over_legacy_cpu_offload():
+    # The API documents memory_mode as overriding cpu_offload when set: fast +
+    # the legacy flag must stay resident, not silently downgrade to offload.
+    plan = plan_diffusion_memory(
+        target = _target(),
+        device_memory = _discrete(80000),
+        model_dense_mib = 4000,
+        runtime_headroom_mib = 2000,
+        requested_mode = MEMORY_MODE_FAST,
+        explicit_offload = True,
+    )
+    assert plan.offload_policy == OFFLOAD_NONE
+    assert not any("explicit cpu_offload" in r for r in plan.reasons)
+
+
 def test_explicit_cpu_offload_ignored_on_cpu_target():
     plan = plan_diffusion_memory(
         target = _target(device = "cpu", backend = "cpu", supports_offload = False),