diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 00476f0fd134..42aacc82ae8b 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1465,14 +1465,20 @@ struct ggml_backend_cuda_context { int curr_stream_no = 0; #ifdef USE_CUDA_GRAPH - // Map from first_node_ptr to cuda_graph - allows multiple graphs per context - // when the computation is split across CPU/GPU (e.g., with --n-cpu-moe) - std::unordered_map> cuda_graphs; + // Map from graph key to cuda_graph - allows multiple graphs per context when the + // computation is split across CPU/GPU (e.g., with --n-cpu-moe), and when the same + // split is called with different tensor shapes (e.g. a speculative verify batch) + // a cuda graph instance is only valid for the shapes it captured, so a caller that + // alternates shapes needs one instance per shape to stay on the graph path. the map is + // bounded only by the time based sweep below, as upstream does: tensor split mode alone + // needs about 2 * n_layers + 1 graphs per device and shape, so any count cap small enough + // to matter evicts graphs that are still in use and re-captures them on every token + std::unordered_map> cuda_graphs; int64_t next_graph_eviction_time = 0; int64_t last_graph_use_time = 0; - ggml_cuda_graph * cuda_graph(const void * first_node_ptr, int64_t time_now = ggml_time_us()) { + ggml_cuda_graph * cuda_graph(uint64_t graph_key, int64_t time_now = ggml_time_us()) { // Give all cached graphs 10s to be reused after an idle period, without updating each graph. if (!cuda_graphs.empty() && time_now - last_graph_use_time >= 10'000'000) { next_graph_eviction_time = time_now + 10'000'000; @@ -1483,7 +1489,7 @@ struct ggml_backend_cuda_context { if (time_now >= next_graph_eviction_time) { next_graph_eviction_time = time_now + 5'000'000; for (auto it = cuda_graphs.begin(); it != cuda_graphs.end(); ) { - if (it->first != first_node_ptr && time_now - it->second->last_used_time >= 10'000'000) { + if (it->first != graph_key && time_now - it->second->last_used_time >= 10'000'000) { it = cuda_graphs.erase(it); } else { ++it; @@ -1491,9 +1497,9 @@ struct ggml_backend_cuda_context { } } - auto it = cuda_graphs.find(first_node_ptr); + auto it = cuda_graphs.find(graph_key); if (it == cuda_graphs.end()) { - it = cuda_graphs.emplace(first_node_ptr, std::make_unique()).first; + it = cuda_graphs.emplace(graph_key, std::make_unique()).first; } it->second->last_used_time = time_now; return it->second.get(); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 06956ce7a3d9..e1910f4f2a6c 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2891,14 +2891,40 @@ static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { return use_cuda_graph; } -static const void * ggml_cuda_graph_get_key(ggml_cgraph * cgraph) { - return cgraph->nodes[0]; +// the key identifies both the split (its first node) and the shapes it was called with. +// a captured cuda graph hard-codes the shapes, so a caller that alternates shapes - a +// speculative verify batch, for example - needs a separate instance per shape. with a +// single key per split, every shape change resets the warmup and no graph is ever used. +// +// this stays O(1) on purpose: walking every node undoes the point of a cuda graph, which is +// to not touch per-node data on the hot path. the first and last node carry the batch +// dimension, which is what changes when a verify batch changes size. a shape this does not +// separate just shares an entry and re-captures, exactly as before, so it can only help. +static uint64_t ggml_cuda_graph_get_key(ggml_cgraph * cgraph) { + if (cgraph->n_nodes == 0) { + return 0; + } + + uint64_t key = (uint64_t) (uintptr_t) cgraph->nodes[0]; + + auto mix = [&key](uint64_t v) { + key = (key ^ v) * 0x100000001b3ull; + }; + + mix(cgraph->n_nodes); + + for (int d = 0; d < GGML_MAX_DIMS; d++) { + mix(cgraph->nodes[0]->ne[d]); + mix(cgraph->nodes[cgraph->n_nodes - 1]->ne[d]); + } + + return key; } static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph) { bool res = false; - const void * graph_key = ggml_cuda_graph_get_key(cgraph); + const uint64_t graph_key = ggml_cuda_graph_get_key(cgraph); ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); if (cgraph->uid != 0 && @@ -4750,7 +4776,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return 0; } -static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, const bool use_cuda_graph, const bool cuda_graph_update_required, const void * graph_key) { +static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, const bool use_cuda_graph, const bool cuda_graph_update_required, uint64_t graph_key) { bool graph_evaluated_or_captured = false; // flag used to determine whether it is an integrated_gpu @@ -4980,7 +5006,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } #ifdef USE_CUDA_GRAPH -static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { +static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, uint64_t graph_key) { ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); if (graph->instances.empty()) { @@ -5003,7 +5029,7 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, bool use_cuda_graph = false; bool cuda_graph_update_required = false; - const void * graph_key = nullptr; + uint64_t graph_key = 0; #ifdef USE_CUDA_GRAPH graph_key = ggml_cuda_graph_get_key(cgraph); @@ -5211,7 +5237,7 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph } #ifdef USE_CUDA_GRAPH - const void * graph_key = ggml_cuda_graph_get_key(cgraph); + const uint64_t graph_key = ggml_cuda_graph_get_key(cgraph); const bool use_cuda_graph = ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); #else const bool use_cuda_graph = false; diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index da17476494ef..186de292c662 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -62,10 +62,13 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( hparams_idx.n_embd_head_k_mla_impl = model.hparams.indexer_head_size; hparams_idx.n_embd_head_v_mla_impl = model.hparams.indexer_head_size; - LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size); + // glm5-next caches the compressor gates beside the keys and they feed a softmax, so a quantized -ctk stays F16 here + const ggml_type type_idx = model.arch == LLM_ARCH_GLM5_NEXT && ggml_is_quantized(type_k) ? GGML_TYPE_F16 : type_k; + + LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells, type = %s\n", __func__, kv_size, ggml_type_name(type_idx)); return new llama_kv_cache( - model, no_alloc, hparams_idx, type_k, type_v, v_trans, offload, unified, + model, no_alloc, hparams_idx, type_idx, type_v, v_trans, offload, unified, kv_size, n_seq_max, n_swa, swa_type, nullptr, filter_idx, nullptr, nullptr, "idx_"); }()) {}