mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-07 06:15:58 -10:00
Hot VRAM resize, opt-in: a segmented expert cache, VRAM <reserve>, POST /v1/vram (#533)
--vram-elastic (the config's "vram_elastic": true) allocates the expert cache as one
reserved address range backed by physical segments (CUDA virtual memory management,
512 MiB each, --vram-segment-mib), so every slot keeps its address. Between
requests, the engine command `VRAM <reserve_mib>` gives the last segments back to
the driver until that much VRAM is free (their experts become CPU misses, the prompt
path's loan moves down and its chunk shrinks to fit), or maps them again and puts
the same experts back when more is free. `VRAM` alone: the start's reserve. The
server's POST /v1/vram {"reserve_mib": N | null} sends it after a running request,
keeps it for an engine that is unloaded, and /v1/status shows the cache's size.
Never automatic. One NVIDIA GPU only (no layer split, helper caches, --peer-device
or the resident low-RAM mode); HIP refuses the flag.
RTX 5070, Q2_0 (4.8 GiB cache): reserve 6000 -> 4.3 GiB freed in 78 ms, decode
44 -> 33 tok/s; grown back in 92 ms, the answers token for token the ones before.
Without the flag the cache is the one cudaMalloc it was (same answers with and
without the flag until a VRAM command).
This commit is contained in:
1 parent
ad0df801ae
commit
89bad9a4ec
8 files changed
+774
-6
No files matched your search
@@ -384,6 +384,10 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
|
||||
target_link_libraries(expert_cache_per_layer_test PRIVATE strata_engine)
|
||||
add_test(NAME expert_cache_per_layer_test COMMAND expert_cache_per_layer_test)
|
||||
set_tests_properties(expert_cache_per_layer_test PROPERTIES SKIP_RETURN_CODE 77 TIMEOUT 60)
|
||||
add_executable(expert_cache_segmented_test tests/core/expert_cache_segmented_test.cpp)
|
||||
target_link_libraries(expert_cache_segmented_test PRIVATE strata_engine)
|
||||
add_test(NAME expert_cache_segmented_test COMMAND expert_cache_segmented_test)
|
||||
set_tests_properties(expert_cache_segmented_test PROPERTIES SKIP_RETURN_CODE 77 TIMEOUT 60)
|
||||
endif()
|
||||
if(STRATA_BUILD_CONVERSATION_TESTS)
|
||||
add_executable(conversation_snapshot_test src/core/conversation_snapshot_test.cpp)
|
||||
|
||||
@@ -414,6 +414,18 @@ the OS file cache, so loading again takes seconds while that RAM is not needed e
|
||||
16 GB with Q2_0 in the low-RAM mode: unloading takes ~0.3 s, and a request to an unloaded model answered after
|
||||
4.6 s (text) or 14.7 s (a picture, image encoder on the CPU).
|
||||
|
||||
**Giving part of the VRAM back while it keeps serving (#533, opt-in, one NVIDIA GPU).** With `"vram_elastic": true`
|
||||
in the config (the engine flag `--vram-elastic`), the expert cache is allocated in 512 MiB segments
|
||||
(`"vram_segment_mib"`), and `POST /v1/vram` with `{"reserve_mib": 8000}` shrinks it between requests until that much
|
||||
VRAM is free for another program; `{"reserve_mib": null}` grows it back towards its full size, keeping the reserve
|
||||
the engine started with (`--vram-reserve-mib`), and `{"reserve_mib": 0}` takes all of it back. A request that is
|
||||
running finishes first. The experts of the segments given back are computed on the CPU, like any expert outside the
|
||||
cache, so answers keep coming, slower; growing back puts the same experts in the same slots. Nothing resizes on its
|
||||
own. Not with a layer split, the helper caches, `--peer-device` or the resident low-RAM mode. Measured on an RTX 5070
|
||||
12 GB with Q2_0 (4.8 GiB cache): `{"reserve_mib": 6000}` took 78 ms and freed 4.3 GiB (the cache keeps 0.5 GiB for
|
||||
the prompt path), decode 44 -> 33 tok/s; growing back took 92 ms and the answers were token for token the ones before
|
||||
the shrink. Without the flag nothing changes (the same answers as without it).
|
||||
|
||||
**Keep what the expert cache learned across restarts (opt-in, engine 0.1.36, #477):** a start fills the GPU's expert
|
||||
cache from the shipped profile, and the adaptive tier (`--adapt-every`) then moves in the experts your requests use.
|
||||
With `"expert_profile_save": "expert-profile-learned.bin"` in `strata-<model>.json` the engine saves that as a
|
||||
|
||||
@@ -93,13 +93,44 @@ public:
|
||||
void close();
|
||||
|
||||
bool valid() const { return base_ != nullptr; }
|
||||
int64_t slots() const { return slots_; }
|
||||
/// The slots usable now: all of them, unless `shrink` gave the tail's VRAM back (#533) - then the ones wholly
|
||||
/// inside the VRAM still backed, always a prefix. The prompt path's loan (the last slots) and every report
|
||||
/// follow this; `full_slots()` is what `open` allocated.
|
||||
int64_t slots() const { return live_slots_; }
|
||||
int64_t full_slots() const { return slots_; }
|
||||
/// Slots actually claimed. Not the same as `slots()` - the cache does not evict, so a run that routes
|
||||
/// fewer distinct experts than there are slots leaves the rest empty.
|
||||
int64_t resident() const { return per_layer_ ? admitted_ : next_free_; }
|
||||
int64_t bytes() const { return off_.empty() ? slots_ * blob_ : (int64_t) off_.back(); }
|
||||
/// The bytes of `slots()` (all of the arena unless shrunk); `full_bytes()`: of every slot.
|
||||
int64_t bytes() const { return slot_end(live_slots_); }
|
||||
int64_t full_bytes() const { return slot_end(slots_); }
|
||||
double gib() const { return (double) bytes() / 1073741824.0; }
|
||||
|
||||
/// **#533: HOT VRAM RESIZE, OPT-IN (--vram-elastic).** With a segment size set before `open`/`open_sized`, the
|
||||
/// arena is one reserved address range backed by physical segments of that size (CUDA virtual memory
|
||||
/// management) instead of one cudaMalloc, so `shrink` can give the tail's VRAM back to the driver - for another
|
||||
/// program - and `grow` can take it again, while every slot keeps its address (the captured graphs, the verify
|
||||
/// plan's pointers and the prompt path's loan all hold addresses into it). 0 (the default): one cudaMalloc, as
|
||||
/// always. CUDA only: a HIP build refuses it at open. The slots past `slots()` after a shrink must hold no
|
||||
/// resident expert and no loan: the CALLER evicts them first (they become CPU misses).
|
||||
void set_segment_bytes(int64_t seg_bytes) { seg_req_ = seg_bytes > 0 ? seg_bytes : 0; }
|
||||
bool segmented() const { return !segs_.empty(); }
|
||||
int64_t segment_bytes() const { return seg_; }
|
||||
/// Bytes of the arena backed by VRAM now (a whole number of segments; the arena when not segmented).
|
||||
int64_t mapped_bytes() const;
|
||||
/// Unmaps every segment past the first `keep_bytes` (rounded UP to a segment boundary): `slots()` becomes the
|
||||
/// slots wholly inside what stays. Waits for the device first. False with `err` when not segmented or a driver
|
||||
/// call failed.
|
||||
bool shrink(int64_t keep_bytes, std::string& err);
|
||||
/// Maps segments again up to `want_bytes` (rounded DOWN to a segment, at most the arena; the last, shorter
|
||||
/// segment only when `want_bytes` covers the arena). Stops at the first segment the driver cannot back (false,
|
||||
/// `err`: what was mapped by then stays). The new slots are empty until the caller fills them.
|
||||
bool grow(int64_t want_bytes, std::string& err);
|
||||
/// The number of leading slots that fit wholly inside the first `bytes` bytes.
|
||||
int64_t slots_within(int64_t bytes) const;
|
||||
/// The bytes the first `n` slots span.
|
||||
int64_t bytes_of(int64_t n) const { return slot_end(n < 0 ? 0 : n > slots_ ? slots_ : n); }
|
||||
|
||||
/// `(layer, expert)` -> slot index, or `kNotResident`. Bounds-checked: a bad layer or expert returns
|
||||
/// `kNotResident` rather than reading whatever is adjacent in the table.
|
||||
int32_t slot_of(int64_t layer, int64_t expert) const;
|
||||
@@ -164,7 +195,17 @@ private:
|
||||
uint8_t* blocking_staging_ = nullptr;
|
||||
std::size_t blocking_staging_bytes_ = 0;
|
||||
#endif
|
||||
int64_t slot_end(int64_t n) const { return off_.empty() ? n * blob_ : (int64_t) off_[(size_t) n]; }
|
||||
bool open_segmented(uint64_t want, std::string& err);
|
||||
void release_segmented();
|
||||
uint8_t* base_ = nullptr;
|
||||
int64_t live_slots_ = 0; ///< #533: slots() - all of them unless shrunk
|
||||
int64_t seg_req_ = 0; ///< #533: the segment size asked for (0: one cudaMalloc)
|
||||
int64_t seg_ = 0; ///< #533: the segment size used (a multiple of the driver's granularity)
|
||||
uint64_t reserved_ = 0; ///< #533: the reserved address range's size
|
||||
std::vector<unsigned long long> segs_; ///< #533: each segment's physical handle (0: unmapped)
|
||||
std::vector<int64_t> seg_size_; ///< #533: each segment's size (the last one may be shorter)
|
||||
int64_t mapped_segs_ = 0; ///< #533: segments [0, mapped_segs_) are backed
|
||||
std::vector<int32_t> residency_; ///< [n_layers * n_expert] -> slot or kNotResident
|
||||
int64_t slots_ = 0;
|
||||
int64_t n_layers_ = 0;
|
||||
|
||||
@@ -444,6 +444,37 @@ class StrataEngine:
|
||||
if len(f) >= 15: # #471 (engine 0.1.36+): the prompt tokens actually read
|
||||
self.last.update(prompt_read=int(f[14]))
|
||||
|
||||
def vram(self, reserve_mib: int | None, timeout: float = 120.0) -> dict:
|
||||
"""#533: `VRAM <reserve_mib>` between requests (the caller holds the service's FIFO): the engine shrinks its
|
||||
expert cache until that much VRAM is free, or grows it back when more is free; None: back to the reserve it
|
||||
started with. -> the engine's figures (expert_slots, expert_cache_mib, vram_free_mib, ...); raises ValueError
|
||||
with the engine's reason (e.g. it was not started with --vram-elastic), EngineDied when it ended."""
|
||||
if not self.alive():
|
||||
raise EngineDied("the engine is not running")
|
||||
self.proc.stdin.write("VRAM" + ("" if reserve_mib is None else f" {int(reserve_mib)}") + "\n")
|
||||
self.proc.stdin.flush()
|
||||
deadline = time.time() + timeout
|
||||
while True:
|
||||
try:
|
||||
line = self.lines.get(timeout=max(0.1, deadline - time.time()))
|
||||
except queue.Empty:
|
||||
raise EngineDied("the engine did not answer the VRAM command") from None
|
||||
if line is None:
|
||||
raise EngineDied("the engine ended")
|
||||
line = line.strip()
|
||||
if line.startswith("ERR"):
|
||||
raise ValueError(line[4:].strip() or "the engine refused the VRAM command")
|
||||
if line.startswith("VRAM "):
|
||||
out = {}
|
||||
for kv in line.split()[1:]:
|
||||
k, _, v = kv.partition("=")
|
||||
out[k] = int(v) if v.lstrip("-").isdigit() else v
|
||||
self.info.update({k: out[k] for k in ("expert_slots", "vram_free_mib") if k in out})
|
||||
self.info["vram"] = out
|
||||
return out
|
||||
if time.time() > deadline:
|
||||
raise EngineDied("the engine did not answer the VRAM command")
|
||||
|
||||
@staticmethod
|
||||
def sampling_keys(sampling: dict) -> str:
|
||||
keys = ""
|
||||
@@ -798,6 +829,13 @@ def engine_args(cfg: dict) -> list[str]:
|
||||
# opt-in: an auto split runs on the first card alone when it holds every profiled expert and the KV
|
||||
if len(gpu_list(cfg)) > 1 and cfg.get("split_skip_if_fits") and "--split-skip-if-fits" not in args:
|
||||
args.append("--split-skip-if-fits")
|
||||
# #533 (opt-in): "vram_elastic": true - the expert cache in segments, so POST /v1/vram can give VRAM back to other
|
||||
# programs and take it back; "vram_segment_mib" sets the segment size (the engine's default: 512)
|
||||
if cfg.get("vram_elastic") is True and "--vram-elastic" not in args:
|
||||
args.append("--vram-elastic")
|
||||
seg = cfg.get("vram_segment_mib")
|
||||
if isinstance(seg, int) and not isinstance(seg, bool) and seg > 0 and "--vram-segment-mib" not in args:
|
||||
args += ["--vram-segment-mib", str(seg)]
|
||||
return learned_profile_args(cfg, args)
|
||||
|
||||
|
||||
@@ -988,6 +1026,7 @@ class Service:
|
||||
self.idle_unload_s = 0
|
||||
self.min_free_vram_mib = 0
|
||||
self.before_load = None
|
||||
self.vram_reserve = None # #533: the last POST /v1/vram reserve (None: the start's)
|
||||
self.reasoning_budget_tokens = 0 # #123: the config's default thinking budget (0: none)
|
||||
self.anthropic_think_unasked = True # #278: "anthropic_thinking": "on_request" -> False
|
||||
self.stop_ids = set(tokenizer.encode(IM_END, parse_special=True) +
|
||||
@@ -1088,6 +1127,35 @@ class Service:
|
||||
"(a minute or two) ...", flush=True)
|
||||
self.engine.restart()
|
||||
print("[strata] the engine is running again", flush=True)
|
||||
if self.vram_reserve is not None and hasattr(self.engine, "vram"): # #533: the reserve asked for last
|
||||
try:
|
||||
self.engine.vram(self.vram_reserve)
|
||||
except (ValueError, EngineDied) as e:
|
||||
print(f"[strata] the VRAM reserve ({self.vram_reserve} MiB) was not applied: {e}", flush=True)
|
||||
|
||||
vram_wait_s = 300.0 # #533: how long POST /v1/vram waits for a running request
|
||||
|
||||
def vram(self, reserve_mib: int | None) -> dict:
|
||||
"""#533, POST /v1/vram: keep `reserve_mib` of VRAM free for other programs (None: the reserve the engine
|
||||
started with), applied between requests - a request that is running finishes first (up to wait_s). Only an
|
||||
engine started with --vram-elastic (the config's "vram_elastic": true) can do it; it never resizes on its own.
|
||||
An unloaded engine applies it when it loads. -> {"status": ..., and the engine's figures}."""
|
||||
if not hasattr(self.engine, "vram"):
|
||||
raise ValueError("this engine cannot resize its VRAM use")
|
||||
if not self.fifo.acquire(timeout=self.vram_wait_s):
|
||||
raise ModelBusy("a request is still running; try again when it has finished")
|
||||
try:
|
||||
if not self.loaded():
|
||||
self.vram_reserve = reserve_mib
|
||||
return {"status": "not loaded", "reserve_mib": reserve_mib,
|
||||
"note": "applied when the model loads"}
|
||||
out = self.engine.vram(reserve_mib)
|
||||
self.vram_reserve = reserve_mib
|
||||
print(f"[strata] VRAM: {out.get('vram_free_mib')} MiB free, expert cache {out.get('expert_cache_mib')} of "
|
||||
f"{out.get('expert_cache_full_mib')} MiB ({out.get('expert_slots')} experts)", flush=True)
|
||||
return {"status": "ok", **out}
|
||||
finally:
|
||||
self.fifo.release()
|
||||
|
||||
def _say_died(self, e: Exception) -> None:
|
||||
"""The server window's line for an engine that died (or was ended, #481) in the middle of a request."""
|
||||
@@ -1326,6 +1394,8 @@ class Service:
|
||||
"activity": {"requests": totals["requests"] + int(busy), "in_flight": int(busy) + int(s.get("queued") or 0),
|
||||
"last_request_at": int(last_at) if last_at else None},
|
||||
"last_timings": last_t,
|
||||
"vram": dict((getattr(self.engine, "info", {}) or {}).get("vram") or {},
|
||||
elastic=bool((getattr(self.engine, "info", {}) or {}).get("vram_elastic"))),
|
||||
"machine": {
|
||||
"at": int(time.time()),
|
||||
"gpu": {"name": static.get("gpu_name"), "used_mib": scaled(hw.get("gpu_mem_used"), 2 ** 20),
|
||||
@@ -2272,6 +2342,17 @@ def make_handler(svc: Service):
|
||||
result = "loaded"
|
||||
self._json(200, {"status": result, **svc.v1_status()})
|
||||
return
|
||||
if path == "/v1/vram": # #533: give VRAM back to other programs, or take it
|
||||
if not self._own_page("the VRAM reserve can be changed"):
|
||||
return
|
||||
r = req.get("reserve_mib")
|
||||
if r is not None and (isinstance(r, bool) or not isinstance(r, int) or r < 0):
|
||||
raise ValueError("reserve_mib: a whole number of MiB (0 or more), or null for the start's")
|
||||
try:
|
||||
self._json(200, svc.vram(r))
|
||||
except EngineDied as e:
|
||||
self._json(503, {"error": {"type": "server_error", "message": str(e)}})
|
||||
return
|
||||
if path in ("/v1/chat/completions", "/v1/messages"):
|
||||
self.record = svc.begin_request(path, req)
|
||||
if path == "/v1/chat/completions":
|
||||
|
||||
@@ -0,0 +1,136 @@
|
||||
"""#533: POST /v1/vram (opt-in hot VRAM resize) - the engine command, the route, the config key, a reserve kept for
|
||||
an unloaded engine and applied when it loads again."""
|
||||
import io
|
||||
import json
|
||||
import queue
|
||||
import unittest
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
from serve.frontend import ChatTemplate
|
||||
from serve.server import ByteTokenizer, EngineDied, Service, StrataEngine, engine_args, serve
|
||||
from serve.test_lifecycle import ResidentEngine
|
||||
|
||||
|
||||
class VramEngine(ResidentEngine):
|
||||
"""The resident test engine with an elastic cache: VRAM n frees what the reserve asks for, None takes the
|
||||
start's (700 MiB)."""
|
||||
def __init__(self, tok, elastic=True):
|
||||
super().__init__(tok)
|
||||
self.elastic, self.calls = elastic, []
|
||||
|
||||
def vram(self, reserve_mib, timeout=120.0):
|
||||
self.calls.append(reserve_mib)
|
||||
if not self.elastic:
|
||||
raise ValueError("VRAM needs an engine started with --vram-elastic (one NVIDIA GPU, --serve)")
|
||||
r = 700 if reserve_mib is None else reserve_mib
|
||||
out = {"reserve_mib": r, "expert_slots": max(100, 4000 - r), "expert_cache_mib": max(500, 6000 - r),
|
||||
"expert_cache_full_mib": 6000, "vram_free_mib": r, "prompt_chunk": 4096}
|
||||
self.info["vram"] = out
|
||||
return out
|
||||
|
||||
|
||||
class EngineCommand(unittest.TestCase):
|
||||
def engine(self, replies):
|
||||
e = StrataEngine("missing-executable", [], lazy=True)
|
||||
e.proc = mock.Mock()
|
||||
e.proc.poll.return_value = None
|
||||
e.proc.stdin = io.StringIO()
|
||||
e.ended = False
|
||||
e.lines = queue.Queue()
|
||||
for r in replies:
|
||||
e.lines.put(r)
|
||||
return e
|
||||
|
||||
def test_the_line_and_the_answer(self):
|
||||
e = self.engine(["strata serve: something else\n",
|
||||
"VRAM reserve_mib=6000 expert_slots=1200 expert_slots_full=3700 expert_cache_mib=1536 "
|
||||
"expert_cache_full_mib=4900 vram_free_mib=6100 prompt_chunk=2048\n"])
|
||||
out = e.vram(6000)
|
||||
self.assertEqual(e.proc.stdin.getvalue(), "VRAM 6000\n")
|
||||
self.assertEqual(out["expert_slots"], 1200)
|
||||
self.assertEqual(out["prompt_chunk"], 2048)
|
||||
self.assertEqual(e.info["expert_slots"], 1200)
|
||||
e = self.engine(["VRAM reserve_mib=700 expert_slots=3700\n"])
|
||||
e.vram(None)
|
||||
self.assertEqual(e.proc.stdin.getvalue(), "VRAM\n")
|
||||
|
||||
def test_refused_and_died(self):
|
||||
e = self.engine(["ERR VRAM needs an engine started with --vram-elastic\n"])
|
||||
with self.assertRaisesRegex(ValueError, "--vram-elastic"):
|
||||
e.vram(100)
|
||||
e = self.engine([None])
|
||||
with self.assertRaises(EngineDied):
|
||||
e.vram(100)
|
||||
|
||||
def test_config_key(self):
|
||||
cfg = {"args": ["--max-context", "8192"]}
|
||||
self.assertEqual(engine_args(cfg), ["--max-context", "8192"]) # off unless asked
|
||||
self.assertEqual(engine_args({**cfg, "vram_elastic": True, "vram_segment_mib": 256}),
|
||||
["--max-context", "8192", "--vram-elastic", "--vram-segment-mib", "256"])
|
||||
self.assertEqual(engine_args({**cfg, "vram_elastic": "yes"}), ["--max-context", "8192"])
|
||||
|
||||
|
||||
class Route(unittest.TestCase):
|
||||
def setUp(self):
|
||||
tok = ByteTokenizer()
|
||||
self.engine = VramEngine(tok)
|
||||
self.svc = Service(self.engine, tok, ChatTemplate(Path(__file__).parent / "chat_template.jinja"))
|
||||
self.httpd = serve(self.svc, port=0)
|
||||
self.base = f"http://127.0.0.1:{self.httpd.server_address[1]}"
|
||||
|
||||
def tearDown(self):
|
||||
self.httpd.shutdown()
|
||||
self.httpd.server_close()
|
||||
|
||||
def request(self, path, body=None, headers=None):
|
||||
req = urllib.request.Request(self.base + path, data=json.dumps(body).encode() if body is not None else None,
|
||||
headers={"Content-Type": "application/json", **(headers or {})})
|
||||
try:
|
||||
response = urllib.request.urlopen(req, timeout=10)
|
||||
except urllib.error.HTTPError as error:
|
||||
response = error
|
||||
with response:
|
||||
return response.status, json.loads(response.read().decode())
|
||||
|
||||
def test_resize_and_status(self):
|
||||
self.engine.loaded = True
|
||||
code, body = self.request("/v1/vram", {"reserve_mib": 4000})
|
||||
self.assertEqual(code, 200, body)
|
||||
self.assertEqual(body["status"], "ok")
|
||||
self.assertEqual(body["expert_cache_mib"], 2000)
|
||||
self.assertEqual(self.request("/v1/status")[1]["vram"]["expert_cache_mib"], 2000)
|
||||
code, body = self.request("/v1/vram", {"reserve_mib": None}) # back to the start's reserve
|
||||
self.assertEqual(code, 200)
|
||||
self.assertEqual(self.engine.calls, [4000, None])
|
||||
|
||||
def test_bad_input_busy_and_foreign_pages(self):
|
||||
self.engine.loaded = True
|
||||
for bad in ({"reserve_mib": -1}, {"reserve_mib": "8 GiB"}, {"reserve_mib": True}, {"reserve_mib": 1.5}):
|
||||
self.assertEqual(self.request("/v1/vram", bad)[0], 400, bad)
|
||||
self.assertEqual(self.request("/v1/vram", {"reserve_mib": 1}, {"Origin": "https://other.example"})[0], 403)
|
||||
self.assertEqual(self.request("/v1/vram", {"reserve_mib": 1}, {"Content-Type": "text/plain"})[0], 415)
|
||||
self.svc.vram_wait_s = 0.2
|
||||
with self.svc.fifo: # a request is running
|
||||
self.assertEqual(self.request("/v1/vram", {"reserve_mib": 1})[0], 409)
|
||||
self.assertEqual(self.engine.calls, [])
|
||||
|
||||
def test_an_engine_without_the_flag_says_so(self):
|
||||
self.engine.loaded, self.engine.elastic = True, False
|
||||
code, body = self.request("/v1/vram", {"reserve_mib": 3000})
|
||||
self.assertEqual(code, 400)
|
||||
self.assertIn("--vram-elastic", body["error"]["message"])
|
||||
|
||||
def test_unloaded_keeps_the_reserve_for_the_next_load(self):
|
||||
code, body = self.request("/v1/vram", {"reserve_mib": 5000})
|
||||
self.assertEqual(code, 200)
|
||||
self.assertEqual(body["status"], "not loaded")
|
||||
self.assertEqual(self.engine.calls, [])
|
||||
self.assertEqual(self.request("/v1/load", {})[0], 200)
|
||||
self.assertEqual(self.engine.calls, [5000])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+245
-2
@@ -2,6 +2,9 @@
|
||||
#include "strata/core/expert_cache.hpp"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#if !defined(STRATA_USE_HIP)
|
||||
#include <cuda.h> // #533: the virtual memory management types (the functions come through the runtime's entry points)
|
||||
#endif
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdio>
|
||||
@@ -142,6 +145,239 @@ bool write_expert_profile(const std::string& path, int64_t n_layers, int64_t n_e
|
||||
|
||||
ExpertCache::~ExpertCache() { close(); }
|
||||
|
||||
// ---- #533: the segmented arena (--vram-elastic). The driver API's virtual memory functions, looked up through the
|
||||
// runtime (no link against the driver library): one address range for the whole arena, backed by physical segments,
|
||||
// and the tail's segments unmapped / mapped again later. Nothing here runs unless a segment size was set.
|
||||
#if !defined(STRATA_USE_HIP)
|
||||
namespace {
|
||||
struct Vmm {
|
||||
CUresult (CUDAAPI* device_get)(CUdevice*, int) = nullptr;
|
||||
CUresult (CUDAAPI* attribute)(int*, CUdevice_attribute, CUdevice) = nullptr;
|
||||
CUresult (CUDAAPI* granularity)(size_t*, const CUmemAllocationProp*, CUmemAllocationGranularity_flags) = nullptr;
|
||||
CUresult (CUDAAPI* reserve)(CUdeviceptr*, size_t, size_t, CUdeviceptr, unsigned long long) = nullptr;
|
||||
CUresult (CUDAAPI* address_free)(CUdeviceptr, size_t) = nullptr;
|
||||
CUresult (CUDAAPI* create)(CUmemGenericAllocationHandle*, size_t, const CUmemAllocationProp*,
|
||||
unsigned long long) = nullptr;
|
||||
CUresult (CUDAAPI* release)(CUmemGenericAllocationHandle) = nullptr;
|
||||
CUresult (CUDAAPI* map)(CUdeviceptr, size_t, size_t, CUmemGenericAllocationHandle, unsigned long long) = nullptr;
|
||||
CUresult (CUDAAPI* unmap)(CUdeviceptr, size_t) = nullptr;
|
||||
CUresult (CUDAAPI* set_access)(CUdeviceptr, size_t, const CUmemAccessDesc*, size_t) = nullptr;
|
||||
bool ok = false;
|
||||
};
|
||||
|
||||
template <class F> bool entry(const char* name, F& f) {
|
||||
void* p = nullptr;
|
||||
cudaDriverEntryPointQueryResult q{};
|
||||
#if CUDART_VERSION >= 12050
|
||||
const cudaError_t e = cudaGetDriverEntryPointByVersion(name, &p, 12000, cudaEnableDefault, &q);
|
||||
#else
|
||||
const cudaError_t e = cudaGetDriverEntryPoint(name, &p, cudaEnableDefault, &q);
|
||||
#endif
|
||||
if (e != cudaSuccess || q != cudaDriverEntryPointSuccess || p == nullptr) {
|
||||
(void) cudaGetLastError();
|
||||
return false;
|
||||
}
|
||||
f = reinterpret_cast<F>(p);
|
||||
return true;
|
||||
}
|
||||
|
||||
const Vmm& vmm() {
|
||||
static const Vmm v = [] {
|
||||
Vmm x;
|
||||
x.ok = entry("cuDeviceGet", x.device_get) && entry("cuDeviceGetAttribute", x.attribute) &&
|
||||
entry("cuMemGetAllocationGranularity", x.granularity) && entry("cuMemAddressReserve", x.reserve) &&
|
||||
entry("cuMemAddressFree", x.address_free) && entry("cuMemCreate", x.create) &&
|
||||
entry("cuMemRelease", x.release) && entry("cuMemMap", x.map) && entry("cuMemUnmap", x.unmap) &&
|
||||
entry("cuMemSetAccess", x.set_access);
|
||||
return x;
|
||||
}();
|
||||
return v;
|
||||
}
|
||||
|
||||
CUmemAllocationProp device_prop(int dev) {
|
||||
CUmemAllocationProp prop{};
|
||||
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
|
||||
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
|
||||
prop.location.id = dev;
|
||||
return prop;
|
||||
}
|
||||
|
||||
// one segment: a physical allocation mapped at `va`, readable and writable by this device
|
||||
bool map_segment(const Vmm& v, int dev, CUdeviceptr va, size_t bytes, unsigned long long& handle) {
|
||||
const CUmemAllocationProp prop = device_prop(dev);
|
||||
CUmemGenericAllocationHandle h = 0;
|
||||
if (v.create(&h, bytes, &prop, 0) != CUDA_SUCCESS) return false;
|
||||
if (v.map(va, bytes, 0, h, 0) != CUDA_SUCCESS) {
|
||||
v.release(h);
|
||||
return false;
|
||||
}
|
||||
CUmemAccessDesc access{};
|
||||
access.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
|
||||
access.location.id = dev;
|
||||
access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
|
||||
if (v.set_access(va, bytes, &access, 1) != CUDA_SUCCESS) {
|
||||
v.unmap(va, bytes);
|
||||
v.release(h);
|
||||
return false;
|
||||
}
|
||||
handle = (unsigned long long) h;
|
||||
return true;
|
||||
}
|
||||
} // namespace
|
||||
#endif
|
||||
|
||||
bool ExpertCache::open_segmented(uint64_t want, std::string& err) {
|
||||
#if defined(STRATA_USE_HIP)
|
||||
(void) want;
|
||||
err = "ExpertCache: --vram-elastic (a segmented expert cache) is CUDA-only for now";
|
||||
return false;
|
||||
#else
|
||||
const Vmm& v = vmm();
|
||||
int dev = 0, supported = 0;
|
||||
CUdevice cu = 0;
|
||||
if (!v.ok || cudaGetDevice(&dev) != cudaSuccess || v.device_get(&cu, dev) != CUDA_SUCCESS ||
|
||||
v.attribute(&supported, CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED, cu) != CUDA_SUCCESS ||
|
||||
!supported) {
|
||||
err = "ExpertCache: --vram-elastic needs the driver's virtual memory management, which this GPU or driver "
|
||||
"does not offer";
|
||||
return false;
|
||||
}
|
||||
const CUmemAllocationProp prop = device_prop(dev);
|
||||
size_t gran = 0;
|
||||
if (v.granularity(&gran, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED) != CUDA_SUCCESS || gran == 0) {
|
||||
err = "ExpertCache: cannot read the driver's allocation granularity";
|
||||
return false;
|
||||
}
|
||||
const uint64_t g = (uint64_t) gran;
|
||||
const uint64_t total = (want + g - 1) / g * g;
|
||||
seg_ = (int64_t) (((uint64_t) seg_req_ + g - 1) / g * g);
|
||||
CUdeviceptr va = 0;
|
||||
if (v.reserve(&va, (size_t) total, 0, 0, 0) != CUDA_SUCCESS) {
|
||||
err = "ExpertCache: cannot reserve the address range of the segmented expert cache";
|
||||
return false;
|
||||
}
|
||||
base_ = reinterpret_cast<uint8_t*>(va);
|
||||
reserved_ = total;
|
||||
for (uint64_t at = 0; at < total; at += (uint64_t) seg_) {
|
||||
segs_.push_back(0);
|
||||
seg_size_.push_back((int64_t) std::min<uint64_t>((uint64_t) seg_, total - at));
|
||||
}
|
||||
for (size_t i = 0; i < segs_.size(); ++i) {
|
||||
if (!map_segment(v, dev, va + (CUdeviceptr) ((uint64_t) i * (uint64_t) seg_), (size_t) seg_size_[i],
|
||||
segs_[i])) {
|
||||
char buf[200];
|
||||
std::snprintf(buf, sizeof buf, "ExpertCache: cudaMalloc failed: segment %zu of %zu (%.2f GiB) of the "
|
||||
"segmented cache could not be allocated", i + 1, segs_.size(),
|
||||
(double) seg_size_[i] / 1073741824.0);
|
||||
err = buf; // "cudaMalloc failed": the auto cache's smaller-retry path reads it as an allocation failure
|
||||
release_segmented();
|
||||
return false;
|
||||
}
|
||||
mapped_segs_ = (int64_t) i + 1;
|
||||
}
|
||||
return true;
|
||||
#endif
|
||||
}
|
||||
|
||||
void ExpertCache::release_segmented() {
|
||||
#if !defined(STRATA_USE_HIP)
|
||||
const Vmm& v = vmm();
|
||||
if (base_ != nullptr) cudaDeviceSynchronize();
|
||||
const CUdeviceptr va = reinterpret_cast<CUdeviceptr>(base_);
|
||||
for (size_t i = 0; i < segs_.size(); ++i)
|
||||
if (segs_[i] != 0) {
|
||||
v.unmap(va + (CUdeviceptr) ((uint64_t) i * (uint64_t) seg_), (size_t) seg_size_[i]);
|
||||
v.release((CUmemGenericAllocationHandle) segs_[i]);
|
||||
}
|
||||
if (base_ != nullptr && reserved_ > 0) v.address_free(va, (size_t) reserved_);
|
||||
#endif
|
||||
segs_.clear();
|
||||
seg_size_.clear();
|
||||
mapped_segs_ = 0;
|
||||
reserved_ = 0;
|
||||
base_ = nullptr;
|
||||
}
|
||||
|
||||
int64_t ExpertCache::mapped_bytes() const {
|
||||
if (segs_.empty()) return base_ != nullptr ? full_bytes() : 0;
|
||||
int64_t b = 0;
|
||||
for (int64_t i = 0; i < mapped_segs_; ++i) b += seg_size_[(size_t) i];
|
||||
return b;
|
||||
}
|
||||
|
||||
int64_t ExpertCache::slots_within(int64_t bytes) const {
|
||||
if (bytes >= full_bytes()) return slots_;
|
||||
if (bytes <= 0) return 0;
|
||||
if (off_.empty()) return blob_ > 0 ? bytes / blob_ : 0;
|
||||
// off_[i + 1] is slot i's end: count the slots whose end is at most `bytes`
|
||||
return (int64_t) (std::upper_bound(off_.begin() + 1, off_.end(), (uint64_t) bytes) - (off_.begin() + 1));
|
||||
}
|
||||
|
||||
bool ExpertCache::shrink(int64_t keep_bytes, std::string& err) {
|
||||
if (segs_.empty()) {
|
||||
err = "the expert cache is not segmented (the engine needs --vram-elastic)";
|
||||
return false;
|
||||
}
|
||||
#if defined(STRATA_USE_HIP)
|
||||
(void) keep_bytes;
|
||||
return false;
|
||||
#else
|
||||
const Vmm& v = vmm();
|
||||
int64_t keep = 0, at = 0; // the segments [0, keep) hold the first keep_bytes
|
||||
while (keep < (int64_t) segs_.size() && at < keep_bytes) at += seg_size_[(size_t) keep++];
|
||||
if (keep >= mapped_segs_) return true;
|
||||
if (cudaDeviceSynchronize() != cudaSuccess) {
|
||||
err = std::string("the device failed before the cache shrank: ") + cudaGetErrorString(cudaGetLastError());
|
||||
return false;
|
||||
}
|
||||
const CUdeviceptr va = reinterpret_cast<CUdeviceptr>(base_);
|
||||
for (int64_t i = mapped_segs_ - 1; i >= keep; --i) {
|
||||
const CUdeviceptr p = va + (CUdeviceptr) ((uint64_t) i * (uint64_t) seg_);
|
||||
if (v.unmap(p, (size_t) seg_size_[(size_t) i]) != CUDA_SUCCESS ||
|
||||
v.release((CUmemGenericAllocationHandle) segs_[(size_t) i]) != CUDA_SUCCESS) {
|
||||
err = "the driver refused to release an expert-cache segment";
|
||||
live_slots_ = slots_within(mapped_bytes());
|
||||
return false;
|
||||
}
|
||||
segs_[(size_t) i] = 0;
|
||||
mapped_segs_ = i;
|
||||
}
|
||||
live_slots_ = slots_within(mapped_bytes());
|
||||
return true;
|
||||
#endif
|
||||
}
|
||||
|
||||
bool ExpertCache::grow(int64_t want_bytes, std::string& err) {
|
||||
if (segs_.empty()) {
|
||||
err = "the expert cache is not segmented (the engine needs --vram-elastic)";
|
||||
return false;
|
||||
}
|
||||
#if defined(STRATA_USE_HIP)
|
||||
(void) want_bytes;
|
||||
return false;
|
||||
#else
|
||||
const Vmm& v = vmm();
|
||||
int dev = 0;
|
||||
cudaGetDevice(&dev);
|
||||
const CUdeviceptr va = reinterpret_cast<CUdeviceptr>(base_);
|
||||
int64_t at = mapped_bytes();
|
||||
while (mapped_segs_ < (int64_t) segs_.size() && at + seg_size_[(size_t) mapped_segs_] <= want_bytes) {
|
||||
const size_t i = (size_t) mapped_segs_;
|
||||
if (!map_segment(v, dev, va + (CUdeviceptr) ((uint64_t) i * (uint64_t) seg_), (size_t) seg_size_[i],
|
||||
segs_[i])) {
|
||||
(void) cudaGetLastError();
|
||||
err = "the driver has no VRAM for another expert-cache segment";
|
||||
live_slots_ = slots_within(mapped_bytes());
|
||||
return false;
|
||||
}
|
||||
at += seg_size_[i];
|
||||
++mapped_segs_;
|
||||
}
|
||||
live_slots_ = slots_within(mapped_bytes());
|
||||
return true;
|
||||
#endif
|
||||
}
|
||||
|
||||
#if defined(STRATA_USE_HIP)
|
||||
bool ExpertCache::ensure_blocking_staging(std::size_t bytes, std::string& err) {
|
||||
if (bytes <= blocking_staging_bytes_) return true;
|
||||
@@ -193,7 +429,9 @@ bool ExpertCache::open(int64_t n_slots, int64_t n_layers, int64_t n_expert, int6
|
||||
}
|
||||
}
|
||||
|
||||
if (cudaMalloc((void**) &base_, (size_t) want) != cudaSuccess) {
|
||||
if (seg_req_ > 0) { // #533: --vram-elastic: physical segments behind one address range (zeroed below)
|
||||
if (!open_segmented(want, err)) return false;
|
||||
} else if (cudaMalloc((void**) &base_, (size_t) want) != cudaSuccess) {
|
||||
base_ = nullptr;
|
||||
char buf[256];
|
||||
std::snprintf(buf, sizeof buf, "ExpertCache: cudaMalloc(%.2f GiB) failed: %s",
|
||||
@@ -211,6 +449,7 @@ bool ExpertCache::open(int64_t n_slots, int64_t n_layers, int64_t n_expert, int6
|
||||
|
||||
residency_.assign((size_t) (n_layers * n_expert), kNotResident);
|
||||
slots_ = n_slots;
|
||||
live_slots_ = n_slots;
|
||||
n_layers_ = n_layers;
|
||||
n_expert_ = n_expert;
|
||||
blob_ = blob_bytes;
|
||||
@@ -247,6 +486,7 @@ bool ExpertCache::open_sized(const std::vector<int64_t>& slot_bytes, int64_t n_l
|
||||
// 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();
|
||||
live_slots_ = slots_;
|
||||
blob_ = mx;
|
||||
off_ = std::move(off);
|
||||
#if defined(STRATA_USE_HIP)
|
||||
@@ -273,12 +513,15 @@ void ExpertCache::close() {
|
||||
blocking_staging_bytes_ = 0;
|
||||
#endif
|
||||
off_.clear();
|
||||
if (base_ != nullptr) {
|
||||
if (!segs_.empty()) {
|
||||
release_segmented();
|
||||
} else if (base_ != nullptr) {
|
||||
cudaFree(base_);
|
||||
base_ = nullptr;
|
||||
}
|
||||
residency_.clear();
|
||||
slots_ = 0;
|
||||
live_slots_ = 0;
|
||||
n_layers_ = 0;
|
||||
n_expert_ = 0;
|
||||
blob_ = 0;
|
||||
|
||||
+154
-2
@@ -340,6 +340,10 @@ struct Options {
|
||||
/// and the KV state, minus this reserve for the graphs, the hit scratch and the head.
|
||||
int vram_reserve_mib = 700;
|
||||
bool vram_reserve_given = false; ///< --vram-reserve-mib on the command line (#496: no smaller automatic reserve)
|
||||
/// #533 (opt-in): the expert cache in physical segments (CUDA virtual memory management), so the serve loop's
|
||||
/// `VRAM <reserve_mib>` command can give part of it back to other programs and take it again. Off: one cudaMalloc.
|
||||
bool vram_elastic = false;
|
||||
int64_t vram_segment_mib = 512;
|
||||
/// Plan v0.3 P5: batched prompt processing in chunks of this many tokens (0 = the token path).
|
||||
int64_t prefill_chunk = 0;
|
||||
/// `--prefill auto`: the largest chunk (up to 8192) whose buffers the expert cache can lend. Every expert a chunk
|
||||
@@ -515,6 +519,9 @@ void usage() {
|
||||
" --conversation-cache-mib N --serve: RAM budget for parked conversations (default 0 = off)\n"
|
||||
" --conversation-cache-slots N --serve: at most N parked conversations (default 4)\n"
|
||||
" --conversation-cache-min-free-mib N --serve: physical RAM floor when parking (default 2560)\n"
|
||||
" --vram-elastic --serve (#533, opt-in, NVIDIA): the expert cache in segments (--vram-segment-mib,\n"
|
||||
" default 512), so the command `VRAM <reserve_mib>` (the server's POST /v1/vram) can\n"
|
||||
" give VRAM back to other programs between requests and take it back later\n"
|
||||
" --prompt-cache-every N --serve: also checkpoint every N fresh prompt tokens (default 16384, 0 = off)\n"
|
||||
" --turn-token ID --serve: the token that opens a chat turn (default 248045, <|im_start|>)\n"
|
||||
" --short-read N --serve: read at most N fresh text tokens through the decode windows instead\n"
|
||||
@@ -1211,6 +1218,8 @@ int main(int argc, char** argv) {
|
||||
}
|
||||
else if (a == "--no-split-rows") o.no_split_rows = true;
|
||||
else if (a == "--no-prefill-borrow") o.no_prefill_borrow = true;
|
||||
else if (a == "--vram-elastic") o.vram_elastic = true;
|
||||
else if (a == "--vram-segment-mib") o.vram_segment_mib = std::atoll(next("--vram-segment-mib"));
|
||||
else if (a == "--prefill-until") o.prefill_until = std::atoll(next("--prefill-until"));
|
||||
else if (a == "--dump-final-r") o.dump_final_r = next("--dump-final-r");
|
||||
else if (a == "--spec") o.spec = std::atoi(next("--spec"));
|
||||
@@ -2952,6 +2961,24 @@ int main(int argc, char** argv) {
|
||||
}
|
||||
o.expert_cache = (int) sized_slots.size();
|
||||
}
|
||||
// #533: --vram-elastic: the cache in physical segments (the VRAM command resizes it between requests). One GPU,
|
||||
// no helper caches, serve mode: anything else keeps the one cudaMalloc, said once.
|
||||
if (o.vram_elastic) {
|
||||
const char* why = !o.serve ? "it works between requests of --serve"
|
||||
: multi_gpu ? "a layer split has a cache per GPU"
|
||||
: std::any_of(o.expert_cache_remote.begin(), o.expert_cache_remote.end(),
|
||||
[](int n) { return n > 0; }) ? "the helper caches on other GPUs"
|
||||
: o.vram_segment_mib < 64 ? "--vram-segment-mib is below 64" : nullptr;
|
||||
#if defined(STRATA_USE_HIP)
|
||||
if (why == nullptr) why = "it is NVIDIA-only for now";
|
||||
#endif
|
||||
if (why != nullptr) {
|
||||
std::fprintf(stderr, "strata generate: --vram-elastic is off: %s\n", why);
|
||||
o.vram_elastic = false;
|
||||
} else {
|
||||
xcache.set_segment_bytes(o.vram_segment_mib << 20);
|
||||
}
|
||||
}
|
||||
if (o.expert_cache > 0) {
|
||||
// keep the first `keep_bytes` of the cache (the profile's hottest experts first); false when nothing is left
|
||||
auto shrink_to = [&](int64_t keep_bytes) -> bool {
|
||||
@@ -4986,7 +5013,7 @@ int main(int argc, char** argv) {
|
||||
"expert_slots_primary=%lld expert_cache_primary_mib=%lld spec=%d "
|
||||
"mtp_max=%d lookup=%d vram_free_mib=%lld cvec=%s arena_mib=%lld pool_workers=%d pcie_frac=%.2f "
|
||||
"spec_min_p=%.2f conversation_cache_mib=%lld conversation_cache_slots=%d "
|
||||
"conversation_cache_min_free_mib=%lld engine=" STRATA_VERSION "\n",
|
||||
"conversation_cache_min_free_mib=%lld vram_elastic=%d engine=" STRATA_VERSION "\n",
|
||||
(long long) o.max_context, o.kv.c_str(),
|
||||
(long long) (g.n_qsa_layers() > 0 && ss.qsa_states[ss.qsa_primary()].kv_mode == 1
|
||||
? ss.qsa_states[ss.qsa_primary()].n_slots * 4 : 0),
|
||||
@@ -4997,7 +5024,7 @@ int main(int argc, char** argv) {
|
||||
(long long) ((o.mmap_experts ? src.resident_bytes() : strata::kernels::cpu::expert_layout().total) >> 20),
|
||||
pool.workers(), o.pcie_frac,
|
||||
o.spec_min_p, (long long) o.conversation_cache_mib, o.conversation_cache_slots,
|
||||
(long long) o.conversation_cache_min_free_mib);
|
||||
(long long) o.conversation_cache_min_free_mib, xcache.segmented() ? 1 : 0);
|
||||
}
|
||||
// issue #29: a request whose heartbeat (tokens, prompt chunks, verify windows) stops for this long is stuck on
|
||||
// a flag nobody will raise - end the engine with where it was, so the server starts it again instead of the
|
||||
@@ -5044,7 +5071,132 @@ int main(int argc, char** argv) {
|
||||
bool mrope_identity = true;
|
||||
std::vector<float> img_rows;
|
||||
std::vector<const float*> row_ptr;
|
||||
// ---- #533: VRAM <reserve_mib>, between requests, only with --vram-elastic. It shrinks the expert cache
|
||||
// until that much VRAM is free for other programs (a game, a CAD session), or grows it back towards its full
|
||||
// size when more than that is free. Never on its own: only this command (the server's POST /v1/vram) moves
|
||||
// it. A shrink gives back the cache's LAST segments: the experts in their slots become CPU misses, as any
|
||||
// expert outside the cache is (no output changes beyond what a smaller cache decodes), the prompt path's
|
||||
// loan moves down to the end of what is left (a smaller chunk when it no longer fits). A grow maps them
|
||||
// again and puts each slot's expert back (or, when the adaptive tier already brought that one back, the
|
||||
// most-routed missing expert of the same layer); the adaptive tier carries on from there.
|
||||
std::vector<std::pair<int32_t, int32_t>> vram_evicted; // (residency index, slot) the shrinks took
|
||||
const int64_t chunk_full = o.prefill_chunk; // the loan's chunk with the whole cache
|
||||
auto relend = [&]() { // the prompt path's loan: the end of the slots left, its chunk as large as fits
|
||||
if (pf_parts.empty() || pf_parts[0].first < 0) return;
|
||||
PfPart& p = pf_parts[0];
|
||||
const int64_t live = xcache.slots();
|
||||
auto fits = [&](int64_t c) {
|
||||
const int64_t k = part_slots(p, c);
|
||||
return k > 0 && k + 128 <= live && !(o.prefill_auto && k * 100 > kAutoLendPct * live);
|
||||
};
|
||||
int64_t c = chunk_full;
|
||||
while (c > 256 && !fits(c)) c = std::max<int64_t>(256, c / 2 / 256 * 256);
|
||||
o.prefill_chunk = c;
|
||||
p.first = (int32_t) std::max<int64_t>(0, live - part_slots(p, c));
|
||||
p.first_now = -1; // laid out again at the next loan
|
||||
};
|
||||
auto vram_command = [&](const std::string& cmd, std::string& e) -> bool {
|
||||
// `VRAM` alone: the reserve the engine started with (--vram-reserve-mib)
|
||||
char* end = nullptr;
|
||||
const bool bare = cmd.find_first_not_of(' ', 4) == std::string::npos;
|
||||
const long long reserve = bare ? (long long) o.vram_reserve_mib : std::strtoll(cmd.c_str() + 4, &end, 10);
|
||||
if (!bare && (end == cmd.c_str() + 4 || reserve < 0)) { e = "expected: VRAM [reserve_mib]"; return false; }
|
||||
if (!xcache.segmented()) {
|
||||
e = "VRAM needs an engine started with --vram-elastic (one NVIDIA GPU, --serve)";
|
||||
return false;
|
||||
}
|
||||
if (src.complement_ready()) {
|
||||
e = "VRAM: not with the resident low-RAM mode (the cache holds experts RAM does not)";
|
||||
return false;
|
||||
}
|
||||
if (peer.valid()) { e = "VRAM: not with --peer-device"; return false; }
|
||||
for (const PfPart& p : pf_parts)
|
||||
if (!p.lent.empty()) { e = "VRAM: a prompt's loan is still out"; return false; }
|
||||
if (!ver.wait_commit(e)) return false;
|
||||
apply_pending(true); // the adaptive tier's swaps in flight land first
|
||||
if (cudaDeviceSynchronize() != cudaSuccess) {
|
||||
e = std::string("VRAM: ") + cudaGetErrorString(cudaGetLastError());
|
||||
return false;
|
||||
}
|
||||
const auto t0 = Clock::now();
|
||||
size_t free_b = 0, total_b = 0;
|
||||
cudaMemGetInfo(&free_b, &total_b);
|
||||
const int64_t want_free = (int64_t) reserve << 20, mapped = xcache.mapped_bytes();
|
||||
const int64_t before = xcache.slots();
|
||||
// what stays at least: the prompt path's smallest loan (256 tokens) and the 128 slots it must leave
|
||||
const int64_t floor_slots = std::min<int64_t>(xcache.full_slots(),
|
||||
128 + (pf_parts.empty() || pf_parts[0].first < 0 ? 0 : part_slots(pf_parts[0], 256)));
|
||||
std::string note;
|
||||
if ((int64_t) free_b < want_free) {
|
||||
const int64_t floor_b = xcache.bytes_of(floor_slots);
|
||||
int64_t keep = mapped - (want_free - (int64_t) free_b);
|
||||
if (keep < floor_b) {
|
||||
keep = floor_b;
|
||||
note = " (the cache keeps its smallest size: the prompt path's buffers)";
|
||||
}
|
||||
if (!xcache.shrink(keep, e)) return false;
|
||||
const int64_t live = xcache.slots();
|
||||
for (size_t i = 0; i < host_res.size(); ++i)
|
||||
if (host_res[i] >= live) {
|
||||
vram_evicted.emplace_back((int32_t) i, host_res[i]);
|
||||
host_res[i] = strata::core::kNotResident;
|
||||
}
|
||||
res_upload();
|
||||
} else if (mapped < xcache.full_bytes()) {
|
||||
std::string gerr;
|
||||
if (!xcache.grow(mapped + ((int64_t) free_b - want_free), gerr)) note = " (" + gerr + ")";
|
||||
const int64_t live = xcache.slots();
|
||||
std::string ferr;
|
||||
std::vector<std::pair<int32_t, int32_t>> keep_out;
|
||||
for (const auto& [i, slot] : vram_evicted) {
|
||||
if (slot >= live) { keep_out.emplace_back(i, slot); continue; }
|
||||
const int64_t layer = i / g.n_expert;
|
||||
int64_t pick = i;
|
||||
if (host_res[(size_t) i] >= 0) { // back already (the adaptive tier): the layer's most-routed miss
|
||||
pick = -1;
|
||||
float best = -1.0f;
|
||||
for (int64_t ex = 0; ex < g.n_expert; ++ex) {
|
||||
const size_t j = (size_t) (layer * g.n_expert + ex);
|
||||
if (host_res[j] >= 0) continue;
|
||||
const float u = drive.d.usage.empty() ? 0.0f : drive.d.usage[j];
|
||||
if (u > best) { best = u; pick = (int64_t) j; }
|
||||
}
|
||||
}
|
||||
if (pick < 0) continue;
|
||||
const uint8_t* b = srcp->blob(layer, pick % g.n_expert);
|
||||
if (b == nullptr || !xcache.fill_slot_queued(slot, b, ferr,
|
||||
(int64_t) strata::kernels::cpu::expert_layout().blob_bytes(layer))) {
|
||||
e = "VRAM: refilling the cache failed: " + ferr;
|
||||
return false;
|
||||
}
|
||||
host_res[(size_t) pick] = slot;
|
||||
}
|
||||
vram_evicted.swap(keep_out);
|
||||
if (!xcache.sync_queued(ferr)) { e = "VRAM: " + ferr; return false; }
|
||||
res_upload();
|
||||
}
|
||||
relend();
|
||||
cudaMemGetInfo(&free_b, &total_b);
|
||||
std::fprintf(stderr, "strata serve: VRAM %lld MiB kept free: the expert cache %lld -> %lld of %lld slots "
|
||||
"(%.2f of %.2f GiB), %lld MiB free now, prompt chunk %lld, in %.0f ms%s\n",
|
||||
reserve, (long long) before, (long long) xcache.slots(), (long long) xcache.full_slots(),
|
||||
(double) xcache.mapped_bytes() / 1073741824.0, (double) xcache.full_bytes() / 1073741824.0,
|
||||
(long long) (free_b >> 20), (long long) o.prefill_chunk,
|
||||
std::chrono::duration<double, std::milli>(Clock::now() - t0).count(), note.c_str());
|
||||
std::printf("VRAM reserve_mib=%lld expert_slots=%lld expert_slots_full=%lld expert_cache_mib=%lld "
|
||||
"expert_cache_full_mib=%lld vram_free_mib=%lld prompt_chunk=%lld\n", reserve,
|
||||
(long long) xcache.slots(), (long long) xcache.full_slots(),
|
||||
(long long) (xcache.mapped_bytes() >> 20), (long long) (xcache.full_bytes() >> 20),
|
||||
(long long) (free_b >> 20), (long long) o.prefill_chunk);
|
||||
return true;
|
||||
};
|
||||
while (next_line(line)) {
|
||||
if (line.rfind("VRAM", 0) == 0) { // #533 (above): between requests, not a request
|
||||
std::string verr;
|
||||
if (!vram_command(line, verr)) std::printf("ERR %s\n", verr.c_str());
|
||||
std::fflush(stdout);
|
||||
continue;
|
||||
}
|
||||
// #477: every --expert-profile-save-every minutes, before the next request (at QUIT: after the loop)
|
||||
if (!heat.empty() && line != "QUIT" && o.expert_profile_save_min > 0 &&
|
||||
Clock::now() - profile_saved_at >= std::chrono::duration<double>(o.expert_profile_save_min * 60.0))
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
// #533: the segmented expert cache (--vram-elastic). A shrink gives the last segments' VRAM back to the driver and
|
||||
// keeps every slot left exactly as it was (same address, same bytes); a grow maps them again at the same addresses.
|
||||
// Uniform and sized slots. Needs a CUDA device with virtual memory management (exits 77 without one).
|
||||
#include "strata/core/expert_cache.hpp"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <cstdio>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
int fail(const char* what, const std::string& err = "") {
|
||||
std::fprintf(stderr, "FAIL %s %s\n", what, err.c_str());
|
||||
return 1;
|
||||
}
|
||||
|
||||
size_t free_vram() {
|
||||
size_t f = 0, t = 0;
|
||||
cudaMemGetInfo(&f, &t);
|
||||
return f;
|
||||
}
|
||||
|
||||
// slot i holds the byte (i * 7 + 1) everywhere
|
||||
bool fill(strata::core::ExpertCache& c, int64_t from, int64_t to, int64_t blob) {
|
||||
std::vector<uint8_t> b((size_t) blob);
|
||||
for (int64_t i = from; i < to; ++i) {
|
||||
std::fill(b.begin(), b.end(), (uint8_t) (i * 7 + 1));
|
||||
if (cudaMemcpy(c.device_slot((int32_t) i), b.data(), (size_t) blob, cudaMemcpyHostToDevice) != cudaSuccess)
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
bool check(const strata::core::ExpertCache& c, int64_t to, int64_t blob) {
|
||||
std::vector<uint8_t> b((size_t) blob);
|
||||
for (int64_t i = 0; i < to; ++i) {
|
||||
if (cudaMemcpy(b.data(), c.device_slot((int32_t) i), (size_t) blob, cudaMemcpyDeviceToHost) != cudaSuccess)
|
||||
return false;
|
||||
for (const uint8_t x : b)
|
||||
if (x != (uint8_t) (i * 7 + 1)) return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
int main() {
|
||||
int n = 0;
|
||||
if (cudaGetDeviceCount(&n) != cudaSuccess || n == 0) {
|
||||
std::puts("no CUDA device: skipped");
|
||||
return 77;
|
||||
}
|
||||
constexpr int64_t kBlob = 1 << 20, kSlots = 96, kSeg = 8ll << 20;
|
||||
for (int sized = 0; sized < 2; ++sized) {
|
||||
strata::core::ExpertCache c;
|
||||
c.set_segment_bytes(kSeg);
|
||||
std::string err;
|
||||
const bool ok = sized ? c.open_sized(std::vector<int64_t>(kSlots, kBlob), 4, 64, err)
|
||||
: c.open(kSlots, 4, 64, kBlob, err);
|
||||
if (!ok) {
|
||||
if (err.find("virtual memory management") != std::string::npos) {
|
||||
std::printf("%s: skipped\n", err.c_str());
|
||||
return 77;
|
||||
}
|
||||
return fail("open", err);
|
||||
}
|
||||
if (!c.segmented() || c.slots() != kSlots || c.mapped_bytes() < kSlots * kBlob) return fail("open sizes");
|
||||
uint8_t* const first = c.device_slot(0);
|
||||
if (!fill(c, 0, kSlots, kBlob)) return fail("fill");
|
||||
const size_t before = free_vram();
|
||||
// keep 30 MiB: rounded up to whole 8 MiB segments = 32 MiB = 32 slots
|
||||
if (!c.shrink(30ll << 20, err)) return fail("shrink", err);
|
||||
if (c.slots() != 32 || c.mapped_bytes() != 32ll << 20 || c.bytes() != 32ll << 20)
|
||||
return fail("shrunk sizes");
|
||||
if (c.full_slots() != kSlots || c.full_bytes() != kSlots * kBlob) return fail("full sizes");
|
||||
const size_t freed = free_vram() - before;
|
||||
if (freed + (8u << 20) < (size_t) (64ll << 20)) { // ~64 MiB back to the driver (WDDM rounds a little)
|
||||
std::fprintf(stderr, "freed %zu bytes\n", freed);
|
||||
return fail("VRAM not given back");
|
||||
}
|
||||
if (c.device_slot(0) != first || !check(c, 32, kBlob)) return fail("slots kept after the shrink");
|
||||
// grow to 70 MiB: whole segments only, 64 MiB = 64 slots
|
||||
if (!c.grow(70ll << 20, err) || c.slots() != 64) return fail("grow part", err);
|
||||
if (!c.grow(1ll << 40, err) || c.slots() != kSlots || c.mapped_bytes() < kSlots * kBlob)
|
||||
return fail("grow all", err);
|
||||
if (c.device_slot(0) != first || !check(c, 32, kBlob)) return fail("slots kept after the grow");
|
||||
if (!fill(c, 32, kSlots, kBlob) || !check(c, kSlots, kBlob)) return fail("the grown slots");
|
||||
if (c.slots_within(0) != 0 || c.slots_within(kBlob * 3 + 5) != 3 || c.bytes_of(5) != 5 * kBlob)
|
||||
return fail("slots_within / bytes_of");
|
||||
c.close();
|
||||
if (c.valid() || c.slots() != 0) return fail("close");
|
||||
}
|
||||
// the default (no segment size): one allocation, shrink refused
|
||||
strata::core::ExpertCache d;
|
||||
std::string err;
|
||||
if (!d.open(8, 2, 4, kBlob, err) || d.segmented() || d.shrink(0, err)) return fail("unsegmented", err);
|
||||
std::puts("expert cache segments: shrink gives VRAM back and keeps the slots, grow maps them again");
|
||||
return 0;
|
||||
}
|
||||
Reference in new issue
Block a user