From 77c02f2861a3a5ab9b5ddd0f31890b256267ae83 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Fri, 9 Oct 2026 17:03:26 +0200 Subject: [PATCH 1/3] fix(qwen4exp): prefix snapshots in system memory beside a discrete GPU Prefix snapshots and concurrent prefix checkpoints were copies in the target's memory, so beside a GPU with its own memory (the R9700 in split mode) their allowance competed with the KV caches. At 131072 tokens with four slots it was 473 MB: captures failed and the prefix cache stopped helping. They now go to system memory through create_snapshot_backend, as the other models' snapshots do, and the server's resident limit bounds them. On the gfx1151 iGPU they stay in its memory, which is the system's, under the planner's allowance as before. Split mode, four slots, 131072 tokens: every capture saved, a 6K-token checkpoint restores in 106 ms, a resent 16K prompt prefills in 0.3 s instead of 10.4 s, and outputs match alone and together (4 of 4 short prompts, 4 of 4 with two 6K prompts). Co-Authored-By: Claude Opus 5.5 --- server/src/qwen4exp/qwen4exp_backend.cpp | 39 +++++++++++++------ server/src/qwen4exp/qwen4exp_backend.h | 3 +- server/src/qwen4exp/qwen4exp_cache.cpp | 11 +++--- server/src/qwen4exp/qwen4exp_cache.h | 12 ++++-- server/src/qwen4exp/qwen4exp_seq_engine.cpp | 13 ++++--- server/src/qwen4exp/qwen4exp_seq_engine.h | 13 ++++--- .../test/unit/test_qwen4exp_prefix_cache.cpp | 23 +++++++---- server/test/unit/test_qwen4exp_seq_engine.cpp | 21 +++++++--- 8 files changed, 87 insertions(+), 48 deletions(-) diff --git a/server/src/qwen4exp/qwen4exp_backend.cpp b/server/src/qwen4exp/qwen4exp_backend.cpp index aa26c5b36..4cdbc1fdf 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" @@ -264,6 +265,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; @@ -305,23 +310,31 @@ bool Qwen4ExpBackend::load_target() { cfg_.max_concurrency > 1 ? std::string() : cfg_.expert_placement_path, expert_budget_bytes_from_env(kQwen4ExpHotExpertBudget))) return false; 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_.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_); + const int fit = qwen4exp_select_chunk(backend_, weights_, cache_, cfg_.max_concurrency, 1, + allowance(2 * (size_t) cfg_.max_concurrency + 1)); 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_); + chunk_ = qwen4exp_select_chunk(backend_, weights_, cache_, 1, 1, allowance(3)); } + 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; @@ -366,7 +379,7 @@ bool Qwen4ExpBackend::start_seq_engine() { } 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 +714,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 +755,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 +816,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..9011c3aae 100644 --- a/server/src/qwen4exp/qwen4exp_backend.h +++ b/server/src/qwen4exp/qwen4exp_backend.h @@ -94,7 +94,8 @@ class Qwen4ExpBackend final : public ModelBackend { // 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..313d2fe94 100644 --- a/server/src/qwen4exp/qwen4exp_cache.cpp +++ b/server/src/qwen4exp/qwen4exp_cache.cpp @@ -352,11 +352,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 +375,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 +390,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..6677f52ee 100644 --- a/server/src/qwen4exp/qwen4exp_cache.h +++ b/server/src/qwen4exp/qwen4exp_cache.h @@ -209,8 +209,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 +224,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_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/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; From 82397fd241ae64cb88e4388afffded6724d89e79 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Fri, 9 Oct 2026 17:28:43 +0200 Subject: [PATCH 2/3] feat(server): --max-ctx auto fits Qwen3.8-Flash-Next's context to the devices `--max-ctx auto` starts from the model's trained context (262,144 tokens) and the qwen4exp backend lowers it, in steps of 4096, to the largest 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. The server serves that context (/props n_ctx). The chunk planner reports what it measured (Qwen4ExpChunkPlan); Qwen4ExpContextFit shrinks a miss by its shortfall, then interpolates between the largest fit and the smallest miss. An unpark keeps the fitted context. The qwen-next-r9700-strix profile uses it. Measured, starting in 15 s each: | Layout | Fitted context | 4 short prompts together vs alone | | --- | --- | --- | | R9700 + Strix Halo, 4 at once | 131,072 | identical, 78.2 tok/s together | | R9700 + Strix Halo, 1 at once | 262,144 | identical | | Strix Halo alone (192 GB), 4 at once | 262,144 | identical, 48.8 tok/s together | Co-Authored-By: Claude Opus 5.5 --- server/docs/QWEN4EXP.md | 16 ++-- server/src/common/model_backend.h | 3 + server/src/placement/placement_config.h | 4 + server/src/qwen4exp/qwen4exp_backend.cpp | 106 +++++++++++++++++------ server/src/qwen4exp/qwen4exp_backend.h | 4 + server/src/qwen4exp/qwen4exp_chunk.h | 37 +++++++- server/src/server/launch_profiles.h | 2 +- server/src/server/server_main.cpp | 21 ++++- server/test/unit/test_launch_policy.cpp | 12 +++ server/test/unit/test_qwen4exp_chunk.cpp | 22 +++++ 10 files changed, 192 insertions(+), 35 deletions(-) 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 4cdbc1fdf..976b287af 100644 --- a/server/src/qwen4exp/qwen4exp_backend.cpp +++ b/server/src/qwen4exp/qwen4exp_backend.cpp @@ -100,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); @@ -201,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. @@ -230,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); @@ -309,27 +325,45 @@ 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; - // 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_.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. - const int fit = qwen4exp_select_chunk(backend_, weights_, cache_, cfg_.max_concurrency, 1, - allowance(2 * (size_t) cfg_.max_concurrency + 1)); - chunk_ = cfg_.chunk > 0 && fit > 0 ? explicit_chunk : fit; - } else if (cfg_.chunk > 0) { - chunk_ = explicit_chunk; - } else { - chunk_ = qwen4exp_select_chunk(backend_, weights_, cache_, 1, 1, allowance(3)); + // --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_); @@ -346,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(); @@ -361,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), diff --git a/server/src/qwen4exp/qwen4exp_backend.h b/server/src/qwen4exp/qwen4exp_backend.h index 9011c3aae..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,6 +92,7 @@ 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. 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/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); From b1fe6be97dc552b069e4c4257e75ec821d33cea2 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Fri, 9 Oct 2026 18:20:04 +0200 Subject: [PATCH 3/3] fix(qwen4exp): split serving keeps one prompt chunk's memory on the target At 131072 tokens with four slots, the split ran out of R9700 memory on a long prompt: a 96K prompt beside three short requests, and also a 96K prompt alone. The chunk planner did not count two kinds of growth: - Every concurrency slot had its own split schedulers, so each slot that ran a prompt kept a prompt chunk's buffers (about 1.9 GB each). The slots now share the first slot's schedulers. One forward runs at a time, which is what the planner counts. - The target's CUDA pool keeps every block it frees, and a prompt chunk's temporaries grow with its position. During a prompt alone, free memory fell by about 270 MB per chunk until the allocation failed. Split prompt chunks now release the pool's cached blocks. A failed graph allocation retries once after releasing them. Split mode, four slots, 131072 tokens: a 96K prompt beside three short requests completes. The R9700 keeps at least 2.15 GB free through the short, long and prefix-cache runs; before, it fell to 0. Outputs match alone and together, and a 16K prompt prefills in 9.6 s, as before. Co-Authored-By: Claude Opus 5.5 --- server/src/qwen4exp/qwen4exp_backend.cpp | 1 + server/src/qwen4exp/qwen4exp_cache.cpp | 1 + server/src/qwen4exp/qwen4exp_cache.h | 5 ++- server/src/qwen4exp/qwen4exp_graph.cpp | 56 +++++++++++++++++------- 4 files changed, 46 insertions(+), 17 deletions(-) diff --git a/server/src/qwen4exp/qwen4exp_backend.cpp b/server/src/qwen4exp/qwen4exp_backend.cpp index 976b287af..2f3df2410 100644 --- a/server/src/qwen4exp/qwen4exp_backend.cpp +++ b/server/src/qwen4exp/qwen4exp_backend.cpp @@ -431,6 +431,7 @@ 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( diff --git a/server/src/qwen4exp/qwen4exp_cache.cpp b/server/src/qwen4exp/qwen4exp_cache.cpp index 313d2fe94..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; diff --git a/server/src/qwen4exp/qwen4exp_cache.h b/server/src/qwen4exp/qwen4exp_cache.h index 6677f52ee..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 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;