core: elastic expert cache via CUDA VMM and on-demand GPU vision lifecycle

CMakeLists.txt

@@ -1,4 +1,4 @@
-# Strata Engine - top-level build (P1.S1).
+# Strata Engine - top-level build (P1.S1).
 #
 # Targets (phase-1 §P1.S1): `strata_artifact` (no CUDA dependency), `tests`.  `strata_plan` and the C++
 # `strata-pack` are added when their sources exist - a target with no sources would make the build graph
@@ -241,7 +241,7 @@ if(STRATA_ENABLE_CUDA)
   add_library(strata_engine STATIC src/core/layer.cpp src/core/session.cpp src/core/expert_source.cpp
     src/core/expert_cache.cpp src/core/native_head.cpp src/core/native_dense.cpp src/core/verify.cpp src/core/mtp.cpp)
   target_include_directories(strata_engine PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
-  target_link_libraries(strata_engine PUBLIC strata_core strata_kernels strata_kernels_cpu CUDA::cudart)
+  target_link_libraries(strata_engine PUBLIC strata_core strata_kernels strata_kernels_cpu CUDA::cudart CUDA::cuda_driver)
 
   # `strata generate` (P2.S6).  The FIRST user-facing program: everything under it is a component, and this is
   # what composes them into a token.  No tokenizer - the prompt arrives as ids - because Gate C1 compares logits

include/strata/core/expert_cache.hpp

@@ -24,10 +24,19 @@
 //
 // The order of work is `Memory/R4-design-note.md` §7: slots and residency first, then the split, then the
 // kernel. This is that first step, and the step it unblocks is the one that can be measured.
+//
+// **ELASTIC CACHE (R4.3).**  When a vision encoder loads onto the GPU it needs ~850 MB of VRAM.  Rather than
+// keeping the cache permanently smaller, the cache is allocated through CUDA Virtual Memory Management
+// (cuMemAddressReserve / cuMemMap / cuMemUnmap, CUDA 10.2+).  The virtual address range is reserved for the
+// MAXIMUM capacity; physical pages are mapped for as many slots as fit.  `shrink(n)` unmaps the tail, freeing
+// physical VRAM for vision.  `grow(n, profile, arena)` maps new pages and refills from the profile.  Because
+// the virtual addresses are STABLE, captured CUDA graphs survive resize without recapture.
+// On GPUs without VMM support (CC < 6.0), the cache falls back to plain cudaMalloc and is not elastic.
 #pragma once
 
 #include <cstdint>
 #include <string>
+#include <utility>
 #include <vector>
 
 namespace strata::core {
@@ -135,6 +144,49 @@ class ExpertCache {
     /// Slots filled so far, for the startup report.
     int64_t fills() const { return fills_; }
 
+    // ---- R4.3: ELASTIC CACHE (CUDA VMM) ----
+
+    /// True if the cache was allocated with CUDA VMM and can be resized without invalidating virtual addresses.
+    bool elastic() const { return vmm_; }
+
+    /// The maximum number of slots the virtual address range was reserved for.  Only meaningful when elastic().
+    int64_t max_slots() const { return max_slots_; }
+
+    /// The total bytes of the full virtual address reservation (max capacity).
+    int64_t max_bytes() const { return max_bytes_; }
+
+    /// Shrink the cache: unmap the tail slots whose byte range exceeds `keep_bytes`, mark them not-resident,
+    /// and free the physical VRAM.  The GPU stream MUST be synchronised before calling this.
+    /// Returns the number of bytes actually freed.  Non-elastic caches return 0.
+    int64_t shrink(int64_t keep_bytes, std::string& err);
+
+    /// Grow the cache: map physical pages to extend the arena up to `new_bytes` (clamped to max_bytes),
+    /// then refill the new slots from the profile.  `profile` is the ranked (layer, expert) list,
+    /// `blob_fn` returns the host blob for a given (layer, expert), `blob_size_fn` returns the blob size
+    /// (0 = uniform blob_).  The GPU stream MUST be synchronised before calling this.
+    /// Returns the number of new slots filled.
+    using BlobFn = const uint8_t* (*)(int32_t layer, int32_t expert, void* user);
+    using BlobSizeFn = int64_t (*)(int32_t layer, void* user);
+    int64_t grow(int64_t new_bytes, const std::vector<std::pair<int32_t, int32_t>>& profile,
+                 BlobFn blob_fn, void* blob_user, BlobSizeFn size_fn, void* size_user, std::string& err);
+
+    /// Query the VMM allocation granularity (2 MiB on most GPUs).  0 if VMM is not available.
+    static size_t vmm_granularity();
+
+    /// Check whether the current GPU supports CUDA VMM.
+    static bool vmm_supported();
+
+    // ---- OPTION B: DYNAMIC ADAPTIVE SWAPPING ----
+
+    /// Evicts whatever expert is currently occupying `slot`, marking it kNotResident.
+    void evict_slot(int32_t slot);
+
+    /// Assigns `slot` to `(layer, expert)` in the residency table.
+    void assign_slot(int32_t slot, int64_t layer, int64_t expert);
+
+    /// Const access to the residency table [n_layers * n_expert].
+    const std::vector<int32_t>& residency_table() const { return residency_; }
+
 private:
     uint8_t* base_ = nullptr;
     std::vector<int32_t> residency_;   ///< [n_layers * n_expert] -> slot or kNotResident
@@ -150,6 +202,21 @@ class ExpertCache {
     std::vector<int32_t> layer_next_;   ///< [n_layers] -> that layer's next free slot
     std::vector<uint64_t> off_;         ///< plan v0.3 P6: slot offsets (slots + 1 entries) when sized
     int64_t admitted_ = 0;
+
+    // ---- R4.3: VMM state ----
+    bool vmm_ = false;                    ///< true if allocated through cuMem* VMM APIs
+    int64_t max_slots_ = 0;               ///< total slots the VA reservation can hold
+    int64_t max_bytes_ = 0;               ///< total bytes of the VA reservation
+    int64_t mapped_bytes_ = 0;            ///< bytes currently backed by physical memory
+
+    // VMM opaque handles (stored as uint64_t to avoid cuda.h in the header)
+    uint64_t vmm_va_ = 0;                 ///< CUdeviceptr of the reserved VA range
+    uint64_t vmm_va_size_ = 0;            ///< size of the VA reservation
+    struct VmmChunk { uint64_t handle; uint64_t offset; uint64_t size; };
+    std::vector<VmmChunk> vmm_chunks_;    ///< physical allocations, in order of mapping offset
+
+    bool open_vmm(int64_t total_bytes, int64_t n_layers, int64_t n_expert, std::string& err);
+    void close_vmm();
 };
 
 }  // namespace strata::core

serve/server.py

@@ -273,6 +273,32 @@ def generate(self, ids, max_new, sampling, cancel, embeddings=None):
                         self._parse_done(line)
                         break
 
+    def resize(self, keep_bytes: int) -> tuple[int, int] | None:
+        """Shrink expert cache to `keep_bytes` so GPU vision has free VRAM."""
+        try:
+            self.proc.stdin.write(f"RESIZE {int(keep_bytes)}\n")
+            self.proc.stdin.flush()
+            line = self.lines.get(timeout=5)
+            if line and line.startswith("OK "):
+                parts = line.split()
+                return int(parts[1]), int(parts[2])
+        except Exception as e:
+            print(f"[strata] resize error: {e}", flush=True)
+        return None
+
+    def grow(self, target_bytes: int = 10_000_000_000) -> tuple[int, int] | None:
+        """Expand expert cache back to `target_bytes` (up to max reserved VRAM)."""
+        try:
+            self.proc.stdin.write(f"GROW {int(target_bytes)}\n")
+            self.proc.stdin.flush()
+            line = self.lines.get(timeout=10)
+            if line and line.startswith("OK "):
+                parts = line.split()
+                return int(parts[1]), int(parts[2])
+        except Exception as e:
+            print(f"[strata] grow error: {e}", flush=True)
+        return None
+
     def close(self):
         try:
             self.proc.stdin.write("QUIT\n")
@@ -283,26 +309,36 @@ def close(self):
 
 
 class Vision:
-    """The resident image encoder: `strata-vision` (llama.cpp mtmd + the mmproj file) reads `ENC <image> <out>`
+    """The image encoder: `strata-vision` (llama.cpp mtmd + the mmproj file) reads `ENC <image> <out>`
     lines and writes each image's embeddings; results are cached by the image's hash, so a conversation that
     sends the same picture again (every turn, with most clients) encodes it once."""
 
     def __init__(self, cfg: dict, log=None, env: dict | None = None):
-        args = [cfg["exe"], "--mmproj", cfg["mmproj"], "--model", cfg["model"]]
+        self.cfg = cfg
+        self.log = log
+        self.env = env
+        self.args = [cfg["exe"], "--mmproj", cfg["mmproj"], "--model", cfg["model"]]
         if cfg.get("gpu"):
-            args.append("--gpu")
+            self.args.append("--gpu")
         if cfg.get("threads"):
-            args += ["--threads", str(cfg["threads"])]
+            self.args += ["--threads", str(cfg["threads"])]
         if cfg.get("max_tokens"):
-            args += ["--max-tokens", str(cfg["max_tokens"])]
+            self.args += ["--max-tokens", str(cfg["max_tokens"])]
         self.dir = Path(tempfile.mkdtemp(prefix="strata-vision-"))
-        self.proc = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=log or subprocess.DEVNULL,
-                                     text=True, encoding="utf-8", bufsize=1, env=env)
-        line = self.proc.stdout.readline()
-        if not line.startswith("READY"):
-            raise RuntimeError("the vision encoder did not start: " + line.strip())
         self.lock = threading.Lock()
         self.cache: dict[str, tuple[Path, int]] = {}
+        self.proc = None
+        if not cfg.get("gpu"):
+            self.proc = subprocess.Popen(self.args, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
+                                         stderr=log or subprocess.DEVNULL,
+                                         text=True, encoding="utf-8", bufsize=1, env=env)
+            line = self.proc.stdout.readline()
+            if not line.startswith("READY"):
+                self.proc.kill()
+                raise RuntimeError("the vision encoder did not start: " + line.strip())
+            print(f"[strata] vision encoder is ready ({line.strip()})", flush=True)
+        else:
+            print("[strata] vision encoder configured (dynamic GPU on-demand)", flush=True)
 
     @staticmethod
     def load(source: str) -> bytes:
@@ -345,19 +381,53 @@ def normalize(data: bytes) -> bytes:
         im.save(out, format="PNG")
         return out.getvalue()
 
-    def encode(self, source: str) -> tuple[Path, int]:
+    def encode(self, source: str, engine=None) -> tuple[Path, int]:
         """-> (embeddings file, number of image tokens)."""
         data = self.normalize(self.load(source))
         key = hashlib.sha256(data).hexdigest()[:32]
         with self.lock:
             if key in self.cache:
                 return self.cache[key]
+            is_gpu = bool(self.cfg.get("gpu"))
+            if is_gpu and engine is not None and hasattr(engine, "resize"):
+                engine.resize(5_400_000_000)
+
+            proc = self.proc
+            should_close = False
+            if proc is None:
+                # On-demand spawn for GPU vision
+                proc = subprocess.Popen(self.args, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
+                                        stderr=self.log or subprocess.DEVNULL,
+                                        text=True, encoding="utf-8", bufsize=1, env=self.env)
+                ready_line = proc.stdout.readline()
+                if not ready_line.startswith("READY"):
+                    proc.kill()
+                    if is_gpu and engine is not None and hasattr(engine, "grow"):
+                        engine.grow(10_000_000_000)
+                    raise RuntimeError("the vision encoder did not start: " + ready_line.strip())
+                should_close = True
+
             img, out = self.dir / f"{key}.img", self.dir / f"{key}.sve"
             img.write_bytes(data)
-            self.proc.stdin.write(f"ENC {img} {out}\n")
-            self.proc.stdin.flush()
-            line = self.proc.stdout.readline().strip()
-            img.unlink(missing_ok=True)
+            try:
+                proc.stdin.write(f"ENC {img} {out}\n")
+                proc.stdin.flush()
+                line = proc.stdout.readline().strip()
+            finally:
+                img.unlink(missing_ok=True)
+                if should_close:
+                    try:
+                        proc.stdin.write("QUIT\n")
+                        proc.stdin.flush()
+                        proc.wait(timeout=5)
+                    except Exception:
+                        try:
+                            proc.kill()
+                        except Exception:
+                            pass
+                if is_gpu and engine is not None and hasattr(engine, "grow"):
+                    engine.grow(10_000_000_000)
+
             if not line.startswith("OK"):
                 raise ValueError("the image could not be read: " + (line[4:] if line.startswith("ERR") else
                                                                      "the vision encoder stopped"))
@@ -368,12 +438,17 @@ def encode(self, source: str) -> tuple[Path, int]:
             return self.cache[key]
 
     def close(self):
-        try:
-            self.proc.stdin.write("QUIT\n")
-            self.proc.stdin.flush()
-            self.proc.wait(timeout=10)
-        except Exception:
-            self.proc.kill()
+        if self.proc:
+            try:
+                self.proc.stdin.write("QUIT\n")
+                self.proc.stdin.flush()
+                self.proc.wait(timeout=5)
+            except Exception:
+                try:
+                    self.proc.kill()
+                except Exception:
+                    pass
+        shutil.rmtree(self.dir, ignore_errors=True)
 
 
 def child_env(cfg: dict) -> dict:
@@ -539,7 +614,7 @@ def prepare(self, messages, tools, kwargs, max_new=None):
             # "reading the prompt" with CPU and GPU busy, for good (reproduced).  So encoding takes its turn in the
             # same FIFO as the requests.
             with self.fifo:
-                encoded = [self.vision.encode(src) for src in images]
+                encoded = [self.vision.encode(src, engine=self.engine) for src in images]
             out, k = [], 0
             for t in ids:                               # one <|image_pad|> per image -> one per image token
                 if t == pad and k < len(encoded):
@@ -1234,6 +1309,8 @@ def main() -> int:
                          "(default: reject with 400, like llama.cpp; also \"fit_max_tokens\": true in the config)")
     ap.add_argument("--api-key", default=os.environ.get("STRATA_API_KEY", ""),
                     help="require this key on /v1/* (Authorization: Bearer ... or x-api-key); also $STRATA_API_KEY")
+    ap.add_argument("--no-vision", action="store_true",
+                    help="disable vision even if the config has a vision block")
     a = ap.parse_args()
     cfg = json.loads(Path(a.config).read_text(encoding="utf-8-sig")) if a.config else {}   # Notepad adds a BOM
     a.host = a.host or cfg.get("host") or "127.0.0.1"   # issue #26: the run scripts pass no --host, the config can
@@ -1266,13 +1343,13 @@ def main() -> int:
         if sampling_defaults:
             pretty = ", ".join(f"{k}={v}" for k, v in sampling_defaults.items())
             print(f"[strata] sampling defaults from the config: {pretty}", flush=True)
-        if cfg.get("vision"):
-            print("loading the vision encoder ...", flush=True)
-            vision = Vision(cfg["vision"], log=open(cfg["log"], "a", encoding="utf-8") if cfg.get("log") else None,
-                            env=env)
         print("loading the model (the first start takes a minute or two) ...", flush=True)
         engine = StrataEngine(cfg["exe"], cfg["args"], cwd=cfg.get("cwd"), log=cfg.get("log"), env=env)
         warn_tight_ram(engine.info.get("arena_mib"))
+        if cfg.get("vision") and not a.no_vision:
+            print("configuring the vision encoder ...", flush=True)
+            vision = Vision(cfg["vision"], log=open(cfg["log"], "a", encoding="utf-8") if cfg.get("log") else None,
+                            env=env)
     else:
         engine, vision, sampling_defaults = MockEngine(tok, a.script), None, {}
     # the model's own chat template (exported with its tokenizer), else the original model's