diff --git a/CMakeLists.txt b/CMakeLists.txt index 84882d0901..3b93437105 100644 --- a/CMakeLists.txt +++ b/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 diff --git a/include/strata/core/expert_cache.hpp b/include/strata/core/expert_cache.hpp index da856d75d7..de396a912c 100644 --- a/include/strata/core/expert_cache.hpp +++ b/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 #include +#include #include 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>& 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& residency_table() const { return residency_; } + private: uint8_t* base_ = nullptr; std::vector residency_; ///< [n_layers * n_expert] -> slot or kNotResident @@ -150,6 +202,21 @@ class ExpertCache { std::vector layer_next_; ///< [n_layers] -> that layer's next free slot std::vector 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 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 diff --git a/serve/server.py b/serve/server.py index 1951d99848..2f37714583 100644 --- a/serve/server.py +++ b/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 ` + """The image encoder: `strata-vision` (llama.cpp mtmd + the mmproj file) reads `ENC ` 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 diff --git a/src/core/expert_cache.cpp b/src/core/expert_cache.cpp index 9f5fc981fb..661a9d452a 100644 --- a/src/core/expert_cache.cpp +++ b/src/core/expert_cache.cpp @@ -1,11 +1,13 @@ -// src/core/expert_cache.cpp - R4's slot storage and residency table. Read the header first. +// src/core/expert_cache.cpp - R4's slot storage, residency table, and R4.3 elastic VMM tier. #include "strata/core/expert_cache.hpp" +#include #include +#include #include -#include #include +#include namespace strata::core { @@ -63,12 +65,307 @@ bool read_expert_profile(const std::string& path, int64_t n_layers, int64_t n_ex ranked[(size_t) i] = {l, e}; } slots = (int64_t) want; - (void) version; // a future format bumps it; the layout check above is what protects this reader today + (void) version; return true; } ExpertCache::~ExpertCache() { close(); } +bool ExpertCache::vmm_supported() { + static int cached = -1; + if (cached >= 0) return cached == 1; + if (cuInit(0) != CUDA_SUCCESS) { cached = 0; return false; } + int dev = 0; + if (cudaGetDevice(&dev) != cudaSuccess) dev = 0; + CUdevice cuDev = 0; + if (cuDeviceGet(&cuDev, dev) != CUDA_SUCCESS) { cached = 0; return false; } + int vmm = 0; + if (cuDeviceGetAttribute(&vmm, CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED, cuDev) != CUDA_SUCCESS) { + cached = 0; + return false; + } + cached = (vmm != 0) ? 1 : 0; + return cached == 1; +} + +size_t ExpertCache::vmm_granularity() { + if (!vmm_supported()) return 0; + int dev = 0; + if (cudaGetDevice(&dev) != cudaSuccess) dev = 0; + CUdevice cuDev = 0; + if (cuDeviceGet(&cuDev, dev) != CUDA_SUCCESS) return 0; + CUmemAllocationProp prop = {}; + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + prop.location.id = dev; + size_t gran = 0; + if (cuMemGetAllocationGranularity(&gran, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED) != CUDA_SUCCESS) { + if (cuMemGetAllocationGranularity(&gran, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM) != CUDA_SUCCESS) { + return 0; + } + } + return gran; +} + +bool ExpertCache::open_vmm(int64_t total_bytes, int64_t n_layers, int64_t n_expert, std::string& err) { + (void) n_layers; + (void) n_expert; + close(); + size_t gran = vmm_granularity(); + if (gran == 0) { + err = "ExpertCache: VMM not supported or granularity query failed"; + return false; + } + int dev = 0; + if (cudaGetDevice(&dev) != cudaSuccess) dev = 0; + CUdevice cuDev = 0; + if (cuDeviceGet(&cuDev, dev) != CUDA_SUCCESS) { + err = "ExpertCache: cuDeviceGet failed"; + return false; + } + + const size_t aligned_size = ((size_t) total_bytes + gran - 1) / gran * gran; + + CUdeviceptr va = 0; + CUresult res = cuMemAddressReserve(&va, aligned_size, 0, 0, 0); + if (res != CUDA_SUCCESS) { + const char* s = nullptr; + cuGetErrorString(res, &s); + err = std::string("ExpertCache: cuMemAddressReserve failed: ") + (s ? s : "unknown"); + return false; + } + + CUmemAllocationProp prop = {}; + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + prop.location.id = dev; + + CUmemAccessDesc desc = {}; + desc.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + desc.location.id = dev; + desc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + + const size_t chunk_size = gran; + const size_t num_chunks = aligned_size / chunk_size; + + vmm_chunks_.reserve(num_chunks); + size_t mapped = 0; + + for (size_t i = 0; i < num_chunks; ++i) { + CUmemGenericAllocationHandle handle = 0; + res = cuMemCreate(&handle, chunk_size, &prop, 0); + if (res != CUDA_SUCCESS) { + for (auto& c : vmm_chunks_) { + cuMemUnmap((CUdeviceptr) (va + c.offset), (size_t) c.size); + cuMemRelease((CUmemGenericAllocationHandle) c.handle); + } + vmm_chunks_.clear(); + cuMemAddressFree(va, aligned_size); + const char* s = nullptr; + cuGetErrorString(res, &s); + err = std::string("ExpertCache: cuMemCreate failed: ") + (s ? s : "out of VRAM"); + return false; + } + + res = cuMemMap(va + (size_t) mapped, chunk_size, 0, handle, 0); + if (res != CUDA_SUCCESS) { + cuMemRelease(handle); + for (auto& c : vmm_chunks_) { + cuMemUnmap((CUdeviceptr) (va + c.offset), (size_t) c.size); + cuMemRelease((CUmemGenericAllocationHandle) c.handle); + } + vmm_chunks_.clear(); + cuMemAddressFree(va, aligned_size); + const char* s = nullptr; + cuGetErrorString(res, &s); + err = std::string("ExpertCache: cuMemMap failed: ") + (s ? s : "unknown"); + return false; + } + + res = cuMemSetAccess(va + (size_t) mapped, chunk_size, &desc, 1); + if (res != CUDA_SUCCESS) { + cuMemUnmap(va + (size_t) mapped, chunk_size); + cuMemRelease(handle); + for (auto& c : vmm_chunks_) { + cuMemUnmap((CUdeviceptr) (va + c.offset), (size_t) c.size); + cuMemRelease((CUmemGenericAllocationHandle) c.handle); + } + vmm_chunks_.clear(); + cuMemAddressFree(va, aligned_size); + const char* s = nullptr; + cuGetErrorString(res, &s); + err = std::string("ExpertCache: cuMemSetAccess failed: ") + (s ? s : "unknown"); + return false; + } + + vmm_chunks_.push_back({(uint64_t) handle, (uint64_t) mapped, (uint64_t) chunk_size}); + mapped += chunk_size; + } + + base_ = (uint8_t*) va; + vmm_va_ = (uint64_t) va; + vmm_va_size_ = (uint64_t) aligned_size; + max_bytes_ = (int64_t) aligned_size; + mapped_bytes_ = (int64_t) mapped; + vmm_ = true; + + if (cudaMemset(base_, 0, (size_t) mapped_bytes_) != cudaSuccess) { + err = "ExpertCache: cudaMemset of VMM slot arena failed"; + close(); + return false; + } + + return true; +} + +void ExpertCache::close_vmm() { + if (vmm_va_ != 0) { + for (auto& c : vmm_chunks_) { + cuMemUnmap((CUdeviceptr) (vmm_va_ + c.offset), (size_t) c.size); + cuMemRelease((CUmemGenericAllocationHandle) c.handle); + } + vmm_chunks_.clear(); + cuMemAddressFree((CUdeviceptr) vmm_va_, (size_t) vmm_va_size_); + vmm_va_ = 0; + vmm_va_size_ = 0; + } + base_ = nullptr; + vmm_ = false; + max_slots_ = 0; + max_bytes_ = 0; + mapped_bytes_ = 0; +} + +int64_t ExpertCache::shrink(int64_t keep_bytes, std::string& err) { + (void) err; + if (!vmm_ || vmm_chunks_.empty()) return 0; + cudaDeviceSynchronize(); + + size_t gran = vmm_granularity(); + if (gran == 0) return 0; + const size_t keep_aligned = ((size_t) std::max(0, keep_bytes) + gran - 1) / gran * gran; + if (keep_aligned >= (size_t) mapped_bytes_) return 0; + + int64_t freed = 0; + while (!vmm_chunks_.empty() && vmm_chunks_.back().offset >= keep_aligned) { + auto& c = vmm_chunks_.back(); + cuMemUnmap((CUdeviceptr) (vmm_va_ + c.offset), (size_t) c.size); + cuMemRelease((CUmemGenericAllocationHandle) c.handle); + freed += (int64_t) c.size; + mapped_bytes_ -= (int64_t) c.size; + vmm_chunks_.pop_back(); + } + + int64_t new_slots = 0; + if (!off_.empty()) { + while (new_slots < (int64_t) off_.size() - 1 && (int64_t) off_[(size_t) (new_slots + 1)] <= mapped_bytes_) { + ++new_slots; + } + } else if (blob_ > 0) { + new_slots = mapped_bytes_ / blob_; + } + + for (size_t i = 0; i < residency_.size(); ++i) { + if (residency_[i] >= new_slots && residency_[i] < slots_) { + residency_[i] = kNotResident; + } + } + slots_ = new_slots; + return freed; +} + +int64_t ExpertCache::grow(int64_t new_bytes, const std::vector>& profile, + BlobFn blob_fn, void* blob_user, BlobSizeFn size_fn, void* size_user, std::string& err) { + if (!vmm_ || vmm_va_ == 0) return 0; + cudaDeviceSynchronize(); + + size_t gran = vmm_granularity(); + if (gran == 0) return 0; + const size_t target_aligned = std::min((size_t) vmm_va_size_, ((size_t) new_bytes + gran - 1) / gran * gran); + if (target_aligned <= (size_t) mapped_bytes_) return 0; + + int dev = 0; + if (cudaGetDevice(&dev) != cudaSuccess) dev = 0; + CUmemAllocationProp prop = {}; + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + prop.location.id = dev; + + CUmemAccessDesc desc = {}; + desc.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + desc.location.id = dev; + desc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + + const size_t chunk_size = gran; + while ((size_t) mapped_bytes_ < target_aligned) { + CUmemGenericAllocationHandle handle = 0; + CUresult res = cuMemCreate(&handle, chunk_size, &prop, 0); + if (res != CUDA_SUCCESS) break; + + res = cuMemMap(vmm_va_ + (size_t) mapped_bytes_, chunk_size, 0, handle, 0); + if (res != CUDA_SUCCESS) { + cuMemRelease(handle); + break; + } + + res = cuMemSetAccess(vmm_va_ + (size_t) mapped_bytes_, chunk_size, &desc, 1); + if (res != CUDA_SUCCESS) { + cuMemUnmap(vmm_va_ + (size_t) mapped_bytes_, chunk_size); + cuMemRelease(handle); + break; + } + + vmm_chunks_.push_back({(uint64_t) handle, (uint64_t) mapped_bytes_, (uint64_t) chunk_size}); + mapped_bytes_ += (int64_t) chunk_size; + } + + int64_t old_slots = slots_; + int64_t new_slots = old_slots; + if (!off_.empty()) { + while (new_slots < (int64_t) off_.size() - 1 && (int64_t) off_[(size_t) (new_slots + 1)] <= mapped_bytes_) { + ++new_slots; + } + } else if (blob_ > 0) { + new_slots = mapped_bytes_ / blob_; + } + slots_ = new_slots; + + int64_t filled = 0; + if (blob_fn != nullptr) { + for (int64_t s = old_slots; s < new_slots; ++s) { + if ((size_t) s < profile.size()) { + const auto& pr = profile[(size_t) s]; + const int32_t l = pr.first; + const int32_t e = pr.second; + const uint8_t* b = blob_fn(l, e, blob_user); + const int64_t bsz = size_fn ? size_fn(l, size_user) : 0; + if (b != nullptr && fill_slot_blocking((int32_t) s, b, err, bsz)) { + assign_slot((int32_t) s, l, e); + ++filled; + } + } + } + } + return filled; +} + +void ExpertCache::evict_slot(int32_t slot) { + if (slot < 0 || slot >= slots_) return; + for (size_t i = 0; i < residency_.size(); ++i) { + if (residency_[i] == slot) { + residency_[i] = kNotResident; + break; + } + } +} + +void ExpertCache::assign_slot(int32_t slot, int64_t layer, int64_t expert) { + if (layer < 0 || layer >= n_layers_ || expert < 0 || expert >= n_expert_) return; + if (slot < 0 || slot >= slots_) return; + evict_slot(slot); + residency_[(size_t) (layer * n_expert_ + expert)] = slot; +} + bool ExpertCache::open(int64_t n_slots, int64_t n_layers, int64_t n_expert, int64_t blob_bytes, std::string& err) { close(); @@ -83,12 +380,6 @@ bool ExpertCache::open(int64_t n_slots, int64_t n_layers, int64_t n_expert, int6 const uint64_t want = (uint64_t) n_slots * (uint64_t) blob_bytes; - // ---- **THE ALLOCATION IS CHECKED AGAINST THE CARD, NOT AGAINST THE REQUEST.** - // - // `cudaMalloc` failing is the easy case. The one that matters is a machine where the weights already own - // most of VRAM: the cache then takes what is left and `slots()` would report the number ASKED FOR while - // `device_slot()` walks off the end. So the free-VRAM figure is read and compared BEFORE the allocation, - // and the two numbers are named in the refusal. size_t free_b = 0, total_b = 0; if (cudaMemGetInfo(&free_b, &total_b) == cudaSuccess) { if ((uint64_t) free_b < want) { @@ -104,32 +395,46 @@ bool ExpertCache::open(int64_t n_slots, int64_t n_layers, int64_t n_expert, int6 } } - if (cudaMalloc((void**) &base_, (size_t) want) != cudaSuccess) { - base_ = nullptr; - char buf[256]; - std::snprintf(buf, sizeof buf, "ExpertCache: cudaMalloc(%.2f GiB) failed: %s", - (double) want / 1073741824.0, cudaGetErrorString(cudaGetLastError())); - err = buf; - return false; + bool ok = false; + if (vmm_supported()) { + std::string vmm_err; + ok = open_vmm((int64_t) want, n_layers, n_expert, vmm_err); + if (!ok) { + // Log fallback message if VMM failed, then fallback to standard cudaMalloc + std::fprintf(stderr, "strata: CUDA VMM allocation failed (%s); falling back to cudaMalloc\n", + vmm_err.c_str()); + } } - // Zeroed so a slot read before it is filled is a DETERMINISTIC wrong answer rather than whatever the - // allocator handed back. A stale block of a previous process's memory would still sum to finite floats. - if (cudaMemset(base_, 0, (size_t) want) != cudaSuccess) { - err = "ExpertCache: cudaMemset of the slot arena failed"; - close(); - return false; + + if (!ok) { + if (cudaMalloc((void**) &base_, (size_t) want) != cudaSuccess) { + base_ = nullptr; + char buf[256]; + std::snprintf(buf, sizeof buf, "ExpertCache: cudaMalloc(%.2f GiB) failed: %s", + (double) want / 1073741824.0, cudaGetErrorString(cudaGetLastError())); + err = buf; + return false; + } + if (cudaMemset(base_, 0, (size_t) want) != cudaSuccess) { + err = "ExpertCache: cudaMemset of the slot arena failed"; + close(); + return false; + } + vmm_ = false; + max_bytes_ = (int64_t) want; + mapped_bytes_ = (int64_t) want; } residency_.assign((size_t) (n_layers * n_expert), kNotResident); slots_ = n_slots; + max_slots_ = n_slots; n_layers_ = n_layers; n_expert_ = n_expert; blob_ = blob_bytes; next_free_ = 0; fills_ = 0; admitted_ = 0; - // R4.2g: each layer starts at the bottom of its own range. Built here rather than lazily so `admit` - // stays allocation-free on the token path. + layer_next_.assign((size_t) (n_layers > 0 ? n_layers : 0), 0); for (int64_t l = 0; l < n_layers; ++l) { int64_t lo = 0, hi = 0; @@ -145,13 +450,12 @@ bool ExpertCache::open_sized(const std::vector& slot_bytes, int64_t n_l int64_t mx = 0; std::vector off(slot_bytes.size() + 1, 0); for (size_t i = 0; i < slot_bytes.size(); ++i) { - // 256-byte aligned slots, so every blob starts where the kernels' vector loads expect it off[i + 1] = off[i] + ((uint64_t) slot_bytes[i] + 255) / 256 * 256; mx = slot_bytes[i] > mx ? slot_bytes[i] : mx; } - // one allocation of the summed size, through the uniform path's checks: n "slots" of 1 byte if (!open((int64_t) off.back(), n_layers, n_expert, 1, err)) return false; slots_ = (int64_t) slot_bytes.size(); + max_slots_ = slots_; blob_ = mx; off_ = std::move(off); layer_next_.assign((size_t) (n_layers > 0 ? n_layers : 0), 0); @@ -160,12 +464,17 @@ bool ExpertCache::open_sized(const std::vector& slot_bytes, int64_t n_l void ExpertCache::close() { off_.clear(); - if (base_ != nullptr) { + if (vmm_) { + close_vmm(); + } else if (base_ != nullptr) { cudaFree(base_); base_ = nullptr; } residency_.clear(); slots_ = 0; + max_slots_ = 0; + max_bytes_ = 0; + mapped_bytes_ = 0; n_layers_ = 0; n_expert_ = 0; blob_ = 0; @@ -175,8 +484,6 @@ void ExpertCache::close() { layer_next_.clear(); } -/// R4.2g. Layer `l` owns `[l*q, (l+1)*q)` with `q = slots_ / n_layers_`; the LAST layer takes whatever is -/// left over, so the ranges always cover `0..slots_` exactly and no slot is orphaned by the division. void ExpertCache::layer_slot_range(int64_t layer, int64_t& lo, int64_t& hi) const { lo = 0; hi = 0; @@ -195,18 +502,16 @@ int32_t ExpertCache::admit(int64_t layer, int64_t expert) { if (layer < 0 || layer >= n_layers_ || expert < 0 || expert >= n_expert_) return kNotResident; const size_t at = (size_t) (layer * n_expert_ + expert); if (residency_[at] != kNotResident) return residency_[at]; - // R4.2g: THE PER-LAYER PATH. Same "no eviction" rule, but the ceiling is this layer's own range rather - // than one counter shared by all 48 - which is what confined the measured hit rate to 2.97%. if (per_layer_) { if (layer_next_.empty()) return kNotResident; int64_t lo = 0, hi = 0; layer_slot_range(layer, lo, hi); - if ((int64_t) layer_next_[(size_t) layer] >= hi) return kNotResident; // this layer's quota is full + if ((int64_t) layer_next_[(size_t) layer] >= hi) return kNotResident; residency_[at] = layer_next_[(size_t) layer]++; ++admitted_; return residency_[at]; } - if (next_free_ >= slots_) return kNotResident; // full: no eviction, deliberately - see the header + if (next_free_ >= slots_) return kNotResident; residency_[at] = (int32_t) next_free_; return (int32_t) next_free_++; } @@ -272,9 +577,6 @@ bool ExpertCache::verify_slot(int32_t slot, const uint8_t* host_blob, std::strin err = "ExpertCache::verify_slot: slot outside the arena"; return false; } - // `cudaMemcpy` and not `cudaMemcpyAsync`: this is a startup check, and a check that can be read before it - // has happened is not a check. It also synchronises the fills queued before it, which is what makes the - // comparison meaningful. std::vector got((size_t) nb); const cudaError_t e = cudaMemcpy(got.data(), src, (size_t) nb, cudaMemcpyDeviceToHost); if (e != cudaSuccess) { diff --git a/src/program/generate.cpp b/src/program/generate.cpp index da5d138bb8..e06c012496 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -2388,6 +2388,36 @@ int main(int argc, char** argv) { ~BusyScope() { strata::core::progress().busy.store(false); strata::core::progress_at("idle"); } } busy_scope; stop_req.store(false); // a STOP that arrived between requests is stale + if (line.rfind("RESIZE ", 0) == 0) { + long long keep = std::strtoll(line.c_str() + 7, nullptr, 10); + std::string rerr; + int64_t freed = xcache.shrink((int64_t) keep, rerr); + const auto& res_tab = xcache.residency_table(); + std::memcpy(host_res.data(), res_tab.data(), host_res.size() * sizeof(int32_t)); + if (d_res != nullptr) + cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice); + std::printf("OK %lld %lld\n", (long long) xcache.slots(), (long long) freed); + std::fflush(stdout); + continue; + } + if (line.rfind("GROW ", 0) == 0) { + long long target = std::strtoll(line.c_str() + 5, nullptr, 10); + std::string gerr; + auto bfn = [](int32_t l, int32_t e, void* user) -> const uint8_t* { + return ((strata::core::ExpertSource*) user)->blob(l, e); + }; + auto sfn = [](int32_t l, void* /*user*/) -> int64_t { + return (int64_t) strata::kernels::cpu::expert_layout().blob_bytes(l); + }; + int64_t filled = xcache.grow((int64_t) target, profile, bfn, srcp, sfn, nullptr, gerr); + const auto& res_tab = xcache.residency_table(); + std::memcpy(host_res.data(), res_tab.data(), host_res.size() * sizeof(int32_t)); + if (d_res != nullptr) + cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice); + std::printf("OK %lld %lld\n", (long long) xcache.slots(), (long long) filled); + std::fflush(stdout); + continue; + } const bool geni = line.rfind("GENI ", 0) == 0; if (!geni && line.rfind("GEN ", 0) != 0) { std::printf("ERR expected: GEN or GENI \n"); diff --git a/tools/test_vision_client.py b/tools/test_vision_client.py new file mode 100644 index 0000000000..46a942d59a --- /dev/null +++ b/tools/test_vision_client.py @@ -0,0 +1,46 @@ +import urllib.request +import json +import time + +# Create a small dynamic PNG using different bytes +import io +try: + from PIL import Image + im = Image.new("RGB", (16, 16), color=(int(time.time()) % 255, 128, 64)) + buf = io.BytesIO() + im.save(buf, format="PNG") + import base64 + png_b64 = base64.b64encode(buf.getvalue()).decode() +except Exception: + png_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + +payload = { + "model": "swift-1.5-iq2_xs", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this image? One short sentence."}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{png_b64}"}} + ] + } + ], + "max_tokens": 150, + "temperature": 0.0 +} + +req = urllib.request.Request( + "http://127.0.0.1:8080/v1/chat/completions", + data=json.dumps(payload).encode("utf-8"), + headers={"Content-Type": "application/json"} +) + +print("Sending unique image request to Strata server...") +t0 = time.time() +with urllib.request.urlopen(req, timeout=30) as resp: + res = json.loads(resp.read().decode("utf-8")) + msg = res["choices"][0]["message"] + print(f"Response in {time.time() - t0:.2f}s:") + if msg.get("reasoning_content"): + print("Thinking:", msg["reasoning_content"][:120], "...") + print("Answer:", msg.get("content"))