diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index d6097649..2c1e3667 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -6238,17 +6238,17 @@ struct ggml_tensor * ggml_mul_multi_add( return result; } -#if defined(_MSC_VER) -#pragma warning(disable: 4244 4267) // possible loss of data -#include -#include -#include -#include -#include -static inline int popcount(uint32_t x) { return __popcnt(x); } -#else -static inline int popcount(uint32_t x) { return __builtin_popcount(x); } -#endif +//#if defined(_MSC_VER) +//#pragma warning(disable: 4244 4267) // possible loss of data +//#include +//#include +//#include +//#include +//#include +//static inline int popcount(uint32_t x) { return __popcnt(x); } +//#else +//static inline int popcount(uint32_t x) { return __builtin_popcount(x); } +//#endif struct ggml_tensor * ggml_hadamard( struct ggml_context * ctx, @@ -6256,8 +6256,16 @@ struct ggml_tensor * ggml_hadamard( int n) { GGML_ASSERT(n > 1); // no point in Hadamard transforms with less than 2 elements - GGML_ASSERT(a->ne[0] % n == 0); - GGML_ASSERT(popcount(n) == 1); // must be a power of 2 + if (a->ne[0] % n != 0) { + fprintf(stderr, "%s: head size %ld is not a multiple of block size %d for tensor %s\n", __func__, a->ne[0], n, a->name); + GGML_ABORT("Fatal error"); + } + if ((n & ~(n-1)) != n) { + fprintf(stderr, "%s: block size %d is not a power of 2 for tensor %s\n", __func__, n, a->name); + GGML_ABORT("Fatal error"); + } + //GGML_ASSERT(a->ne[0] % n == 0); + //GGML_ASSERT(popcount(n) == 1); // must be a power of 2 struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, a->ne); diff --git a/src/graphs/build_gemma4.cpp b/src/graphs/build_gemma4.cpp index 355db0c4..64853ea6 100644 --- a/src/graphs/build_gemma4.cpp +++ b/src/graphs/build_gemma4.cpp @@ -301,14 +301,18 @@ static ggml_cgraph * build_gemma4_graph_parallel(llm_build_context & llm, llama_ const int64_t n_head_kv = wk->splits[id]->ne[1] / n_embd_head_k; if (cparams.k_cache_hadamard) { - Qcur = ggml_hadamard(ctx0, Qcur, n_embd_head_k); - Kcur = ggml_hadamard(ctx0, Kcur, n_embd_head_k); - cb(Qcur, "Qcur_h", il_cb); - cb(Kcur, "Kcur_h", il_cb); + if (int block_size = lctx.model.hadamard_size_k(il); block_size > 0) { + Qcur = ggml_hadamard(ctx0, Qcur, block_size); + Kcur = ggml_hadamard(ctx0, Kcur, block_size); + cb(Qcur, "Qcur_h", il_cb); + cb(Kcur, "Kcur_h", il_cb); + } } if (cparams.v_cache_hadamard) { - Vcur = ggml_hadamard(ctx0, Vcur, n_embd_head_v); - cb(Vcur, "Vcur_h", il_cb); + if (int block_size = lctx.model.hadamard_size_v(il); block_size > 0) { + Vcur = ggml_hadamard(ctx0, Vcur, block_size); + cb(Vcur, "Vcur_h", il_cb); + } } GGML_ASSERT(kv_self.size == cparams.n_ctx); @@ -357,8 +361,10 @@ static ggml_cgraph * build_gemma4_graph_parallel(llm_build_context & llm, llama_ cb(cur, "fa", il_cb); cur->op_params[4] = n_swa; if (cparams.v_cache_hadamard) { - cur = ggml_hadamard(ctx0, cur, n_embd_head_v); - cb(cur, "fa_h", il_cb); + if (int block_size = lctx.model.hadamard_size_v(il); block_size > 0) { + cur = ggml_hadamard(ctx0, cur, block_size); + cb(cur, "fa_h", il_cb); + } } cur = ggml_reshape_2d(ctx0, cur, wo->splits[id]->ne[0], n_tokens); if (il == hparams.n_layer-1 && inp_out_ids) { diff --git a/src/graphs/build_minimaxm2.cpp b/src/graphs/build_minimaxm2.cpp index 07497c7c..8c632616 100644 --- a/src/graphs/build_minimaxm2.cpp +++ b/src/graphs/build_minimaxm2.cpp @@ -103,17 +103,21 @@ ggml_cgraph* llm_build_context::build_minimaxm2() { cb(Kcur, "Kcur_roped", il_id); if (cparams.k_cache_hadamard) { - Qcur = ggml_hadamard(ctx0, Qcur, n_embd_head_k); - Kcur = ggml_hadamard(ctx0, Kcur, n_embd_head_k); - cb(Qcur, "Qcur_hadamard", il_id); - cb(Kcur, "Kcur_hadamard", il_id); + if (int block_size = lctx.model.hadamard_size_k(il); block_size > 0) { + Qcur = ggml_hadamard(ctx0, Qcur, block_size); + Kcur = ggml_hadamard(ctx0, Kcur, block_size); + cb(Qcur, "Qcur_hadamard", il_id); + cb(Kcur, "Kcur_hadamard", il_id); + } } ggml_build_forward_expand(gf, Qcur); ggml_build_forward_expand(gf, Kcur); if (cparams.v_cache_hadamard) { - Vcur = ggml_hadamard(ctx0, Vcur, n_embd_head_v); - cb(Vcur, "Vcur_hadamard", il_id); - ggml_build_forward_expand(gf, Vcur); + if (int block_size = lctx.model.hadamard_size_v(il); block_size > 0) { + Vcur = ggml_hadamard(ctx0, Vcur, block_size); + cb(Vcur, "Vcur_hadamard", il_id); + ggml_build_forward_expand(gf, Vcur); + } } // Store K, V in KV cache @@ -150,8 +154,10 @@ ggml_cgraph* llm_build_context::build_minimaxm2() { cb(cur, "fa", il_id); if (cparams.v_cache_hadamard) { - cur = ggml_hadamard(ctx0, cur, n_embd_head_v); - cb(cur, "fa_h", il_id); + if (int block_size = lctx.model.hadamard_size_v(il); block_size > 0) { + cur = ggml_hadamard(ctx0, cur, block_size); + cb(cur, "fa_h", il_id); + } } cur = ggml_reshape_2d(ctx0, cur, wo->splits[id]->ne[0], n_tokens); diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index 04fe1911..54a649ab 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -1635,8 +1635,10 @@ static ggml_tensor * llm_build_kqv( //ggml_flash_attn_ext_set_prec(cur, GGML_PREC_F32); if (cparams.v_cache_hadamard) { - cur = ggml_hadamard(ctx, cur, n_embd_head_v); - cb(cur, "fa_h", il); + if (int block_size = lctx.model.hadamard_size_v(il); block_size > 0) { + cur = ggml_hadamard(ctx, cur, block_size); + cb(cur, "fa_h", il); + } } cur = ggml_reshape_2d(ctx, cur, n_embd_head_v*n_head, n_tokens); } else { @@ -1802,15 +1804,19 @@ ggml_tensor * llm_build_context::llm_build_kv( const llama_cparams & cparams = lctx.cparams; if (cparams.k_cache_hadamard) { - q_cur = ggml_hadamard(ctx, q_cur, hparams.n_embd_head_k(il)); - if (k_cur) { - k_cur = ggml_hadamard(ctx, k_cur, hparams.n_embd_head_k(il)); - cb(k_cur, "Kcur_hadamard", il); + if (int block_size = lctx.model.hadamard_size_k(il); block_size > 0) { + q_cur = ggml_hadamard(ctx, q_cur, block_size); + if (k_cur) { + k_cur = ggml_hadamard(ctx, k_cur, block_size); + cb(k_cur, "Kcur_hadamard", il); + } + cb(q_cur, "Qcur_hadamard", il); } - cb(q_cur, "Qcur_hadamard", il); } if (cparams.v_cache_hadamard && v_cur) { - v_cur = ggml_hadamard(ctx, v_cur, hparams.n_embd_head_v(il)); + if (int block_size = lctx.model.hadamard_size_v(il); block_size > 0) { + v_cur = ggml_hadamard(ctx, v_cur, block_size); + } } // these nodes are added to the graph together so that they are not reordered @@ -2649,14 +2655,18 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens cb(Qcur, "Qcur_temp_scaled", il_cb); } if (cparams.k_cache_hadamard) { - Qcur = ggml_hadamard(ctx0, Qcur, hparams.n_embd_head_k(il)); - Kcur = ggml_hadamard(ctx0, Kcur, hparams.n_embd_head_k(il)); - cb(Qcur, "Qcur_hadamard", il_cb); - cb(Kcur, "Kcur_hadamard", il_cb); + if (int block_size = lctx.model.hadamard_size_k(il); block_size > 0) { + Qcur = ggml_hadamard(ctx0, Qcur, block_size); + Kcur = ggml_hadamard(ctx0, Kcur, block_size); + cb(Qcur, "Qcur_hadamard", il_cb); + cb(Kcur, "Kcur_hadamard", il_cb); + } } if (cparams.v_cache_hadamard) { - Vcur = ggml_hadamard(ctx0, Vcur, hparams.n_embd_head_v(il)); - cb(Vcur, "Vcur_hadamard", il_cb); + if (int block_size = lctx.model.hadamard_size_v(il); block_size > 0) { + Vcur = ggml_hadamard(ctx0, Vcur, block_size); + cb(Vcur, "Vcur_hadamard", il_cb); + } } ggml_build_forward_expand(gf, Qcur); ggml_build_forward_expand(gf, Kcur); @@ -2732,8 +2742,10 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens } if (cparams.v_cache_hadamard) { - cur = ggml_hadamard(ctx0, cur, n_embd_head_v); - cb(cur, "flash_attn_h", il_cb); + if (int block_size = lctx.model.hadamard_size_v(il); block_size > 0) { + cur = ggml_hadamard(ctx0, cur, block_size); + cb(cur, "flash_attn_h", il_cb); + } } if (model.layers[il].wqkv_gate) { diff --git a/src/llama-model.h b/src/llama-model.h index 347187c6..5ff084fe 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -519,6 +519,26 @@ struct llama_model { return arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_MISTRAL4; } + static inline int hadamard_size(int head_size) { + if ((head_size & ~(head_size - 1)) == head_size) return head_size; + // Note: we do not include 32 as an option because the CUDA Hadamard implementation + // does not hcurrently andle a block size of 32. + for (int i = 512; i >= 64; i >>= 1) { + if (head_size % i == 0) return i; + } + return 0; + } + + inline int hadamard_size_k(int il) const { + if (is_mla_model()) return 64; + return hadamard_size(hparams.n_embd_head_k(il)); + } + + inline int hadamard_size_v(int il) const { + if (is_mla_model()) return 64; + return hadamard_size(hparams.n_embd_head_v(il)); + } + size_t cache_size(int il, ggml_type type_k, ggml_type type_v, uint32_t kv_size, int mla_attn, int n_seq_max, bool flash_attn) const; void set_tensor_overrides(const llama_model_params& params); diff --git a/src/llama.cpp b/src/llama.cpp index a12066e9..67e05a47 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -798,7 +798,7 @@ static bool llama_kv_cache_init( } } - bool is_mla_attn = model.arch == LLM_ARCH_DEEPSEEK2 || model.arch == LLM_ARCH_GLM_DSA || model.arch == LLM_ARCH_MISTRAL4; + bool is_mla_attn = model.is_mla_model(); bool split_cache = false; bool replicate_mla = false; @@ -2111,7 +2111,7 @@ static void llm_load_print_meta(llama_model_loader & ml, llama_model & model) { // general kv LLAMA_LOG_INFO("%s: general.name = %s\n", __func__, model.name.c_str()); - if (model.arch == LLM_ARCH_DEEPSEEK2 || model.arch == LLM_ARCH_GLM_DSA || model.arch == LLM_ARCH_MISTRAL4) { + if (model.is_mla_model()) { LLAMA_LOG_INFO("%s: n_layer_dense_lead = %d\n", __func__, hparams.n_layer_dense_lead); LLAMA_LOG_INFO("%s: n_lora_q = %d\n", __func__, hparams.n_lora_q); LLAMA_LOG_INFO("%s: n_lora_kv = %d\n", __func__, hparams.n_lora_kv); @@ -2241,7 +2241,7 @@ static void llm_requantize_output_tensor(llama_model & model, ggml_type new_type } static void llm_prepare_mla(llama_model & model, int mla) { - if (model.arch != LLM_ARCH_DEEPSEEK2 && model.arch != LLM_ARCH_GLM_DSA && model.arch != LLM_ARCH_MISTRAL4) return; + if (!model.is_mla_model()) return; const auto& hparams = model.hparams; const int n_layer = model.layers.size(); int n_to_compute = 0; @@ -2815,7 +2815,7 @@ static void llm_prepare_mla(llama_model & model, int mla) { // skips the runtime cache_nope un-Hadamard. Math identity by H^T H = I. static void llm_apply_khad_pretransform(llama_model & model) { if (model.khad_pretransformed) return; - if (model.arch != LLM_ARCH_DEEPSEEK2 && model.arch != LLM_ARCH_GLM_DSA && model.arch != LLM_ARCH_MISTRAL4) return; + if (!model.is_mla_model()) return; // High-enough bpw to survive one quant->F32->H->quant roundtrip within PPL noise. // Cliff is ~2.7 bpw: IQ3_XXS (3.06) sits at +0.05 noise edge; IQ2_XS (2.31) drifts +0.20. @@ -3066,7 +3066,7 @@ static std::pair, double> get_layer_sizes(const llama_model_ ggml_tensor * wkv_b = nullptr; }; std::vector mla_tensors; - bool has_mla = model.arch == LLM_ARCH_DEEPSEEK2 || model.arch == LLM_ARCH_GLM_DSA || model.arch == LLM_ARCH_MISTRAL4; + bool has_mla = model.is_mla_model(); if (has_mla) { mla_tensors.resize(n_layer); } @@ -3443,7 +3443,7 @@ static bool llm_load_tensors( model.main_gpu = device_count - 1; } - if (model.arch == LLM_ARCH_DEEPSEEK2 || model.arch == LLM_ARCH_GLM_DSA || model.arch == LLM_ARCH_MISTRAL4) { + if (model.is_mla_model()) { if (model.n_gpu_layers > 0 && model.n_gpu_layers < model.hparams.n_layer && mla_attn != 3) { LLAMA_LOG_WARN("=============================================================================\n"); LLAMA_LOG_WARN("MLA models with ngl < n_layer and split mode graph do not work with mla = %d\n", mla_attn); @@ -3918,7 +3918,7 @@ static bool llm_load_tensors( } } - if ((model.arch == LLM_ARCH_DEEPSEEK2 || model.arch == LLM_ARCH_GLM_DSA || model.arch == LLM_ARCH_MISTRAL4)) { + if (model.is_mla_model()) { // -sm graph/attn needs wk_b->extra populated; run prepare even under dry-run. const bool graph_mode = (model.split_mode == LLAMA_SPLIT_MODE_GRAPH || model.split_mode == LLAMA_SPLIT_MODE_ATTN); @@ -5889,7 +5889,7 @@ static void llama_kv_cache_defrag_internal(struct llama_context & lctx) { } static bool get_can_shift(struct llama_context & lctx) { - bool no_shift = lctx.model.arch == LLM_ARCH_DEEPSEEK2 || lctx.model.arch == LLM_ARCH_GLM_DSA; // not supported due to MLA + bool no_shift = lctx.model.is_mla_model(); no_shift = no_shift || lctx.model.hparams.rope_type == LLAMA_ROPE_TYPE_IMROPE; return !no_shift; } @@ -6648,11 +6648,31 @@ struct llama_context * llama_init_from_model( LLAMA_LOG_WARN("%s: there is no point in Hadamard transforms with not quantized K-cache. Turning K-cache Hadamard off\n", __func__); params.k_cache_hadamard = false; } + if (params.k_cache_hadamard) { + int nok = 0; + for (int il = 0; il < model->hparams.n_layer; ++il) { + if (model->hadamard_size_k(il) > 0) ++nok; + } + if (nok == 0) { + LLAMA_LOG_WARN("%s: no layer allows a K-head Hadamard transform. Turning K-cache Hadamard off\n", __func__); + params.k_cache_hadamard = false; + } + } if (params.v_cache_hadamard && !ggml_is_quantized(params.type_v)) { LLAMA_LOG_WARN("%s: there is no point in Hadamard transforms with not quantized V-cache. Turning V-cache Hadamard off\n", __func__); params.v_cache_hadamard = false; } + if (params.v_cache_hadamard) { + int nok = 0; + for (int il = 0; il < model->hparams.n_layer; ++il) { + if (model->hadamard_size_v(il) > 0) ++nok; + } + if (nok == 0) { + LLAMA_LOG_WARN("%s: no layer allows a V-head Hadamard transform. Turning V-cache Hadamard off\n", __func__); + params.v_cache_hadamard = false; + } + } llama_context * ctx = new llama_context(*model); @@ -6777,7 +6797,7 @@ struct llama_context * llama_init_from_model( params.seed = time(NULL); } - if (model->arch != LLM_ARCH_DEEPSEEK2 && model->arch != LLM_ARCH_GLM_DSA && model->arch != LLM_ARCH_MISTRAL4 && cparams.mla_attn != 0) { + if (!model->is_mla_model() && cparams.mla_attn != 0) { cparams.mla_attn = 0; } else { if (model->n_gpu_layers > 0 && model->n_gpu_layers < model->hparams.n_layer && cparams.mla_attn != 3) { @@ -6810,7 +6830,7 @@ struct llama_context * llama_init_from_model( LLAMA_LOG_INFO("%s: n_batch = %u\n", __func__, cparams.n_batch); LLAMA_LOG_INFO("%s: n_ubatch = %u\n", __func__, cparams.n_ubatch); LLAMA_LOG_INFO("%s: flash_attn = %d\n", __func__, cparams.flash_attn); - if (model->arch == LLM_ARCH_DEEPSEEK2 || model->arch == LLM_ARCH_GLM_DSA || model->arch == LLM_ARCH_MISTRAL4) { + if (model->is_mla_model()) { LLAMA_LOG_INFO("%s: mla_attn = %d\n", __func__, cparams.mla_attn); } LLAMA_LOG_INFO("%s: attn_max_b = %d\n", __func__, cparams.attn_max_batch);