nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added split mode graph for Command-R/R+ models. (#1491)

This commit is contained in:
gapeleon 2026-03-23 18:10:41 +11:00 committed by GitHub
parent 87e4b9260b
commit 716ecd6457
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 58 additions and 78 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

101
src/llama-build-context.cpp
View File

@ -1544,6 +1544,7 @@ static ggml_tensor * llm_build_kqv(
|| model.arch == LLM_ARCH_GPTNEOX
|| model.arch == LLM_ARCH_QWEN2
|| model.arch == LLM_ARCH_COHERE2
|| model.arch == LLM_ARCH_COMMAND_R
|| model.arch == LLM_ARCH_GLM4
|| model.arch == LLM_ARCH_GLM4_MOE
|| model.arch == LLM_ARCH_MIMO2;
@ -1669,7 +1670,7 @@ static ggml_tensor * llm_build_kqv(
auto q_i = ggml_view_3d(ctx, q, q->ne[0], q->ne[1], this_ne12, q->nb[1], q->nb[2], q->nb[2]*i12);
auto kq_i = ggml_mul_mat(ctx, k_i, q_i);
if (model.arch == LLM_ARCH_PHI2 || model.arch == LLM_ARCH_PHI3 || model.arch == LLM_ARCH_GPTNEOX || model.arch == LLM_ARCH_QWEN2 ||
model.arch == LLM_ARCH_COHERE2 || model.arch == LLM_ARCH_GLM4 || model.arch == LLM_ARCH_GLM4_MOE) {
model.arch == LLM_ARCH_COHERE2 || model.arch == LLM_ARCH_COMMAND_R || model.arch == LLM_ARCH_GLM4 || model.arch == LLM_ARCH_GLM4_MOE) {
ggml_mul_mat_set_prec(kq_i, GGML_PREC_F32);
}
if (model.arch == LLM_ARCH_GROK) {
@ -1934,14 +1935,16 @@ std::tuple<ggml_tensor*, ggml_tensor*, ggml_tensor*> llm_build_context::llm_buil
auto [Q, K, V] = llm_build_mul_mat_qkv(gf, cur, wq, bq, wk, bk, wv, bv, attention_scale, il, add_graph_split);
auto Qcur = ggml_reshape_3d(ctx0, Q, n_embd_head_k, Q->ne[0]/n_embd_head_k, n_tokens);
// Command-R/R+ uses LayerNorm (not RMSNorm) for per-head Q/K normalisation
const auto qk_norm_type = (model.arch == LLM_ARCH_COMMAND_R) ? LLM_NORM : LLM_NORM_RMS;
if (q_norm) {
Qcur = llm_build_norm(ctx0, Qcur, hparams, q_norm, NULL, LLM_NORM_RMS, cb, il);
Qcur = llm_build_norm(ctx0, Qcur, hparams, q_norm, NULL, qk_norm_type, cb, il);
cb(Qcur, "Qcur_normed", il);
}
auto Kcur = ggml_reshape_3d(ctx0, K, n_embd_head_k, K->ne[0]/n_embd_head_k, n_tokens);
if (k_norm) {
Kcur = llm_build_norm(ctx0, Kcur, hparams, k_norm, NULL, LLM_NORM_RMS, cb, il);
Kcur = llm_build_norm(ctx0, Kcur, hparams, k_norm, NULL, qk_norm_type, cb, il);
cb(Kcur, "Kcur_normed", il);
}
auto Vcur = V;
@ -6242,79 +6245,32 @@ ggml_cgraph * llm_build_context::build_command_r() {
for (int il = 0; il < n_layer; ++il) {
// norm
cur = llm_build_norm(ctx0, inpL, hparams, model.layers[il].attn_norm, NULL, LLM_NORM, cb, il);
cb(cur, "attn_norm", il);
struct ggml_tensor * ffn_inp = cur;
// self-attention (norm applied inside; handles graph-parallel split path automatically)
auto attn_out = build_std_attention(gf, model.layers[il].attn_norm, inpL, inp_pos,
nullptr, nullptr, KQ_mask, nullptr, nullptr,
1.0f / sqrtf(float(n_embd_head)), 0.f,
0, il, /*do_rope=*/true, /*add_graph_split=*/true,
/*add_input=*/true, /*is_norm=*/true, /*is_multi=*/false);
cb(attn_out, "attn_out", il);
// self-attention
{
auto [Qcur, Kcur, Vcur] = llm_build_mul_mat_qkv(gf, cur, model.layers[il].wq, model.layers[il].bq,
model.layers[il].wk, model.layers[il].bk,
model.layers[il].wv, model.layers[il].bv, 0.f, il);
if (model.layers[il].attn_q_norm) {
Qcur = ggml_view_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens,
ggml_element_size(Qcur) * n_embd_head,
ggml_element_size(Qcur) * n_embd_head * n_head,
0);
cb(Qcur, "Qcur", il);
Kcur = ggml_view_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens,
ggml_element_size(Kcur) * n_embd_head,
ggml_element_size(Kcur) * n_embd_head * n_head_kv,
0);
cb(Kcur, "Kcur", il);
Qcur = llm_build_norm(ctx0, Qcur, hparams, model.layers[il].attn_q_norm, NULL, LLM_NORM, cb, il);
cb(Qcur, "Qcur", il);
Kcur = llm_build_norm(ctx0, Kcur, hparams, model.layers[il].attn_k_norm, NULL, LLM_NORM, cb, il);
cb(Kcur, "Kcur", il);
}
Qcur = ggml_rope_ext(
ctx0, ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens), inp_pos, nullptr,
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, ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens), inp_pos, nullptr,
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, n_tokens, kv_head, n_kv, 1.0f/sqrtf(float(n_embd_head)), cb, il);
}
if (il == n_layer - 1) {
if (il == n_layer - 1 && n_tokens > 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);
attn_out = ggml_get_rows(ctx0, attn_out, inp_out_ids);
inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
}
struct ggml_tensor * attn_out = cur;
attn_out->op_params[3] = 1; // turn off attention reduce; the FFN reduce will cover both
// feed-forward network
{
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);
}
// feed-forward network (norm applied inside; graph-parallel when ffn tensors are split)
cur = llm_build_ffn(ctx0, lctx, model.layers[il].attn_norm, inpL,
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, gf, /*add_input=*/false, /*is_norm=*/true, attn_out);
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);
@ -6327,13 +6283,13 @@ ggml_cgraph * llm_build_context::build_command_r() {
cur = llm_build_norm(ctx0, cur, hparams, model.output_norm, NULL, LLM_NORM, cb, -1);
cb(cur, "result_norm", -1);
// lm_head
cur = llm_build_lora_mm(lctx, ctx0, model.output, cur);
if (f_logit_scale) {
cur = ggml_scale(ctx0, cur, f_logit_scale);
cb(cur, "result_norm_scaled", -1);
}
// lm_head
cur = llm_build_lora_mm(lctx, ctx0, model.output, cur);
cb(cur, "result_output", -1);
ggml_build_forward_expand(gf, cur);
@ -10029,6 +9985,7 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens
|| model.arch == LLM_ARCH_GPTNEOX
|| model.arch == LLM_ARCH_QWEN2
|| model.arch == LLM_ARCH_COHERE2
|| model.arch == LLM_ARCH_COMMAND_R
|| model.arch == LLM_ARCH_GLM4
// || model.arch == LLM_ARCH_GLM4_MOE
|| model.arch == LLM_ARCH_MIMO2;

34
src/llama-load-tensors.cpp
View File

@ -2060,16 +2060,15 @@ bool create_tensors_helper::create_command_r_tensors(const LLM_TN & tn) {
}
for (int i = 0; i < n_layer; ++i) {
ggml_context * ctx_layer = ctx_for_layer(i);
ggml_context * ctx_split = ctx_for_layer_split(i);
auto & layer = model.layers[i];
layer.attn_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd});
layer.attn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
if (n_layer >= 64){
layer.attn_q_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k, n_head});
layer.attn_k_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k, n_head_kv});
layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k, n_head}, 0);
layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k, n_head_kv}, 0);
}
layer.wq = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd});
@ -4006,6 +4005,13 @@ bool create_tensors_helper::create_tensors() {
auto tt = ggml_internal_get_type_traits(layer.wo->type);
if (tt.blck_size > granularity_vo) granularity_vo = tt.blck_size;
GGML_ASSERT(granularity_vo % hparams.n_embd_head_v == 0);
// Command-R: align KQ split to wo's block size so wq row
// counts remain valid after splitting.
if (model.arch == LLM_ARCH_COMMAND_R) {
if (tt.blck_size > granularity_kq && layer.wq->ne[1] % tt.blck_size == 0) {
granularity_kq = tt.blck_size;
}
}
}
auto split_vo = create_split(layer.wo->ne[0], granularity_vo, cur_splits, mem_used); //, true);
auto split_kq = create_split(layer.wq->ne[1], granularity_kq, cur_splits, mem_used); //, true);
@ -4027,7 +4033,14 @@ bool create_tensors_helper::create_tensors() {
} else {
prepare_split_tensors(1, ctx_split, layer.wq, layer.split_wq, split_kq, mem_used);
if (layer.attn_q_norm) {
prepare_split_tensors(-1, ctx_split, layer.attn_q_norm, layer.split_q_norm, split_kq, mem_used);
if (layer.attn_q_norm->ne[1] > 1) {
// 2D per-head norm (e.g., Command-R+): split along the Q-head dimension
auto split_q_heads = split_kq;
for (auto & s : split_q_heads) s /= hparams.n_embd_head_k;
prepare_split_tensors(1, ctx_split, layer.attn_q_norm, layer.split_q_norm, split_q_heads, mem_used);
} else {
prepare_split_tensors(-1, ctx_split, layer.attn_q_norm, layer.split_q_norm, split_kq, mem_used);
}
}
if (layer.bq) {
prepare_split_tensors(0, ctx_split, layer.bq, layer.split_bq, split_kq, mem_used);
@ -4076,7 +4089,16 @@ bool create_tensors_helper::create_tensors() {
prepare_split_tensors(0, ctx_split, layer.bk, layer.split_bk, split_kq, mem_used);
}
if (layer.attn_k_norm) {
prepare_split_tensors(-1, ctx_split, layer.attn_k_norm, layer.split_k_norm, split_kq, mem_used);
if (layer.attn_k_norm->ne[1] > 1) {
// 2D per-head norm (e.g., Command-R+): split along the KV-head dimension
// split_kq has already been divided by gqa_ratio, so values are in
// (n_embd_head_k * n_head_kv) units; divide again to get head units
auto split_k_heads = split_kq;
for (auto & s : split_k_heads) s /= hparams.n_embd_head_k;
prepare_split_tensors(1, ctx_split, layer.attn_k_norm, layer.split_k_norm, split_k_heads, mem_used);
} else {
prepare_split_tensors(-1, ctx_split, layer.attn_k_norm, layer.split_k_norm, split_kq, mem_used);
}
}
}
prepare_split_tensors(1, ctx_split, layer.wv, layer.split_wv, split_vo, mem_used);

1
src/llama.cpp
View File

@ -1977,6 +1977,7 @@ static bool is_model_split_supported(const llama_model & model) {
LLM_ARCH_QWEN3MOE,
LLM_ARCH_GLM4_MOE,
LLM_ARCH_MISTRAL3,
LLM_ARCH_COMMAND_R,
LLM_ARCH_COHERE2,
LLM_ARCH_MIMO2,
LLM_ARCH_QWEN3,