longcat: cache prompt embeddings, so the encoder need not load (v0.2.81)

The UMT5-XXL text encoder is ~11 GB and turns a prompt into a few MB of
embeddings that are a pure function of the prompt and the encoder. v0.2.76 made
that run once per request instead of once per segment; this makes it once per
prompt, ever.

The saved GPU pass is not the point — it is seconds. The point is that a cached
prompt needs no encoder AT ALL, so with the cache on the encoder is built
lazily: ~11 GB of host RAM and about a minute of load that a repeated prompt
never spends. That is the answer to "can we keep everything in VRAM": not by
quantising the encoder (v0.2.79 measured int8 resident at 24.0 GB on a 24 GB
card, and NF4 at 22.1 GB with less headroom than a 29-frame clip already uses),
but by not having an encoder resident at all.

Off by default — embeddings are derived data, but they are also several MB each
under the user's cache directory, so growing that is opted into:

    "prompt_cache": "off",          # off | disk
    "prompt_cache_dir": "",         # blank = <coderai home>/prompt_cache/longcat
    "prompt_cache_max_gb": 0,       # 0 = 2 GB, least-recently-used evicted

per model or server-wide, validated at startup rather than after a seven-minute
load, settable from the models page, and in api_model_configure's whitelist so a
save from that page does not silently drop it.

An entry served to the wrong pipeline would be a silent quality change, which is
worse than a miss, so the key covers the checkpoint, the dtype, the encoder's
quantisation and the call's own arguments — an NF4 encoder cannot be served a
bf16 entry. `device` is deliberately excluded: it moves the embeddings without
changing them, and including it would split one entry into one per device.

Every failure is a miss, never an error. Entries are written to a .tmp and
renamed, read with weights_only=True (a cache directory is not a trusted source
for a pickle), deleted when unreadable, and a cache that cannot be written warns
once and generation carries on without it.

_LazyEncoder is deliberately not a transparent proxy: upstream guards every use
with `if self.text_encoder is not None`, so a proxy that looked like a module
would be moved to the card by pipe.to() and offloaded by _place_pipeline,
building the very weights it exists to avoid. pipe.text_encoder stays None and
the holder sits beside it in _state, which is why both paths already do the
right thing.

Found while wiring it: a lazily realised encoder can be quantised, and the
offload swap would have shuttled bitsandbytes weights to the host. The swap
guard now checks the model, not just the offload mode.

/health reports the cache and whether the encoder was ever built — laziness is
invisible otherwise — and GET /prompt-cache returns its stats.

No new dependencies. Unverified on the card: the round trip, eviction, corrupt
entries and the key are covered by 54 tests against real tensors, but no
generation has run through it yet.
Co-Authored-By: Claude Opus 5 <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.80"
+__version__ = "0.2.81"
 __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.

codai/admin/routes.py

@@ -2945,7 +2945,10 @@ async def api_model_configure(request: Request, username: str = Depends(require_
                 
                     # LongCat component quantisation, block-sparse attention, the base
                     # checkpoint an avatar family borrows from, and the CP tile.
-                    "text_encoder_quant", "bsa", "base_model", "cp_split_hw",):
+                    "text_encoder_quant", "bsa", "base_model", "cp_split_hw",
+                    # The prompt-embedding cache: an encode outlives the request that
+                    # paid for it, and the text encoder is then built only on a miss.
+                    "prompt_cache", "prompt_cache_dir", "prompt_cache_max_gb",):
         if key in data:
             entry[key] = data[key]
     # Work-parallel fan-out over engines and nodes (codai/cluster/fanout.py).

codai/admin/templates/models.html

@@ -719,6 +719,15 @@ window.__DEFAULT_WHISPER_SERVER_PATH__ = {{ default_whisper_server_path|tojson }
               <option value="auto">auto — on when triton is available</option>
             </select></div>
         </div>
+        <div style="display:grid;grid-template-columns:1fr 1fr 1fr;gap:.6rem;margin-top:.5rem">
+          <div class="form-row" style="margin:0"><label class="form-label">Prompt cache</label>
+            <select id="cfg-longcat-prompt-cache" class="form-input">
+              <option value="">off — encode every request</option>
+              <option value="disk">disk — encode once ever, encoder built only on a miss</option>
+            </select></div>
+          <div class="form-row" style="margin:0"><label class="form-label">Prompt cache dir <span class="muted">(blank = ~/.coderai)</span></label><input type="text" id="cfg-longcat-prompt-cache-dir" class="form-input" placeholder="~/.coderai/prompt_cache/longcat"></div>
+          <div class="form-row" style="margin:0"><label class="form-label">Prompt cache cap <span class="muted">(GB)</span></label><input type="number" id="cfg-longcat-prompt-cache-max" class="form-input" min="0" step="0.5" placeholder="2"></div>
+        </div>
         <div style="display:grid;grid-template-columns:1fr 1fr 1fr;gap:.6rem;margin-top:.5rem">
           <div class="form-row" style="margin:0"><label class="form-label">Segments <span class="muted">(blank = 1)</span></label><input type="number" id="cfg-longcat-segments" class="form-input" min="1" placeholder="1"></div>
           <div class="form-row" style="margin:0"><label class="form-label">Conditioning frames</label><input type="number" id="cfg-longcat-cond-frames" class="form-input" min="1" placeholder="13"></div>
@@ -4586,6 +4595,9 @@ function _populateCluster(m, s){
   set('cfg-longcat-offload', s.offload_strategy || '');
   set('cfg-longcat-te-quant', s.text_encoder_quant || '');
   set('cfg-longcat-bsa', s.bsa || '');
+  set('cfg-longcat-prompt-cache', s.prompt_cache || '');
+  set('cfg-longcat-prompt-cache-dir', s.prompt_cache_dir || '');
+  set('cfg-longcat-prompt-cache-max', s.prompt_cache_max_gb || '');
   set('cfg-longcat-base-model', s.base_model || '');
   set('cfg-longcat-cp-split', s.cp_split_hw || '');
   set('cfg-longcat-segments', s.num_segments || '');
@@ -4633,6 +4645,9 @@ function _collectCluster(){
     out.offload_strategy = val('cfg-longcat-offload') || '';
     out.text_encoder_quant = val('cfg-longcat-te-quant') || '';
     out.bsa = val('cfg-longcat-bsa') || '';
+    out.prompt_cache = val('cfg-longcat-prompt-cache') || '';
+    out.prompt_cache_dir = val('cfg-longcat-prompt-cache-dir') || '';
+    out.prompt_cache_max_gb = num('cfg-longcat-prompt-cache-max');
     out.base_model = val('cfg-longcat-base-model') || '';
     out.cp_split_hw = val('cfg-longcat-cp-split') || '';
     out.num_segments = num('cfg-longcat-segments');

codai/api/longcat_worker.py

@@ -47,6 +47,11 @@ Per-model config keys (models.json entry, all optional)::
     "bsa": "off",                  # block-sparse attention: on | off | auto
     "base_model": "",              # avatar families borrow tokenizer/text_encoder/vae
     "text_encoder_quant": "none",   # none | int8 | nf4 | fp4 (keeps it resident)
+    "prompt_cache": "off",         # off | disk — cache prompt embeddings between
+                                   #   requests; with it on the ~11 GB text encoder is
+                                   #   built only when a prompt misses
+    "prompt_cache_dir": "",        # blank = <coderai home>/prompt_cache/longcat
+    "prompt_cache_max_gb": 0,      # 0 = the service default (2 GB), LRU-evicted
     "cp_split_hw": "",             # context-parallel tile, e.g. "1x2"
     "used_vram_gb": 0,             # what to reserve before starting; 0 = measured
     "gpu_device": 0
@@ -508,6 +513,28 @@ def ensure_service(model_path: str, config: dict = None,
                    or "").strip()
         if te_q:
             cmd += ["--text-encoder-quant", te_q]
+        # The prompt cache makes an encode survive the request that paid for it, which
+        # is also what lets the encoder be built lazily. Per model, then server-wide:
+        # a checkpoint whose prompts never repeat gains nothing but the disk.
+        pc = str(config.get("prompt_cache") or
+                 (getattr(sec, "prompt_cache", "") if sec is not None else "")
+                 or "").strip()
+        if pc:
+            cmd += ["--prompt-cache", pc]
+        pc_dir = str(config.get("prompt_cache_dir") or
+                     (getattr(sec, "prompt_cache_dir", "") if sec is not None else "")
+                     or "").strip()
+        if pc_dir:
+            cmd += ["--prompt-cache-dir", pc_dir]
+        pc_max = (config.get("prompt_cache_max_gb")
+                  or (getattr(sec, "prompt_cache_max_gb", 0) if sec is not None else 0))
+        try:
+            if float(pc_max or 0) > 0:
+                cmd += ["--prompt-cache-max-gb", str(float(pc_max))]
+        except (TypeError, ValueError):
+            # A bad value is the service's to reject with a message naming the key,
+            # rather than this silently dropping it.
+            cmd += ["--prompt-cache-max-gb", str(pc_max)]
         if sec is not None and str(getattr(sec, "extra_args", "") or "").strip():
             import shlex
             cmd += shlex.split(str(sec.extra_args))

tools/longcat_common.py

@@ -374,6 +374,95 @@ def component_quant_bpe(value: str):
     return COMPONENT_QUANTS.get(str(value or "none").strip().lower())
 
 
+# ── prompt-embedding cache ────────────────────────────────────────────────────
+#
+# The UMT5-XXL text encoder is ~11 GB and runs ONCE per request, to turn a prompt into
+# a few MB of embeddings. v0.2.76 stopped re-running it per segment; caching the result
+# on DISK stops re-running it per request, and — because the entry is keyed on
+# everything that can change the answer — lets the encoder be loaded LAZILY: a request
+# whose prompt is already cached never builds it at all, which is ~11 GB of host RAM
+# and a minute of load time that nothing has to spend.
+#
+# Off by default. Embeddings are derived data, but they are also several MB each and
+# live under the user's cache directory, so growing that is something to opt into.
+PROMPT_CACHE_MODES = ("off", "disk")
+
+# Each entry is prompt_embeds + mask + the negative pair, 512x4096 at bf16 ≈ 4 MB per
+# side. A cap rather than a count, because a long prompt list is the case that runs away.
+DEFAULT_PROMPT_CACHE_MAX_GB = 2.0
+
+# Bumped when the stored tuple's shape or the key's inputs change, so a cache written by
+# an older build is ignored rather than misread.
+PROMPT_CACHE_FORMAT = 1
+
+
+def prompt_cache_problems(mode: str, max_gb=None) -> list:
+    """Reasons this prompt-cache configuration cannot be used, or []."""
+    problems = []
+    m = str(mode or "off").strip().lower()
+    if m not in PROMPT_CACHE_MODES:
+        problems.append(f"prompt_cache must be one of {sorted(PROMPT_CACHE_MODES)}; "
+                        f"got {mode!r}")
+    if max_gb is not None and str(max_gb).strip() != "":
+        try:
+            if float(max_gb) <= 0:
+                problems.append(f"prompt_cache_max_gb must be positive; got {max_gb!r}")
+        except (TypeError, ValueError):
+            problems.append(f"prompt_cache_max_gb must be a number; got {max_gb!r}")
+    return problems
+
+
+def prompt_cache_key(*, checkpoint: str, dtype: str, text_encoder_quant: str,
+                     args: tuple = (), kwargs: dict = None) -> str:
+    """A stable filename-safe key for one encode_prompt call.
+
+    Keyed on everything that changes the EMBEDDINGS and nothing that does not. The
+    checkpoint and the encoder's quantisation are in here because an NF4 encoder does
+    not produce the bf16 encoder's numbers -- reusing one for the other would be a
+    silent quality change, which is worse than a cache miss. ``device`` is deliberately
+    excluded: it moves the result, it does not alter it.
+    """
+    import hashlib
+    import json as _json
+
+    def _plain(v):
+        """Reduce an argument to something JSON can hold, deterministically."""
+        if isinstance(v, (str, int, float, bool)) or v is None:
+            return v
+        if isinstance(v, (list, tuple)):
+            return [_plain(x) for x in v]
+        if isinstance(v, set):
+            return sorted(_plain(x) for x in v)
+        if isinstance(v, dict):
+            return {str(k): _plain(v[k]) for k in sorted(v, key=str)}
+        # A torch.device, a torch.dtype, anything else: its repr is stable within a
+        # build, and `device` is dropped before we get here.
+        return repr(v)
+
+    kwargs = dict(kwargs or {})
+    kwargs.pop("device", None)
+    payload = {
+        "format": PROMPT_CACHE_FORMAT,
+        # basename, not the full path: the same checkpoint reached through a symlink or
+        # a differently-mounted hfcache is the same encoder.
+        "checkpoint": os.path.basename(os.path.normpath(str(checkpoint or ""))),
+        "dtype": str(dtype or ""),
+        "text_encoder_quant": str(text_encoder_quant or "none").strip().lower(),
+        "args": _plain(list(args)),
+        "kwargs": _plain(kwargs),
+    }
+    blob = _json.dumps(payload, sort_keys=True, separators=(",", ":"))
+    return hashlib.sha256(blob.encode("utf-8")).hexdigest()
+
+
+def prompt_cache_dir(configured: str = "") -> str:
+    """Where the prompt cache lives. Under the coderai cache root unless overridden."""
+    if str(configured or "").strip():
+        return os.path.expanduser(str(configured).strip())
+    root = os.environ.get("CODERAI_HOME") or os.path.join("~", ".coderai")
+    return os.path.expanduser(os.path.join(root, "prompt_cache", "longcat"))
+
+
 def dit_subdir(family: str = "") -> str:
     """Which subfolder holds the bf16 transformer for this family."""
     return AVATAR_DIT_SUBDIR if family else BASE_DIT_SUBDIR