feat(qwen4exp): fit the context to the devices, prefix snapshots in system memory

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 |

server/src/common/model_backend.h

@@ -158,6 +158,9 @@ struct ModelBackend {
     // Print the "[<arch>-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

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 {

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<std::pair<int, Plan>> 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<Qwen4ExpSeqEngine>(
         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;

server/src/qwen4exp/qwen4exp_backend.h

@@ -23,6 +23,8 @@
 
 namespace luce::common {
 
+struct Qwen4ExpChunkPlan;
+
 struct Qwen4ExpBackendConfig {
     std::string     model_path;
     std::optional<std::string> 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<Qwen4ExpSnapshot, kMaxSlots> snapshots_;
     int live_slot_ = -1;
     std::vector<int32_t> tokens_;