longcat: place the lazily realised text encoder (v0.2.81 fixup)

A deferred encoder is not there when _place_pipeline runs, so pipe.to("cuda")
has nothing to move. Under offload that was already right — the encoder belongs
on the host and the wrapper brings it over for the encode — but with
offload_strategy "" the encoder was built on the CPU and left there, while
_get_t5_prompt_embeds sends it inputs on the card. The first cache miss would
have died on a device mismatch.

Placement cannot be re-derived on a miss: by then the offload argument is long
gone. So _place_pipeline records where a realised encoder belongs, in each of
its branches, and realise() replays it.

A quantised encoder is exempt — device_map put it on the card as it loaded and
bitsandbytes weights are not meant to be moved afterwards. That check now lives
in _is_quantised_module, which _encoder_is_quantised delegates to, because a
lazily realised encoder has to be asked the same question BEFORE it is attached
to a pipeline and the two answers must not drift.

test_quantisation_is_detected_from_the_model_not_the_config followed the names
into the helper and now asserts the delegation too, so neither function can
start consulting the config.

Found by re-reading the committed diff, not by a test: nothing here runs the
real pipeline, which is exactly the gap that let it through.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

tests/test_longcat_component_quant.py

@@ -105,11 +105,19 @@ def test_a_quantised_encoder_is_never_offloaded():
 
 
 def test_quantisation_is_detected_from_the_model_not_the_config():
-    """A pipeline loaded quantised must not be offloaded by a config that changed since."""
+    """A pipeline loaded quantised must not be offloaded by a config that changed since.
+
+    The bitsandbytes module names live in _is_quantised_module, which _encoder_is_quantised
+    delegates to: a lazily realised encoder has to be asked the same question before it is
+    attached to a pipeline, and the two answers must not drift. Both are checked here so
+    neither can start consulting the config.
+    """
     src = SERVICE.read_text(encoding="utf-8")
-    body = src[src.index("def _encoder_is_quantised"):src.index("def _quant_config")]
+    body = src[src.index("def _is_quantised_module"):src.index("def _quant_config")]
     assert "Linear4bit" in body and "Linear8bitLt" in body
     assert "text_encoder_quant" not in body
+    delegating = src[src.index("def _encoder_is_quantised"):src.index("def _quant_config")]
+    assert "_is_quantised_module(" in delegating
 
 
 def test_missing_bitsandbytes_is_a_clear_error():

tests/test_longcat_prompt_cache.py

@@ -265,6 +265,76 @@ def test_realising_twice_builds_once(SVC):
     assert built == [1] and pipe.text_encoder == "encoder"
 
 
+def test_a_realised_encoder_lands_where_the_placement_already_happened(SVC):
+    """pipe.to("cuda") runs at LOAD time, when a deferred encoder is not there to be
+    moved. Without replaying that, the weights sit on the host while
+    _get_t5_prompt_embeds sends them inputs on the card, and the first miss dies on a
+    device mismatch — the bug that laziness introduces if placement is not remembered.
+    """
+    moved = []
+
+    class _Encoder:
+        def to(self, device, **kw):
+            moved.append(device)
+            return self
+
+    class _Pipe:
+        text_encoder = None
+
+    lazy = SVC._LazyEncoder(_Encoder)
+    lazy.device = "cuda"
+    lazy.realise(_Pipe())
+    assert moved == ["cuda"]
+
+
+def test_under_offload_it_is_built_on_the_host_where_it_belongs(SVC):
+    """The wrapper brings it over for the encode and takes it back. Building it on the
+    card would be the 24.8 GB peak that offload exists to prevent."""
+    moved = []
+
+    class _Encoder:
+        def to(self, device, **kw):
+            moved.append(device)
+            return self
+
+    class _Pipe:
+        text_encoder = None
+
+    lazy = SVC._LazyEncoder(_Encoder)      # cpu is the default
+    assert lazy.device == "cpu"
+    lazy.realise(_Pipe())
+    assert moved == []
+
+
+def test_a_quantised_encoder_is_not_moved_after_it_is_built(SVC):
+    """device_map put it on the card as it loaded; bitsandbytes weights are not meant
+    to be moved afterwards."""
+    moved = []
+
+    class _Encoder:
+        is_loaded_in_4bit = True
+
+        def to(self, device, **kw):
+            moved.append(device)
+            return self
+
+    class _Pipe:
+        text_encoder = None
+
+    lazy = SVC._LazyEncoder(_Encoder)
+    lazy.device = "cuda"
+    lazy.realise(_Pipe())
+    assert moved == []
+
+
+def test_placement_is_recorded_by_the_function_that_knows_the_mode(SVC):
+    """It cannot be re-derived on a miss: by then the offload argument is long gone."""
+    src = SERVICE.read_text(encoding="utf-8")
+    place = src[src.index("def _place_pipeline("):src.index("def _load_quantized_base_dit(")]
+    assert place.count("_lazy_goes_to(") == 4, "each branch must say where it belongs"
+    assert '_lazy_goes_to("cpu")' in place and '_lazy_goes_to("cuda")' in place
+
+
 def test_it_is_not_a_transparent_proxy(SVC):
     """Upstream guards every use with `if self.text_encoder is not None`. A proxy that
     looked like a module would be moved to the card by pipe.to() and offloaded by

tools/longcat_service.py

@@ -356,13 +356,12 @@ def _attention_kwargs():
         "attention has no fallback path")
 
 
-def _encoder_is_quantised(pipe) -> bool:
-    """Whether the text encoder carries bitsandbytes weights.
+def _is_quantised_module(encoder) -> bool:
+    """Whether this module carries bitsandbytes weights.
 
-    Asked of the MODEL rather than the config, so a pipeline that was loaded quantised
-    cannot be offloaded by a config that has since changed.
+    Its own function because a lazily realised encoder has to be asked the same
+    question before it is attached to a pipeline, and the two answers must not drift.
     """
-    encoder = getattr(pipe, "text_encoder", None)
     if encoder is None:
         return False
     if getattr(encoder, "is_quantized", False) or getattr(encoder, "is_loaded_in_4bit", False) \
@@ -375,6 +374,15 @@ def _encoder_is_quantised(pipe) -> bool:
                for m in walk())
 
 
+def _encoder_is_quantised(pipe) -> bool:
+    """Whether the pipeline's text encoder carries bitsandbytes weights.
+
+    Asked of the MODEL rather than the config, so a pipeline that was loaded quantised
+    cannot be offloaded by a config that has since changed.
+    """
+    return _is_quantised_module(getattr(pipe, "text_encoder", None))
+
+
 def _quant_config(kind: str, dtype):
     """A BitsAndBytesConfig for a component, or None to load at the pipeline dtype.
 
@@ -587,14 +595,27 @@ class _LazyEncoder:
         self._build = build
         self._lock = threading.Lock()
         self.realised = False
+        # Where a realised encoder belongs, set by _place_pipeline once it knows the
+        # offload mode. It CANNOT be inferred here: placement happens at load time,
+        # when this holder is all there is, so the decision has to be remembered rather
+        # than re-derived on a miss.
+        self.device = "cpu"
 
     def realise(self, pipe):
-        """Build the encoder and attach it to the pipeline. Idempotent."""
+        """Build the encoder, place it, and attach it to the pipeline. Idempotent."""
         with self._lock:
             if pipe.text_encoder is not None:
                 return pipe.text_encoder
             log("prompt cache miss: building the text encoder now (deferred at load)")
-            pipe.text_encoder = self._build()
+            encoder = self._build()
+            # pipe.to("cuda") already ran, at load time, when there was no encoder to
+            # move. Without this the weights sit on the host while _get_t5_prompt_embeds
+            # sends it inputs on the card, and the first miss dies on a device mismatch.
+            # A quantised encoder is exempt: device_map put it on the card as it loaded,
+            # and bitsandbytes weights are not meant to be moved afterwards.
+            if self.device != "cpu" and not _is_quantised_module(encoder):
+                encoder = encoder.to(self.device, non_blocking=True)
+            pipe.text_encoder = encoder
             self.realised = True
             return pipe.text_encoder
 
@@ -768,6 +789,12 @@ def _place_pipeline(pipe, offload):
         log("WARNING: no CUDA device visible — this will be extremely slow")
         return pipe
 
+    def _lazy_goes_to(device):
+        """Remember where a lazily realised encoder belongs. No-op without one."""
+        lazy = _state.get("lazy_encoder")
+        if lazy is not None:
+            lazy.device = device
+
     # A quantised encoder is already on the card and must stay there: bitsandbytes
     # quantises as the weights land on CUDA and is not built to shuttle them back. It
     # is also small enough not to need to — which is the whole point of quantising it.
@@ -782,12 +809,16 @@ def _place_pipeline(pipe, offload):
         if pipe.vae is not None:
             pipe.vae = pipe.vae.to("cuda", non_blocking=True)
         pipe.device = "cuda"
+        _lazy_goes_to("cuda")
         _wrap_encode_prompt(pipe, swap=False, cache=_state.get("prompt_cache"),
                             lazy=_state.get("lazy_encoder"))
         return pipe
 
     if offload not in ("model", "sequential"):
         pipe.to("cuda")
+        # A deferred encoder was not here for that to(), so record where it belongs.
+        # Resident means resident: it goes to the card when a miss builds it.
+        _lazy_goes_to("cuda")
         # Still worth caching: 12 segments would otherwise run 12 identical encodes.
         _wrap_encode_prompt(pipe, swap=False, cache=_state.get("prompt_cache"),
                             lazy=_state.get("lazy_encoder"))
@@ -811,6 +842,9 @@ def _place_pipeline(pipe, offload):
         pipe.text_encoder = encoder
     if encoder is not None:
         pipe.text_encoder = encoder.to("cpu")
+    # Under offload the encoder lives on the host and the wrapper brings it over for
+    # the encode, so a deferred one is built there too — which is also the default.
+    _lazy_goes_to("cpu")
 
     _wrap_encode_prompt(pipe, swap=True, cache=_state.get("prompt_cache"),
                         lazy=_state.get("lazy_encoder"))