longcat: encode the prompt once per request, not once per segment (v0.2.76)

_run_segments calls generate_t2v once PER SEGMENT, and every one of those calls
self.encode_prompt. A one-minute video at 15 fps is 12 segments, so that was 12
identical encodes of the same prompt — and under offload each one also moved the
13.6 GB DiT off the card and the ~11 GB text encoder on and back again. About
35 GB over PCIe per segment, twelve times, for an answer that cannot change.

The embeddings are a few MB, so they are now cached on the arguments that
produce them, bounded to the last four. Segment 2 onward never touches the
encoder: for the rest of a long generation the card holds nothing but the DiT
and the VAE. Caching applies without offload too, where it saves the redundant
forwards if not the transfers.

The wrapper passes its arguments through verbatim instead of naming the first
two — naming them baked in an assumption about upstream's signature that the
avatar pipeline need not share, and it collided on the first test that used a
differently-shaped one.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0185zEBiy3KWB37xNGSrxL7w

codai/__init__.py

@@ -16,7 +16,7 @@
 
 # Canonical product version for CoderAI — single source of truth. Both the API
 # metadata and the admin web UI read from here.
-__version__ = "0.2.75"
+__version__ = "0.2.76"
 __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.

tests/test_longcat_placement.py

@@ -276,3 +276,91 @@ def test_both_load_paths_get_the_backend(service):
     body = src[src.index("def load_pipeline"):src.index("def _cp_split_hw")]
     assert body.count("_dit_kwargs()") == 2
     assert "_attention_kwargs()" in src[src.index("def _dit_kwargs"):]
+
+
+# ------------------------------------------------- one encode per request, not per segment
+# _run_segments calls generate_t2v once PER SEGMENT and each calls self.encode_prompt.
+# A one-minute video at 15 fps is 12 segments, so that was 12 identical encodes — and
+# under offload, 12 round trips of ~35 GB over PCIe for an answer that cannot change.
+
+def test_the_same_prompt_is_encoded_once(service):
+    pipe = FakePipeline()
+    service._place_pipeline(pipe, "model")
+
+    for _ in range(12):                       # twelve segments, one prompt
+        out = pipe.encode_prompt("a boxer", device="cuda")
+
+    assert pipe.encode_calls == ["a boxer"], "encoded once per request, not per segment"
+    assert out == ("embeds", "mask")
+
+
+def test_the_encoder_stays_off_the_card_after_the_first_segment(service):
+    """The point of caching here: for the rest of a long generation the card holds
+    nothing but the DiT and the VAE."""
+    pipe = FakePipeline()
+    service._place_pipeline(pipe, "model")
+    pipe.encode_prompt("a boxer", device="cuda")
+    visits_after_first = len([d for d in pipe.text_encoder.history if d == "cuda"])
+
+    for _ in range(11):
+        pipe.encode_prompt("a boxer", device="cuda")
+
+    assert len([d for d in pipe.text_encoder.history if d == "cuda"]) == visits_after_first
+    assert pipe.dit.device == "cuda"
+
+
+def test_a_different_prompt_is_encoded_again(service):
+    pipe = FakePipeline()
+    service._place_pipeline(pipe, "model")
+
+    pipe.encode_prompt("a boxer", device="cuda")
+    pipe.encode_prompt("a dancer", device="cuda")
+    pipe.encode_prompt("a boxer", device="cuda")
+
+    assert pipe.encode_calls == ["a boxer", "a dancer"]
+
+
+def test_differing_keyword_arguments_are_not_conflated(service):
+    """Same text, different guidance or sequence length is a different encode."""
+    pipe = FakePipeline()
+    service._place_pipeline(pipe, "model")
+
+    pipe.encode_prompt("a boxer", device="cuda", max_sequence_length=512)
+    pipe.encode_prompt("a boxer", device="cuda", max_sequence_length=226)
+
+    assert len(pipe.encode_calls) == 2
+
+
+def test_the_cache_is_bounded(service):
+    """A long-lived service must not accumulate embeddings for every prompt it sees."""
+    pipe = FakePipeline()
+    service._place_pipeline(pipe, "model")
+
+    for i in range(10):
+        pipe.encode_prompt(f"prompt {i}", device="cuda")
+    pipe.encode_prompt("prompt 0", device="cuda")      # evicted by now
+
+    assert len(pipe.encode_calls) == 11
+
+
+def test_caching_also_applies_without_offload(service):
+    """Fully resident still runs 12 redundant encodes otherwise."""
+    pipe = FakePipeline()
+    service._place_pipeline(pipe, None)
+
+    for _ in range(5):
+        pipe.encode_prompt("a boxer", device="cuda")
+
+    assert pipe.encode_calls == ["a boxer"]
+    assert pipe.text_encoder.device == "cuda", "nothing was offloaded"
+
+
+def test_an_unhashable_argument_still_encodes(service):
+    """A dict or tensor among the arguments must not break the call, just skip the cache."""
+    pipe = FakePipeline()
+    service._place_pipeline(pipe, "model")
+
+    pipe.encode_prompt("a boxer", device="cuda", extra={"weights": [1, 2]})
+    pipe.encode_prompt("a boxer", device="cuda", extra={"weights": [1, 2]})
+
+    assert len(pipe.encode_calls) == 2

tools/longcat_service.py

@@ -34,6 +34,7 @@ Endpoints:
 """
 
 import argparse
+import collections
 import contextlib
 import base64
 import importlib.util
@@ -322,6 +323,80 @@ def _attention_kwargs():
 
 
 @contextlib.contextmanager
+def _wrap_encode_prompt(pipe, swap: bool):
+    """Encode a prompt at most once per request, and (when offloading) swap for it.
+
+    _run_segments calls generate_t2v once PER SEGMENT, and every one of those calls
+    self.encode_prompt — so a one-minute video at 15 fps is 12 segments and was 12
+    identical encodes. Under offload each of those also moved the 13.6 GB DiT off the
+    card and the ~11 GB encoder on and back: about 35 GB over PCIe per segment, for an
+    answer that cannot change, since the prompt is the same every time.
+
+    The embeddings are a few MB, so they are cached on the arguments that produce them.
+    Segment 2 onward then never touches the encoder at all, which means the card holds
+    nothing but the DiT and VAE for the whole of a long generation.
+    """
+    import gc
+    import torch
+
+    inner = pipe.encode_prompt      # the bound method, before it is shadowed below
+    cache = collections.OrderedDict()
+
+    def _move_dit(device):
+        """The DiT and the LoRA networks hanging off it, the way upstream's to() does."""
+        if pipe.dit is None:
+            return
+        pipe.dit = pipe.dit.to(device, non_blocking=True)
+        if getattr(pipe.dit, "lora_dict", None):
+            for network in pipe.dit.lora_dict.values():
+                for lora in network.loras:
+                    lora.to(device, non_blocking=True)
+
+    def encode_prompt(*args, **kwargs):
+        # Pass through verbatim: naming the first parameters here would bake in an
+        # assumption about upstream's signature that the avatar pipeline need not share.
+        def _hashable(v):
+            return tuple(v) if isinstance(v, (list, set)) else v
+        try:
+            key = (tuple(_hashable(a) for a in args),
+                   tuple(sorted((k, _hashable(v)) for k, v in kwargs.items())))
+            hash(key)
+            if key in cache:
+                cache.move_to_end(key)
+                return cache[key]
+        except TypeError:
+            key = None          # an unhashable argument: encode, do not cache
+
+        if swap:
+            # They take TURNS. The INT8 DiT is 13.6 GB resident and the encoder is
+            # ~11 GB: bringing the encoder over while the DiT stays is 24.8 GB on a
+            # 24 GB card, which is the OOM that model offload exists to prevent.
+            _move_dit("cpu")
+            gc.collect()
+            torch.cuda.empty_cache()
+            pipe.text_encoder = pipe.text_encoder.to("cuda", non_blocking=True)
+        try:
+            out = inner(*args, **kwargs)
+        finally:
+            if swap:
+                # finally, not just on success: leaving 11 GB stranded after a failed
+                # encode would make the retry the thing that OOMs, and a DiT left on
+                # the host would run every later segment at PCIe speed.
+                pipe.text_encoder = pipe.text_encoder.to("cpu")
+                gc.collect()
+                torch.cuda.empty_cache()
+                _move_dit("cuda")
+        if key is not None:
+            cache[key] = out
+            while len(cache) > 4:       # a few MB each; bounded for a long-lived service
+                cache.popitem(last=False)
+        return out
+
+    # Set on the INSTANCE: all four generate_* methods call self.encode_prompt.
+    pipe.encode_prompt = encode_prompt
+    return pipe
+
+
 def _bsa_for(pipe, ctx):
     """Run this pass with block-sparse attention only if the geometry allows it.
 
@@ -380,6 +455,8 @@ def _place_pipeline(pipe, offload):
 
     if offload not in ("model", "sequential"):
         pipe.to("cuda")
+        # Still worth caching: 12 segments would otherwise run 12 identical encodes.
+        _wrap_encode_prompt(pipe, swap=False)
         log("the whole pipeline is resident on the card")
         return pipe
     if offload == "sequential":
@@ -401,40 +478,7 @@ def _place_pipeline(pipe, offload):
     if encoder is not None:
         pipe.text_encoder = encoder.to("cpu")
 
-    inner = pipe.encode_prompt      # the bound method, before it is shadowed below
-
-    def _move_dit(device):
-        """The DiT and the LoRA networks hanging off it, the way upstream's to() does."""
-        if pipe.dit is None:
-            return
-        pipe.dit = pipe.dit.to(device, non_blocking=True)
-        if getattr(pipe.dit, "lora_dict", None):
-            for network in pipe.dit.lora_dict.values():
-                for lora in network.loras:
-                    lora.to(device, non_blocking=True)
-
-    def encode_prompt(*args, **kwargs):
-        # They take TURNS. The INT8 DiT is 13.6 GB resident and the encoder is ~11 GB:
-        # bringing the encoder over while the DiT stays is 24.8 GB on a 24 GB card,
-        # which is the OOM that model offload exists to prevent. Peak is max(DiT,
-        # encoder), not their sum. The swap costs a round trip over PCIe, once per
-        # generation, against a denoising loop measured in minutes.
-        _move_dit("cpu")
-        gc.collect()
-        torch.cuda.empty_cache()
-        pipe.text_encoder = pipe.text_encoder.to("cuda", non_blocking=True)
-        try:
-            return inner(*args, **kwargs)
-        finally:
-            # finally, not just on success: leaving 11 GB stranded after a failed
-            # encode would make the retry the thing that OOMs.
-            pipe.text_encoder = pipe.text_encoder.to("cpu")
-            gc.collect()
-            torch.cuda.empty_cache()
-            _move_dit("cuda")
-
-    # Set on the INSTANCE: all four generate_* methods call self.encode_prompt.
-    pipe.encode_prompt = encode_prompt
+    _wrap_encode_prompt(pipe, swap=True)
     log("offload: the DiT and VAE are resident; the text encoder (~11 GB) and the DiT "
         "swap places for encode_prompt so only one of them is ever on the card")
     return pipe