From 328c3ff5e0ce7e7f808a583df00f6e7cb35baead Mon Sep 17 00:00:00 2001 From: Iwan Kawrakow Date: Fri, 12 Dec 2025 14:59:33 +0000 Subject: [PATCH] This works and TG is descent, but PP is low --- ggml/src/ggml-cuda/norm.cu | 26 ++++++---- src/llama-build-context.cpp | 97 ++++++++++++++++--------------------- src/llama-build-context.h | 5 +- src/llama-load-tensors.cpp | 9 +--- src/llama.cpp | 1 + 5 files changed, 63 insertions(+), 75 deletions(-) diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index c4619e60..9791b0f9 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -1,14 +1,14 @@ #include "norm.cuh" -template -static __global__ void norm_f32(const float * x, float * dst, const int ncols, const float eps) { +template +static __global__ void norm_f32(const T * x, T * dst, const int ncols, const float eps) { const int row = blockIdx.x*blockDim.y + threadIdx.y; const int tid = threadIdx.x; float2 mean_var = make_float2(0.f, 0.f); for (int col = tid; col < ncols; col += block_size) { - const float xi = x[row*ncols + col]; + const float xi = (float)x[row*ncols + col]; mean_var.x += xi; mean_var.y += xi * xi; } @@ -32,7 +32,7 @@ static __global__ void norm_f32(const float * x, float * dst, const int ncols, c const float inv_std = rsqrtf(var + eps); for (int col = tid; col < ncols; col += block_size) { - dst[row*ncols + col] = (x[row*ncols + col] - mean) * inv_std; + dst[row*ncols + col] = (T)(((float)x[row*ncols + col] - mean) * inv_std); } } @@ -261,14 +261,15 @@ static __global__ void fused_rms_norm_f32_nc( } } -static void norm_f32_cuda(const float * x, float * dst, const int ncols, const int nrows, const float eps, cudaStream_t stream) { +template +static void norm_f32_cuda(const T * x, T * dst, const int ncols, const int nrows, const float eps, cudaStream_t stream) { GGML_ASSERT(ncols % WARP_SIZE == 0); if (ncols < 1024) { const dim3 block_dims(WARP_SIZE, 1, 1); - norm_f32<<>>(x, dst, ncols, eps); + norm_f32<<>>(x, dst, ncols, eps); } else { const dim3 block_dims(1024, 1, 1); - norm_f32<1024><<>>(x, dst, ncols, eps); + norm_f32<1024, T><<>>(x, dst, ncols, eps); } } @@ -364,8 +365,9 @@ void ggml_cuda_op_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ASSERT(ggml_is_contiguous(src0)); - GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT( dst->type == GGML_TYPE_F32); + GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); + GGML_ASSERT( dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16); + GGML_ASSERT(src0->type == dst->type); const int64_t ne00 = src0->ne[0]; const int64_t nrows = ggml_nrows(src0); @@ -373,7 +375,11 @@ void ggml_cuda_op_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { float eps; memcpy(&eps, dst->op_params, sizeof(float)); - norm_f32_cuda(src0_d, dst_d, ne00, nrows, eps, stream); + if (dst->type == GGML_TYPE_F32) { + norm_f32_cuda(src0_d, dst_d, ne00, nrows, eps, stream); + } else { + norm_f32_cuda((const half *)src0_d, (half *)dst_d, ne00, nrows, eps, stream); + } } void ggml_cuda_op_group_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index bf7a141c..01a5e2a2 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -1224,10 +1224,10 @@ llm_expert_gating_func_type gating_op, } else { cur = routed_out; } - if (cur->ne[1] >= 32) { - cur = ggml_cast(ctx, cur, GGML_TYPE_F16); - cb(cur, "ffn_out_f16", il_cb); - } + //if (cur->ne[1] >= 32) { + // cur = ggml_cast(ctx, cur, GGML_TYPE_F16); + // cb(cur, "ffn_out_f16", il_cb); + //} ggml_build_forward_expand(graph, routed_out); results.push_back(cur); } @@ -1743,7 +1743,8 @@ ggml_cgraph * llm_build_context::build_llama() { // self-attention if (use_rope) { - cur = build_std_attention(gf, inpL, inp_pos, nullptr, this_KQ_mask, nullptr, nullptr, kq_scale, hparams.f_attention_scale, this_n_swa, il); + cur = build_std_attention(gf, model.layers[il].attn_norm, inpL, inp_pos, nullptr, + this_KQ_mask, nullptr, nullptr, kq_scale, hparams.f_attention_scale, this_n_swa, il); } else { @@ -1935,7 +1936,8 @@ ggml_cgraph * llm_build_context::build_mistral3() { auto rope_factors = build_rope_factors(il); - cur = build_std_attention(gf, inpL, inp_pos, rope_factors, KQ_mask, nullptr, inp_attn_scale, kq_scale, hparams.f_attention_scale, 0, il); + cur = build_std_attention(gf, model.layers[il].attn_norm, inpL, inp_pos, rope_factors, KQ_mask, + nullptr, inp_attn_scale, kq_scale, hparams.f_attention_scale, 0, il); if (il == n_layer - 1 && inp_out_ids) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); @@ -3927,7 +3929,7 @@ ggml_cgraph * llm_build_context::build_qwen3moe() { //cur = llm_build_norm(ctx0, inpL, hparams, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, cb, il); //cb(cur, "attn_norm", il); - cur = build_std_attention(gf, inpL, inp_pos, nullptr, KQ_mask, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), 0.0f, 0, il); + cur = build_std_attention(gf, model.layers[il].attn_norm, inpL, inp_pos, nullptr, KQ_mask, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), 0.0f, 0, il); if (il == n_layer - 1) { // skip computing output for unused tokens @@ -6806,7 +6808,7 @@ ggml_cgraph * llm_build_context::build_glm4_moe() { // self-attention if (rope_cache == nullptr) { - cur = build_std_attention(gf, inpL, inp_pos, nullptr, KQ_mask, nullptr, nullptr, kq_scale, 0.0f, 0, il); + cur = build_std_attention(gf, model.layers[il].attn_norm, inpL, inp_pos, nullptr, KQ_mask, nullptr, nullptr, kq_scale, 0.0f, 0, il); } else { // Pre-attention norm cur = llm_build_norm(ctx0, inpL, hparams, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, cb, il); @@ -7223,54 +7225,28 @@ ggml_cgraph * llm_build_context::build_cohere2() { struct ggml_tensor * ffn_inp = cur; // self-attention - { - // rope freq factors for 128k context - struct ggml_tensor * rope_factors = build_rope_factors(il); + cur = build_std_attention(gf, nullptr, cur, inp_pos, nullptr, KQ_mask_l, nullptr, nullptr, 1.0f / sqrtf(float(n_embd_head)), 0.f, + is_sliding ? hparams.n_swa : 0, il, is_sliding); - auto [Qcur, Kcur, Vcur] = llm_build_mul_mat_qkv(gf, cur, - model.layers[il].wqkv, model.layers[il].bqkv, - model.layers[il].wqk, model.layers[il].bqk, - model.layers[il].wq, model.layers[il].bq, - model.layers[il].wk, model.layers[il].bk, - model.layers[il].wv, model.layers[il].bv, nullptr, nullptr, 0.f, il); - - if (is_sliding) { - Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, - beta_fast, beta_slow); - cb(Qcur, "Qcur", il); - - Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, - rope_factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, - attn_factor, beta_fast, beta_slow); - cb(Kcur, "Kcur", il); - }; - - cur = llm_build_kv(ctx0, lctx, kv_self, gf, model.layers[il].wo, model.layers[il].bo, Kcur, Vcur, Qcur, - KQ_mask_l, n_tokens, kv_head, n_kv, 1.0f / sqrtf(float(n_embd_head)), cb, il, nullptr, - is_sliding ? hparams.n_swa : 0); - } + cur = ggml_add(ctx0, cur, inpL); + cb(cur, "attn_out", il); if (il == n_layer - 1) { // skip computing output for unused tokens struct ggml_tensor * inp_out_ids = build_inp_out_ids(); cur = ggml_get_rows(ctx0, cur, inp_out_ids); - inpL = ggml_get_rows(ctx0, inpL, inp_out_ids); ffn_inp = ggml_get_rows(ctx0, ffn_inp, inp_out_ids); } - struct ggml_tensor * attn_out = cur; + auto attn_out = cur; // feed-forward network - { - cur = llm_build_ffn(ctx0, lctx, nullptr, ffn_inp, model.layers[il].ffn_up, NULL, NULL, model.layers[il].ffn_gate, + cur = llm_build_ffn(ctx0, lctx, nullptr, ffn_inp, model.layers[il].ffn_up, NULL, NULL, model.layers[il].ffn_gate, NULL, NULL, model.layers[il].ffn_down, NULL, NULL, NULL, LLM_FFN_SILU, LLM_FFN_PAR, - cb, il); - cb(cur, "ffn_out", il); - } + cb, il, gf); + cb(cur, "ffn_out", il); // add together residual + FFN + self-attention - cur = ggml_add(ctx0, cur, inpL); cur = ggml_add(ctx0, cur, attn_out); cur = lctx.cvec.apply_to(ctx0, cur, il); cb(cur, "l_out", il); @@ -7280,6 +7256,9 @@ ggml_cgraph * llm_build_context::build_cohere2() { } cur = inpL; + if (cur->type != GGML_TYPE_F32) { + cur = ggml_cast(ctx0, cur, GGML_TYPE_F32); + } cur = llm_build_norm(ctx0, cur, hparams, model.output_norm, NULL, LLM_NORM, cb, -1); cb(cur, "result_norm", -1); @@ -9308,12 +9287,14 @@ ggml_cgraph * llm_build_context::llama_build_graph( return result; } -ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tensor * input, ggml_tensor * inp_pos, ggml_tensor * rope_factors_in, - ggml_tensor * KQ_mask, ggml_tensor * sinks, ggml_tensor * inp_attn_scale, float KQ_scale, float f_attn_scale, int n_swa, int il) { +ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tensor * the_attn_norm, + ggml_tensor * input, ggml_tensor * inp_pos, ggml_tensor * rope_factors_in, + ggml_tensor * KQ_mask, ggml_tensor * sinks, ggml_tensor * inp_attn_scale, float KQ_scale, float f_attn_scale, + int n_swa, int il, bool do_rope) { if (!model.layers[il].wqkv && !model.layers[il].wqk && cparams.flash_attn && model.layers[il].wq->extra && model.layers[il].wk->extra && model.layers[il].wv->extra && model.layers[il].wo->extra) { if (kv_self.k_l[il]->extra && kv_self.v_l[il]->extra) { - ggml_split_tensor_t * attn_norm = model.layers[il].attn_norm ? (ggml_split_tensor_t *)model.layers[il].attn_norm->extra : nullptr; + ggml_split_tensor_t * attn_norm = the_attn_norm ? (ggml_split_tensor_t *)the_attn_norm->extra : nullptr; auto wq = (ggml_split_tensor_t *)model.layers[il].wq->extra; auto wk = (ggml_split_tensor_t *)model.layers[il].wk->extra; auto wv = (ggml_split_tensor_t *)model.layers[il].wv->extra; @@ -9374,10 +9355,12 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens auto extra = (ggml_split_tensor_t *)model.layers[il].rope_freqs->extra; rope_factors = extra->splits[id]; } - Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow); - Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, rope_factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow); + if (do_rope) { + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, rope_factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + } cb(Qcur, "Qcur", il_cb); cb(Kcur, "Kcur", il_cb); if (inp_attn_scale) { @@ -9477,9 +9460,9 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens cur = ggml_add(ctx0, cur, bo->splits[id]); cb(cur, "kqv_wo_biased", il_cb); } - if (cur->ne[1] >= 32) { - cur = ggml_cast(ctx0, cur, GGML_TYPE_F16); - } + //if (cur->ne[1] >= 32) { + // cur = ggml_cast(ctx0, cur, GGML_TYPE_F16); + //} ggml_build_forward_expand(gf, cur); attn.push_back(cur); } @@ -9497,8 +9480,8 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens } auto cur = input; - if (model.layers[il].attn_norm) { - cur = llm_build_norm(ctx0, cur, hparams, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, cb, il); + if (the_attn_norm) { + cur = llm_build_norm(ctx0, cur, hparams, the_attn_norm, NULL, LLM_NORM_RMS, cb, il); cb(cur, "attn_norm", il); } @@ -9508,10 +9491,12 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens model.layers[il].wq, model.layers[il].bq, model.layers[il].wk, model.layers[il].bk, model.layers[il].wv, model.layers[il].bv, model.layers[il].attn_q_norm, model.layers[il].attn_k_norm, f_attn_scale, il); - Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors_in, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + if (do_rope) { + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors_in, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - Kcur = ggml_rope_ext( ctx0, Kcur, inp_pos, rope_factors_in, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + Kcur = ggml_rope_ext( ctx0, Kcur, inp_pos, rope_factors_in, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + } cb(Qcur, "Qcur", il); cb(Kcur, "Kcur", il); diff --git a/src/llama-build-context.h b/src/llama-build-context.h index c5a27654..18b820a2 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -407,7 +407,8 @@ llm_expert_gating_func_type gating_op, static ggml_cgraph * llama_build_graph(llama_context & lctx, const llama_batch & batch, bool worst_case); - ggml_tensor * build_std_attention(ggml_cgraph * gf, ggml_tensor * cur, ggml_tensor * inp_pos, ggml_tensor * rope_factors, - ggml_tensor * KQ_mask, ggml_tensor * sinks, ggml_tensor * inp_attn_scale, float KQ_scale, float f_attn_scale, int n_swa, int il); + ggml_tensor * build_std_attention(ggml_cgraph * gf, ggml_tensor * attn_norm, ggml_tensor * cur, ggml_tensor * inp_pos, ggml_tensor * rope_factors, + ggml_tensor * KQ_mask, ggml_tensor * sinks, ggml_tensor * inp_attn_scale, float KQ_scale, float f_attn_scale, + int n_swa, int il, bool do_rope = true); }; diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index 2f8794cc..db2d5695 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -2259,18 +2259,13 @@ bool create_tensors_helper::create_chatglm_tensors(const LLM_TN & tn) { bool create_tensors_helper::create_cohere2_tensors(const LLM_TN & tn) { LOADING_PRELUDE - model.tok_embd = create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0); - // output - model.output_norm = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); - // init output from the input tok embed - model.output = create_tensor(ctx_output_split, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, - llama_model_loader::TENSOR_DUPLICATED); + create_embd_output(tn, n_embd, n_vocab, true, false); //true); for (int i = 0; i < n_layer; ++i) { auto & layer = model.layers[i]; - ggml_context * ctx_layer = ctx_for_layer(i); ggml_context * ctx_split = ctx_for_layer_split(i); + ggml_context * ctx_layer = ctx_for_layer(i); layer.attn_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0); diff --git a/src/llama.cpp b/src/llama.cpp index 53dba75a..6b8b2887 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -1729,6 +1729,7 @@ static bool is_model_split_supported(const llama_model & model) { LLM_ARCH_QWEN3MOE, LLM_ARCH_GLM4_MOE, LLM_ARCH_MISTRAL3, + LLM_ARCH_COHERE2, }; auto it = k_supported.find(model.arch); return it != k_supported.end();