Report the compute dtype on fp8-cast encoders and inject the pre-cast TE on the dense fast path

Two more findings from the hosted-TE GPU smokes:

- Module.dtype reports the first floating parameter, which after the
  layerwise fp8 cast is the fp8 STORAGE dtype. Flux2 derives its prompt
  embed and latent dtypes from encoder.dtype and feeds them to
  randn_tensor, which has no fp8 kernel, so ANY flux.2 load with
  text_encoder_quant=fp8 crashed at generation (pre-existing, runtime
  cast included). The cast now swaps in a subclass whose dtype property
  reports the compute dtype; forward behaviour is unchanged.
- The dense transformer_quant fast path assembles companions through
  _assemble_pipe, which never received the pre-cast TE injection, so the
  hosted encoder engaged on full-pipeline and GGUF builds but not on the
  fast path. Threaded through like the other two branches.

Verified live on B200: qwen-image (full pipeline), flux.2-dev (GGUF picker
with int8 DiT prequant), ltx-2 (video backend) all engage the hosted TE,
render non-black, and report text_encoder_quant=fp8 truthfully.

studio/backend/core/inference/diffusion.py

@@ -1380,6 +1380,7 @@ class DiffusionBackend:
                             prequant_path = transformer_prequant_path,
                             allow_dense_fallback = dense_fallback_allowed,
                             lora_specs = loras,
+                            text_encoder_quant = text_encoder_quant,
                         )
                     except Exception as exc:  # noqa: BLE001 — fall back to the GGUF build
                         logger.warning(
@@ -1812,6 +1813,7 @@ class DiffusionBackend:
         base_local_dir: Optional[str] = None,
         allow_dense_fallback: bool = True,
         lora_specs: Optional[list[tuple[str, float]]] = None,
+        text_encoder_quant: Optional[str] = None,
     ) -> tuple[Any, str]:
         """Build the opt-in fast pipeline and return ``(pipe, engaged_scheme)``.
 
@@ -1867,7 +1869,7 @@ class DiffusionBackend:
                 if transformer is not None:
                     pipe = self._assemble_pipe(
                         pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir,
-                        fam = fam,
+                        fam = fam, te_quant_mode = text_encoder_quant, target = target,
                     )
                     return pipe, scheme
 
@@ -1882,7 +1884,8 @@ class DiffusionBackend:
             base, subfolder = "transformer", torch_dtype = dtype, token = hf_token
         )
         pipe = self._assemble_pipe(
-            pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir, fam = fam
+            pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir,
+            fam = fam, te_quant_mode = text_encoder_quant, target = target,
         )
         if lora_specs:
             # Bake the adapters BEFORE quantize_: peft injects its wrappers on the dense
@@ -1929,6 +1932,8 @@ class DiffusionBackend:
         device: str,
         base_local_dir: Optional[str] = None,
         fam: Optional[DiffusionFamily] = None,
+        te_quant_mode: Optional[str] = None,
+        target: Any = None,
     ) -> Any:
         """Assemble the diffusers pipeline around ``transformer`` and place it on ``device``
         (a no-op for an already-placed pre-quantized transformer; it moves the companions)."""
@@ -1948,6 +1953,20 @@ class DiffusionBackend:
             # The repo ships no Llama text_encoder_4; assemble it from the open mirror
             # (diffusion_hidream.py) exactly like the full-pipeline load branch.
             pipe_kwargs.update(hidream_te4_kwargs(dtype, hf_token))
+        # Same pre-cast TE injection as the full-pipeline and GGUF branches: the dense
+        # fast path supplies only the transformer, so the companion TE is the big download.
+        if target is not None:
+            pipe_kwargs.update(
+                te_prequant_pipe_kwargs(
+                    fam,
+                    base,
+                    te_quant_mode = te_quant_mode,
+                    target = target,
+                    dtype = dtype,
+                    hf_token = hf_token,
+                    logger = logger,
+                )
+            )
         pipe = pipeline_cls.from_pretrained(base_local_dir or base, **pipe_kwargs)
         pipe.to(device)
         return pipe

studio/backend/core/inference/diffusion_precision.py

@@ -288,6 +288,22 @@ def _cast_fp8(encoder: Any, target: Any) -> None:
         skip_modules_classes = (torch.nn.Embedding,),
     )
 
+    # Module.dtype reports the first floating parameter, which is now fp8 STORAGE; pipelines
+    # derive tensor dtypes from encoder.dtype (Flux2 casts prompt embeds to it and feeds the
+    # result to randn_tensor, which has no fp8 kernel; VLM pipelines cast pixel_values to it,
+    # racing the upcast hooks). The encoder computes in target.dtype, so report that.
+    compute_dtype = getattr(target, "dtype", None)
+    if compute_dtype is not None and not getattr(encoder, "_unsloth_te_dtype_override", False):
+        cls = type(encoder)
+        encoder.__class__ = type(
+            cls.__name__,
+            (cls,),
+            {
+                "dtype": property(lambda self, _d = compute_dtype: _d),
+                "_unsloth_te_dtype_override": True,
+            },
+        )
+
 
 def _has_layerwise_hooks(encoder: Any) -> bool:
     """True when any submodule already carries the diffusers layerwise-casting hook."""

studio/backend/tests/test_diffusion_te_prequant.py

@@ -263,6 +263,39 @@ def test_hosted_te_prequant_entries():
     ) == "LTX-2-text_encoder-FP8.pt"
 
 
+def test_assemble_pipe_injects_precast_te(monkeypatch):
+    """The dense transformer_quant fast path assembles companions through _assemble_pipe,
+    which must inject the hosted pre-cast TE like the full-pipeline and GGUF branches."""
+    import core.inference.diffusion as dif
+
+    seen: dict = {}
+
+    class FakePipe:
+        def to(self, device):
+            return self
+
+    class FakePipelineCls:
+        @staticmethod
+        def from_pretrained(base, **kw):
+            seen.update(kw)
+            return FakePipe()
+
+    monkeypatch.setattr(
+        dif, "te_prequant_pipe_kwargs", lambda *a, **k: {"text_encoder": "PRECAST"}
+    )
+    dif.DiffusionBackend._assemble_pipe(
+        FakePipelineCls, "org/base", "TR", None, None, "cpu", None,
+        fam = None, te_quant_mode = "fp8", target = object(),
+    )
+    assert seen["text_encoder"] == "PRECAST"
+    seen.clear()
+    # No target (defensive default) keeps the assembly unchanged.
+    dif.DiffusionBackend._assemble_pipe(
+        FakePipelineCls, "org/base", "TR", None, None, "cpu", None, fam = None,
+    )
+    assert "text_encoder" not in seen
+
+
 def test_cast_fp8_is_idempotent_on_precast_encoder():
     """A pre-cast encoder arrives with the layerwise hooks installed; the runtime re-apply in
     quantize_text_encoders must be a no-op (re-registering the hook name raises, which made
@@ -275,8 +308,13 @@ def test_cast_fp8_is_idempotent_on_precast_encoder():
     enc = torch.nn.Sequential(torch.nn.Linear(64, 64), torch.nn.LayerNorm(64))
     _cast_fp8(enc, target)
     assert enc[0].weight.dtype == torch.float8_e4m3fn
+    # Module.dtype must report the COMPUTE dtype: pipelines derive tensor dtypes from it
+    # (Flux2 feeds it to randn_tensor, which has no fp8 kernel).
+    assert enc.dtype == torch.bfloat16
+    assert isinstance(enc, torch.nn.Sequential)
     _cast_fp8(enc, target)  # must not raise
     assert enc[0].weight.dtype == torch.float8_e4m3fn
+    assert enc.dtype == torch.bfloat16
 
 
 def test_builder_metadata_survives_weights_only_load(tmp_path):