diff --git a/server/docs/QWEN4EXP.md b/server/docs/QWEN4EXP.md index b78085a1c..c72db8ac9 100644 --- a/server/docs/QWEN4EXP.md +++ b/server/docs/QWEN4EXP.md @@ -103,7 +103,9 @@ UD-Q4_K_XL) on the expert device: ``` The profile sets `--target-device hip:0` (the R9700, gfx1201), -`--expert-device hip:1` (the Strix Halo, gfx1151) and `--max-ctx 131072`. +`--expert-device hip:1` (the Strix Halo, gfx1151) and `--max-ctx auto`: the +model's 262,144-token context, lowered in steps of 4096 to what the R9700 +holds beside the dense weights and every slot's cache. Both GPUs run one scheduler: prompt chunks pipeline over 2 to 4 streams so the Strix runs one stream's experts while the R9700 runs another's attention. @@ -141,9 +143,11 @@ a 16K prompt at 44 tok/s. Greedy output is the same with adaptive MTP, a fixed `--verify-width` or MTP off (`--verify-width 1`): every generated token runs its hot experts on the R9700 with the same kernel. -Split mode serves one sequence: `--max-concurrency` above 1 is refused with -`--expert-device`. The prefix cache works as on the Strix; the snapshot -allowance comes out of the R9700's memory after an 8192-row prompt chunk. +With `--max-concurrency 2..4` split mode serves several requests at once. +Prefix snapshots live in system memory beside the R9700, so its memory holds +context; the server's resident limit bounds them. On the Strix alone they stay +in its memory, and `--max-ctx auto` keeps room for one full-context snapshot +per slot. ## Support @@ -154,9 +158,9 @@ allowance comes out of the R9700's memory after an 8192-row prompt chunk. | Attention block ratios other than 4, contexts of 2^24 tokens or more | refused at load | | gfx1151 kernels: MMB bf16 and Q8_0 -> F16 WMMA GEMMs, fused HC / GDN / PLE, M-RoPE into the flash-attention layout | done | | Chat template, reasoning effort, thinking budget, `preserve_thinking`, `sampling_no_thinking` | done | -| Concurrent serving (`--max-concurrency 2..4`) | supported on one GPU, exact independent slots with full per-slot caches; MTP and the prefix cache are single-slot only (MTP is switched off); more than 4, or split mode, is refused | +| Concurrent serving (`--max-concurrency 2..4`) | supported on one GPU and in split mode, exact independent slots with full per-slot caches; MTP drafts for a request running alone; the prefix cache serves every slot; more than 4 is refused | | MTP speculative decoding | sidecar discovered automatically; adaptive k=1..7 by default (code 16K / 64K 32.4 / 28.2 tok/s, counting 43.5), output identical to MTP off; `--verify-width 1` disables, `2..8` selects fixed k=1..7 | -| Prefix cache | done: snapshots at chat cut points, restored on hits; a 64K agent turn's first token in ~3.5 s instead of ~65 s (131K context). Prefill keeps a 4096-row chunk first (8192 in split mode) and snapshots get the remaining memory, so at 262K context with MTP long prefixes do not fit: use `--max-ctx 131072` or less for agent workloads | +| Prefix cache | done: snapshots at chat cut points, restored on hits; a 64K agent turn's first token in ~3.5 s instead of ~65 s (131K context). Prefill keeps a 4096-row chunk first (8192 in split mode). Snapshots share the Strix's memory with the caches (`--max-ctx auto` leaves one full-context snapshot per slot) and live in system memory beside a discrete GPU | | Layer split | refused | | Split mode (`--expert-device`, `--expert-placement`) | routed experts on a second GPU, pipelined prompt chunks, hot experts on the target for decode and verify; measured on R9700 + Strix Halo | | Other GPUs | generic paths; kernels, defaults and quality gates are tuned and measured on gfx1151, and on gfx1201 as the split-mode target | diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 9856b9e64..296ec68c4 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -158,6 +158,9 @@ struct ModelBackend { // Print the "[-daemon] ready ..." banner on stdout. virtual void print_ready_banner() const = 0; virtual int prefill_chunk_size() const { return 0; } + // The context one sequence can hold, from a backend that can fit it to + // its devices (DevicePlacement::fit_ctx); 0 from one that cannot. + virtual int context_limit() const { return 0; } // ── Park / unpark ──────────────────────────────────────────────── // Backend decides which resources to release/restore. Returns true on diff --git a/server/src/placement/placement_config.h b/server/src/placement/placement_config.h index a2574a4f2..8ea71038d 100644 --- a/server/src/placement/placement_config.h +++ b/server/src/placement/placement_config.h @@ -51,6 +51,10 @@ struct DevicePlacement { bool peer_access = false; // enable CUDA/HIP peer access between GPUs int max_ctx = 8192; // max KV cache context length + // --max-ctx auto: max_ctx is the model's trained context, and a backend + // that sizes its own cache lowers it to what its devices hold + // (ModelBackend::context_limit reports the result). + bool fit_ctx = false; bool is_multi_device() const { return layer_split_gpus.size() > 1; } bool is_layer_split() const { diff --git a/server/src/qwen4exp/qwen4exp_backend.cpp b/server/src/qwen4exp/qwen4exp_backend.cpp index aa26c5b36..2f3df2410 100644 --- a/server/src/qwen4exp/qwen4exp_backend.cpp +++ b/server/src/qwen4exp/qwen4exp_backend.cpp @@ -4,6 +4,7 @@ #include "common/peer_access.h" #include "common/sampler.h" +#include "common/snapshot_backend.h" #include "ggml-cuda.h" @@ -99,7 +100,7 @@ static bool load_qwen4exp_hot_experts(ggml_backend_t backend, Qwen4ExpWeights & } int qwen4exp_select_chunk(ggml_backend_t backend, const Qwen4ExpWeights & w, - Qwen4ExpCache & cache, int slots, int resident_slots, size_t * snapshot_budget) { + Qwen4ExpCache & cache, int slots, int resident_slots, size_t * snapshot_budget, Qwen4ExpChunkPlan * plan) { if (slots < 1 || resident_slots < 1 || resident_slots > slots) return 0; size_t free_device = 0, total = 0, free_host = 0; ggml_backend_cuda_get_device_memory(ggml_backend_cuda_get_device_id(backend), &free_device, &total); @@ -200,7 +201,7 @@ int qwen4exp_select_chunk(ggml_backend_t backend, const Qwen4ExpWeights & w, : w.hot ? kQwen4ExpSplitChunkFloor : kQwen4ExpSplitMaxChunk; struct Plan { size_t graph = 0, ring = 0, host = 0, scratch = 0; }; std::vector> plans; // one per probed chunk size - const int chunk = qwen4exp_fit_chunk(cache.max_ctx, available, fixed, [&](int n) { + auto chunk_workspace = [&](int n) { size_t graph = 0, inputs = 0, mask = 0, host = runtime_host, scratch = runtime_scratch; // End-of-context QSA workspace, and the largest dense span before QSA. // Include an unaligned context tail, which can fall back to dense FA. @@ -229,7 +230,23 @@ int qwen4exp_select_chunk(ggml_backend_t backend, const Qwen4ExpWeights & w, std::fprintf(stderr, "[qwen4exp] chunk-plan rows=%d graph=%zu ring=%zu host=%zu scratch=%zu required=%zu available=%zu\n", n, graph, slots * (inputs + mask), host, scratch, fixed + workspace, available); return workspace; - }, snapshot_budget, w.expert_backend ? kQwen4ExpSplitChunkFloor : 4096, max_rows); + }; + const int floor_rows = w.expert_backend ? kQwen4ExpSplitChunkFloor : 4096; + if (plan) { + const size_t floor = chunk_workspace(std::min({floor_rows, max_rows, cache.max_ctx})); + plan->available = available; + plan->required = floor == SIZE_MAX ? SIZE_MAX : fixed + floor + available / 10; + // Every slot's attention and indexer rows; only this cache adds the MTP draft layer's. + size_t rows = 0; + for (const auto * group : {&cache.attn_k, &cache.attn_v, &cache.indexer_raw, &cache.indexer_k}) { + for (const ggml_tensor * t : *group) if (t) rows += ggml_nbytes(t); + } + const size_t mtp_rows = (cache.mtp_k ? ggml_nbytes(cache.mtp_k) : 0) + (cache.mtp_v ? ggml_nbytes(cache.mtp_v) : 0); + plan->position_bytes = ((size_t) slots * rows + mtp_rows) / cache.max_ctx + 1; + if (plan->required > available) return 0; // the context fit moves on without searching smaller chunks + } + const int chunk = qwen4exp_fit_chunk(cache.max_ctx, available, fixed, chunk_workspace, snapshot_budget, + floor_rows, max_rows); std::fprintf(stderr, "[qwen4exp] chunk-auto ctx=%d slots=%d resident=%d chunk=%d state=%zu fixed=%zu available=%zu headroom=%zu\n", cache.max_ctx, slots, resident_slots, chunk, state, fixed + (snapshot_budget ? *snapshot_budget : 0), available, available / 10); @@ -264,6 +281,10 @@ bool Qwen4ExpBackend::init() { cfg_.device.gpu); return false; } + // Prefix snapshots stay in the gfx1151 iGPU's memory, which is the system's. Beside a GPU with its own memory + // they go to system memory (create_snapshot_backend), which leaves that memory to the KV caches. + snap_backend_ = ggml_backend_cuda_qwen4exp_supported(backend_) ? backend_ : create_snapshot_backend(backend_); + if (!snap_backend_) return false; if (!cfg_.expert_placement_path.empty() && !cfg_.expert_device) { std::fprintf(stderr, "[qwen4exp] --expert-placement needs --expert-device\n"); return false; @@ -304,24 +325,50 @@ bool Qwen4ExpBackend::load_target() { if (!load_qwen4exp_hot_experts(backend_, weights_, cfg_.max_concurrency > 1 ? std::string() : cfg_.expert_placement_path, expert_budget_bytes_from_env(kQwen4ExpHotExpertBudget))) return false; - snapshot_budget_ = SIZE_MAX; - const int explicit_chunk = weights_.expert_backend ? std::min(cfg_.chunk, kQwen4ExpSplitMaxChunk) : cfg_.chunk; - if (cfg_.max_concurrency > 1) { - // Prefix checkpoints: each slot's restore point and capture in flight - // plus one shared head, the server's concurrent prefix budget. The - // planner caps the prefill granule; an explicit --chunk sets it, and the - // allowance is measured either way. - snapshot_budget_ = (2 * (size_t) cfg_.max_concurrency + 1) * snapshot_bytes_estimate(cache_.max_ctx); - const int fit = qwen4exp_select_chunk(backend_, weights_, cache_, cfg_.max_concurrency, 1, &snapshot_budget_); - chunk_ = cfg_.chunk > 0 && fit > 0 ? explicit_chunk : fit; - std::fprintf(stderr, "[qwen4exp] prefix checkpoint allowance=%zu bytes\n", snapshot_budget_); - } else if (cfg_.chunk > 0) { - chunk_ = explicit_chunk; - } else { - snapshot_budget_ = 3 * snapshot_bytes_estimate(cache_.max_ctx); - chunk_ = qwen4exp_select_chunk(backend_, weights_, cache_, 1, 1, &snapshot_budget_); - std::fprintf(stderr, "[qwen4exp] prefix snapshot allowance=%zu bytes\n", snapshot_budget_); + // --max-ctx auto: the largest context, up to the trained one, at which every slot's cache, the resident graphs, a + // floor-sized prompt chunk and, with snapshots on the target, one full snapshot per slot fit + // (Qwen4ExpContextFit). The context found stays: an unpark sizes itself for it. + for (Qwen4ExpContextFit search;;) { + Qwen4ExpChunkPlan plan; + chunk_ = plan_chunk(cfg_.device.fit_ctx ? &plan : nullptr); + if (!cfg_.device.fit_ctx) break; + const int ctx = cache_.max_ctx; + const size_t snapshots = snap_backend_ == backend_ + ? (size_t) cfg_.max_concurrency * snapshot_bytes_estimate(ctx) : 0; + const size_t position = plan.position_bytes + snapshots / ctx; + const bool fit = plan.required <= plan.available && snapshots <= plan.available - plan.required; + if (fit) { + search.fits = ctx; + search.spare = plan.available - plan.required - snapshots; + } else { + search.misses = ctx; + search.missing = plan.required == SIZE_MAX ? (size_t) ctx / 4 * position + : plan.required + snapshots - plan.available; + } + int next = search.next(position); + if (!next && !search.fits) { + std::fprintf(stderr, "[qwen4exp] no context fits beside the weights\n"); + return false; + } + if (!next) { + cfg_.device.fit_ctx = false; + std::fprintf(stderr, "[qwen4exp] context fitted to the devices: %d tokens, %d at once\n", + search.fits, cfg_.max_concurrency); + if (fit) break; + next = search.fits; // the last plan missed: plan the fit again + } + free_qwen4exp_cache(cache_); + cache_ = {}; + cfg_.device.max_ctx = next; + if (!create_main_cache()) { + std::fprintf(stderr, "[qwen4exp] cache creation failed\n"); + return false; + } } + if (snapshot_budget_ != SIZE_MAX) + std::fprintf(stderr, "[qwen4exp] prefix snapshot allowance=%zu bytes\n", snapshot_budget_); + else if (snap_backend_ != backend_) + std::fprintf(stderr, "[qwen4exp] prefix snapshots in system memory\n"); if (chunk_ <= 0) { std::fprintf(stderr, "[qwen4exp] insufficient prefill memory at the configured context\n"); return false; @@ -333,6 +380,27 @@ bool Qwen4ExpBackend::load_target() { return true; } +// The prompt chunk, measured against free memory unless --chunk sets it (concurrent serving measures either way, +// to size the prefix allowance; so does a context fit, which reads `plan`). +int Qwen4ExpBackend::plan_chunk(Qwen4ExpChunkPlan * plan) { + snapshot_budget_ = SIZE_MAX; + // Snapshots in the target's memory get an allowance the planner sizes and captures respect. In system memory + // the server's resident limit bounds them. + auto allowance = [this](size_t snapshots) -> size_t * { + if (snap_backend_ != backend_) return nullptr; + snapshot_budget_ = snapshots * snapshot_bytes_estimate(cache_.max_ctx); + return &snapshot_budget_; + }; + const int explicit_chunk = weights_.expert_backend ? std::min(cfg_.chunk, kQwen4ExpSplitMaxChunk) : cfg_.chunk; + if (cfg_.chunk > 0 && cfg_.max_concurrency == 1 && !plan) return explicit_chunk; + // Prefix checkpoints with concurrent slots: each slot's restore point and capture in flight plus one shared + // head, the server's concurrent prefix budget. One sequence keeps three. + const int slots = cfg_.max_concurrency; + const int fit = qwen4exp_select_chunk(backend_, weights_, cache_, slots, 1, + allowance(slots > 1 ? 2 * (size_t) slots + 1 : 3), plan); + return cfg_.chunk > 0 && fit > 0 ? explicit_chunk : fit; +} + void Qwen4ExpBackend::release_target() { for (int i = 0; i < kMaxSlots; ++i) snapshot_free(i); tokens_.clear(); logits_.clear(); @@ -348,7 +416,8 @@ void Qwen4ExpBackend::release_target() { // Slot 0 of the concurrency engine; with concurrent slots its recurrent state is // the first slab of the slot states the multi-slot decode graph batches over. bool Qwen4ExpBackend::create_main_cache() { - if (cfg_.max_concurrency > 1 && + // The slot states do not depend on the context, so a context fit's rebuilt cache keeps them. + if (cfg_.max_concurrency > 1 && !slot_states_.buf && !create_qwen4exp_slot_states(backend_, weights_, cfg_.max_concurrency, slot_states_)) return false; return create_qwen4exp_cache(backend_, weights_, cfg_.device.max_ctx, cache_, /*mtp=*/true, cfg_.verify_width == 0 ? QWEN4EXP_MTP_MAX_DRAFT : std::max(1, cfg_.verify_width - 1), @@ -362,11 +431,12 @@ bool Qwen4ExpBackend::start_seq_engine() { for (Qwen4ExpCache & cache : seq_caches_) { if (!create_qwen4exp_cache(backend_, weights_, cfg_.device.max_ctx, cache, false, 1, &slot_states_, (int) caches.size())) return false; + cache.split_owner = &cache_; caches.push_back(&cache); } seq_engine_ = std::make_unique( backend_, weights_, std::move(caches), cfg_.device.max_ctx, chunk_, snapshot_budget_, - cfg_.verify_width, &mtp_width_, &mtp_costs_); + cfg_.verify_width, &mtp_width_, &mtp_costs_, snap_backend_); std::fprintf(stderr, "[qwen4exp-seq] independent-slot engine enabled: %d full F16 caches, ctx=%d, chunk=%d\n", cfg_.max_concurrency, cfg_.device.max_ctx, chunk_); @@ -701,10 +771,10 @@ bool Qwen4ExpBackend::snapshot_save_replacing(int slot, int source) { } if (!snapshot_fits(slot)) return false; auto & s = snapshots_[slot]; - if (!save_qwen4exp_snapshot(backend_, cache_, s)) return false; + if (!save_qwen4exp_snapshot(backend_, snap_backend_, cache_, s)) return false; s.tokens = tokens_; s.logits = logits_; - std::fprintf(stderr, "[qwen4exp-snap] slot=%d pos=%d device_bytes=%zu\n", + std::fprintf(stderr, "[qwen4exp-snap] slot=%d pos=%d bytes=%zu\n", slot, s.cur_pos, ggml_backend_buffer_get_size(s.buf)); return true; } @@ -742,12 +812,12 @@ int Qwen4ExpBackend::snapshot_cur_pos(int slot) const { } size_t Qwen4ExpBackend::snapshot_bytes_estimate(int tokens) const { - if (!backend_ || !cache_.buf) return 0; + if (!snap_backend_ || !cache_.buf) return 0; const int n = std::clamp(tokens, 0, cache_.max_ctx); if (!n) return 0; size_t host = 0; - const size_t device = qwen4exp_snapshot_bytes(backend_, cache_, n, &host); - return device + host + (size_t(n) + weights_.n_vocab + + const size_t copies = qwen4exp_snapshot_bytes(snap_backend_, cache_, n, &host); + return copies + host + (size_t(n) + weights_.n_vocab + std::max(0, weights_.ple_ngram_size - 1)) * sizeof(int32_t); } @@ -803,6 +873,8 @@ void Qwen4ExpBackend::free_drafter() {} void Qwen4ExpBackend::shutdown() { if (backend_) release_target(); + free_snapshot_backend(snap_backend_, backend_); + snap_backend_ = nullptr; if (weights_.expert_backend) { ggml_backend_free(weights_.expert_backend); weights_.expert_backend = nullptr; diff --git a/server/src/qwen4exp/qwen4exp_backend.h b/server/src/qwen4exp/qwen4exp_backend.h index 3402efa22..7f132fa69 100644 --- a/server/src/qwen4exp/qwen4exp_backend.h +++ b/server/src/qwen4exp/qwen4exp_backend.h @@ -23,6 +23,8 @@ namespace luce::common { +struct Qwen4ExpChunkPlan; + struct Qwen4ExpBackendConfig { std::string model_path; std::optional draft_path; // MTP sidecar; absent = auto-discover @@ -49,6 +51,7 @@ class Qwen4ExpBackend final : public ModelBackend { // ModelBackend interface void print_ready_banner() const override; int prefill_chunk_size() const override { return chunk_; } + int context_limit() const override { return cache_.max_ctx; } bool park(ParkTarget target) override; bool unpark(ParkTarget target) override; @@ -89,12 +92,14 @@ class Qwen4ExpBackend final : public ModelBackend { bool snapshot_save_replacing(int slot, int source); bool snapshot_fits(int slot, int replaced = -1) const; bool create_main_cache(); + int plan_chunk(Qwen4ExpChunkPlan * plan = nullptr); bool start_seq_engine(); // no-op at --max-concurrency 1 // Weights, caches, hot experts, the prefill chunk and the prefix allowance, then the slot engine: init and // unpark share it, so an unpark sizes itself for the memory free at that time. release_target() undoes it. bool load_target(); void release_target(); - size_t snapshot_budget_ = SIZE_MAX; // auto chunk reserves and enforces this allowance + ggml_backend_t snap_backend_ = nullptr; // prefix snapshot storage: the target, or system memory (see init) + size_t snapshot_budget_ = SIZE_MAX; // with snapshots on the target, auto chunk reserves and enforces this allowance std::array snapshots_; int live_slot_ = -1; std::vector tokens_; diff --git a/server/src/qwen4exp/qwen4exp_cache.cpp b/server/src/qwen4exp/qwen4exp_cache.cpp index 993dc98d9..d0e5bb7b1 100644 --- a/server/src/qwen4exp/qwen4exp_cache.cpp +++ b/server/src/qwen4exp/qwen4exp_cache.cpp @@ -268,6 +268,7 @@ void free_qwen4exp_cache(Qwen4ExpCache & c) { if (c.split_sched) { ggml_backend_sched_free(c.split_sched); c.split_sched = nullptr; } if (c.split_sched_short) { ggml_backend_sched_free(c.split_sched_short); c.split_sched_short = nullptr; } if (c.split_cpu) { ggml_backend_free(c.split_cpu); c.split_cpu = nullptr; } + c.split_owner = nullptr; if (c.input_ring.buf) { ggml_backend_buffer_free(c.input_ring.buf); c.input_ring.buf = nullptr; @@ -352,11 +353,11 @@ void snapshot_strips(const Qwen4ExpCache & c, int pos, int blocks, bool mtp, F v } } // namespace -size_t qwen4exp_snapshot_bytes(ggml_backend_t backend, const Qwen4ExpCache & c, int tokens, size_t * host_bytes) { +size_t qwen4exp_snapshot_bytes(ggml_backend_t store, const Qwen4ExpCache & c, int tokens, size_t * host_bytes) { if (host_bytes) *host_bytes = 0; const int pos = std::clamp(tokens, 0, c.max_ctx); if (!pos) return 0; - const size_t alignment = ggml_backend_buft_get_alignment(ggml_backend_get_default_buffer_type(backend)); + const size_t alignment = ggml_backend_buft_get_alignment(ggml_backend_get_default_buffer_type(store)); size_t bytes = 0, count = 0; // The supported QSA layout pools four rows per block. Dense prefill can // have fewer valid blocks; this bound also covers its later completion. @@ -375,7 +376,8 @@ void free_qwen4exp_snapshot(Qwen4ExpSnapshot & s) { s = {}; } -bool save_qwen4exp_snapshot(ggml_backend_t backend, const Qwen4ExpCache & c, Qwen4ExpSnapshot & s) { +bool save_qwen4exp_snapshot(ggml_backend_t backend, ggml_backend_t store, const Qwen4ExpCache & c, + Qwen4ExpSnapshot & s) { free_qwen4exp_snapshot(s); if (c.cur_pos <= 0 || (c.mtp_k && c.mtp_prev_pos != c.cur_pos - 1)) return false; size_t count = 0; @@ -389,9 +391,9 @@ bool save_qwen4exp_snapshot(ggml_backend_t backend, const Qwen4ExpCache & c, Qwe auto * copy = ggml_new_tensor_1d(s.ctx, t->type, n); s.strips.emplace_back(live, copy); }); - s.buf = ggml_backend_alloc_ctx_tensors(s.ctx, backend); + s.buf = ggml_backend_alloc_ctx_tensors(s.ctx, store); if (!s.buf) { free_qwen4exp_snapshot(s); return false; } - for (auto [live, copy] : s.strips) ggml_backend_tensor_copy_async(backend, backend, live, copy); + for (auto [live, copy] : s.strips) ggml_backend_tensor_copy_async(backend, store, live, copy); ggml_backend_synchronize(backend); s.cur_pos = c.cur_pos; s.indexer_blocks = c.indexer_blocks; diff --git a/server/src/qwen4exp/qwen4exp_cache.h b/server/src/qwen4exp/qwen4exp_cache.h index f801e407b..e1d7c9da4 100644 --- a/server/src/qwen4exp/qwen4exp_cache.h +++ b/server/src/qwen4exp/qwen4exp_cache.h @@ -186,7 +186,10 @@ struct Qwen4ExpCache { std::array mtp_rank_workspace; // Split mode (Qwen4ExpWeights::expert_backend): schedulers over the target, - // the expert device and the CPU, reused across forwards. + // the expert device and the CPU, reused across forwards. Concurrency slots + // use the first slot's (split_owner): one forward runs at a time, so one + // prompt chunk's buffers serve every slot, as the chunk planner counts them. + Qwen4ExpCache * split_owner = nullptr; // not owned; null: this cache's own ggml_backend_sched_t split_sched = nullptr; // prompt chunks ggml_backend_sched_t split_sched_short = nullptr; // decode and short batches uint64_t split_short_gen = 0; // bumped on every short-scheduler allocation @@ -209,8 +212,10 @@ void clear_qwen4exp_batched_decode_workspace(Qwen4ExpBatchedDecodeWorkspace & wo // cur_pos. KV is left intact: the next sequence overwrites it from position 0. void reset_qwen4exp_state(ggml_backend_t backend, Qwen4ExpCache & c); -// Prefix-sized device copies. Live strip views remain valid until the cache is -// freed; callers must release snapshots first. No verify scratch is retained. +// Prefix-sized copies, in the memory of the `store` backend they were saved to: +// the device's own, or system memory (create_snapshot_backend). Live strip views +// remain valid until the cache is freed; callers must release snapshots first. +// No verify scratch is retained. struct Qwen4ExpSnapshot { ggml_context * ctx = nullptr; ggml_backend_buffer_t buf = nullptr; @@ -222,9 +227,11 @@ struct Qwen4ExpSnapshot { std::vector logits; }; -size_t qwen4exp_snapshot_bytes(ggml_backend_t backend, const Qwen4ExpCache & c, int tokens, +size_t qwen4exp_snapshot_bytes(ggml_backend_t store, const Qwen4ExpCache & c, int tokens, size_t * host_bytes = nullptr); -bool save_qwen4exp_snapshot(ggml_backend_t backend, const Qwen4ExpCache & c, Qwen4ExpSnapshot & s); +// Copy `c`, which lives on `backend`, into a snapshot allocated on `store`. +bool save_qwen4exp_snapshot(ggml_backend_t backend, ggml_backend_t store, const Qwen4ExpCache & c, + Qwen4ExpSnapshot & s); // Restore into `c`, the cache `s` was saved from or another one with the same trunk layout (a // concurrency slot resuming another slot's checkpoint). The MTP draft layer's state comes along // when both caches have the layer; a cache that has it, restored from a snapshot without it, keeps diff --git a/server/src/qwen4exp/qwen4exp_chunk.h b/server/src/qwen4exp/qwen4exp_chunk.h index 97f88290f..714a924f8 100644 --- a/server/src/qwen4exp/qwen4exp_chunk.h +++ b/server/src/qwen4exp/qwen4exp_chunk.h @@ -59,6 +59,40 @@ int qwen4exp_fit_chunk(int max_ctx, size_t available, size_t fixed, Measure meas return 0; } +// --max-ctx auto: the next context to try after a plan at `ctx` lacks `missing` bytes, when one context position +// takes `position_bytes` across every slot. At least one step lower, in multiples of `step`; 0 when none is left. +inline int qwen4exp_shrink_context(int ctx, size_t missing, size_t position_bytes, int step = 4096) { + const int64_t fewer = position_bytes ? (int64_t) std::min((missing + position_bytes - 1) / position_bytes, + (size_t) ctx) + : ctx; + const int64_t next = std::min((int64_t) ctx - step, ((int64_t) ctx - fewer) / step * step); + return next >= step ? (int) next : 0; +} + +// --max-ctx auto: the search for the largest context that fits. +struct Qwen4ExpContextFit { + int fits = 0, misses = 0; // the largest context that fit and the smallest that did not (0: none yet) + size_t spare = 0, missing = 0; // the bytes the fit left over and the miss lacked + + // The next context to plan; 0 once settled (the answer is `fits`, none when it is 0). Until a context fits, a + // miss shrinks by its shortfall, which overshoots because the graph workspaces shrink with the context too. + // Then interpolation between the fit and the miss, kept strictly between them, settles to a `step` multiple. + int next(size_t position_bytes, int step = 4096) const { + if (!fits) return qwen4exp_shrink_context(misses, missing, position_bytes, step); + if (!misses || misses - fits <= step) return 0; + const double t = (double) spare / ((double) spare + (double) missing); + const int64_t guess = (int64_t) (fits + t * (misses - fits)) / step * step; + return (int) std::clamp(guess, (int64_t) fits + step, (int64_t) misses - step); + } +}; + +// What a chunk plan measured, for fitting the context to the device. +struct Qwen4ExpChunkPlan { + size_t available = 0; // device memory free beside the weights and the resident cache + size_t required = 0; // the other slots' caches, resident graphs, a floor-sized chunk and the headroom + size_t position_bytes = 0; // one context position across every slot's cache +}; + struct Qwen4ExpWeights; struct Qwen4ExpCache; // Call after weights and resident_slots caches are allocated, before prefill. @@ -68,6 +102,7 @@ struct Qwen4ExpCache; // that cannot fit). // The caller must enforce the returned allowance on later snapshot captures. int qwen4exp_select_chunk(ggml_backend_t backend, const Qwen4ExpWeights & w, - Qwen4ExpCache & cache, int slots = 1, int resident_slots = 1, size_t * snapshot_budget = nullptr); + Qwen4ExpCache & cache, int slots = 1, int resident_slots = 1, size_t * snapshot_budget = nullptr, + Qwen4ExpChunkPlan * plan = nullptr); } // namespace luce::common diff --git a/server/src/qwen4exp/qwen4exp_graph.cpp b/server/src/qwen4exp/qwen4exp_graph.cpp index 083393842..98e602387 100644 --- a/server/src/qwen4exp/qwen4exp_graph.cpp +++ b/server/src/qwen4exp/qwen4exp_graph.cpp @@ -1812,13 +1812,17 @@ static ggml_context * graph_context(Qwen4ExpDecodeWorkspace * pool, ggml_backend // Split mode: one scheduler for decode, verify and short batches, one for prompt chunks; a verify width keeps its // own short-batch scheduler (and allocation) beside its graph (`pool`). The chunk planner measures the live // split-input copies through ggml_backend_sched_reserve_size. +// The cache whose split schedulers `c` uses (Qwen4ExpCache::split_owner). +static Qwen4ExpCache & split_owner(Qwen4ExpCache & c) { return c.split_owner ? *c.split_owner : c; } + static ggml_backend_sched_t split_scheduler(ggml_backend_t backend, const Qwen4ExpWeights & w, Qwen4ExpCache & cache, Qwen4ExpDecodeWorkspace * pool, bool verify, bool short_batch) { + Qwen4ExpCache & owner = split_owner(cache); ggml_backend_sched_t & sched = verify && pool ? pool->sched - : short_batch ? cache.split_sched_short : cache.split_sched; + : short_batch ? owner.split_sched_short : owner.split_sched; if (!sched) { - if (!cache.split_cpu) cache.split_cpu = ggml_backend_cpu_init(); - ggml_backend_t backends[3] = { backend, w.expert_backend, cache.split_cpu }; + if (!owner.split_cpu) owner.split_cpu = ggml_backend_cpu_init(); + ggml_backend_t backends[3] = { backend, w.expert_backend, owner.split_cpu }; // Both are parallel: per-copy events make a split wait for its // cross-device inputs on the GPU; a single-copy scheduler instead // synchronizes the receiving device from the host before every such @@ -1887,7 +1891,7 @@ static void measure_split_graph(ggml_backend_t backend, const Qwen4ExpWeights & ggml_cgraph * gf, const std::vector & expert_nodes, int64_t T, int pos0, bool verify, Qwen4ExpGraphMemory & measure) { ggml_backend_sched_t sched = split_scheduler(backend, w, cache, nullptr, verify, T <= kSplitShortRows); - const ggml_backend_t devices[3] = { backend, w.expert_backend, cache.split_cpu }; + const ggml_backend_t devices[3] = { backend, w.expert_backend, split_owner(cache).split_cpu }; const ggml_tensor * unplaced = nullptr; for (int i = 0; i < ggml_graph_n_nodes(gf) && !unplaced; ++i) { const ggml_tensor * node = ggml_graph_node(gf, i); @@ -1907,7 +1911,7 @@ static void measure_split_graph(ggml_backend_t backend, const Qwen4ExpWeights & size_t sizes[3] = {}; ggml_backend_sched_reserve_size(sched, gf, sizes); ggml_backend_sched_reset(sched); - if (T <= kSplitShortRows) ++cache.split_short_gen; // a retained stable graph must allocate again + if (T <= kSplitShortRows) ++split_owner(cache).split_short_gen; // a retained stable graph must allocate again measure.graph = sizes[0]; } @@ -1947,23 +1951,36 @@ static bool place_inputs_in_ring(ggml_backend_t backend, const Qwen4ExpWeights & return true; } +// The CUDA pools keep every block they free, so the growing temporaries of a long prompt can hold the memory a graph +// allocation needs. Release the devices' cached blocks (ggml_backend_cuda_trim_pool); true when any came back. +static bool release_cached_pools(ggml_backend_t backend, const Qwen4ExpWeights & w) { + size_t freed = ggml_backend_cuda_trim_pool(backend); + if (w.expert_backend) freed += ggml_backend_cuda_trim_pool(w.expert_backend); + if (freed) std::fprintf(stderr, "[qwen4exp] released %zu MiB of cached pool memory for a graph\n", freed >> 20); + return freed > 0; +} + // Allocate the graph: on the split scheduler (weights decide placement, see qwen4exp_pin_split), or with the -// workspace's retained allocator (`pool`), or a fresh one (`galloc`, the caller frees it). On failure the caller -// frees its context unless the workspace owns it. +// workspace's retained allocator (`pool`), or a fresh one (`galloc`, the caller frees it). A failed allocation +// retries once with the cached pool memory released. On failure the caller frees its context unless the workspace +// owns it. static bool allocate_graph(ggml_backend_t backend, const Qwen4ExpWeights & w, Qwen4ExpCache & cache, Qwen4ExpDecodeWorkspace * pool, bool verify, ggml_cgraph * gf, const std::vector & expert_nodes, int64_t T, int64_t kv_len, ggml_backend_sched_t & split_sched, ggml_gallocr_t & galloc) { if (w.expert_backend) { split_sched = split_scheduler(backend, w, cache, pool, verify, T <= kSplitShortRows); - ggml_backend_sched_reset(split_sched); - qwen4exp_pin_split(split_sched, gf, backend, w.expert_backend, expert_nodes); - if (!ggml_backend_sched_alloc_graph(split_sched, gf)) { + auto alloc = [&] { + ggml_backend_sched_reset(split_sched); + qwen4exp_pin_split(split_sched, gf, backend, w.expert_backend, expert_nodes); + return ggml_backend_sched_alloc_graph(split_sched, gf); + }; + if (!alloc() && !(release_cached_pools(backend, w) && alloc())) { std::fprintf(stderr, "[qwen4exp] split graph alloc failed (T=%lld kv_len=%lld)\n", (long long) T, (long long) kv_len); return false; } - if (split_sched == cache.split_sched_short) ++cache.split_short_gen; + if (split_sched == split_owner(cache).split_sched_short) ++split_owner(cache).split_short_gen; return true; } if (pool) { @@ -1981,8 +1998,10 @@ static bool allocate_graph(ggml_backend_t backend, const Qwen4ExpWeights & w, Qw // T=1 graphs keep the same broad shape but advancing KV views can change // lifetimes. Recompute assignments while retaining the allocator buffers; // reusing the old index-wise plan produced incorrect tokens. - const bool reserve_ok = !pool || !pool->planned || ggml_gallocr_reserve(galloc, gf); - if (!reserve_ok || !ggml_gallocr_alloc_graph(galloc, gf)) { + auto alloc = [&] { + return (!pool || !pool->planned || ggml_gallocr_reserve(galloc, gf)) && ggml_gallocr_alloc_graph(galloc, gf); + }; + if (!alloc() && !(release_cached_pools(backend, w) && alloc())) { std::fprintf(stderr, "[qwen4exp] graph alloc failed (T=%lld kv_len=%lld)\n", (long long) T, (long long) kv_len); if (pool) { @@ -2159,11 +2178,13 @@ static Qwen4ExpForwardResult forward_impl(ggml_backend_t backend, // allocation: the scheduler rewrote its cross-device sources, so it is never split again. if (use_stable_graph && stable_ws.gf && (stable_ws.qsa_blocks >= 0) == stable_qsa && (stable_ws.hidden != nullptr) == (out_hidden != nullptr) && - (!split || stable_ws.sched || stable_ws.split_gen == cache.split_short_gen) && + (!split || stable_ws.sched || stable_ws.split_gen == split_owner(cache).split_short_gen) && (verify ? stable_ws.row_spans == row_spans : kv_len <= stable_ws.kv_bucket && stable_ws.next_pos == pos0) && (!stable_qsa || (stable_ws.kv_bucket == stable_kv_bucket && stable_ws.qsa_budget == w.indexer_top_k / 4 && cache.indexer_blocks == pos0 / 4))) { - if (!run_stable_graph(stable_in, stable_ws, cache.split_sched_short, verify, out_logits, out_hidden)) return res; + if (!run_stable_graph(stable_in, stable_ws, split_owner(cache).split_sched_short, verify, out_logits, out_hidden)) { + return res; + } stable_ws.next_pos = (int) kv_len; cache.cur_pos = (int) kv_len; ++stable_ws.replays; @@ -2311,7 +2332,7 @@ static Qwen4ExpForwardResult forward_impl(ggml_backend_t backend, if (use_stable_graph) { stable_ws.expert_nodes = expert_nodes; - stable_ws.split_gen = cache.split_short_gen; + stable_ws.split_gen = split_owner(cache).split_short_gen; stable_ws.gf = gf; stable_ws.inp_emb = inp_emb; stable_ws.positions = positions; @@ -2356,6 +2377,9 @@ static Qwen4ExpForwardResult forward_impl(ggml_backend_t backend, } return res; } + // A split prompt chunk's temporaries grow with its position, and the target's pool keeps every block it frees: + // a long prompt would leave each chunk's behind until the device runs out. Release them after the chunk. + if (split && T > kSplitShortRows) ggml_backend_cuda_trim_pool(backend); cache.ple_prev = inputs->ple_prev; // Commit only after the graph computed: a failed compute must not mark blocks the kernel never pooled. cache.cur_pos = (int) kv_len; diff --git a/server/src/qwen4exp/qwen4exp_seq_engine.cpp b/server/src/qwen4exp/qwen4exp_seq_engine.cpp index 0a8fce8b9..5c62b3cff 100644 --- a/server/src/qwen4exp/qwen4exp_seq_engine.cpp +++ b/server/src/qwen4exp/qwen4exp_seq_engine.cpp @@ -31,8 +31,9 @@ Qwen4ExpSeqEngine::Qwen4ExpSeqEngine( ggml_backend_t backend, const Qwen4ExpWeights & weights, std::vector caches, int max_ctx, int prefill_granule, size_t prefix_allowance, int verify_width, AdaptiveSpecWidth * mtp_width, - SpecWidthCostMemory * mtp_costs) - : backend_(backend), weights_(weights), caches_(std::move(caches)), + SpecWidthCostMemory * mtp_costs, ggml_backend_t snapshot_store) + : backend_(backend), store_(snapshot_store ? snapshot_store : backend), weights_(weights), + caches_(std::move(caches)), pool_(pool_blocks(max_ctx, caches_.size()), (uint32_t)caches_.size(), 256), slots_(pool_, max_ctx), @@ -108,9 +109,9 @@ SeqEngine::AdmitResult Qwen4ExpSeqEngine::admit_cold( size_t Qwen4ExpSeqEngine::estimate_prefix_store_bytes(int tokens) const { if (caches_.empty() || !caches_[0] || tokens <= 0) return 0; size_t host = 0; - const size_t device = qwen4exp_snapshot_bytes(backend_, *caches_[0], tokens, &host); - if (!device) return 0; - return device + host + (size_t(tokens) + std::max(0, weights_.ple_ngram_size - 1)) * sizeof(int32_t); + const size_t copies = qwen4exp_snapshot_bytes(store_, *caches_[0], tokens, &host); + if (!copies) return 0; + return copies + host + (size_t(tokens) + std::max(0, weights_.ple_ngram_size - 1)) * sizeof(int32_t); } int Qwen4ExpSeqEngine::checkpoint_index(PrefixStoreRef checkpoint) const { @@ -321,7 +322,7 @@ PrefixStoreEvent Qwen4ExpSeqEngine::capture_prefix(int slot, PrefixCaptureTicket return event; } Qwen4ExpSnapshot & s = checkpoints_[(size_t)index]; - if (!save_qwen4exp_snapshot(backend_, *caches_[(size_t)slot], s)) { + if (!save_qwen4exp_snapshot(backend_, store_, *caches_[(size_t)slot], s)) { free_qwen4exp_snapshot(s); event.elapsed_us = elapsed_us(); event.error = "qwen4exp prefix capture failed"; diff --git a/server/src/qwen4exp/qwen4exp_seq_engine.h b/server/src/qwen4exp/qwen4exp_seq_engine.h index cf8fe4311..380af4fe2 100644 --- a/server/src/qwen4exp/qwen4exp_seq_engine.h +++ b/server/src/qwen4exp/qwen4exp_seq_engine.h @@ -23,11 +23,11 @@ namespace luce::common { // other slots decode, a step runs one granule in total, so a long prompt never // holds the live streams for more than one. // -// Prefix checkpoints are prefix-sized device copies of a slot cache -// (qwen4exp snapshots, the same payload the single-slot path keeps), -// restorable into any slot of the same layout. The scheduler owns their -// identities (1..kPrefixCheckpoints) and LRU policy; the engine enforces the -// backend's byte allowance on capture. +// Prefix checkpoints are prefix-sized copies of a slot cache (qwen4exp +// snapshots, the same payload the single-slot path keeps) on `snapshot_store`, +// the device or system memory, restorable into any slot of the same layout. +// The scheduler owns their identities (1..kPrefixCheckpoints) and LRU policy; +// the engine enforces the backend's byte allowance on capture. // // The first slot's cache carries the MTP draft layer and verify rollback state // (one set: its rollback buffers are close to 1 GiB). Admission takes the lowest @@ -49,7 +49,7 @@ class Qwen4ExpSeqEngine final : public SeqEngine { std::vector caches, int max_ctx, int prefill_granule = 512, size_t prefix_allowance = 0, int verify_width = 1, AdaptiveSpecWidth * mtp_width = nullptr, - SpecWidthCostMemory * mtp_costs = nullptr); + SpecWidthCostMemory * mtp_costs = nullptr, ggml_backend_t snapshot_store = nullptr); ~Qwen4ExpSeqEngine() override; int slot_count() const override { return slots_.slot_count(); } @@ -93,6 +93,7 @@ class Qwen4ExpSeqEngine final : public SeqEngine { StepResult mtp_step(const StepInput & input); ggml_backend_t backend_; + ggml_backend_t store_; // checkpoint memory: snapshot_store, else the backend const Qwen4ExpWeights & weights_; std::vector caches_; PagedKvPool pool_; diff --git a/server/src/server/launch_profiles.h b/server/src/server/launch_profiles.h index f7a07ef9a..dcdfde49b 100644 --- a/server/src/server/launch_profiles.h +++ b/server/src/server/launch_profiles.h @@ -237,7 +237,7 @@ inline const std::vector & launch_profiles() { { {"--target-device", "hip:0"}, {"--expert-device", "hip:1"}, - {"--max-ctx", "131072"}, + {"--max-ctx", "auto"}, }, {}, }, diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index b31a4e6d8..fc208c666 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -125,7 +125,9 @@ static void print_usage(const char * prog) { " --mmproj-device hip:N Run the DS4V image encoder on another GPU (one-GPU layout)\n" " --port Listen port (default: 8080)\n" " --host Bind address (default: 0.0.0.0)\n" - " --max-ctx Max context length (default: 8192)\n" + " --max-ctx Max context length (default: 8192). auto: the model's\n" + " trained context, lowered to what the devices hold\n" + " (Qwen3.8-Flash-Next)\n" " --max-tokens Default max output tokens (legacy alias for\n" " --default-max-tokens; loses to --default-max-tokens\n" " when both are passed)\n" @@ -458,7 +460,8 @@ static int parse_model_options(int argc, char ** argv, ModelOptions & model, } else if (std::strcmp(argv[i], "--host") == 0 && i + 1 < argc) { sconfig.host = argv[++i]; } else if (std::strcmp(argv[i], "--max-ctx") == 0 && i + 1 < argc) { - int v = std::atoi(argv[++i]); + bargs.device.fit_ctx = std::strcmp(argv[++i], "auto") == 0; + int v = bargs.device.fit_ctx ? 0 : std::atoi(argv[i]); sconfig.max_ctx = v; bargs.device.max_ctx = v; } else if (std::strcmp(argv[i], "--max-tokens") == 0 && i + 1 < argc) { @@ -983,6 +986,16 @@ static int parse_model_options(int argc, char ** argv, ModelOptions & model, if (bargs.max_concurrency > 1 && inspect_gguf_model_info(bargs.model_path.c_str()).arch != "qwen4exp") bargs.paged_attention = true; + // --max-ctx auto starts from the trained context; the backend lowers it to what its devices hold. + if (bargs.device.fit_ctx) { + const GgufMetadata meta = read_gguf_metadata(bargs.model_path, false); + if (meta.general_architecture != "qwen4exp" || meta.context_length <= 0) { + std::fprintf(stderr, "[server] --max-ctx auto needs a model whose backend fits its context to the " + "devices (Qwen3.8-Flash-Next); give --max-ctx a token count\n"); + return 2; + } + bargs.device.max_ctx = meta.context_length; + } if (sconfig.decode_kv_offload_bytes && sconfig.decode_kv_offload_bytes != kAutoKvOffloadBytes && bargs.max_concurrency <= 1) { std::fprintf(stderr, "[server] --decode-kv-offload-mb requires --max-concurrency greater than 1\n"); @@ -1486,6 +1499,10 @@ static int load_model(ModelOptions & model, LoadedModel & loaded, bool multi_mod return 1; } ModelBackend * backend = backend_owner.get(); + // --max-ctx auto: serve the context the backend fitted to its devices. + if (backend_placement.target.fit_ctx && backend->context_limit() > 0) { + sconfig.max_ctx = backend->context_limit(); + } // Cross-check the capability table against the backend that was actually // built. arch_supports_remote_draft() admitted this launch from the arch // string alone; if the two ever disagree the table is stale, and failing diff --git a/server/test/unit/test_launch_policy.cpp b/server/test/unit/test_launch_policy.cpp index c8d2b27f7..83e775349 100644 --- a/server/test/unit/test_launch_policy.cpp +++ b/server/test/unit/test_launch_policy.cpp @@ -6,6 +6,7 @@ #include "placement/device_select.h" #include "server/launch_profiles.h" +#include #include #include #include @@ -254,6 +255,16 @@ struct LaunchPolicyFixture : CommonFixture { const std::vector merged = launch_profile_args(*profile, given); CHECK(!contains_flag(merged, "--draft-device")); } + + void test_qwen_next_split_profile_fits_its_context() { + // The split fits the context to the R9700; a number on the command line still wins. + const LaunchProfile * profile = find_launch_profile("qwen-next-r9700-strix"); + CHECK(profile != nullptr); + const std::vector args = launch_profile_args(*profile, {}); + const auto max_ctx = std::find(args.begin(), args.end(), "--max-ctx"); + CHECK(max_ctx != args.end() && max_ctx + 1 != args.end() && *(max_ctx + 1) == "auto"); + CHECK(!contains_flag(launch_profile_args(*profile, {"--max-ctx", "65536"}), "--max-ctx")); + } }; } // namespace @@ -271,4 +282,5 @@ TEST_CASE(LaunchPolicyFixture, launch_policy_suite) { test_profile_replaces_documented_recipe(); test_ds41_profile_is_the_lucebox_recipe(); test_ds41_gorgon_profile_is_the_measured_recipe(); + test_qwen_next_split_profile_fits_its_context(); } diff --git a/server/test/unit/test_qwen4exp_chunk.cpp b/server/test/unit/test_qwen4exp_chunk.cpp index 5967b8986..6a499ec29 100644 --- a/server/test/unit/test_qwen4exp_chunk.cpp +++ b/server/test/unit/test_qwen4exp_chunk.cpp @@ -2,7 +2,29 @@ #include "qwen4exp/qwen4exp_chunk.h" #include using luce::common::qwen4exp_fit_chunk; +using luce::common::Qwen4ExpContextFit; +using luce::common::qwen4exp_shrink_context; int main() { + // --max-ctx auto: step down by the shortfall over a position's bytes, in 4096 multiples, at least one step. + assert(qwen4exp_shrink_context(262144, 10000000000ULL, 138400) == 188416); + assert(qwen4exp_shrink_context(135168, 300000000ULL, 138400) == 131072); + assert(qwen4exp_shrink_context(262144, 1, 1000) == 258048); + assert(qwen4exp_shrink_context(8192, 1000000000000ULL, 1) == 0); + assert(qwen4exp_shrink_context(4096, 1, 1) == 0); + assert(qwen4exp_shrink_context(131072, 0, 0) == 0); + // Shrink until a context fits, then interpolate between it and the smallest miss. + auto fit = [](int fits, size_t spare, int misses, size_t missing) { + Qwen4ExpContextFit f; + f.fits = fits; f.spare = spare; f.misses = misses; f.missing = missing; + return f.next(138400); + }; + assert(fit(0, 0, 262144, 10000000000ULL) == 188416); + assert(fit(0, 0, 4096, 1000000000ULL) == 0); + assert(fit(262144, 1, 0, 0) == 0); + assert(fit(122880, 1000000000ULL, 262144, 9000000000ULL) == 135168); + assert(fit(122880, 0, 262144, 9000000000ULL) == 126976); // at least one step above the fit + assert(fit(122880, 5000000000ULL, 131072, 100000000ULL) == 126976); // at least one step below the miss + assert(fit(131072, 100000000ULL, 135168, 200000000ULL) == 0); auto workspace = [](int n) { return size_t(n) * 10; }; assert(qwen4exp_fit_chunk(10000, 100000, 0, workspace) == 8960); assert(qwen4exp_fit_chunk(10000, 100000, 20000, workspace) == 6912); diff --git a/server/test/unit/test_qwen4exp_prefix_cache.cpp b/server/test/unit/test_qwen4exp_prefix_cache.cpp index 358ab7d6f..a20a929b4 100644 --- a/server/test/unit/test_qwen4exp_prefix_cache.cpp +++ b/server/test/unit/test_qwen4exp_prefix_cache.cpp @@ -2,6 +2,7 @@ #include "qwen4exp/qwen4exp_backend.h" #include "qwen4exp/qwen4exp_chunk.h" #include "qwen4exp/qwen4exp_graph.h" +#include "common/snapshot_backend.h" #include "server/tokenizer.h" #include "ggml-backend-impl.h" #include "ggml-cpu.h" @@ -144,7 +145,7 @@ struct Qwen4ExpPrefixTest { static void budget() { Qwen4ExpBackend b({}); - b.backend_ = ggml_backend_cpu_init(); + b.backend_ = b.snap_backend_ = ggml_backend_cpu_init(); CHECK(b.backend_); auto & c = b.cache_; c.ctx = ggml_init({8 * ggml_tensor_overhead(), nullptr, true}); @@ -199,8 +200,9 @@ static bool same_tokens(const GenerateResult & warm, const GenerateResult & cold return warm.ok() && cold.ok() && warm.tokens == cold.tokens; } -static void copies(ggml_backend_t backend) { - CHECK(backend); +// `store`: the backend's own memory, or system memory beside a GPU. +static void copies(ggml_backend_t backend, ggml_backend_t store) { + CHECK(backend && store); Qwen4ExpCache c; c.max_ctx = 63; // packed KV has a padded fourth key beyond the logical limit c.ctx = ggml_init({64 * ggml_tensor_overhead(), nullptr, true}); @@ -225,9 +227,10 @@ static void copies(ggml_backend_t backend) { } c.cur_pos = pos; c.indexer_blocks = pos / 4; c.mtp_prev_pos = pos - 1; c.kv_bucket_base = 512; c.ple_prev = {11, 12}; - CHECK(save_qwen4exp_snapshot(backend, c, s)); + CHECK(save_qwen4exp_snapshot(backend, store, c, s)); + CHECK(ggml_backend_buffer_get_type(s.buf) == ggml_backend_get_default_buffer_type(store)); size_t host = 0; - CHECK(ggml_backend_buffer_get_size(s.buf) <= qwen4exp_snapshot_bytes(backend, c, pos, &host)); + CHECK(ggml_backend_buffer_get_size(s.buf) <= qwen4exp_snapshot_bytes(store, c, pos, &host)); CHECK(host >= ggml_get_mem_size(s.ctx) + s.strips.capacity() * sizeof(s.strips[0])); for (auto [live, copy] : s.strips) CHECK(bytes(live) == bytes(copy)); // All states, including QSA raw incomplete blocks and MTP carry, must @@ -251,10 +254,11 @@ static void copies(ggml_backend_t backend) { } } --c.mtp_prev_pos; - CHECK(!save_qwen4exp_snapshot(backend, c, s)); // incomplete MTP is not publishable + CHECK(!save_qwen4exp_snapshot(backend, store, c, s)); // incomplete MTP is not publishable free_qwen4exp_snapshot(s); free_qwen4exp_snapshot(s); free_qwen4exp_cache(c); - std::puts("PASS prefix snapshot bytes, metadata, overwrite/reset, MTP validity"); + std::printf("PASS prefix snapshot bytes, metadata, overwrite/reset, MTP validity (%s memory)\n", + store == backend ? "device" : "system"); } static void memory(ggml_backend_t backend) { @@ -494,7 +498,10 @@ int main(int argc, char ** argv) { Qwen4ExpPrefixTest::budget(); auto backend = argc > 1 ? ggml_backend_cuda_init(0) : ggml_backend_cpu_init(); CHECK(backend); - copies(backend); + copies(backend, backend); + const ggml_backend_t store = create_snapshot_backend(backend); + if (store != backend) copies(backend, store); + free_snapshot_backend(store, backend); memory(backend); ggml_backend_free(backend); // test_qwen4exp_prefix_cache MODEL [width] [chunk] [ctx] [expert_gpu] [placement.csv] diff --git a/server/test/unit/test_qwen4exp_seq_engine.cpp b/server/test/unit/test_qwen4exp_seq_engine.cpp index 0f8ff92a3..a5012c4a7 100644 --- a/server/test/unit/test_qwen4exp_seq_engine.cpp +++ b/server/test/unit/test_qwen4exp_seq_engine.cpp @@ -6,6 +6,7 @@ #include "common/concurrency/seq_engine.h" #include "seq_engine_contract.h" #include "common/sampler.h" +#include "common/snapshot_backend.h" #include "server/tokenizer.h" #include "ggml-cuda.h" #include "qwen4exp_test_state.h" @@ -555,12 +556,20 @@ int main(int argc, char ** argv) { if (mtp_ok) { std::vector served = ptrs; served[0] = &mtp_cache; - Qwen4ExpSeqEngine engine(backend, weights, served, ctx, 512, size_t(8) << 30, /*verify_width=*/0); - const bool prefix_ok = run_prefix_store(engine, weights); - std::printf("[qwen4exp-seq] prefix-store=%s\n", prefix_ok ? "PASS" : "FAIL"); - const bool cuts_ok = run_prefix_cuts(engine, weights); - std::printf("[qwen4exp-seq] prefix-cuts=%s\n", cuts_ok ? "PASS" : "FAIL"); - ok = ok && prefix_ok && cuts_ok; + // Checkpoints in the device's memory, and in system memory as beside a discrete GPU. + const ggml_backend_t system = create_snapshot_backend(backend); + for (ggml_backend_t store : {backend, system}) { + Qwen4ExpSeqEngine engine(backend, weights, served, ctx, 512, size_t(8) << 30, /*verify_width=*/0, + nullptr, nullptr, store); + const char * memory = store == backend ? "device" : "system"; + const bool prefix_ok = run_prefix_store(engine, weights); + std::printf("[qwen4exp-seq] prefix-store(%s)=%s\n", memory, prefix_ok ? "PASS" : "FAIL"); + const bool cuts_ok = run_prefix_cuts(engine, weights); + std::printf("[qwen4exp-seq] prefix-cuts(%s)=%s\n", memory, cuts_ok ? "PASS" : "FAIL"); + ok = ok && prefix_ok && cuts_ok; + if (system == backend) break; + } + free_snapshot_backend(system, backend); } else { std::fprintf(stderr, "[qwen4exp-seq] MTP cache creation failed\n"); ok = false;