ideogram-4: build the FP8 text encoder at target dtype, size it as bf16-resident, optimize both DiTs

- The FP8 Qwen3-VL text encoder was constructed at the process fp32 default before the
  dequantized bf16 weights are copied in. That ~8B-param fp32 scaffold peaks ~2x on host RAM
  (loading FIRST, before the DiTs), so a 64 GB host can OOM. Build it at the target dtype under
  set_default_dtype, mirroring the DiT loader; rotary inv_freq is still computed in explicit fp32.

- The auto-policy memory table listed the text encoder at 8.8 GB, its FP8 on-disk size, while the
  DiTs were doubled to their bf16-resident sizes. The loader dequantizes the encoder to bf16 too
  (~16.3 GB), so the entry understated the resident footprint by ~7.5 GB and could let the planner
  pick a resident placement that OOMs. Size it as bf16-resident.

- Speed (regional compile, QKV fuse) and the attention backend only touched pipe.transformer, so
  ideogram-4's second denoiser (unconditional_transformer, run every step for dual-branch CFG)
  stayed eager/native while status reported the optimization as engaged. Iterate every denoiser DiT
  (mirroring the offload path) so both experts are optimized. Guarded on attr presence, so single-DiT
  families are unchanged.

studio/backend/core/inference/diffusion_attention.py

@@ -244,13 +244,27 @@ def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> Non
                 )
 
 
+def _attention_dits(pipe: Any) -> list:
+    """Every DiT the denoise loop runs each step: the primary ``transformer`` plus a second
+    expert some families carry (Ideogram's ``unconditional_transformer`` for its dual-branch
+    CFG, an MoE ``transformer_2``). The attention backend must be set on ALL of them, else the
+    second DiT keeps the native default while status reports the requested kernel as engaged."""
+    dits: list = []
+    for attr in ("transformer", "transformer_2", "unconditional_transformer"):
+        m = getattr(pipe, attr, None)
+        if m is not None and m not in dits:
+            dits.append(m)
+    return dits
+
+
 def apply_attention_backend(
     pipe: Any,
     backend: Optional[str],
     *,
     logger: Any = None,
 ) -> Optional[str]:
-    """Set ``backend`` on ``pipe.transformer`` via the diffusers dispatcher.
+    """Set ``backend`` on EVERY denoiser DiT (``pipe.transformer`` plus a second expert such as
+    Ideogram's ``unconditional_transformer``) via the diffusers dispatcher.
 
     Returns the backend actually engaged, or None when left at the native default (either
     because ``backend`` was None or because the requested kernel was unavailable -> graceful
@@ -261,28 +275,32 @@ def apply_attention_backend(
     defaults to None). So a load that wants native must restore it explicitly: otherwise it
     silently inherits a backend an earlier load pinned (e.g. cuDNN under a speed profile),
     breaking the bit-identical/``off`` guarantee. Best-effort throughout."""
-    transformer = getattr(pipe, "transformer", None)
-    fn = getattr(transformer, "set_attention_backend", None)
-    if not callable(fn):
+    setters = [s for s in (getattr(t, "set_attention_backend", None) for t in _attention_dits(pipe))
+               if callable(s)]
+    if not setters:
         return None
     if backend is not None:
         _ensure_attention_backend_installed(backend, logger)
-        try:
-            fn(backend)
-            # set_attention_backend also pins the backend in diffusers' process-wide
-            # registry. This transformer's own processors keep it locally (their
-            # _attention_backend is now explicit), so reset the global default back to
-            # native -- otherwise a later component whose processors are unconfigured
-            # (backend None) silently inherits this kernel.
+        engaged = False
+        for fn in setters:
+            try:
+                fn(backend)
+                engaged = True
+            except Exception as exc:  # noqa: BLE001 — unavailable kernel -> restore native below
+                _warn(logger, backend, exc)
+        if engaged:
+            # set_attention_backend also pins the backend in diffusers' process-wide registry.
+            # Each DiT's own processors now keep it locally (their _attention_backend is now
+            # explicit), so reset the global default back to native ONCE -- otherwise a later
+            # component whose processors are unconfigured (backend None) inherits this kernel.
             _reset_global_backend_to_native(logger)
             if logger is not None:
                 logger.info("diffusion.attention: backend=%s", backend)
             return backend
-        except Exception as exc:  # noqa: BLE001 — unavailable kernel -> restore native below
-            _warn(logger, backend, exc)
-    # No backend requested, or the requested one failed: pin the native default so a stale
-    # process-wide backend from a previous load can't leak into this one.
-    _restore_native_backend(fn, logger)
+    # No backend requested, or every set failed: pin the native default so a stale process-wide
+    # backend from a previous load can't leak into this one. Fresh DiTs follow the process-wide
+    # backend, so one reset via any DiT's setter covers them all.
+    _restore_native_backend(setters[0], logger)
     return None
 
 

studio/backend/core/inference/diffusion_auto_policy.py

@@ -61,10 +61,11 @@ _FAMILY_BF16_GB: dict[str, tuple[float, float, float]] = {
     "krea-2": (26.3, 8.9, 0.5),
     # Two ~9.3B DiTs (the conditional transformer PLUS the separate
     # unconditional_transformer driving Ideogram's dual-branch CFG), both resident
-    # for every generation, and a Qwen3-VL text encoder. The vendor repo stores the
-    # DiTs as raw float8 (9.29 GB each); these are the bf16-resident sizes after the
-    # loader's dtype cast, per this table's contract.
-    "ideogram-4": (37.2, 8.8, 0.2),
+    # for every generation, and a Qwen3-VL text encoder. The vendor repo stores the DiTs
+    # AND the text encoder as raw float8 (9.29 GB per DiT, 8.8 GB encoder); these are the
+    # bf16-resident sizes after the loader's dtype cast, per this table's contract, so the
+    # encoder doubles to ~16.3 GB just like each DiT (37.2 = 2 x 18.6).
+    "ideogram-4": (37.2, 16.3, 0.2),
 }
 
 # Base-repo overrides for families whose picker offers multiple sizes under one family

studio/backend/core/inference/diffusion_ideogram4.py

@@ -308,8 +308,16 @@ def load_ideogram4_text_encoder(
 
     # Construct normally (so __init__ computes the non-persistent rotary inv_freq
     # buffers the checkpoint omits) then copy the dequantized weights in with
-    # assign=False. Host RAM is ample, so the transient dense init is fine.
-    model = Qwen3VLModel(config).to(dtype)
+    # assign=False. Build at the target dtype (mirrors the DiT loader below): this ~8B-param
+    # Qwen3-VL scaffold is ~2x at the process fp32 default (~33 GB vs ~16 GB) and loads FIRST
+    # on host RAM, so the fp32 transient can OOM a 64 GB host. rotary inv_freq is computed in
+    # explicit fp32 in __init__, so a bf16 default leaves it correct.
+    default_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    try:
+        model = Qwen3VLModel(config).to(dtype)
+    finally:
+        torch.set_default_dtype(default_dtype)
     missing, unexpected = model.load_state_dict(state_dict, strict = False)
     real_missing = [k for k in missing if not k.endswith("inv_freq")]
     if real_missing or unexpected:

studio/backend/core/inference/diffusion_speed.py

@@ -272,6 +272,20 @@ def _vae_channels_last(pipe: Any, logger: Any) -> bool:
         return False
 
 
+def _denoiser_dits(pipe: Any) -> list:
+    """Every DiT the denoise loop runs each step: the primary ``transformer`` plus a second
+    expert some families carry (Ideogram's ``unconditional_transformer`` for its dual-branch
+    CFG, an MoE ``transformer_2``). Speed / attention optims must reach ALL of them -- mirroring
+    the offload path (diffusion_memory streams the same set) -- else the second DiT runs
+    eager / native for every generation while status over-reports the optimisation as engaged."""
+    dits: list = []
+    for attr in ("transformer", "transformer_2", "unconditional_transformer"):
+        m = getattr(pipe, attr, None)
+        if m is not None and m not in dits:
+            dits.append(m)
+    return dits
+
+
 def _compile_repeated_blocks(
     pipe: Any,
     logger: Any,
@@ -280,9 +294,8 @@ def _compile_repeated_blocks(
     cache_active: bool = False,
     offload_active: bool = False,
 ) -> bool:
-    transformer = getattr(pipe, "transformer", None)
-    fn = getattr(transformer, "compile_repeated_blocks", None)
-    if not callable(fn):
+    dits = [t for t in _denoiser_dits(pipe) if callable(getattr(t, "compile_repeated_blocks", None))]
+    if not dits:
         return False
     # default: mode="default" + dynamic=True -- fast cold start, robust to resolution
     # changes (no recompile). max: mode="max-autotune-no-cudagraphs" + dynamic=False --
@@ -319,11 +332,19 @@ def _compile_repeated_blocks(
             for _limit_attr in ("recompile_limit", "cache_size_limit"):  # name varies by torch ver
                 if hasattr(dynamo_cfg, _limit_attr):
                     setattr(dynamo_cfg, _limit_attr, max(getattr(dynamo_cfg, _limit_attr) or 0, 64))
-        fn(**kwargs)
-        return True
     except Exception as exc:  # noqa: BLE001 — optimisation only
         _warn(logger, "compile_repeated_blocks", exc)
         return False
+    # Compile every denoiser DiT (a dual-DiT family such as Ideogram runs both each step); a
+    # per-DiT failure degrades that one to eager without dropping the others.
+    engaged = False
+    for transformer in dits:
+        try:
+            transformer.compile_repeated_blocks(**kwargs)
+            engaged = True
+        except Exception as exc:  # noqa: BLE001 — optimisation only
+            _warn(logger, "compile_repeated_blocks", exc)
+    return engaged
 
 
 def _enable_cudnn_benchmark(logger: Any) -> bool:
@@ -399,16 +420,26 @@ def _enable_fp16_accumulation(
 
 
 def _fuse_qkv(pipe: Any, logger: Any) -> bool:
-    for owner in (pipe, getattr(pipe, "transformer", None)):
-        fn = getattr(owner, "fuse_qkv_projections", None)
-        if callable(fn):
+    # Prefer the pipe-level fuse (it covers every component the pipe knows about); else fuse each
+    # denoiser DiT directly so a dual-DiT family (Ideogram) fuses BOTH experts, not just the first.
+    fn = getattr(pipe, "fuse_qkv_projections", None)
+    if callable(fn):
+        try:
+            fn()
+            return True
+        except Exception as exc:  # noqa: BLE001 — optimisation only
+            _warn(logger, "fuse_qkv_projections", exc)
+            return False
+    engaged = False
+    for transformer in _denoiser_dits(pipe):
+        tfn = getattr(transformer, "fuse_qkv_projections", None)
+        if callable(tfn):
             try:
-                fn()
-                return True
+                tfn()
+                engaged = True
             except Exception as exc:  # noqa: BLE001 — optimisation only
                 _warn(logger, "fuse_qkv_projections", exc)
-                return False
-    return False
+    return engaged
 
 
 def _warn(logger: Any, what: str, exc: Exception) -> None:

studio/backend/tests/test_diffusion_attention.py

@@ -177,6 +177,17 @@ def test_apply_sets_backend():
     assert engaged == "_native_cudnn" and t.set_to == "_native_cudnn"
 
 
+def test_apply_sets_backend_on_both_dits():
+    # A dual-DiT family (Ideogram) runs transformer + unconditional_transformer each step, so the
+    # backend must be set on BOTH; otherwise the second DiT keeps the native default while status
+    # reports the requested kernel as engaged.
+    t1, t2 = _FakeTransformer(), _FakeTransformer()
+    pipe = types.SimpleNamespace(transformer = t1, unconditional_transformer = t2)
+    engaged = apply_attention_backend(pipe, "_native_cudnn")
+    assert engaged == "_native_cudnn"
+    assert t1.set_to == "_native_cudnn" and t2.set_to == "_native_cudnn"
+
+
 def test_apply_falls_back_on_unavailable_kernel(monkeypatch):
     # an unavailable kernel must not fail the load -> returns None (diffusers default).
     monkeypatch.setattr(att, "_active_attention_backend", lambda: "native")

studio/backend/tests/test_diffusion_speed.py

@@ -177,6 +177,7 @@ class _Pipe:
         *,
         with_compile = False,
         with_fuse = False,
+        with_second_dit = False,
     ) -> None:
         self.vae = types.SimpleNamespace(mem_format = None, to = self._vae_to)
         self.transformer = types.SimpleNamespace()
@@ -186,6 +187,12 @@ class _Pipe:
             self.fuse_qkv_projections = self._fuse
         self.compiled = False
         self.fused = False
+        # A dual-DiT family (Ideogram) carries a second denoiser expert that runs every step.
+        self.second_compiled = False
+        if with_second_dit:
+            self.unconditional_transformer = types.SimpleNamespace()
+            if with_compile:
+                self.unconditional_transformer.compile_repeated_blocks = self._compile2
 
     def _vae_to(self, *, memory_format):
         self.vae.mem_format = memory_format
@@ -194,6 +201,9 @@ class _Pipe:
         self.compiled = True
         self.compile_kwargs = kwargs
 
+    def _compile2(self, **kwargs):
+        self.second_compiled = True
+
     def _fuse(self):
         self.fused = True
 
@@ -218,6 +228,20 @@ def test_speed_off_applies_nothing(monkeypatch):
     assert torch.backends.cudnn.benchmark is False
 
 
+def test_speed_compiles_both_dits_for_dual_dit_family(monkeypatch):
+    # A dual-DiT family (Ideogram: transformer + unconditional_transformer) runs BOTH DiTs each
+    # denoise step, so the regional block compile must engage on both, not just the first --
+    # otherwise the second DiT runs eager while status reports compile as engaged.
+    _stub_torch(monkeypatch)
+    _stub_gguf_accel(monkeypatch)
+    pipe = _Pipe(with_compile = True, with_second_dit = True)
+    applied = apply_speed_optims(
+        pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT
+    )
+    assert applied["compiled"] is True
+    assert pipe.compiled is True and pipe.second_compiled is True
+
+
 def test_speed_default_dense_falls_back_to_regional_compile(monkeypatch):
     # A DENSE model has no GGUF dequant to compile, so `default` falls back to the
     # regional block compile (its only compile lever) -- and no GGUF accelerators.