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/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", { 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]); +}