feat(qwen4exp): 262K context for four requests in split mode, KV of some layers on the Strix

server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu

@@ -1227,6 +1227,10 @@ static bool ggml_backend_cuda_buffer_cpy_tensor(ggml_backend_buffer_t buffer, co
     if (ggml_backend_buffer_is_cuda(src->buffer)) {
         ggml_backend_cuda_buffer_context * src_ctx = (ggml_backend_cuda_buffer_context *)src->buffer->context;
         ggml_backend_cuda_buffer_context * dst_ctx = (ggml_backend_cuda_buffer_context *)dst->buffer->context;
+        // The per-thread stream is the current device's. On ROCm a peer copy queued on another device's stream
+        // completes without writing (an R9700 to a Strix Halo iGPU), so queue it on the source's, as
+        // ggml_backend_cuda_cpy_tensor_async does.
+        ggml_cuda_set_device(src_ctx->device);
         if (src_ctx->device == dst_ctx->device) {
             CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(src), cudaMemcpyDeviceToDevice, cudaStreamPerThread));
         } else {

server/docs/QWEN4EXP.md

@@ -104,8 +104,11 @@ 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 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.
+model's 262,144-token context. When the R9700 cannot hold every slot's cache
+beside the dense weights, the last full-attention layers keep their K/V and
+indexer rows on the Strix and run their attention there, as few layers as
+cover the shortfall. The context is lowered in steps of 4096 only once the
+Strix cannot take more.
 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.
 
@@ -144,6 +147,10 @@ a 16K prompt at 44 tok/s. Greedy output is the same with adaptive MTP, a fixed
 its hot experts on the R9700 with the same kernel.
 
 With `--max-concurrency 2..4` split mode serves several requests at once.
+Four at once get the full 262,144 tokens with 7 of the 12 attention layers'
+rows on the Strix: a request alone decodes at about 46 tok/s and four
+together at 69 tok/s, against 59 and 78 at 131,072 tokens with every row on
+the R9700.
 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

server/src/qwen4exp/qwen4exp_backend.cpp

@@ -99,6 +99,36 @@ static bool load_qwen4exp_hot_experts(ggml_backend_t backend, Qwen4ExpWeights &
     return true;
 }
 
+// Split mode: the last `layers` full-attention layers run their attention on the expert device, beside their cache
+// rows (Qwen4ExpCache::remote_layers), so their projection, norm and indexer weights move there. Layers moved before
+// stay; the target's copies stay in its weights buffer.
+static bool move_qwen4exp_attention(Qwen4ExpWeights & w, int layers) {
+    const ggml_backend_buffer_type_t expert = ggml_backend_get_default_buffer_type(w.expert_backend);
+    std::vector<ggml_tensor **> fields;
+    for (int il = w.n_layer - 1; il >= 0 && layers > 0; --il) {
+        Qwen4ExpLayer & L = w.layers[(size_t) il];
+        if (!L.is_full_attention) continue;
+        --layers;
+        for (ggml_tensor ** f : {&L.wq, &L.wk, &L.wv, &L.wo, &L.q_norm, &L.k_norm,
+                                 &L.indexer_q_proj, &L.indexer_k_proj, &L.indexer_q_norm, &L.indexer_k_norm}) {
+            if (*f && ggml_backend_buffer_get_type((*f)->buffer) != expert) fields.push_back(f);
+        }
+    }
+    if (fields.empty()) return true;
+    ggml_context * ctx = w.derived.open(fields.size());
+    std::vector<ggml_tensor *> moved;
+    for (ggml_tensor ** f : fields) moved.push_back(ctx ? ggml_dup_tensor(ctx, *f) : nullptr);
+    if (!ctx || !w.derived.commit(w.expert_backend)) return false;
+    std::vector<uint8_t> bytes;
+    for (size_t i = 0; i < fields.size(); ++i) {   // through the host: the two GPUs need not reach each other
+        bytes.resize(ggml_nbytes(*fields[i]));
+        ggml_backend_tensor_get(*fields[i], bytes.data(), 0, bytes.size());
+        ggml_backend_tensor_set(moved[i], bytes.data(), 0, bytes.size());
+        *fields[i] = moved[i];
+    }
+    return true;
+}
+
 int qwen4exp_select_chunk(ggml_backend_t backend, const Qwen4ExpWeights & w,
         Qwen4ExpCache & cache, int slots, int resident_slots, size_t * snapshot_budget, Qwen4ExpChunkPlan * plan) {
     if (slots < 1 || resident_slots < 1 || resident_slots > slots) return 0;
@@ -236,13 +266,15 @@ int qwen4exp_select_chunk(ggml_backend_t backend, const Qwen4ExpWeights & w,
         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;
+        // Every slot's attention and indexer rows on the target; only this cache adds the MTP draft layer's.
+        size_t rows = 0, layer = 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);
+            for (const ggml_tensor * t : *group) if (t && t->buffer == cache.buf) rows += ggml_nbytes(t);
+            if (!group->empty() && group->front()) layer += ggml_nbytes(group->front());
         }
         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;
+        plan->layer_bytes = (size_t) slots * layer;
         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,
@@ -325,9 +357,15 @@ 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;
-    // --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.
+    if (remote_layers_ && !move_qwen4exp_attention(weights_, remote_layers_)) {
+        std::fprintf(stderr, "[qwen4exp] moving attention to the expert device failed\n");
+        return false;
+    }
+    // --max-ctx auto: the trained context, or the largest that fits (Qwen4ExpContextFit) with every slot's cache, the
+    // resident graphs, a floor-sized prompt chunk and, with snapshots on the target, one full snapshot per slot. In
+    // split mode a miss first moves the fewest full-attention layers' cache rows and attention to the expert device
+    // that cover it, while that device holds them; the context shrinks once none can move. The layout found stays:
+    // an unpark sizes itself for it.
     for (Qwen4ExpContextFit search;;) {
         Qwen4ExpChunkPlan plan;
         chunk_ = plan_chunk(cfg_.device.fit_ctx ? &plan : nullptr);
@@ -337,23 +375,33 @@ bool Qwen4ExpBackend::load_target() {
             ? (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;
+        const size_t missing = fit ? 0 : plan.required == SIZE_MAX ? (size_t) ctx / 4 * position
+                                                                    : plan.required + snapshots - plan.available;
+        int next = 0;
         if (fit) {
             search.fits = ctx;
             search.spare = plan.available - plan.required - snapshots;
+        } else if (const int remote = remote_layers_for(plan, missing); remote > remote_layers_) {
+            remote_layers_ = remote;
+            if (!move_qwen4exp_attention(weights_, remote_layers_)) {
+                std::fprintf(stderr, "[qwen4exp] moving attention to the expert device failed\n");
+                return false;
+            }
+            next = ctx;   // plan the same context with the rows moved
         } else {
             search.misses = ctx;
-            search.missing = plan.required == SIZE_MAX ? (size_t) ctx / 4 * position
-                                                       : plan.required + snapshots - plan.available;
+            search.missing = missing;
         }
-        int next = search.next(position);
+        if (!next) 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);
+            std::fprintf(stderr, "[qwen4exp] context fitted to the devices: %d tokens, %d at once, %d of %zu "
+                         "attention layers on the expert device\n", search.fits, cfg_.max_concurrency, remote_layers_,
+                         cache_.full_layer_ids.size());
             if (fit) break;
             next = search.fits;   // the last plan missed: plan the fit again
         }
@@ -421,7 +469,23 @@ bool Qwen4ExpBackend::create_main_cache() {
         !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),
-        slot_states_.ctx ? &slot_states_ : nullptr, 0);
+        slot_states_.ctx ? &slot_states_ : nullptr, 0, remote_layers_);
+}
+
+// --max-ctx auto in split mode: the full-attention layers to keep on the expert device so a plan that lacks
+// `missing` bytes fits, when that device holds every slot's rows of them beside its own headroom (else the layers
+// already there).
+int Qwen4ExpBackend::remote_layers_for(const Qwen4ExpChunkPlan & plan, size_t missing) const {
+    const int n_full = (int) cache_.full_layer_ids.size();
+    if (!weights_.expert_backend || !plan.layer_bytes || plan.required == SIZE_MAX || remote_layers_ >= n_full)
+        return remote_layers_;
+    const int remote = std::min<int>(n_full, remote_layers_ + (int) ((missing + plan.layer_bytes - 1) / plan.layer_bytes));
+    size_t free_device = 0, total = 0, free_host = 0;
+    ggml_backend_cuda_get_device_memory(ggml_backend_cuda_get_device_id(weights_.expert_backend), &free_device, &total);
+    ggml_backend_dev_memory(ggml_backend_get_device(weights_.expert_backend), &free_host, &total);
+    const size_t available = std::min(free_device, free_host) +
+                             (cache_.remote_buf ? ggml_backend_buffer_get_size(cache_.remote_buf) : 0);
+    return (size_t) remote * plan.layer_bytes <= available - available / 10 ? remote : remote_layers_;
 }
 
 bool Qwen4ExpBackend::start_seq_engine() {
@@ -430,7 +494,7 @@ bool Qwen4ExpBackend::start_seq_engine() {
     std::vector<Qwen4ExpCache *> caches{&cache_};
     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;
+                                   &slot_states_, (int) caches.size(), remote_layers_)) return false;
         cache.split_owner = &cache_;
         caches.push_back(&cache);
     }

server/src/qwen4exp/qwen4exp_backend.h

@@ -93,12 +93,14 @@ class Qwen4ExpBackend final : public ModelBackend {
     bool snapshot_fits(int slot, int replaced = -1) const;
     bool create_main_cache();
     int plan_chunk(Qwen4ExpChunkPlan * plan = nullptr);
+    int remote_layers_for(const Qwen4ExpChunkPlan & plan, size_t missing) const;
     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();
     ggml_backend_t snap_backend_ = nullptr; // prefix snapshot storage: the target, or system memory (see init)
+    int remote_layers_ = 0;   // split mode: full-attention layers whose cache rows and attention run on the expert device
     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;

server/src/qwen4exp/qwen4exp_cache.cpp

@@ -49,7 +49,7 @@ void free_qwen4exp_slot_states(Qwen4ExpSlotStates & s) {
 
 bool create_qwen4exp_cache(ggml_backend_t backend, const Qwen4ExpWeights & w,
                            int max_ctx, Qwen4ExpCache & out, bool mtp, int mtp_draft,
-                           const Qwen4ExpSlotStates * slot_states, int state_slot) {
+                           const Qwen4ExpSlotStates * slot_states, int state_slot, int remote_layers) {
     const Qwen4ExpCudaScope profile(w.gfx1151);
     // The QSA cell-id kernel is exact for positions below 2^24.
     if (max_ctx <= 0 || max_ctx >= (1 << 24)) {
@@ -115,22 +115,31 @@ bool create_qwen4exp_cache(ggml_backend_t backend, const Qwen4ExpWeights & w,
         out.ple_conv_state[i] = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, ple_hist, hc_dim);
     }
 
+    if (remote_layers < 0 || remote_layers > (int) n_full || (remote_layers > 0 && !w.expert_backend)) return false;
+    out.remote_layers = remote_layers;
+    out.remote_backend = remote_layers > 0 ? w.expert_backend : nullptr;
+    if (remote_layers > 0) {
+        out.remote_ctx = ggml_init({4 * (size_t) remote_layers * ggml_tensor_overhead(), nullptr, true});
+        if (!out.remote_ctx) return false;
+    }
+
     // Packed attention reads groups of four keys even for an unaligned
     // logical context limit. Padding is masked by the causal cell IDs.
     const int64_t align = w.qsa ? 4 : 1;
     const int64_t kv_capacity = (static_cast<int64_t>(max_ctx) + align - 1) / align * align;
     for (size_t i = 0; i < n_full; ++i) {
-        out.attn_k[i] = ggml_new_tensor_3d(out.ctx, kv_type,
+        ggml_context * rows = out.remote((int) i) ? out.remote_ctx : out.ctx;
+        out.attn_k[i] = ggml_new_tensor_3d(rows, kv_type,
             w.n_embd_head_k, kv_capacity, w.n_head_kv);
-        out.attn_v[i] = ggml_new_tensor_3d(out.ctx, kv_type,
+        out.attn_v[i] = ggml_new_tensor_3d(rows, kv_type,
             w.n_embd_head_v, kv_capacity, w.n_head_kv);
         const int il = out.full_layer_ids[i];
         const int ratio = il < (int) w.compress_ratios.size() ? w.compress_ratios[il] : 0;
         if (w.indexer_head_size > 0 && ratio > 0) {
             const int64_t max_blocks = (static_cast<int64_t>(max_ctx) + ratio - 1) / ratio;
-            out.indexer_k[i] = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32,
+            out.indexer_k[i] = ggml_new_tensor_2d(rows, GGML_TYPE_F32,
                 w.indexer_head_size, max_blocks + 1);
-            out.indexer_raw[i] = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32,
+            out.indexer_raw[i] = ggml_new_tensor_2d(rows, GGML_TYPE_F32,
                 w.indexer_head_size, max_ctx);
         }
     }
@@ -192,9 +201,9 @@ bool create_qwen4exp_cache(ggml_backend_t backend, const Qwen4ExpWeights & w,
     }
 
     out.buf = ggml_backend_alloc_ctx_tensors(out.ctx, backend);
-    if (!out.buf) {
-        ggml_free(out.ctx);
-        out.ctx = nullptr;
+    if (out.remote_ctx) out.remote_buf = ggml_backend_alloc_ctx_tensors(out.remote_ctx, w.expert_backend);
+    if (!out.buf || (out.remote_ctx && !out.remote_buf)) {
+        free_qwen4exp_cache(out);
         return false;
     }
 
@@ -277,6 +286,10 @@ void free_qwen4exp_cache(Qwen4ExpCache & c) {
     }
     if (c.buf) { ggml_backend_buffer_free(c.buf); c.buf = nullptr; }
     if (c.ctx) { ggml_free(c.ctx); c.ctx = nullptr; }
+    if (c.remote_buf) { ggml_backend_buffer_free(c.remote_buf); c.remote_buf = nullptr; }
+    if (c.remote_ctx) { ggml_free(c.remote_ctx); c.remote_ctx = nullptr; }
+    c.remote_layers = 0;
+    c.remote_backend = nullptr;
     c.slot_states = nullptr;
     c.state_slot = -1;
     c.attn_k.clear();
@@ -393,6 +406,7 @@ bool save_qwen4exp_snapshot(ggml_backend_t backend, ggml_backend_t store, const
     });
     s.buf = ggml_backend_alloc_ctx_tensors(s.ctx, store);
     if (!s.buf) { free_qwen4exp_snapshot(s); return false; }
+    if (c.remote_backend) ggml_backend_synchronize(c.remote_backend);   // its layers' last K/V writes
     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;
@@ -423,10 +437,12 @@ bool restore_qwen4exp_snapshot(ggml_backend_t backend, const Qwen4ExpSnapshot &
     ggml_context * views = ggml_init({target.size() * ggml_tensor_overhead(), nullptr, true});
     if (!views) return false;
     ggml_backend_synchronize(backend); // finish rollback before clearing its source/destination buffer
+    if (c.remote_backend) ggml_backend_synchronize(c.remote_backend);
     clear_position_bound_graphs(c);
     // Clear masked suffixes too: stable QSA scores the entire bucket. Stacked slot
     // states are not in c.buf; the snapshot rewrites this slot's slabs in full.
     ggml_backend_buffer_clear(c.buf, 0);
+    if (c.remote_buf) ggml_backend_buffer_clear(c.remote_buf, 0);
     for (size_t i = 0; i < target.size(); ++i) {
         ggml_tensor * copy = s.strips[i].second;
         ggml_tensor * live = ggml_view_1d(views, target[i].t, ggml_nelements(copy), target[i].off);

server/src/qwen4exp/qwen4exp_cache.h

@@ -115,6 +115,13 @@ void free_qwen4exp_slot_states(Qwen4ExpSlotStates & s);
 struct Qwen4ExpCache {
     ggml_context *        ctx     = nullptr;
     ggml_backend_buffer_t buf     = nullptr;
+    // Split mode: the K/V and indexer rows of the last `remote_layers` full-attention layers live on the expert
+    // device, where those layers' attention runs (see Qwen4ExpLayerBuilder::on_cache_device).
+    int                   remote_layers  = 0;
+    ggml_backend_t        remote_backend = nullptr;   // not owned
+    ggml_context *        remote_ctx     = nullptr;
+    ggml_backend_buffer_t remote_buf     = nullptr;
+    bool remote(int fi) const { return fi >= (int) full_layer_ids.size() - remote_layers; }
 
     int       max_ctx  = 0;
     int       cur_pos  = 0;
@@ -198,10 +205,12 @@ struct Qwen4ExpCache {
 
 // `mtp` adds the MTP draft layer's K/V and the verify rollback state (needs a loaded sidecar).
 // `slot_states`: the recurrent state is slab `state_slot` of those (allocated by the caller).
+// `remote_layers`: the last that many full-attention layers keep their rows on w.expert_backend.
 bool create_qwen4exp_cache(ggml_backend_t backend, const Qwen4ExpWeights & w,
                            int max_ctx, Qwen4ExpCache & out, bool mtp = false,
                            int mtp_draft = 1, // allocate the explicit draft cap once
-                           const Qwen4ExpSlotStates * slot_states = nullptr, int state_slot = -1);
+                           const Qwen4ExpSlotStates * slot_states = nullptr, int state_slot = -1,
+                           int remote_layers = 0);
 
 void free_qwen4exp_cache(Qwen4ExpCache & c);