From adbe3caaf258805a41003876aaeaf27f32e52ee6 Mon Sep 17 00:00:00 2001 From: Helldez <74499579+Helldez@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:34:35 +0200 Subject: [PATCH] feat(moe): stream split multi-shard ggufs natively + DeepSeek V4 Flash recipe (#144) * feat(moe): stream split multi-shard ggufs natively + DeepSeek V4 Flash recipe Hugging Face rejects single files above 50 GB, so every large model ships as -00001-of-0000N.gguf shards; until now the streamer assumed one file, forcing a merge with double the disk. gguf_offsets now fans the first shard out to the whole set and resolves every tensor to (shard, offset); the expert streamer and the dense loader open one positioned reader per shard and route each read by the tensor's shard index. Pass the first shard, exactly as llama.cpp takes it; a missing sibling fails the load with the shard named. Add the deepseek4 recipe row: V3.2-style routing (256 routed experts, a per-expert bias like lfm2moe, an always-on shared expert that stays resident) over the standard split expert suffixes. The V4 compressed-attention machinery is dense-side llama.cpp code, invisible to the streaming seam. The byte-identity gates gain a 4-shard qwen3moe fixture (metadata-only first shard, the layout large quants actually use); make-tiny-moe.py learns --split-max-tensors. All gates pass, split included. * fix(moe): cache auto must budget for the anon dense conversion The auto budget read MemAvailable while the dense weights were still reclaimable page cache, then dense-weights=anon converted them into buffers the kernel cannot take back: the same bytes planned twice. Latent since the anon policy shipped (dense sets were 2-3 GiB and explicit budgets were the benched path); DeepSeek V4 Flash's 6.5 GiB dense set turned it into a device-taking overcommit on first load. The budget now deducts the pending conversion and says so in the log. * fix(moe): review pass on the multi-shard path Three defects the split rewrite introduced, none of which the gates could see: - The shard index rode in an int8_t, so a model past 127 shards wrapped to a negative index into the reader vector. The bounds check could never catch it: it validated the untruncated value. Widened to int16_t, which covers the whole -%05d-of-%05d filename space. - DenseWeights::warm() reused one flag as both the inner loop condition and the partial-warm report, so the first shard that failed to open silently skipped the warm-up of every later shard. Per-shard condition, sticky report. - The dense readers stayed allocated for the session after read_anonymous had copied and rebound every tensor: fds and a per-lane bounce buffer per shard, sitting next to a cache counting every MiB. Released at the end of init. Also: the streaming banner read O_DIRECT off shard 0, which under the small-first-shard layout is metadata only and too short to verify, so it could claim a mode the shards carrying experts had not got. It now reports the weakest of the readers. * build: the engine version says 0.19.0, like the changelog does The version is declared in CMakeLists.txt and reported by `--version` and by the run-parameter preamble of every metrics CSV, so a committed benchmark file names the engine that produced it. This release section landed while the number stayed at 0.18.0, which would have stamped the wrong engine on every CSV this branch produces, defeating the one purpose the string has. --- CHANGELOG.md | 18 ++++ CMakeLists.txt | 2 +- README.md | 8 +- core/src/engine/session.cpp | 4 +- core/src/moe/arch_registry.cpp | 8 ++ core/src/moe/dense_weights.cpp | 149 +++++++++++++++----------- core/src/moe/dense_weights.h | 50 +++++---- core/src/moe/expert_stream_source.cpp | 94 +++++++++++----- core/src/moe/expert_stream_source.h | 28 +++-- core/src/moe/gguf_offsets.cpp | 76 +++++++++++-- core/src/moe/gguf_offsets.h | 26 ++++- docs/architecture.md | 3 +- docs/limitations.md | 2 +- scripts/make-tiny-moe.py | 28 ++++- tests/CMakeLists.txt | 16 +++ 15 files changed, 374 insertions(+), 138 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 35897b2..3c23704 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,24 @@ All notable changes to this project are documented here. The format follows [Keep a Changelog](https://keepachangelog.com/), and the project aims to follow Semantic Versioning. +## [0.19.0] - 2026-08-01 + +### Added +- **Split (multi-shard) ggufs stream natively.** Hugging Face rejects single files above 50 GB, so + every large model ships as `-00001-of-0000N.gguf` shards — and until now the streamer assumed one + file, forcing a merge with double the disk. `gguf_offsets` now fans the first shard out to the + whole set and resolves every tensor to its (shard, offset); the expert streamer and the dense + loader open one positioned reader per shard and route each read by the tensor's shard index. Pass + the first shard, as with llama.cpp itself; a missing sibling fails the load with the shard named. + The byte-identity gates gained a 4-shard fixture (metadata-only first shard, the layout large + quants actually use) proving streamed == resident across shard boundaries. +- **DeepSeek V4 Flash (`deepseek4`) recipe.** V3.2-style expert routing — 256 routed experts with a + per-expert bias (the `lfm2moe` pattern) plus an always-on shared expert that stays resident — over + the standard split expert suffixes, so streaming is one registry row. The V4 attention machinery + (compressed sparse attention, the lightning indexer and its dedicated KV cache) is dense-side + llama.cpp code, invisible to the streaming seam. Requires a llama.cpp with DeepSeek V4 support + (the pinned submodule has it). + ## [0.18.1] - 2026-07-29 ### Fixed diff --git a/CMakeLists.txt b/CMakeLists.txt index d4f6b17..0967d5f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -19,7 +19,7 @@ cmake_minimum_required(VERSION 3.21) # The version is declared here and nowhere else in the build: the engine reports it (bmoe-cli # --version, and the run-parameter preamble of every metrics CSV), so a committed benchmark file # says which engine produced it. Keep it in step with CHANGELOG.md and the app's versionName. -project(bigmoeonedge VERSION 0.18.0 LANGUAGES CXX) +project(bigmoeonedge VERSION 0.19.0 LANGUAGES CXX) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) diff --git a/README.md b/README.md index ea91d2a..93efeee 100644 --- a/README.md +++ b/README.md @@ -90,9 +90,10 @@ shortcut: the downloader takes any direct gguf URL, so any model from the finishes, pick the model and chat. The telemetry panel shows tok/s and the compute-vs-flash split live, and every streaming knob below is in Settings. -The flagship gpt-oss-120b ships from Hugging Face as two shards, so it needs a one-time merge -and a manual copy to the device: steps in the -[Android example README](examples/android/README.md#gpt-oss-120b). +Models above Hugging Face's 50 GB per-file limit (gpt-oss-120b, DeepSeek V4 Flash) ship as +multi-shard ggufs. The engine reads split models natively — point it at the first shard +(`-00001-of-...`) with the siblings alongside; no merge step. gpt-oss-120b still needs a manual +copy to the device: steps in the [Android example README](examples/android/README.md#gpt-oss-120b). ## Features @@ -149,6 +150,7 @@ Defaults are the measured winning recipe for a model near RAM. | `gemma4` | Gemma 4 MoE (e.g. 26B-A4B) | Fused expert layout, handled by its registry row | | `gpt-oss` | OpenAI gpt-oss-20b / 120b | Purely routed; MXFP4 weights stream unchanged | | `lfm2moe` | Liquid AI LFM2 / LFM2.5 MoE (e.g. 8B-A1B) | Hybrid conv/attention stack with leading dense blocks; those stay resident | +| `deepseek4` | DeepSeek V4 Flash (284B-A13B) | V3.2-style routing (256 experts + shared); compressed attention is dense-side; ships multi-shard | Adding an architecture is one row in the registry; expert counts and layouts are discovered from the model file at runtime, so nothing about a specific model is hardcoded in the streaming path. diff --git a/core/src/engine/session.cpp b/core/src/engine/session.cpp index 9176a71..9457136 100644 --- a/core/src/engine/session.cpp +++ b/core/src/engine/session.cpp @@ -454,6 +454,7 @@ std::unique_ptr Session::open(const SessionConfig & cfg, auto it = offs.off_by_name.find(t->name); if (it == offs.off_by_name.end()) return fail(std::string("no gguf offset for tensor ") + t->name); L.proj[p].file_off = it->second; + L.proj[p].file_idx = offs.file_by_name.at(t->name); // same parse as the offset, so present const int ne2 = (int) t->ne[2]; if (n_expert == 0) n_expert = ne2; @@ -487,12 +488,13 @@ std::unique_ptr Session::open(const SessionConfig & cfg, d.tensor = kv.second; d.file_off = off->second; d.size = sz->second; + d.file_idx = offs.file_by_name.at(name); dense.push_back(d); } im.source.set_dense_tensors(std::move(dense)); } - if (!im.source.init(cfg.model_path, n_expert, std::move(layers), cfg.moe)) + if (!im.source.init(offs.shard_paths, n_expert, std::move(layers), cfg.moe)) return fail("expert stream source init failed"); im.hook->set_source(&im.source); diff --git a/core/src/moe/arch_registry.cpp b/core/src/moe/arch_registry.cpp index 457748e..08f97f1 100644 --- a/core/src/moe/arch_registry.cpp +++ b/core/src/moe/arch_registry.cpp @@ -41,6 +41,14 @@ static const MoeRecipe k_recipes[] = { // node the hook reads (ffn_moe_topk). Both lower the streamed fraction relative to a purely // routed model — see docs/limitations.md. {"lfm2moe", {"ffn_gate_exps", "ffn_up_exps", "ffn_down_exps"}}, + // deepseek4 (DeepSeek V4 Flash, 284B-A13B) reuses the V3.2 MoE routing — 256 routed experts, + // top-k with a per-expert bias (exp_probs_b, the lfm2moe pattern) plus one always-on shared + // expert (ffn_*_shexp) that matches no suffix and stays mmap-resident. The experts name the + // standard split suffixes, so streaming is one row. The V4 attention novelties (compressed + // sparse attention, the lightning indexer, its dedicated KV cache) are dense-side machinery + // inside llama.cpp and invisible to the streaming seam. Models this size ship as multi-shard + // ggufs; the streamer resolves each expert tensor to its (shard, offset) — see gguf_offsets. + {"deepseek4", {"ffn_gate_exps", "ffn_up_exps", "ffn_down_exps"}}, }; static const int k_n_recipes = (int) (sizeof(k_recipes) / sizeof(k_recipes[0])); diff --git a/core/src/moe/dense_weights.cpp b/core/src/moe/dense_weights.cpp index 76bf691..6f74caa 100644 --- a/core/src/moe/dense_weights.cpp +++ b/core/src/moe/dense_weights.cpp @@ -28,17 +28,20 @@ DenseWeights::byte_ranges(std::vector> expert_rang } bool DenseWeights::init(DenseWeightsMode mode, - const std::string & path, + const std::vector & paths, size_t align, - std::vector> ranges, + std::vector>> ranges, std::vector tensors) { mode_ = mode; - path_ = path; + paths_ = paths; align_ = align ? align : 4096; ranges_ = std::move(ranges); tensors_ = std::move(tensors); - const size_t slash = path_.find_last_of("/\\"); - basename_ = slash == std::string::npos ? path_ : path_.substr(slash + 1); + basenames_.clear(); + for (const std::string & p : paths_) { + const size_t slash = p.find_last_of("/\\"); + basenames_.push_back(slash == std::string::npos ? p : p.substr(slash + 1)); + } if (mode_ == DenseWeightsMode::Anonymous || mode_ == DenseWeightsMode::Pinned) { if (tensors_.empty()) return true; // nothing captured to rebind — behave as Mmap @@ -47,11 +50,19 @@ bool DenseWeights::init(DenseWeightsMode mode, "platform does not provide (Android only)\n"); return false; } - // A single-lane reader with a bounce large enough for our chunk; O_DIRECT independent of the - // expert stream. Sized to the largest tensor is unnecessary — we read in bounded chunks. + // Single-lane readers (one per shard) with a bounce large enough for our chunk; O_DIRECT + // independent of the expert stream. Sized to the largest tensor is unnecessary — we read in + // bounded chunks. const size_t chunk = 8ull << 20; - if (!reader_.open(path_, 1, /*direct=*/true, align_, chunk + 2 * align_)) return false; + for (const std::string & p : paths_) { + readers_.push_back(std::unique_ptr(new FileReader())); + if (!readers_.back()->open(p, 1, /*direct=*/true, align_, chunk + 2 * align_)) return false; + } if (!read_anonymous(align_)) return false; + // The tensors are copied and rebound; nothing reads through these again. Their fds and + // per-lane bounce buffers would otherwise sit allocated for the whole session, next to + // the expert cache that is counting every MiB. + readers_.clear(); drop_mmap_copies(pio::vm_page()); } else if (mode_ == DenseWeightsMode::Warmed) { warm(); @@ -79,6 +90,10 @@ bool DenseWeights::read_anonymous(size_t align) { if (pinned) pinned_.reserve(tensors_.size()); for (const DenseTensorRef & d : tensors_) { if (!d.tensor || d.size == 0) continue; + if (d.file_idx < 0 || d.file_idx >= (int) readers_.size()) { + std::fprintf(stderr, "bmoe: dense tensor points at shard %d of %zu\n", d.file_idx, readers_.size()); + return false; + } void * buf = nullptr; if (pinned) { pio::PinnedAlloc pa; @@ -101,7 +116,7 @@ bool DenseWeights::read_anonymous(size_t align) { buf_sz_.push_back((size_t) d.size); for (uint64_t done = 0; done < d.size;) { const uint64_t n = std::min(chunk, d.size - done); - if (reader_.read(0, (char *) buf + done, d.file_off + done, n) < 0) return false; + if (readers_[(size_t) d.file_idx]->read(0, (char *) buf + done, d.file_off + done, n) < 0) return false; done += n; } d.tensor->data = buf; // rebind the model weight onto its private copy @@ -112,6 +127,25 @@ bool DenseWeights::read_anonymous(size_t align) { return true; } +// Per-shard VMA resolution for every consumer that must turn (file_idx, offset) into an address. +// llama.cpp maps each shard of a split model separately, so the lookup is per shard basename. +void DenseWeights::resolve_vmas() { + if (vmas_tried_) return; + vmas_tried_ = true; + vmas_.resize(basenames_.size()); + for (size_t s = 0; s < basenames_.size(); ++s) + pio::file_mapped_regions(basenames_[s].c_str(), vmas_[s]); +} + +const char * DenseWeights::addr_of(int file_idx, uint64_t off) const { + if (file_idx < 0 || file_idx >= (int) vmas_.size()) return nullptr; + for (const auto & v : vmas_[file_idx]) { + const uint64_t span = (uint64_t) (v.end - v.start); + if (off >= v.file_offset && off < v.file_offset + span) return (const char *) v.start + (off - v.file_offset); + } + return nullptr; +} + // Hand the mmap copies back. The capture warm-up decode faulted these dense pages in mmap-resident, // and read_anonymous has just copied them into anon buffers and rebound every tensor — so the file- // backed pages are referenced by nobody. Left alone they sit resident until reclaim, doubling the @@ -119,19 +153,11 @@ bool DenseWeights::read_anonymous(size_t align) { // read-only mapping, so nothing is lost and nothing will refault the range. Best-effort — needs // /proc/self/maps to turn a file offset into an address; where that is unreadable the pages stay. void DenseWeights::drop_mmap_copies(size_t page) { - std::vector vmas; - if (!pio::file_mapped_regions(basename_.c_str(), vmas) || vmas.empty()) return; - auto addr_of = [&](uint64_t off) -> char * { - for (const auto & v : vmas) { - const uint64_t span = (uint64_t) (v.end - v.start); - if (off >= v.file_offset && off < v.file_offset + span) return (char *) v.start + (off - v.file_offset); - } - return nullptr; - }; + resolve_vmas(); uint64_t dropped = 0; for (const DenseTensorRef & d : tensors_) { if (!d.tensor || d.size == 0) continue; - char * a = addr_of(d.file_off); + const char * a = addr_of(d.file_idx, d.file_off); if (!a) continue; // Align INWARD to whole pages (start up, end down), so a page shared with an adjacent tensor // that stays mmap-resident — an expert slice, or the next dense tensor — is never dropped. @@ -149,30 +175,34 @@ void DenseWeights::drop_mmap_copies(size_t page) { // ── Warmed: one sequential buffered sweep over the dense ranges to populate the page cache ── void DenseWeights::warm() { - pio::fd_t fd = pio::open_read(path_.c_str(), false); - if (!pio::fd_ok(fd)) return; const size_t chunk = 8ull << 20; void * buf = pio::alloc_aligned(align_, chunk); - if (!buf) { - pio::close_fd(fd); - return; - } + if (!buf) return; const auto t0 = clock_t_::now(); uint64_t warmed = 0; - bool ok = true; - for (const auto & r : ranges_) { - for (uint64_t a = r.first; a < r.second && ok;) { - const long long got = pio::pread_at(fd, buf, (size_t) std::min(chunk, r.second - a), a); - if (got <= 0) { - ok = false; // best-effort: those pages just stay cold - break; - } - a += (uint64_t) got; - warmed += (uint64_t) got; + bool all_ok = true; // sticky, for the report only: a failed shard must not abort the others + for (size_t s = 0; s < paths_.size() && s < ranges_.size(); ++s) { + pio::fd_t fd = pio::open_read(paths_[s].c_str(), false); + if (!pio::fd_ok(fd)) { + all_ok = false; // best-effort: that shard's pages just stay cold + continue; } + bool shard_ok = true; + for (const auto & r : ranges_[s]) { + for (uint64_t a = r.first; a < r.second && shard_ok;) { + const long long got = pio::pread_at(fd, buf, (size_t) std::min(chunk, r.second - a), a); + if (got <= 0) { + shard_ok = all_ok = false; + break; + } + a += (uint64_t) got; + warmed += (uint64_t) got; + } + } + pio::close_fd(fd); } + const bool ok = all_ok; pio::aligned_free(buf); - pio::close_fd(fd); const double s = std::chrono::duration(clock_t_::now() - t0).count(); std::fprintf(stderr, "bmoe: dense warm-up — %llu MiB in %.1f s%s\n", (unsigned long long) (warmed >> 20), s, ok ? "" : " (partial)"); @@ -214,46 +244,39 @@ void DenseWeights::sample_anon(size_t page) { resident_frac_ = sampled ? (double) resident / (double) sampled : -1.0; } -// Mmap/Warmed: the weights are one mmap of the gguf; a dense file offset becomes an address through -// the VMA that backs it (/proc/self/maps), which an app may read, being its own. Probe sample_pages -// points spread evenly across the dense byte ranges, one page each. +// Mmap/Warmed: the weights are llama.cpp's mmaps of the gguf shards; a dense (shard, offset) becomes +// an address through the VMA that backs it (/proc/self/maps), which an app may read, being its own. +// Probe sample_pages points spread evenly across the dense bytes of ALL shards, one page each, so the +// fraction keeps meaning "of the whole dense set" whatever the shard layout. void DenseWeights::sample_mmap(size_t page) { if (ranges_.empty()) return; - if (!vmas_tried_) { // resolve the mmap once; llama.cpp has mapped the file by the first decode - vmas_tried_ = true; - pio::file_mapped_regions(basename_.c_str(), vmas_); - } - if (vmas_.empty()) return; // maps unreadable → leave resident_frac_ at -1 + resolve_vmas(); uint64_t total = 0; - for (const auto & r : ranges_) - total += r.second - r.first; + for (const auto & shard : ranges_) + for (const auto & r : shard) + total += r.second - r.first; if (total == 0) return; - auto addr_of = [&](uint64_t off) -> const char * { - for (const auto & v : vmas_) { - const uint64_t span = (uint64_t) (v.end - v.start); - if (off >= v.file_offset && off < v.file_offset + span) - return (const char *) v.start + (off - v.file_offset); - } - return nullptr; - }; - size_t sampled = 0, resident = 0; for (int k = 0; k < sample_pages; ++k) { uint64_t target = (total * (uint64_t) k) / (uint64_t) sample_pages; + int fi = -1; // found shard; a dense offset of 0 is valid, so the sentinel is the index uint64_t off = 0; - for (const auto & r : ranges_) { - const uint64_t len = r.second - r.first; - if (target < len) { - off = r.first + target; - break; + for (size_t s = 0; s < ranges_.size() && fi < 0; ++s) { + for (const auto & r : ranges_[s]) { + const uint64_t len = r.second - r.first; + if (target < len) { + fi = (int) s; + off = r.first + target; + break; + } + target -= len; } - target -= len; } // Align the probe DOWN to its page: vm_resident_sample counts only pages fully inside the // range, so a single page handed in at an arbitrary offset would clip to nothing. - if (const char * a = addr_of(off)) { + if (const char * a = fi >= 0 ? addr_of(fi, off) : nullptr) { const char * pg = (const char *) ((uintptr_t) a & ~(uintptr_t) (page - 1)); pio::vm_resident_sample(pg, page, &sampled, &resident); } @@ -271,7 +294,7 @@ void DenseWeights::shutdown() { bases_.clear(); buf_sz_.clear(); tensors_.clear(); - reader_.close(); + readers_.clear(); } } // namespace bmoe diff --git a/core/src/moe/dense_weights.h b/core/src/moe/dense_weights.h index ab90f20..b519167 100644 --- a/core/src/moe/dense_weights.h +++ b/core/src/moe/dense_weights.h @@ -7,10 +7,13 @@ // sensor that reports how much of the dense set the kernel still has — the half the expert cache's // own residency sensor is blind to. // -// It reads through its OWN FileReader, so its O_DIRECT choice is independent of the expert stream's: +// It reads through its OWN FileReaders, so its O_DIRECT choice is independent of the expert stream's: // the dense set may be pulled cache-bypassing while the experts are not, or the reverse. There is one // definition of "which bytes are dense" here (byte_ranges), shared by the warm sweep, the anonymous // read, and the sensor — the three used to compute it separately. +// +// A split model hands in one path (and one set of dense ranges) per shard file; every consumer walks +// the set. A single-file model is the one-entry case of the same shape, not a separate path. #pragma once #include "../io/file_reader.h" @@ -18,6 +21,7 @@ #include "bmoe/config.h" // DenseWeightsMode #include +#include #include #include #include @@ -27,11 +31,12 @@ struct ggml_tensor; namespace bmoe { // One dense weight tensor: where its bytes live in the gguf, and the tensor whose ->data we rebind. -// Read whole and contiguous (`size` bytes from `file_off`), unlike an expert slice. +// Read whole and contiguous (`size` bytes from `file_off` in shard `file_idx`), unlike an expert slice. struct DenseTensorRef { ggml_tensor * tensor = nullptr; uint64_t file_off = 0; uint64_t size = 0; + int file_idx = 0; // which shard file holds the bytes (0 for a single-file model) }; class DenseWeights { @@ -42,20 +47,20 @@ public: DenseWeights(const DenseWeights &) = delete; DenseWeights & operator=(const DenseWeights &) = delete; - // The dense byte ranges: the complement of `expert_ranges` within [0, file_size). Sorts the - // input; the result is the gaps between expert ranges plus the trailing tail. The single source - // of truth every consumer here shares. + // The dense byte ranges OF ONE FILE: the complement of `expert_ranges` within [0, file_size). + // Sorts the input; the result is the gaps between expert ranges plus the trailing tail. The + // single source of truth every consumer here shares; called once per shard. static std::vector> byte_ranges(std::vector> expert_ranges, uint64_t file_size); - // Apply `mode` for the model at `path` (dense `ranges` precomputed via byte_ranges; `tensors` - // needed only for Anonymous). `align` is the O_DIRECT block size. Runs once at load, before the - // streamer's workers start (Anonymous rebinds tensor->data on the caller's thread). Returns false - // only on a hard Anonymous failure (alloc/read); Mmap and Warmed cannot fail. + // Apply `mode` for the model files `paths` (per-shard dense `ranges` precomputed via byte_ranges; + // `tensors` needed only for Anonymous/Pinned). `align` is the O_DIRECT block size. Runs once at + // load, before the streamer's workers start (Anonymous rebinds tensor->data on the caller's + // thread). Returns false only on a hard Anonymous failure (alloc/read); Mmap and Warmed cannot fail. bool init(DenseWeightsMode mode, - const std::string & path, + const std::vector & paths, size_t align, - std::vector> ranges, + std::vector>> ranges, std::vector tensors); // Sample how much of the dense set the kernel still has in RAM (mincore), setting resident_frac(). @@ -74,26 +79,29 @@ private: void drop_mmap_copies(size_t page); void sample_anon(size_t page); void sample_mmap(size_t page); + void resolve_vmas(); // per-shard /proc/self/maps lookup, once + const char * addr_of(int file_idx, uint64_t off) const; DenseWeightsMode mode_ = DenseWeightsMode::Warmed; - std::string path_; - std::string basename_; + std::vector paths_; // one per shard, in shard order + std::vector basenames_; // what /proc/self/maps entries are matched against size_t align_ = 4096; - // Anonymous/Pinned: our own reader, the tensors we read, and the buffers backing them. `bases_` - // is what the sensor probes and is filled by both modes; `bufs_` and `pinned_` are the two - // release lists, exactly one of which is populated for a given run. - FileReader reader_; + // Anonymous/Pinned: our own readers (one per shard; FileReader is not movable, hence the + // unique_ptr), the tensors we read, and the buffers backing them. `bases_` is what the sensor + // probes and is filled by both modes; `bufs_` and `pinned_` are the two release lists, exactly + // one of which is populated for a given run. + std::vector> readers_; std::vector tensors_; std::vector bases_; std::vector bufs_; std::vector pinned_; std::vector buf_sz_; - // Mmap/Warmed: the dense byte ranges and the mmap VMAs that back them, resolved once from - // /proc/self/maps for the sensor. - std::vector> ranges_; - std::vector vmas_; + // Mmap/Warmed: the dense byte ranges and the mmap VMAs that back them, per shard, resolved once + // from /proc/self/maps for the sensor (llama.cpp maps every shard of a split model). + std::vector>> ranges_; + std::vector> vmas_; bool vmas_tried_ = false; double resident_frac_ = -1.0; // last sampled dense residency; -1 = never/unmeasured diff --git a/core/src/moe/expert_stream_source.cpp b/core/src/moe/expert_stream_source.cpp index 3ce0fbd..91ae337 100644 --- a/core/src/moe/expert_stream_source.cpp +++ b/core/src/moe/expert_stream_source.cpp @@ -25,7 +25,7 @@ ExpertStreamSource::~ExpertStreamSource() { } // ── init: allocate buffers, rebind expert tensors, start the read pool ────────────── -bool ExpertStreamSource::init(const std::string & gguf_path, +bool ExpertStreamSource::init(const std::vector & shard_paths, int n_expert, std::vector layers, const MoeStreamConfig & cfg) { @@ -34,6 +34,10 @@ bool ExpertStreamSource::init(const std::string & gguf_path, std::fprintf(stderr, "bmoe: expert streaming needs a MoE model (n_expert=%d)\n", n_expert); return false; } + if (shard_paths.empty()) { + std::fprintf(stderr, "bmoe: expert streaming got no model file\n"); + return false; + } n_expert_ = n_expert; layers_ = std::move(layers); @@ -77,7 +81,24 @@ bool ExpertStreamSource::init(const std::string & gguf_path, const size_t hard_cap = ceil > 0 ? std::min(ceil, total_expert_bytes_) : total_expert_bytes_; const size_t min_budget = std::min((size_t) MoeStreamConfig::cache_min_mb * 1024ull * 1024ull, hard_cap); - const uint64_t avail = pio::mem_available_bytes(); + uint64_t avail = pio::mem_available_bytes(); + // Under anon/ahwb the dense weights are ABOUT to move from reclaimable page cache into + // buffers the kernel cannot take back (dense_.init runs below). At this moment the kernel + // still counts those pages as available, so a budget sized from the raw number plans the + // dense set twice and overcommits by its whole size — enough to take the device down on a + // model whose dense set is large (DeepSeek V4's is ~6.5 GiB). Budget as if the conversion + // had already happened. + if (cfg.dense_weights == DenseWeightsMode::Anonymous || cfg.dense_weights == DenseWeightsMode::Pinned) { + uint64_t dense_pending = 0; + for (const DenseTensorRef & d : dense_tensors_) + if (d.tensor) dense_pending += d.size; + const uint64_t deduct = std::min(avail, dense_pending); + if (deduct > 0) { + avail -= deduct; + std::fprintf(stderr, "bmoe: cache auto — reserving %llu MiB for the dense-weight conversion\n", + (unsigned long long) (deduct / (1024 * 1024))); + } + } if (avail == 0) { cache_max_ = min_budget; std::fprintf(stderr, "bmoe: cache auto — available memory unknown, using the %zu MiB floor\n", @@ -151,12 +172,25 @@ bool ExpertStreamSource::init(const std::string & gguf_path, clookups_ = 0; } - // Read pool: the reader owns a private fd + bounce per lane so concurrent preads never contend, - // and the O_DIRECT request + its verify/fallback. The dense-weights loader opens its own reader, - // so its cache-bypass choice is independent of this one. + // Read pool: one reader per shard file, each owning a private fd + bounce per lane so concurrent + // preads never contend, and the O_DIRECT request + its verify/fallback (per file — a shard set + // could in principle straddle storage with different O_DIRECT behaviour). The dense-weights + // loader opens its own readers, so its cache-bypass choice is independent of this one. const size_t max_slice = max_full_any / (size_t) n_expert_; const size_t bounce_cap = max_slice + 2 * align_; - if (!reader_.open(gguf_path, io_threads_, cfg.o_direct, align_, bounce_cap)) return false; + for (const std::string & sp : shard_paths) { + readers_.push_back(std::unique_ptr(new FileReader())); + if (!readers_.back()->open(sp, io_threads_, cfg.o_direct, align_, bounce_cap)) return false; + } + for (const LayerExperts & L : layers_) { + if (!L.bound) continue; + for (int p = 0; p < MoeRecipe::max_exps; ++p) + if (L.proj[p].nb2 && (L.proj[p].file_idx < 0 || L.proj[p].file_idx >= (int) readers_.size())) { + std::fprintf(stderr, "bmoe: expert tensor points at shard %d of %zu\n", L.proj[p].file_idx, + readers_.size()); + return false; + } + } seen_.assign(n_expert_, 0); jobs_.reserve((size_t) n_expert_ * MoeRecipe::max_exps); @@ -195,16 +229,18 @@ bool ExpertStreamSource::init(const std::string & gguf_path, // are the complement of the expert ranges in the file, computed once here and shared by the warm // sweep and the sensor. { - std::vector> exp; + std::vector>> exp(readers_.size()); for (const LayerExperts & L : layers_) { if (!L.bound) continue; for (int p = 0; p < MoeRecipe::max_exps; ++p) { const uint64_t sz = (uint64_t) L.proj[p].nb2 * (uint64_t) n_expert_; - if (sz) exp.push_back({L.proj[p].file_off, L.proj[p].file_off + sz}); + if (sz) exp[L.proj[p].file_idx].push_back({L.proj[p].file_off, L.proj[p].file_off + sz}); } } - auto ranges = DenseWeights::byte_ranges(std::move(exp), reader_.file_size()); - if (!dense_.init(cfg.dense_weights, gguf_path, align_, std::move(ranges), std::move(dense_tensors_))) { + std::vector>> ranges(readers_.size()); + for (size_t s = 0; s < readers_.size(); ++s) + ranges[s] = DenseWeights::byte_ranges(std::move(exp[s]), readers_[s]->file_size()); + if (!dense_.init(cfg.dense_weights, shard_paths, align_, std::move(ranges), std::move(dense_tensors_))) { std::fprintf(stderr, "bmoe: dense-weights init failed\n"); return false; } @@ -218,8 +254,13 @@ bool ExpertStreamSource::init(const std::string & gguf_path, for (int lane = first_worker_lane; lane < io_threads_; ++lane) io_pool_.emplace_back(&ExpertStreamSource::io_worker, this, lane); - std::fprintf(stderr, "bmoe: expert streaming ON n_expert=%d o_direct=%d io_threads=%d cache=%zu MiB\n", n_expert_, - (int) reader_.direct(), io_threads_, cache_max_ >> 20); + // Each shard verified O_DIRECT for itself, so report the weakest: a metadata-only first shard + // is too short to verify at all and would flatter the number for the shards carrying experts. + bool all_direct = true; + for (const auto & r : readers_) + all_direct = all_direct && r->direct(); + std::fprintf(stderr, "bmoe: expert streaming ON n_expert=%d o_direct=%d io_threads=%d cache=%zu MiB shards=%zu\n", + n_expert_, (int) all_direct, io_threads_, cache_max_ >> 20, readers_.size()); return true; } @@ -230,7 +271,7 @@ bool ExpertStreamSource::init(const std::string & gguf_path, bool ExpertStreamSource::read_slice(int lane, const IoJob & j) { if (j.nbytes == 0) return true; const auto t0 = clock_t_::now(); - const long long window = reader_.read(lane, j.dst, j.off, j.nbytes); + const long long window = readers_[(size_t) j.file]->read(lane, j.dst, j.off, j.nbytes); if (window < 0) return false; if (io_trace_on_) { @@ -373,8 +414,8 @@ void ExpertStreamSource::prefetch(int il, const int32_t * ids, int n_ids) { ok = false; break; } - staged.push_back( - {dst, L.proj[p].file_off + (uint64_t) e * slice, slice, id, e, (int16_t) il, (int8_t) p, 1}); + staged.push_back({dst, L.proj[p].file_off + (uint64_t) e * slice, slice, id, (int16_t) L.proj[p].file_idx, + e, (int16_t) il, (int8_t) p, 1}); ++njobs; } if (!ok) { @@ -713,7 +754,7 @@ bool ExpertStreamSource::load_layer(int il, const int32_t * ids, int n_ids) { const uint64_t slice = L.proj[p].nb2; if (slice == 0) continue; // absent slot in a fused layout jobs_.push_back({(char *) slot_[p] + (uint64_t) e * slice, L.proj[p].file_off + (uint64_t) e * slice, - slice, -1, e, (int16_t) il, (int8_t) p, 0}); + slice, -1, (int16_t) L.proj[p].file_idx, e, (int16_t) il, (int8_t) p, 0}); } return true; } @@ -727,7 +768,7 @@ bool ExpertStreamSource::load_layer(int il, const int32_t * ids, int n_ids) { const uint64_t slice = L.proj[p].nb2; if (slice == 0) continue; // absent slot in a fused layout jobs_.push_back({(char *) lbuf_[p][il] + (uint64_t) e * slice, L.proj[p].file_off + (uint64_t) e * slice, - slice, -1, e, (int16_t) il, (int8_t) p, 0}); + slice, -1, (int16_t) L.proj[p].file_idx, e, (int16_t) il, (int8_t) p, 0}); } return true; }; @@ -876,7 +917,7 @@ bool ExpertStreamSource::load_layer_async(int il, const int32_t * ids, int n_ids for (int e : staged_) { const int32_t flag = (int32_t) ((size_t) p * (size_t) n_expert_ + (size_t) e); jobs_.push_back({(char *) slot_[p] + (uint64_t) e * slice, L.proj[p].file_off + (uint64_t) e * slice, - slice, flag, e, (int16_t) il, (int8_t) p, 0}); + slice, flag, (int16_t) L.proj[p].file_idx, e, (int16_t) il, (int8_t) p, 0}); } } } else { @@ -920,8 +961,8 @@ bool ExpertStreamSource::load_layer_async(int il, const int32_t * ids, int n_ids if (!seen_[e]) continue; // cache hit, already marked ready const int32_t flag = (int32_t) ((size_t) p * (size_t) n_expert_ + (size_t) e); jobs_.push_back({(char *) lbuf_[p][il] + (uint64_t) e * slice, - L.proj[p].file_off + (uint64_t) e * slice, slice, flag, e, (int16_t) il, (int8_t) p, - 0}); + L.proj[p].file_off + (uint64_t) e * slice, slice, flag, (int16_t) L.proj[p].file_idx, + e, (int16_t) il, (int8_t) p, 0}); } }; if (!two_wave) { @@ -1094,11 +1135,16 @@ void ExpertStreamSource::enable_overlap_hook() { IExpertSource::Stats ExpertStreamSource::stats() const { Stats s; - s.read_bytes = (uint64_t) reader_.read_bytes(); + long long rd_bytes = 0, rd_syscall_ns = 0; + for (const auto & r : readers_) { + rd_bytes += r->read_bytes(); + rd_syscall_ns += r->syscall_ns(); + } + s.read_bytes = (uint64_t) rd_bytes; // Serial: read_ns_ is the wall time the caller was blocked in the read phase. Overlap: the - // caller never blocks, so "I/O time" is instead the summed lane-busy time (the reader's syscall + // caller never blocks, so "I/O time" is instead the summed lane-busy time (the readers' syscall // ns), which the runtime reports as lane-busy per token rather than as a slice of wall time. - s.read_seconds = (overlap_ ? reader_.syscall_ns() : read_ns_.load()) / 1e9; + s.read_seconds = (overlap_ ? rd_syscall_ns : read_ns_.load()) / 1e9; s.mgmt_seconds = mgmt_ns_.load() / 1e9; s.spec_read_bytes = (uint64_t) spec_read_bytes_.load(); s.spec_experts = spec_experts_.load(); @@ -1162,7 +1208,7 @@ void ExpertStreamSource::shutdown() { // contract as the slot and LRU buffers freed above. dense_.shutdown(); dense_tensors_.clear(); - reader_.close(); // closes the lane fds and frees the bounces + readers_.clear(); // closes every shard's lane fds and frees the bounces jobs_.clear(); layers_.clear(); active_ = false; diff --git a/core/src/moe/expert_stream_source.h b/core/src/moe/expert_stream_source.h index b88a64c..e997068 100644 --- a/core/src/moe/expert_stream_source.h +++ b/core/src/moe/expert_stream_source.h @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -38,8 +39,9 @@ namespace bmoe { // One expert weight tensor to rebind, with where its data lives in the gguf. struct ExpertTensorRef { ggml_tensor * tensor = nullptr; // persistent weight tensor whose ->data we rebind - uint64_t file_off = 0; // absolute file offset of the tensor's data + uint64_t file_off = 0; // byte offset of the tensor's data within its shard uint64_t nb2 = 0; // bytes per expert (== tensor->nb[2]) + int file_idx = 0; // which shard file holds the bytes (0 for a single-file model) }; // The expert weight tensors of one MoE layer, one per recipe suffix slot. The split @@ -59,10 +61,13 @@ public: ExpertStreamSource & operator=(const ExpertStreamSource &) = delete; // Build buffers, rebind the bound layers' expert tensors onto them, and start the - // I/O pool. `layers` is indexed by layer id (unbound entries are skipped). - // Returns false on any allocation/open failure or an inconsistent tensor. - bool - init(const std::string & gguf_path, int n_expert, std::vector layers, const MoeStreamConfig & cfg); + // I/O pool. `layers` is indexed by layer id (unbound entries are skipped); each tensor's + // file_idx indexes `shard_paths` (a single-file model passes one path). Returns false on + // any allocation/open failure or an inconsistent tensor. + bool init(const std::vector & shard_paths, + int n_expert, + std::vector layers, + const MoeStreamConfig & cfg); // Supply the dense (non-expert) weight tensors the DenseWeights policy may need (only the // Anonymous mode reads+rebinds them). Call BEFORE init, which hands them to the dense module. @@ -109,6 +114,10 @@ private: uint64_t off = 0; uint64_t nbytes = 0; int32_t flag = -1; // overlap: index into ready_ to publish on completion; -1 = serial + // Which shard reader serves this read. Wide enough for the whole filename format + // (-%05d-of-%05d): an int8_t would wrap a 200-shard model into a negative index that the + // init-time bounds check, which sees the untruncated value, could never catch. + int16_t file = 0; // Which (layer, expert, projection) this read serves. Known at every enqueue site and // otherwise thrown away; carried so the I/O trace can attribute a read without the read // path having to guess. Inert unless the trace is on. @@ -187,10 +196,11 @@ private: int n_expert_ = 0; size_t align_ = 4096; - // The positioned reader that owns the fd pool, bounces and O_DIRECT decision. Expert slices are - // read through it; the dense-weights loader constructs its own, so their O_DIRECT choices are - // independent (see docs/architecture.md). file_size() replaces the old fsize_ member. - FileReader reader_; + // The positioned readers that own the fd pools, bounces and O_DIRECT decision — one per shard + // file, in shard order (a single-file model has exactly one). Expert slices are read through + // them; the dense-weights loader constructs its own, so their O_DIRECT choices are independent + // (see docs/architecture.md). FileReader is not movable, hence the unique_ptr. + std::vector> readers_; std::vector layers_; diff --git a/core/src/moe/gguf_offsets.cpp b/core/src/moe/gguf_offsets.cpp index c21dfa8..92876f1 100644 --- a/core/src/moe/gguf_offsets.cpp +++ b/core/src/moe/gguf_offsets.cpp @@ -2,6 +2,7 @@ #include "gguf.h" +#include #include namespace bmoe { @@ -44,17 +45,15 @@ gguf_context * open_meta(const char * path) { return gguf_init_from_file(path, gp); } -void fill_offsets(const gguf_context * gctx, GgufOffsets & out) { +void fill_offsets(const gguf_context * gctx, int file_idx, GgufOffsets & out) { const uint64_t data_off = (uint64_t) gguf_get_data_offset(gctx); const int64_t n = gguf_get_n_tensors(gctx); - out.off_by_name.reserve((size_t) n); - out.size_by_name.reserve((size_t) n); for (int64_t i = 0; i < n; ++i) { const char * name = gguf_get_tensor_name(gctx, i); out.off_by_name[name] = data_off + (uint64_t) gguf_get_tensor_offset(gctx, i); out.size_by_name[name] = (uint64_t) gguf_get_tensor_size(gctx, i); + out.file_by_name[name] = file_idx; } - out.ok = true; } void fill_model_info(const gguf_context * gctx, GgufModelInfo & out) { @@ -71,12 +70,75 @@ void fill_model_info(const gguf_context * gctx, GgufModelInfo & out) { out.ok = true; } +// The llama.cpp split filename convention: `-%05d-of-%05d.gguf`, 1-based. If +// `path` matches, returns the shard count and sets `prefix` (everything before the +// shard number); the caller rebuilds each sibling with make_shard_path. Returns 0 for +// a plain single-file path. The FIRST shard must be the one asked for — the same rule +// llama.cpp applies — so downstream shard indices line up with llama.cpp's own mmaps. +int split_count_from_path(const std::string & path, std::string & prefix) { + // -DDDDD-of-DDDDD.gguf → 19 chars after the separating dash + static const size_t tail_len = 5 + 4 + 5 + 5; // "DDDDD" "-of-" "DDDDD" ".gguf" + if (path.size() < tail_len + 1) return 0; + const size_t tail = path.size() - tail_len; + auto five_digits = [&](size_t at) { + for (size_t i = at; i < at + 5; ++i) + if (path[i] < '0' || path[i] > '9') return false; + return true; + }; + if (path[tail - 1] != '-' || !five_digits(tail) || path.compare(tail + 5, 4, "-of-") != 0 || + !five_digits(tail + 9) || path.compare(tail + 14, 5, ".gguf") != 0) + return 0; + const int no = std::stoi(path.substr(tail, 5)); + const int count = std::stoi(path.substr(tail + 9, 5)); + if (count <= 1) return 0; // a -00001-of-00001 file is just a single file + if (no != 1) { + std::fprintf(stderr, "bmoe: %s is shard %d of a split model — pass the first shard (-00001-of-)\n", + path.c_str(), no); + return -1; + } + prefix = path.substr(0, tail - 1); + return count; +} + +std::string make_shard_path(const std::string & prefix, int no, int count) { + char tail[32]; + std::snprintf(tail, sizeof(tail), "-%05d-of-%05d.gguf", no, count); + return prefix + tail; +} + +// Offsets for `path` and, if it is the first file of a split set, for every sibling +// shard. Any missing or unparsable shard fails the whole read: a partial map would +// surface later as "no gguf offset for tensor X", which points at the wrong culprit. +bool fill_all_offsets(const char * path, GgufOffsets & out, gguf_context * first) { + std::string prefix; + const int count = split_count_from_path(path, prefix); + if (count < 0) return false; + + fill_offsets(first, 0, out); + out.shard_paths.push_back(path); + + for (int no = 2; no <= count; ++no) { + const std::string sp = make_shard_path(prefix, no, count); + gguf_context * gctx = open_meta(sp.c_str()); + if (!gctx) { + std::fprintf(stderr, "bmoe: split model is missing shard %s\n", sp.c_str()); + out = GgufOffsets{}; + return false; + } + fill_offsets(gctx, no - 1, out); + out.shard_paths.push_back(sp); + gguf_free(gctx); + } + out.ok = true; + return true; +} + } // namespace GgufOffsets read_gguf_offsets(const char * path) { GgufOffsets out; if (gguf_context * gctx = open_meta(path)) { - fill_offsets(gctx, out); + fill_all_offsets(path, out, gctx); gguf_free(gctx); } return out; @@ -94,10 +156,10 @@ GgufModelInfo read_gguf_model_info(const char * path) { GgufMeta read_gguf_meta(const char * path) { GgufMeta out; if (gguf_context * gctx = open_meta(path)) { - fill_offsets(gctx, out.offsets); + const bool offs_ok = fill_all_offsets(path, out.offsets, gctx); fill_model_info(gctx, out.info); gguf_free(gctx); - out.ok = true; + out.ok = offs_ok; } return out; } diff --git a/core/src/moe/gguf_offsets.h b/core/src/moe/gguf_offsets.h index 3d0d8c3..135ad08 100644 --- a/core/src/moe/gguf_offsets.h +++ b/core/src/moe/gguf_offsets.h @@ -2,26 +2,42 @@ // API. The expert streamer needs these offsets to pread individual expert slices; it // gets the tensor pointers (to rebind) from the graph, and the file layout from here. // The two are matched by tensor name. +// +// A model may ship as several shard files (`model-00001-of-00003.gguf`, the llama.cpp +// split convention — Hugging Face rejects single files above 50 GB, so every large model +// arrives this way). llama.cpp loads the whole set when handed the first shard; here the +// same first-shard path fans out to one parse per shard, and every tensor resolves to +// (shard index, offset in that shard) instead of an offset in one file. #pragma once #include #include #include +#include namespace bmoe { struct GgufOffsets { - // tensor name -> absolute file offset of its data (data_offset + per-tensor offset) + // tensor name -> byte offset of its data (data_offset + per-tensor offset) WITHIN ITS SHARD std::unordered_map off_by_name; // tensor name -> its data size in bytes. The streamer does not need this (an expert slice // is strided by nb2, read from the graph's tensor), but the route trace does: it is the // only way to say how many bytes of a layer are dense, i.e. left mmap-resident. std::unordered_map size_by_name; + // tensor name -> index into shard_paths. Always filled; 0 for every tensor of a + // single-file model, so consumers index shard_paths unconditionally. + std::unordered_map file_by_name; + // The files, in shard order. A single-file model is the degenerate case: one entry, + // the path that was asked for. Never empty when ok. + std::vector shard_paths; + // Per shard, the file's total data size is not knowable from the gguf header alone; + // consumers that need it (dense byte ranges) take it from the opened file instead. bool ok = false; }; // Parse `path` with no_alloc (metadata only, no tensor data read into RAM) and collect -// the offset of every tensor. Returns ok=false if the file cannot be opened as gguf. +// the offset of every tensor — across every shard when `path` is the first file of a +// split set. Returns ok=false if any file cannot be opened as gguf or a shard is missing. GgufOffsets read_gguf_offsets(const char * path); // The handful of model metadata needed BEFORE the model is loaded: the architecture (to @@ -37,10 +53,12 @@ struct GgufModelInfo { // Peek `path`'s metadata (no_alloc, no tensor bytes) for the fields above. Returns ok=false // if the file cannot be opened as gguf; a present file with a missing key leaves that field 0. +// Metadata lives in the first shard, so this never needs the rest of a split set. GgufModelInfo read_gguf_model_info(const char * path); -// Both of the above from ONE parse. gguf_init_from_file walks the whole KV section even with -// no_alloc, so a caller that needs offsets and model info should not pay that walk twice. +// Both of the above from ONE parse of the first file. gguf_init_from_file walks the whole +// KV section even with no_alloc, so a caller that needs offsets and model info should not +// pay that walk twice. Further shards of a split set are parsed once each for offsets. struct GgufMeta { GgufOffsets offsets; GgufModelInfo info; diff --git a/docs/architecture.md b/docs/architecture.md index 2360305..3e690f1 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -21,7 +21,8 @@ core/ src/ io/ platform_io — O_DIRECT reads + reserve/commit/evict VM, cross-platform file_reader — pooled positioned reader, per-consumer O_DIRECT - moe/ gguf_offsets, arch_registry, expert_stream_source, router_hook + moe/ gguf_offsets (tensor → (shard, offset), split ggufs included), arch_registry, + expert_stream_source (one reader per shard), router_hook dense_weights — non-expert weight policy + the residency sensor engine/ session — composition + the generation loop (open/generate/close) runtime — the one-shot run() wrapper over a Session diff --git a/docs/limitations.md b/docs/limitations.md index 89030dd..328edeb 100644 --- a/docs/limitations.md +++ b/docs/limitations.md @@ -30,7 +30,7 @@ serial path, and only a single ~25-line hook (with an explicit sunset) for the o GPU offload of the streamed experts is not supported (the dense parts can still use the GPU). Decode is flash-I/O-bound anyway, so this is rarely the bottleneck. - **Shared experts stay resident.** Architectures with an always-on shared expert (e.g. - `gemma4`) stream the routed experts but keep the shared expert — and any dense layers — + `gemma4`, `deepseek4`) stream the routed experts but keep the shared expert — and any dense layers — resident (in the page cache, or in the engine's own buffers under `--dense-weights anon`), so the streamed fraction (and the memory saving) is smaller than for a purely routed model like `qwen3moe`. The same applies to architectures whose first blocks are dense by design diff --git a/scripts/make-tiny-moe.py b/scripts/make-tiny-moe.py index 07e0b43..0d8992c 100644 --- a/scripts/make-tiny-moe.py +++ b/scripts/make-tiny-moe.py @@ -87,11 +87,29 @@ def add_attn_tensors(w, p, s): # --- qwen3moe: split expert layout ------------------------------------------------- -def build_qwen3moe(out): +def make_writer(out, arch, split_max_tensors): + """A plain writer, or a sharding one when --split-max-tensors is set. + + Sharded output mirrors how real >50 GB models arrive from Hugging Face: the writer + emits -%05d-of-%05d.gguf siblings, with a metadata-only first shard + (small_first_shard, the layout unsloth ships). The byte-identity gates then prove the + multi-shard streaming path against the same tensors the single-file fixture uses. + """ + if not split_max_tensors: + return gguf.GGUFWriter(out, arch) + try: + return gguf.GGUFWriter(out, arch, + split_max_tensors=split_max_tensors, + small_first_shard=True) + except TypeError: + raise SystemExit("this gguf package cannot write split files: pip install -U gguf") + + +def build_qwen3moe(out, split_max_tensors=0): tokens, scores, toktypes = build_vocab() n_vocab = len(tokens) - w = gguf.GGUFWriter(out, "qwen3moe") + w = make_writer(out, "qwen3moe", split_max_tensors) w.add_name("tiny-moe") w.add_context_length(N_CTX) w.add_embedding_length(N_EMBD) @@ -223,12 +241,16 @@ def main(): ap = argparse.ArgumentParser() ap.add_argument("--arch", choices=["qwen3moe", "gemma4"], default="qwen3moe") ap.add_argument("--out", default="tiny-moe.gguf") + ap.add_argument("--split-max-tensors", type=int, default=0, + help="emit a sharded gguf (N tensors per shard, metadata-only first shard)") args = ap.parse_args() if args.arch == "gemma4": + if args.split_max_tensors: + raise SystemExit("--split-max-tensors is exercised via the qwen3moe fixture only") build_gemma4(args.out) else: - build_qwen3moe(args.out) + build_qwen3moe(args.out, args.split_max_tensors) if __name__ == "__main__": diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 8d154ca..8e213ad 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -53,6 +53,22 @@ if(PYTHON3) add_test(NAME moe_gates_${spec} COMMAND bmoe_moe_gates ${_model}) set_tests_properties(moe_gates_${spec} PROPERTIES FIXTURES_REQUIRED ${_fixture}) endforeach() + + # The same qwen3moe fixture, written as a 4-shard split gguf (metadata-only first shard — + # the layout large models actually ship in, since Hugging Face rejects files above 50 GB). + # 51 tensors at 20 per shard is deterministic: -00001-of-00004 .. -00004-of-00004. The gate + # binary gets the FIRST shard, exactly as a user would pass it; the streamer must resolve + # every expert tensor to its (shard, offset) and produce byte-identical output. + add_test( + NAME moe_make_tiny_model_split + COMMAND ${PYTHON3} ${CMAKE_SOURCE_DIR}/scripts/make-tiny-moe.py --arch qwen3moe + --split-max-tensors 20 --out "${CMAKE_CURRENT_BINARY_DIR}/tiny-moe-split.gguf" + ) + set_tests_properties(moe_make_tiny_model_split PROPERTIES FIXTURES_SETUP TINY_MOE_MODEL_split) + + add_test(NAME moe_gates_split + COMMAND bmoe_moe_gates "${CMAKE_CURRENT_BINARY_DIR}/tiny-moe-split-00001-of-00004.gguf") + set_tests_properties(moe_gates_split PROPERTIES FIXTURES_REQUIRED TINY_MOE_MODEL_split) else() message(WARNING "python3 not found: MoE gates will not run (models cannot be generated)") endif()