From abe19e9a51a6591775561f1482e963fe418588c4 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Fri, 9 Oct 2026 20:52:20 +0200 Subject: [PATCH 1/2] fix(ggml-cuda): copy a tensor between two devices on the source's stream ggml_backend_cuda_buffer_cpy_tensor queued its peer copy on the per-thread stream of whichever device was current. On ROCm a copy from an R9700 to a Strix Halo iGPU queued on the destination's stream completes without writing, so the scheduler's copies of graph inputs to the other device (row indices, positions, masks) arrived as zeros. It now sets the source device first, as ggml_backend_cuda_cpy_tensor_async queues its copy on the source's stream. Split serving never sent a graph input to the expert device until the next change, which runs some attention layers there. test_cuda_copy_batch gains cross_device_tensor_copy (skipped below two devices). On an R9700 + Strix Halo it reads 0 0 0 0 without the fix and 19 20 21 22 with it. Co-Authored-By: Claude Opus 5.5 --- .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 4 +++ server/test/unit/test_cuda_copy_batch.cpp | 34 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index 35bc9cda6..2721a37b0 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/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 { diff --git a/server/test/unit/test_cuda_copy_batch.cpp b/server/test/unit/test_cuda_copy_batch.cpp index 4fabae97a..8294954d0 100644 --- a/server/test/unit/test_cuda_copy_batch.cpp +++ b/server/test/unit/test_cuda_copy_batch.cpp @@ -218,3 +218,37 @@ TEST_CASE(CudaCopyBatchFixture, graph_copy_long_run) { REQUIRE(runs >= 2); } } + +// ggml_backend_tensor_copy between two devices, the last call having selected the destination (as the scheduler's +// copy of a graph input to another device does): the destination must hold the source. +TEST_CASE(CudaCopyBatchFixture, cross_device_tensor_copy) { + if (ggml_backend_cuda_get_device_count() < 2) { + SKIP("needs two CUDA/HIP devices"); + } + ggml_backend_t gpus[2] = { ggml_backend_cuda_init(0), ggml_backend_cuda_init(1) }; + REQUIRE(gpus[0] != nullptr); + REQUIRE(gpus[1] != nullptr); + ggml_init_params ip{}; + ip.mem_size = ggml_tensor_overhead(); + ip.no_alloc = true; + ggml_context * ctx[2] = { ggml_init(ip), ggml_init(ip) }; + ggml_tensor * src = ggml_new_tensor_1d(ctx[0], GGML_TYPE_I32, 4); + ggml_tensor * dst = ggml_new_tensor_1d(ctx[1], GGML_TYPE_I32, 4); + ggml_backend_buffer_t buf[2] = { ggml_backend_alloc_ctx_tensors(ctx[0], gpus[0]), + ggml_backend_alloc_ctx_tensors(ctx[1], gpus[1]) }; + REQUIRE(buf[0] != nullptr); + REQUIRE(buf[1] != nullptr); + const int32_t rows[4] = { 19, 20, 21, 22 }, zero[4] = {}; + ggml_backend_tensor_set(src, rows, 0, sizeof(rows)); + ggml_backend_tensor_set(dst, zero, 0, sizeof(zero)); // selects the destination device + ggml_backend_tensor_copy(src, dst); + int32_t got[4] = {}; + ggml_backend_tensor_get(dst, got, 0, sizeof(got)); + std::printf("[cuda-copy-batch] cross-device copy: %d %d %d %d\n", got[0], got[1], got[2], got[3]); + for (int i = 0; i < 2; ++i) { + ggml_backend_buffer_free(buf[i]); + ggml_free(ctx[i]); + ggml_backend_free(gpus[i]); + } + for (int i = 0; i < 4; ++i) REQUIRE(got[i] == rows[i]); +} From cb6116c6526e97895ae521817883f28fa83899ec Mon Sep 17 00:00:00 2001 From: mrciffa Date: Fri, 9 Oct 2026 21:08:08 +0200 Subject: [PATCH 2/2] feat(qwen4exp): split serving keeps attention layers' KV on the Strix when the R9700 is full `--max-ctx auto` lowered the context once the R9700 could not hold every slot's cache beside the dense weights: four requests at once got 131,072 tokens. A miss now first moves the last full-attention layers' K/V and indexer rows to the expert device, as few layers as cover the shortfall while that device holds them, and runs those layers' attention there (input projections through the output projection). The context shrinks only once no more layers can move. - Qwen4ExpCache::remote_layers: those rows live in a second buffer on the expert device; snapshots save and restore them with the rest. - move_qwen4exp_attention moves the layers' attention and indexer weights. - Qwen4ExpLayerBuilder::on_cache_device pins the nodes a layer's attention adds to the expert device (single-sequence, verify, pipelined prompt and multi-slot builders); qwen4exp_pin_split places them. - The chunk plan counts the rows on the target only and reports one layer's rows across slots (Qwen4ExpChunkPlan::layer_bytes). - The QSA prefix bootstrap pools each layer on the device holding its rows. R9700 + Strix Halo, Manage's split launcher, four slots: 262,144 tokens with 7 of 12 layers moved, 3.9 GiB left on the R9700. Outputs repeat run to run and match alone and together (4 of 4); a resent 16K prompt prefills in 0.25 s instead of 11.0 s with the same reply; three facts hidden in a 158,927-token prompt are all recalled; a 158,848-token prompt completes beside three short requests. A request alone decodes at 45.9 tok/s and four together at 68.6, against 59.3 and 78.3 at 131,072 tokens with every row on the R9700. Co-Authored-By: Claude Opus 5.5 --- server/docs/QWEN4EXP.md | 11 +- server/src/qwen4exp/qwen4exp_backend.cpp | 90 +++++++++++++--- server/src/qwen4exp/qwen4exp_backend.h | 2 + server/src/qwen4exp/qwen4exp_cache.cpp | 32 ++++-- server/src/qwen4exp/qwen4exp_cache.h | 11 +- server/src/qwen4exp/qwen4exp_chunk.h | 3 +- server/src/qwen4exp/qwen4exp_graph.cpp | 128 ++++++++++++++--------- server/src/server/launch_profiles.h | 3 +- 8 files changed, 204 insertions(+), 76 deletions(-) diff --git a/server/docs/QWEN4EXP.md b/server/docs/QWEN4EXP.md index c72db8ac9..762568a2a 100644 --- a/server/docs/QWEN4EXP.md +++ b/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 diff --git a/server/src/qwen4exp/qwen4exp_backend.cpp b/server/src/qwen4exp/qwen4exp_backend.cpp index 2f3df2410..45b7150cd 100644 --- a/server/src/qwen4exp/qwen4exp_backend.cpp +++ b/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 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 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 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(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 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); } diff --git a/server/src/qwen4exp/qwen4exp_backend.h b/server/src/qwen4exp/qwen4exp_backend.h index 7f132fa69..1e5346a56 100644 --- a/server/src/qwen4exp/qwen4exp_backend.h +++ b/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 snapshots_; int live_slot_ = -1; diff --git a/server/src/qwen4exp/qwen4exp_cache.cpp b/server/src/qwen4exp/qwen4exp_cache.cpp index d0e5bb7b1..3c8935cec 100644 --- a/server/src/qwen4exp/qwen4exp_cache.cpp +++ b/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(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(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); diff --git a/server/src/qwen4exp/qwen4exp_cache.h b/server/src/qwen4exp/qwen4exp_cache.h index e1d7c9da4..ab9f9ce48 100644 --- a/server/src/qwen4exp/qwen4exp_cache.h +++ b/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); diff --git a/server/src/qwen4exp/qwen4exp_chunk.h b/server/src/qwen4exp/qwen4exp_chunk.h index 714a924f8..e95f270a8 100644 --- a/server/src/qwen4exp/qwen4exp_chunk.h +++ b/server/src/qwen4exp/qwen4exp_chunk.h @@ -90,7 +90,8 @@ struct Qwen4ExpContextFit { 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 + size_t position_bytes = 0; // one context position across every slot's cache on the device + size_t layer_bytes = 0; // one full-attention layer's rows across every slot's cache }; struct Qwen4ExpWeights; diff --git a/server/src/qwen4exp/qwen4exp_graph.cpp b/server/src/qwen4exp/qwen4exp_graph.cpp index 98e602387..96b31be78 100644 --- a/server/src/qwen4exp/qwen4exp_graph.cpp +++ b/server/src/qwen4exp/qwen4exp_graph.cpp @@ -1292,14 +1292,12 @@ Qwen4ExpInputs qwen4exp_prepare_inputs(const Qwen4ExpWeights & w, return res; } -// Split mode placement: the routed experts on the expert device, every other -// node and the graph inputs on the target. Left unpinned, the scheduler places -// a node after its graph-order neighbour, often an expert split, and keeps graph -// inputs on the CPU, copying them to each split that reads them; a strided view -// of an input (MTP prefill's per-block positions) cannot be copied as one span. +// Split mode placement: the routed experts, and the attention of layers whose cache rows live there, on the expert +// device; every other node and the graph inputs on the target. Left unpinned, the scheduler places a node after its +// graph-order neighbour, often an expert split, and keeps graph inputs on the CPU, copying them to each split that +// reads them; a strided view of an input (MTP prefill's per-block positions) cannot be copied as one span. static void qwen4exp_pin_split(ggml_backend_sched_t sched, ggml_cgraph * gf, ggml_backend_t target, ggml_backend_t expert, const std::vector & expert_nodes) { - for (ggml_tensor * t : expert_nodes) ggml_backend_sched_set_tensor_backend(sched, t, expert); const std::unordered_set on_expert(expert_nodes.begin(), expert_nodes.end()); for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { ggml_tensor * node = ggml_graph_node(gf, i); @@ -1312,7 +1310,9 @@ static void qwen4exp_pin_split(ggml_backend_sched_t sched, ggml_cgraph * gf, ggm } const bool view = node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE; - if (!view && !on_expert.count(node) && ggml_backend_supports_op(target, node)) { + if (on_expert.count(node)) { + ggml_backend_sched_set_tensor_backend(sched, node, expert); + } else if (!view && ggml_backend_supports_op(target, node)) { ggml_backend_sched_set_tensor_backend(sched, node, target); } } @@ -1373,6 +1373,24 @@ struct Qwen4ExpLayerBuilder { // Layer il + 1 starts with PLE, so layer il's output cannot be fused with its attention norm. bool next_is_ple(int il) const { return il + 1 < w.n_layer && w.layers[il + 1].is_ple && has_ple; } + // Full-attention layer `fi`'s attention, from its input projections through the output projection: with its + // cache rows on the expert device (Qwen4ExpCache::remote), every node `build` adds runs there, beside those rows + // and the layer's weights (move_qwen4exp_attention); only the input and the output cross. + template ggml_tensor * on_cache_device(int fi, Build build) { + if (!cache.remote(fi)) return build(); + ggml_tensor * last = nullptr; + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t; t = ggml_get_next_tensor(ctx, t)) last = t; + ggml_tensor * out = build(); + for (ggml_tensor * t = last ? ggml_get_next_tensor(ctx, last) : ggml_get_first_tensor(ctx); t; + t = ggml_get_next_tensor(ctx, t)) { + // Views stay with their source, as qwen4exp_pin_split leaves them; inputs stay on the target. + const bool view = t->op == GGML_OP_NONE || t->op == GGML_OP_VIEW || t->op == GGML_OP_RESHAPE || + t->op == GGML_OP_PERMUTE || t->op == GGML_OP_TRANSPOSE; + if (!view && !(t->flags & GGML_TENSOR_FLAG_INPUT)) expert_nodes.push_back(t); + } + return out; + } + // The HC mix into layer il's attention input, from the previous layer's fused norm when xn_next holds it // (consumed). `inject` receives the mix's injection. ggml_tensor * attn_input(int il, ggml_tensor * res_hc, ggml_tensor *& xn_next, ggml_tensor *& inject) { @@ -1397,12 +1415,14 @@ struct Qwen4ExpLayerBuilder { ggml_tensor * cur = attn_input(il, res_hc, xn_next, inject); if (L.is_full_attention) { const int fi = full_idx[il]; - return build_full_attn(ctx, gf, cur, L, w, - cache.attn_k[fi], cache.attn_v[fi], cache.indexer_k[fi], - w.qsa ? cache.indexer_raw[fi] : nullptr, - b.positions, b.mask, kv_row, b.kv_len, b.pos0, - il < (int) w.compress_ratios.size() ? w.compress_ratios[il] : 0, - b.n_pooled, b.qsa, rows, qsa_ws); + return on_cache_device(fi, [&] { + return build_full_attn(ctx, gf, cur, L, w, + cache.attn_k[fi], cache.attn_v[fi], cache.indexer_k[fi], + w.qsa ? cache.indexer_raw[fi] : nullptr, + b.positions, b.mask, kv_row, b.kv_len, b.pos0, + il < (int) w.compress_ratios.size() ? w.compress_ratios[il] : 0, + b.n_pooled, b.qsa, rows, qsa_ws); + }); } const int li = lin_idx[il]; return build_linear_attn(ctx, gf, cur, L, w, cache.ssm_state[li], cache.conv_state[li], @@ -1760,24 +1780,30 @@ static bool run_stable_graph(Qwen4ExpStableInputs & in, Qwen4ExpDecodeWorkspace } // Dense decode may not have maintained the pooled QSA prefix: pool the already-cached complete blocks before pos0 -// once. The stable graph handles this forward's own blocks. +// once, on the device holding each layer's rows (Qwen4ExpCache::remote). The stable graph handles this forward's own +// blocks. static bool bootstrap_qsa_prefix(ggml_backend_t backend, const Qwen4ExpWeights & w, Qwen4ExpCache & cache, int pos0) { - ggml_context * bc = ggml_init({8 * 1024 * 1024, nullptr, true}); - if (!bc) return false; - ggml_cgraph * bg = ggml_new_graph_custom(bc, 8192, false); - for (size_t fi = 0; fi < cache.full_layer_ids.size(); ++fi) { - const auto & L = w.layers[cache.full_layer_ids[fi]]; - qsa_pooled_keys(bc, bg, L, w, cache.indexer_k[fi], cache.indexer_raw[fi], nullptr, - pos0, 4, cache.indexer_blocks, pos0 / 4); - } - ggml_gallocr_t ba = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - const bool ok = ba && ggml_gallocr_alloc_graph(ba, bg) && - ggml_backend_graph_compute(backend, bg) == GGML_STATUS_SUCCESS; - ggml_backend_cuda_graph_invalidate_range(backend, ggml_get_mem_buffer(bc), ggml_get_mem_size(bc)); - if (ba) ggml_gallocr_free(ba); - ggml_free(bc); - if (ok) cache.indexer_blocks = pos0 / 4; - return ok; + for (ggml_backend_t device : {backend, cache.remote_backend}) { + if (!device) continue; + ggml_context * bc = ggml_init({8 * 1024 * 1024, nullptr, true}); + if (!bc) return false; + ggml_cgraph * bg = ggml_new_graph_custom(bc, 8192, false); + for (size_t fi = 0; fi < cache.full_layer_ids.size(); ++fi) { + if (cache.remote((int) fi) != (device != backend)) continue; + const auto & L = w.layers[cache.full_layer_ids[fi]]; + qsa_pooled_keys(bc, bg, L, w, cache.indexer_k[fi], cache.indexer_raw[fi], nullptr, + pos0, 4, cache.indexer_blocks, pos0 / 4); + } + ggml_gallocr_t ba = ggml_graph_n_nodes(bg) ? ggml_gallocr_new(ggml_backend_get_default_buffer_type(device)) : nullptr; + const bool ok = !ggml_graph_n_nodes(bg) || (ba && ggml_gallocr_alloc_graph(ba, bg) && + ggml_backend_graph_compute(device, bg) == GGML_STATUS_SUCCESS); + ggml_backend_cuda_graph_invalidate_range(device, ggml_get_mem_buffer(bc), ggml_get_mem_size(bc)); + if (ba) ggml_gallocr_free(ba); + ggml_free(bc); + if (!ok) return false; + } + cache.indexer_blocks = pos0 / 4; + return true; } // The graph's metadata context: a decode or verify workspace's own (`pool`), reset for a rebuild, or a fresh one. @@ -2904,27 +2930,29 @@ static bool build_batched_decode_nodes(ggml_context * ctx, const Qwen4ExpWeights ggml_tensor * cur = lb.attn_input(il, G.res_hc, G.xn_next, G.inject); if (L.is_full_attention) { - ggml_tensor * qfull = mm(ctx, L.wq, cur); - ggml_tensor * kraw = mm(ctx, L.wk, cur); - ggml_tensor * vraw = mm(ctx, L.wv, cur); - if (ws.kv_rows.empty()) { - // One mask and K/V row per slot and one M-RoPE position input per group, shared by - // every attention layer (a split scheduler tracks at most GGML_SCHED_MAX_SPLIT_INPUTS). - for (int s = 0; s < n_slots; ++s) { - const int64_t span = spans[(size_t) s]; - ws.masks.push_back(span > 0 ? ggml_new_tensor_2d(ctx, GGML_TYPE_F16, span, 1) : nullptr); - ws.kv_rows.push_back(ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1)); - if (ws.masks.back()) ggml_set_input(ws.masks.back()); - ggml_set_input(ws.kv_rows.back()); + cur = lb.on_cache_device(lb.full_idx[il], [&] { + ggml_tensor * qfull = mm(ctx, L.wq, cur); + ggml_tensor * kraw = mm(ctx, L.wk, cur); + ggml_tensor * vraw = mm(ctx, L.wv, cur); + if (ws.kv_rows.empty()) { + // One mask and K/V row per slot and one M-RoPE position input per group, shared by + // every attention layer (a split scheduler tracks at most GGML_SCHED_MAX_SPLIT_INPUTS). + for (int s = 0; s < n_slots; ++s) { + const int64_t span = spans[(size_t) s]; + ws.masks.push_back(span > 0 ? ggml_new_tensor_2d(ctx, GGML_TYPE_F16, span, 1) : nullptr); + ws.kv_rows.push_back(ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1)); + if (ws.masks.back()) ggml_set_input(ws.masks.back()); + ggml_set_input(ws.kv_rows.back()); + } } - } - if (!G.positions) { - G.positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 4 * n); - ggml_set_input(G.positions); - } - cur = mm(ctx, L.wo, build_full_attn_slots(ctx, gf, qfull, kraw, vraw, cur, G.positions, - ws.masks.data() + G.s0, ws.kv_rows.data() + G.s0, spans.data() + G.s0, L, w, - caches + G.s0, lb.full_idx[il], use_qsa, qsa.data() + G.s0)); + if (!G.positions) { + G.positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 4 * n); + ggml_set_input(G.positions); + } + return mm(ctx, L.wo, build_full_attn_slots(ctx, gf, qfull, kraw, vraw, cur, G.positions, + ws.masks.data() + G.s0, ws.kv_rows.data() + G.s0, spans.data() + G.s0, L, w, + caches + G.s0, lb.full_idx[il], use_qsa, qsa.data() + G.s0)); + }); } else { // Dense input projections are batched; with stacked slot states the conv and // recurrence run once for the group, else each slot's sees one row and its own state. diff --git a/server/src/server/launch_profiles.h b/server/src/server/launch_profiles.h index dcdfde49b..16cc1f878 100644 --- a/server/src/server/launch_profiles.h +++ b/server/src/server/launch_profiles.h @@ -231,7 +231,8 @@ inline const std::vector & launch_profiles() { { "qwen-next-r9700-strix", "Qwen3.8-Flash-Next with dense work, KV cache and the MTP drafter on " - "an R9700 (gfx1201) and the routed experts on Strix Halo (gfx1151); " + "an R9700 (gfx1201) and the routed experts on Strix Halo (gfx1151), " + "which also keeps some attention layers' KV cache when the R9700 is full; " "add --expert-placement with the box's routing stats to also run " "the most-routed experts on the R9700", {