diff --git a/app/serve.cpp b/app/serve.cpp index 2020305f..ee62a728 100644 --- a/app/serve.cpp +++ b/app/serve.cpp @@ -62,6 +62,10 @@ #include "gguf_meta.h" +#ifdef ONEBIT_MOE +#include "gguf_index.h" +#endif + #ifdef ONEBIT_NPU #include "unified.h" #endif @@ -125,7 +129,8 @@ struct Options { int prefill_min_tokens = 0; bool lean = false; std::string mtp; - int moe_slots = 0; // --moe-slots N: stream routed experts from the model file, N held in RAM + int moe_slots = 0; // --moe-slots N: stream routed experts from the model file, N held in RAM (-1: auto) + int moe_prefetch = 0; // --moe-prefetch N: gate-ahead prefetch N MoE layers ahead (0: off) std::string moe_subst; // --moe-subst R: a resident expert scoring at least R x a missing one takes its place std::string mmproj; // a vision (or audio) projector: image parts in chat messages int mtp_max = 0; @@ -533,6 +538,52 @@ struct Launch { }; // The backend process for one device: its command line and environment. +#ifdef ONEBIT_MOE +uint64_t mem_available() { + FILE* f = std::fopen("/proc/meminfo", "r"); + if (!f) return 0; + char line[256]; + unsigned long long kb = 0; + while (std::fgets(line, sizeof line, f)) + if (std::sscanf(line, "MemAvailable: %llu kB", &kb) == 1) break; + std::fclose(f); + return uint64_t(kb) << 10; +} + +// --moe-slots auto: as many experts as fit in free memory beside the tensors the GPU holds, keeping +// a reserve (8 GiB or 15% of free memory) for the KV cache, buffers and everything else. The +// streamer gives every MoE layer an equal share, so the cache costs slots x the mean expert size. +int auto_moe_slots(const std::string& model) { + const auto idx = moe::GgufIndex::open(model); + if (idx.experts.empty()) throw std::runtime_error("--moe-slots: " + model + " has no routed experts"); + double mean = 0; + for (const auto& [layer, parts] : idx.experts) mean += double(idx.expert_bytes(layer)); + mean /= double(idx.experts.size()); + uint64_t all = 0, experts = 0, mapped = 0; + for (const auto& t : idx.tensors) { + all += t.bytes; + if (t.name.find("_exps.") != std::string::npos) experts += t.bytes; + else if (t.name.rfind("per_layer_token_embd", 0) == 0) mapped += t.bytes; // stays file-backed + } + const double gib = double(1ull << 30); + const uint64_t others = all - experts - mapped, avail = mem_available(); + const uint64_t reserve = std::max(8ull << 30, avail / 100 * 15); + const int64_t total = int64_t(idx.n_expert) * int64_t(idx.experts.size()); + const int64_t floor = 64 * int64_t(idx.experts.size()); // decode batches stream only above ~8 x top-k per layer + const double room = double(avail) - double(others) - double(reserve); + int64_t n = room > 0 ? int64_t(room / mean) : 0; + n = std::min(n, total); + if (n < floor) { + std::fprintf(stderr, "1bit serve: --moe-slots auto: only %.1f GiB free for experts; using %lld slots (%.1f GiB) anyway\n", + std::max(0.0, room) / gib, (long long) floor, floor * mean / gib); + n = floor; + } + std::fprintf(stderr, "1bit serve: --moe-slots auto: %lld of %lld experts (%.1f GiB); %.1f GiB free, %.1f GiB of other tensors on the GPU\n", + (long long) n, (long long) total, n * mean / gib, avail / gib, others / gib); + return int(n); +} +#endif + Launch launch_for(const Options& o, const std::string& device, int child_port) { Launch l; l.device = device; @@ -583,13 +634,23 @@ Launch launch_for(const Options& o, const std::string& device, int child_port) { env.push_back("ONEBIT_PREFILL_DEVICE=HRX0"); if (o.prefill_min_tokens > 0) env.push_back("ONEBIT_PREFILL_MIN_TOKENS=" + std::to_string(o.prefill_min_tokens)); } - if (o.moe_slots > 0) { + if (o.moe_slots != 0) { // routed experts stay in the file (memory-mapped, never read whole); N of them are cached in // Vulkan buffers and streamed in as the router asks (docs/moe-streaming.md) if (device != "vulkan" || split || fork_arch) throw std::runtime_error("--moe-slots works with --device vulkan, without --prefill-device"); env.push_back("ONEBIT_MOE_FILE=" + std::filesystem::absolute(o.model).string()); - env.push_back("ONEBIT_MOE_SLOTS=" + std::to_string(o.moe_slots)); + int slots = o.moe_slots; + if (slots < 0) { +#ifdef ONEBIT_MOE + slots = auto_moe_slots(o.model); +#else + throw std::runtime_error("--moe-slots auto needs a Linux build (ONEBIT_MOE)"); +#endif + } + env.push_back("ONEBIT_MOE_SLOTS=" + std::to_string(slots)); + // off unless asked: where the drive is the limit, wrong guesses cost more than right ones save + env.push_back("ONEBIT_MOE_PREFETCH=" + std::to_string(o.moe_prefetch)); if (!o.moe_subst.empty()) env.push_back("ONEBIT_MOE_SUBST=" + o.moe_subst); // the experts stay file-backed: without these, llama.cpp copies every one into RAM // (a pinned Vulkan host buffer, or a CPU repack), which is what streaming avoids @@ -987,8 +1048,9 @@ void usage(FILE* out) { " [--prefill-device hrx] [--prefill-min-tokens N] (with --device vulkan)\n" " [--lean] ROCmFPX formats: ROCmFP4 on vulkan, ROCmI4 with --device rocm\n" " [--mtp HEAD.gguf] [--mtp-max N] [--mtp-p-min P] multi-token prediction (vulkan, hrx, rocm)\n" - " [--moe-slots N] [--moe-subst R] stream MoE experts from the file, N held in RAM;\n" - " R: resident experts stand in for missing ones (vulkan; docs/moe-streaming.md)\n" + " [--moe-slots N|auto] [--moe-subst R] [--moe-prefetch N] stream MoE experts from the file,\n" + " N held in RAM (auto: what fits); R: resident experts stand in for missing\n" + " ones; prefetch N layers ahead (default 0) (vulkan; docs/moe-streaming.md)\n" " [--mmproj MMPROJ.gguf] images in chat messages (vulkan, hrx, rocm)\n" " [--parallel N] N requests decoded together (continuous batching)\n" " [--adaptive] [--adaptive-at N] Vulkan (+MTP) for N in flight (default 1), ROCm batches the rest\n" @@ -1028,7 +1090,11 @@ int run_serve(int argc, char** argv) { else if (a == "--prefill-min-tokens") o.prefill_min_tokens = std::stoi(next()); else if (a == "--lean") o.lean = true; else if (a == "--mtp") o.mtp = next(); - else if (a == "--moe-slots") o.moe_slots = std::stoi(next()); + else if (a == "--moe-slots") { + const std::string v = next(); + o.moe_slots = v == "auto" ? -1 : std::stoi(v); + } + else if (a == "--moe-prefetch") o.moe_prefetch = std::stoi(next()); else if (a == "--moe-subst") o.moe_subst = next(); else if (a == "--mmproj") o.mmproj = next(); else if (a == "--mtp-max") o.mtp_max = std::stoi(next()); diff --git a/docs/moe-streaming.md b/docs/moe-streaming.md index 9e18f42b..7e23dd89 100644 --- a/docs/moe-streaming.md +++ b/docs/moe-streaming.md @@ -612,6 +612,11 @@ the change is within the noise (four rounds, 6.7 against 7.0 tok/s). ## Next +0. **Prompts through the slots.** A batch above `ONEBIT_MOE_MAX_BATCH` (8) runs its experts on the + CPU from the mapped file: Flash-Next prompts went at 0.4-0.8 tok/s through `serve`. Candidates: + cap llama-server's ubatch at the streamed batch size, or stream larger batches when the + layer's slots hold all their experts. (Measuring it was interrupted: the Flash-Next files were + removed from the box on 2026-09-26.) 1. **Fewer host round trips.** The remap per layer caps streamed decode at about 57 tok/s on Coder-30B ("Where streamed decode spends its time"). 6,144 slots (every expert) runs erratically on the shared box. diff --git a/docs/serve.md b/docs/serve.md index 8d9ae4ba..84645924 100644 --- a/docs/serve.md +++ b/docs/serve.md @@ -27,7 +27,7 @@ OpenAI client. [--llama-server PATH] [--zinc PATH] [--ds4 PATH] [--ssd-streaming] [--hrx-libhsa PATH] [--mlx-server PATH] [--prefill-device hrx] [--prefill-min-tokens N] [--lean] [--mtp HEAD.gguf] [--mtp-max N] [--mtp-p-min P] [--mmproj MMPROJ.gguf] - [--moe-slots N] [--moe-subst R] + [--moe-slots N|auto] [--moe-subst R] [--moe-prefetch N] [--parallel N] [--adaptive] [--adaptive-at N] [--embed MODEL.gguf] [--rerank MODEL.gguf] [--npu-opt KEY=VALUE ...] @@ -167,6 +167,17 @@ Qwen3.8-Flash-Next UD-Q4_K_XL (111 GB) runs this way from 18-32 GiB of GPU memor exact at 4,608-9,216 slots. On that model use `--moe-subst 0.9` or none: 0.5 costs a KL divergence of 0.07 there. +`--moe-slots auto` sizes the cache from free memory: `MemAvailable`, less the other tensors the +GPU holds and a reserve (8 GiB or 15%), divided by the model's mean expert size. On a box with +40 GiB free it gave Flash-Next 9,331 of 24,576 experts (27.2 GiB) beside 5.1 GiB of other +tensors, and `serve` held 145 MiB of anonymous memory. Other processes can take that memory +back later; on a shared box, give a number instead. The gate-ahead prefetch is off unless +`--moe-prefetch N` asks for it (it doubled Flash-Next's decode time). + +Prompts are slow when streaming: a batch larger than 8 tokens runs the experts on the CPU +from the mapped file (Flash-Next: 0.4-0.8 tok/s of prompt). Keep prompts short, or see +docs/moe-streaming.md, Next. + ## Many requests at once (`--parallel`) `--parallel N` gives llama-server N slots (`-np N`): requests that arrive together decode