Public/koboldcpp
1
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-09-18 16:55:14 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'upstream' into concedo_experimental

Browse Source 
# Conflicts:
#	README.md
#	docs/backend/SYCL.md
#	ggml/CMakeLists.txt
#	ggml/src/ggml-cpu/repack.cpp
#	ggml/src/ggml-sycl/backend.hpp
#	ggml/src/ggml-sycl/common.hpp
#	ggml/src/ggml-sycl/convert.cpp
#	ggml/src/ggml-sycl/cpy.cpp
#	ggml/src/ggml-sycl/cpy.hpp
#	ggml/src/ggml-sycl/dequantize.hpp
#	ggml/src/ggml-sycl/fattn.cpp
#	ggml/src/ggml-sycl/fattn.hpp
#	ggml/src/ggml-sycl/ggml-sycl.cpp
#	ggml/src/ggml-sycl/mmvq.cpp
#	ggml/src/ggml-sycl/norm.cpp
#	ggml/src/ggml-sycl/norm.hpp
#	ggml/src/ggml-sycl/vecdotq.hpp
#	ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
#	ggml/src/ggml-webgpu/ggml-webgpu.cpp
#	ggml/src/ggml-webgpu/wgsl-shaders/common_decls.tmpl
#	ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl
#	ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_quant_staging.tmpl
#	ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_tile.wgsl
#	ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_vec_split.wgsl
#	ggml/src/ggml-zendnn/CMakeLists.txt
#	ggml/src/ggml-zendnn/ggml-zendnn.cpp
#	scripts/sync-ggml.last
#	tests/CMakeLists.txt
#	tests/snapshots/qwen3.6-27b.schema
#	tests/test-autorelease.cpp
#	tests/test-backend-ops.cpp
#	tests/test-backend-sampler.cpp
#	tests/test-model-load-cancel.cpp
#	tests/test-quant-type-selection.cpp
This commit is contained in:
Concedo
2026-07-31 19:44:27 +08:00
parent 98926f27c1 000547513f
commit 4a69c10078
47 changed files with 985 additions and 630 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-context.cpp
+33 -10
View File
@@ -482,6 +482,9 @@ llama_context::llama_context(
}
llama_context::~llama_context() {
// wait for any pending asynchronous copies into the output buffers before they are freed
synchronize();
if (!model.hparams.no_alloc) {
for (size_t i = 0; i < backend_ptrs.size(); ++i) {
ggml_backend_t backend = backend_ptrs[i];
@@ -1427,13 +1430,17 @@ int llama_context::encode(const llama_batch & batch_inp) {
// micro-batching is not possible for non-causal encoding, so we process the batch in a single shot
GGML_ASSERT(cparams.n_ubatch >= n_tokens && "encoder requires n_ubatch >= n_tokens");
// TODO: this clear of the buffer can easily be forgotten - need something better
// sync first so any in-flight async copies into embd_seq complete before it is freed
if (!embd_seq.empty()) {
synchronize();
}
embd_seq.clear();
if (t_compute_start_us == 0) {
t_compute_start_us = ggml_time_us();
}
// TODO: this clear of the buffer can easily be forgotten - need something better
embd_seq.clear();
sched_reserve();
n_queued_tokens += n_tokens;
@@ -1772,13 +1779,18 @@ int llama_context::decode(const llama_batch & batch_inp) {
// GGML_ASSERT((cparams.causal_attn || cparams.n_ubatch >= n_tokens_all) && "non-causal attention requires n_ubatch >= n_tokens");
// TODO: this clear of the buffer can easily be forgotten - need something better
// sync first so any in-flight async copies into embd_seq complete before it is freed
if (!embd_seq.empty()) {
synchronize();
}
embd_seq.clear();
if (t_compute_start_us == 0) {
t_compute_start_us = ggml_time_us();
}
n_queued_tokens += n_tokens_all;
// TODO: this clear of the buffer can easily be forgotten - need something better
embd_seq.clear();
output_swaps.clear();
sched_reserve();
@@ -3551,6 +3563,22 @@ llama_context * llama_init_from_model(
}
}
if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4) && params.type_k != params.type_v) {
LLAMA_LOG_ERROR("%s: model does not support different K (%s) and V (%s) cache types\n", __func__, ggml_type_name(params.type_k), ggml_type_name(params.type_v));
return nullptr;
}
if (ggml_is_quantized(params.type_v) && params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_ENABLED) {
if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO) {
LLAMA_LOG_INFO("%s: enabling flash_attn since it is required for quantized V cache\n", __func__);
params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
}
if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_DISABLED) {
LLAMA_LOG_ERROR("%s: quantized V cache requires flash_attn to be enabled\n", __func__);
return nullptr;
}
}
if (params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED && ggml_is_quantized(params.type_k)) {
const uint32_t blck_size = ggml_blck_size(params.type_k);
for (uint32_t il = 0; il < model->hparams.n_layer(); ++il) {
@@ -3573,11 +3601,6 @@ llama_context * llama_init_from_model(
}
}
if (ggml_is_quantized(params.type_v) && params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_DISABLED) {
LLAMA_LOG_ERROR("%s: V cache quantization requires flash_attn\n", __func__);
return nullptr;
}
if (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED &&
params.pooling_type != model->hparams.pooling_type) {
//user-specified pooling-type is different from the model default
src/llama-model.cpp
+3 -1
View File
@@ -2385,7 +2385,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
}
if ((arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_HY_V3 || arch == LLM_ARCH_GLM_DSA) && hparams.n_layer_nextn > 0) {
if ((arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_HY_V3 || arch == LLM_ARCH_GLM_DSA ||
arch == LLM_ARCH_MIMO2) &&
hparams.n_layer_nextn > 0) {
if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) {
filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
} else {
src/llama-vocab.cpp
+16 -1
View File
@@ -2817,7 +2817,14 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
if (suppress_idx != -1) {
const int n = gguf_get_arr_n(ctx, suppress_idx);
const int32_t * data = (const int32_t *) gguf_get_arr_data(ctx, suppress_idx);
suppress_tokens.assign(data, data + n);
// drop out-of-range ids
suppress_tokens.reserve(n);
for (int i = 0; i < n; ++i) {
const int32_t id = data[i];
if (id >= 0 && id < (int) id_to_token.size()) {
suppress_tokens.push_back(id);
}
}
}
}
@@ -4506,6 +4513,14 @@ bool llama_vocab_get_add_sep(const struct llama_vocab * vocab) {
return vocab->get_add_sep();
}
const llama_token * llama_vocab_get_suppress_tokens(const struct llama_vocab * vocab, int32_t * n_suppress_tokens) {
const std::vector<llama_token> & tokens = vocab->get_suppress_tokens();
if (n_suppress_tokens) {
*n_suppress_tokens = (int32_t) tokens.size();
}
return tokens.data();
}
llama_token llama_vocab_fim_pre(const struct llama_vocab * vocab) {
return vocab->token_fim_pre();
}
src/models/gemma4.cpp
-37
View File
@@ -142,33 +142,6 @@ std::unique_ptr<llm_graph_context> llama_model_gemma4::build_arch_graph(const ll
// idx * x->ne[0] * x->ne[1] * ggml_element_size(x));
// }
// TODO @ngxson : maybe improve this in the future
class llm_graph_input_logits_bias : public llm_graph_input_i {
public:
llm_graph_input_logits_bias(const llama_vocab & vocab) {
arr.resize(vocab.n_tokens(), 0.0f);
for (llama_token id : vocab.get_suppress_tokens()) {
if (0 <= id && id < (int32_t)vocab.n_tokens()) {
arr[id] = -INFINITY;
}
}
}
virtual ~llm_graph_input_logits_bias() = default;
void set_input(const llama_ubatch * /*ubatch*/) override {
const int64_t n_vocab = arr.size();
ggml_backend_tensor_set(logits_bias, arr.data(), 0, n_vocab*ggml_element_size(logits_bias));
}
bool can_reuse(const llm_graph_params & /*params*/) override {
return true;
}
ggml_tensor * logits_bias = nullptr; // F32 [n_vocab]
std::vector<float> arr;
};
llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_params & params) :
llm_graph_context(params),
model(model),
@@ -429,16 +402,6 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping);
}
// apply logits bias if needed (e.g. for gemma4_unified patch)
// this is to mirror the suppress_tokens patch on transformers, to avoid model from outputing <image|> and <audio|> tokens (which is a known issue related to the checkpoint)
// TODO: maybe handle this inside the sampling system in the future
if (!model.vocab.get_suppress_tokens().empty()) {
auto inp_bias = std::make_unique<llm_graph_input_logits_bias>(model.vocab);
inp_bias->logits_bias = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, inp_bias->arr.size());
cur = ggml_add(ctx0, cur, inp_bias->logits_bias);
res->add_input(std::move(inp_bias));
}
cb(cur, "result_output", -1);
res->t_logits = cur;
src/models/mimo2.cpp
+181 -21
View File
@@ -25,9 +25,13 @@ void llama_model_mimo2::load_arch_hparams(llama_model_loader & ml) {
}
}
void llama_model_mimo2::load_arch_tensors(llama_model_loader &) {
void llama_model_mimo2::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
// output
@@ -40,41 +44,46 @@ void llama_model_mimo2::load_arch_tensors(llama_model_loader &) {
uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(i);
uint32_t n_head = hparams.n_head(i);
// NextN/MTP layers (the last n_nextn blocks) are preserved but disabled pending support
const bool is_nextn = i >= n_layer;
const int skip = is_nextn ? TENSOR_SKIP : 0;
const int flags = is_nextn ? mtp_flags : 0;
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, skip);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_v * n_head, n_embd }, skip);
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, flags);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_v * n_head, n_embd }, flags);
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, skip);
layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, TENSOR_NOT_REQUIRED | skip);
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags);
layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, TENSOR_NOT_REQUIRED | flags);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, skip);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags);
// non-MoE branch
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED | skip);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, TENSOR_NOT_REQUIRED | skip);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED | skip);
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED | flags);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, TENSOR_NOT_REQUIRED | flags);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, TENSOR_NOT_REQUIRED | flags);
// MoE branch
int64_t n_ff_exp = hparams.n_ff_exp;
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, TENSOR_NOT_REQUIRED | skip);
layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED | skip);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, TENSOR_NOT_REQUIRED | skip);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED | skip);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED | skip);
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, TENSOR_NOT_REQUIRED | flags);
layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED | flags);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, TENSOR_NOT_REQUIRED | flags);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, TENSOR_NOT_REQUIRED | flags);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED | flags);
if (is_nextn) {
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, skip);
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, skip);
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, skip);
layer.layer_out_norm = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", i), {n_embd}, skip);
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, flags);
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, flags);
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, flags);
layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED | flags);
layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED | flags);
layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED | flags);
layer.layer_out_norm = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED | flags);
}
}
}
std::unique_ptr<llm_graph_context> llama_model_mimo2::build_arch_graph(const llm_graph_params & params) const {
if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
return std::make_unique<graph_mtp>(*this, params);
}
return std::make_unique<graph>(*this, params);
}
@@ -89,6 +98,8 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param
ggml_tensor * inp_out_ids = build_inp_out_ids();
const float v_scale = hparams.f_attn_value_scale;
const bool emit_h_nextn = cparams.embeddings_nextn;
const bool crop_last_layer = inp_out_ids && (!emit_h_nextn || cparams.embeddings_nextn_masked);
for (int il = 0; il < n_layer; ++il) {
ggml_tensor * inpSA = inpL;
@@ -168,7 +179,7 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param
}
}
if (il == n_layer - 1 && inp_out_ids) {
if (il == n_layer - 1 && crop_last_layer) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -218,6 +229,15 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param
cur = inpL;
if (emit_h_nextn) {
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
if (!cparams.embeddings_nextn_masked && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
}
cur = build_norm(cur,
model.output_norm, NULL,
LLM_NORM_RMS, -1);
@@ -233,3 +253,143 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param
ggml_build_forward_expand(gf, cur);
}
// Mirrors MiMo's appended NextN block: normalize and fuse token and hidden inputs, run the decoder block,
// expose its pre-head-norm state to the next draft step, then apply the shared output norm and LM head.
// Converted checkpoints may store that shared norm as layer_out_norm, so it remains in the fallback chain.
llama_model_mimo2::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
: llm_graph_context(params) {
GGML_ASSERT(hparams.n_layer_nextn > 0 && "MIMO2 MTP requires n_layer_nextn > 0");
const int il = hparams.n_layer() + cparams.nextn_layer_offset;
GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
"nextn_layer_offset out of range [0, n_layer_nextn)");
const auto & layer = model.layers[il];
GGML_ASSERT(layer.nextn.eh_proj && "MIMO2 MTP block missing nextn.eh_proj");
GGML_ASSERT(layer.nextn.enorm && "MIMO2 MTP block missing nextn.enorm");
GGML_ASSERT(layer.nextn.hnorm && "MIMO2 MTP block missing nextn.hnorm");
GGML_ASSERT(layer.wqkv && "MIMO2 MTP requires fused attn_qkv");
const uint32_t n_head_l = hparams.n_head(il);
const uint32_t n_head_kv_l = hparams.n_head_kv(il);
const float freq_base_l = model.get_rope_freq_base(cparams, il);
const float freq_scale_l = model.get_rope_freq_scale(cparams, il);
const float v_scale = hparams.f_attn_value_scale;
auto inp = std::make_unique<llm_graph_input_embd>(hparams.n_embd);
inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
ggml_set_input(inp->tokens);
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
ggml_set_input(inp->embd);
ggml_set_name(inp->embd, "mtp_h_input");
ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
ggml_tensor * h_input = inp->embd;
ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
cb(tok_embd, "mtp_tok_embd", il);
res->add_input(std::move(inp));
ggml_tensor * inp_pos = build_inp_pos();
ggml_tensor * inp_out_ids = build_inp_out_ids();
auto * inp_attn = build_attn_inp_kv_iswa();
ggml_tensor * h_norm = build_norm(h_input, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
cb(h_norm, "mtp_hnorm", il);
ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
cb(e_norm, "mtp_enorm", il);
ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);
cb(concat, "mtp_concat", il);
ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
cb(cur, "mtp_eh_proj", il);
ggml_tensor * inpSA = cur;
cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "mtp_attn_norm", il);
ggml_tensor * qkv = build_lora_mm(layer.wqkv, cur, layer.wqkv_s);
cb(qkv, "mtp_wqkv", il);
const size_t row_k = ggml_row_size(qkv->type, n_embd_head_k);
const size_t row_v = ggml_row_size(qkv->type, n_embd_head_v);
const size_t row_full = qkv->nb[1];
const size_t k_off = row_k * n_head_l;
const size_t v_off = k_off + row_k * n_head_kv_l;
ggml_tensor * Qcur = ggml_view_3d(ctx0, qkv, n_embd_head_k, n_head_l, n_tokens, row_k, row_full, 0);
ggml_tensor * Kcur = ggml_view_3d(ctx0, qkv, n_embd_head_k, n_head_kv_l, n_tokens, row_k, row_full, k_off);
ggml_tensor * Vcur = ggml_view_3d(ctx0, qkv, n_embd_head_v, n_head_kv_l, n_tokens, row_v, row_full, v_off);
Qcur = ggml_rope_ext(
ctx0, Qcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
ext_factor, attn_factor, beta_fast, beta_slow);
Kcur = ggml_rope_ext(
ctx0, Kcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
ext_factor, attn_factor, beta_fast, beta_slow);
cb(Qcur, "mtp_Qcur", il);
cb(Kcur, "mtp_Kcur", il);
cb(Vcur, "mtp_Vcur", il);
cur = build_attn(inp_attn,
layer.wo, nullptr, layer.wo_s,
Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr,
1.0f / sqrtf(float(n_embd_head_k)), il);
cb(cur, "mtp_attn_out", il);
if (v_scale) {
cur = ggml_scale(ctx0, cur, v_scale);
cb(cur, "mtp_attn_out_scaled", il);
}
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
cb(ffn_inp, "mtp_ffn_inp", il);
cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "mtp_ffn_norm", il);
GGML_ASSERT(layer.ffn_gate && layer.ffn_down && layer.ffn_up && "MIMO2 MTP requires dense FFN tensors");
cur = build_ffn(cur,
layer.ffn_up, layer.ffn_up_b, nullptr,
layer.ffn_gate, layer.ffn_gate_b, nullptr,
layer.ffn_down, layer.ffn_down_b, nullptr,
nullptr,
LLM_FFN_SILU, LLM_FFN_PAR, il);
cb(cur, "mtp_ffn_out", il);
cur = ggml_add(ctx0, cur, ffn_inp);
cb(cur, "mtp_post_ffn", il);
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
ggml_tensor * head_norm_w = layer.nextn.shared_head_norm
? layer.nextn.shared_head_norm
: (layer.layer_out_norm ? layer.layer_out_norm : model.output_norm);
GGML_ASSERT(head_norm_w && "MIMO2 MTP missing head norm fallback");
cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);
cb(cur, "mtp_shared_head_norm", -1);
ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
GGML_ASSERT(head_w && "MIMO2 MTP missing LM head fallback");
cur = build_lora_mm(head_w, cur, head_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
src/models/minimax-m3.cpp
+21 -72
View File
@@ -2,7 +2,6 @@
#include "llama-kv-cache.h"
#include <cmath>
#include <vector>
#include <algorithm>
#include <cstdint>
// MiniMax-M3: MiniMax-M2 style GQA (per-head QK-norm, partial rotary) with
@@ -126,68 +125,6 @@ public:
int64_t nblk;
};
// pooled score of a block with no visible token: -inf from the mask, or -FLT_MAX from the
// max-pool identity when every element of the block is -inf
static inline bool msa_score_masked(float x) { return x <= -1e30f; }
// MSA block selection (batch regime)
// CPU custom op, the token-level expansion and the combination with the causal mask happen on the GPU.
static void msa_block_mask_op(struct ggml_tensor * dst, int ith, int nth, void * userdata) {
const struct ggml_tensor * bs = dst->src[0];
const struct ggml_tensor * bias = dst->src[1];
const msa_params * p = (const msa_params *) userdata;
const int nblk = (int) bs->ne[0];
const int Hd = (int) bs->ne[1];
const int S = (int) bs->ne[2];
GGML_ASSERT(bs->type == GGML_TYPE_F32 && ggml_is_contiguous(bs));
GGML_ASSERT(bias->type == GGML_TYPE_F32 && ggml_is_contiguous(bias));
GGML_ASSERT(dst->type == GGML_TYPE_F16 && ggml_is_contiguous(dst));
GGML_ASSERT(dst->ne[0] == nblk && dst->ne[1] == S && dst->ne[2] == Hd);
GGML_ASSERT(bias->ne[0] == nblk && bias->ne[1] == S);
const int topk = p->topk_blocks < nblk ? p->topk_blocks : nblk;
const ggml_fp16_t f16_zero = ggml_fp32_to_fp16(0.0f);
const ggml_fp16_t f16_ninf = ggml_fp32_to_fp16(-INFINITY);
std::vector<float> rank(nblk);
std::vector<char> valid(nblk);
std::vector<int> ord(nblk);
ggml_fp16_t * out = (ggml_fp16_t *) dst->data;
for (int i = ith; i < S; i += nth) {
const float * bias_col = (const float *) bias->data + (size_t) i * nblk;
for (int h = 0; h < Hd; ++h) {
const float * bs_col = (const float *) bs->data + ((size_t) i * Hd + h) * nblk;
for (int bk = 0; bk < nblk; ++bk) {
// a block is selectable if it has a visible token or is locally forced
valid[bk] = !msa_score_masked(bs_col[bk]) || bias_col[bk] > 0.0f;
rank [bk] = bias_col[bk] > 0.0f ? bias_col[bk] : bs_col[bk];
ord [bk] = bk;
}
std::partial_sort(ord.begin(), ord.begin() + topk, ord.end(),
[&](int a, int b) { return rank[a] > rank[b]; });
ggml_fp16_t * dst_col = out + ((size_t) h * S + i) * nblk;
for (int bk = 0; bk < nblk; ++bk) {
dst_col[bk] = f16_ninf;
}
for (int t = 0; t < topk; ++t) {
const int bk = ord[t];
if (!valid[bk]) {
break; // sorted desc: first invalid -> fewer than topk selectable blocks
}
dst_col[bk] = f16_zero;
}
}
}
}
// One FA call for all GQA groups (and at multi-stream decode, all streams) by mapping them onto the FA sequence dim (ne[3])
ggml_tensor * llama_model_minimax_m3::graph::build_attn_msa_fa(
ggml_tensor * q_cur, // [D, HQ, T]
@@ -433,8 +370,8 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
msa_mf->nb[1], msa_mf->nb[1], st*msa_mf->nb[3]);
ggml_tensor * km_s = ggml_view_3d(ctx0, msa_kqm, n_kv, n_tps, 1,
msa_kqm->nb[1], msa_kqm->nb[3], st*msa_kqm->nb[3]);
ggml_tensor * bias_s = ggml_view_2d(ctx0, msa_loc->bias, nblk, n_tps,
msa_loc->bias->nb[1], st*n_tps*msa_loc->bias->nb[1]);
ggml_tensor * bias_s = ggml_view_3d(ctx0, msa_loc->bias, nblk, 1, n_tps,
msa_loc->bias->nb[1], msa_loc->bias->nb[1], st*n_tps*msa_loc->bias->nb[1]);
ggml_tensor * q_s = ggml_view_3d(ctx0, Qcur, D, n_head, n_tps,
Qcur->nb[1], Qcur->nb[2], st*n_tps*Qcur->nb[2]);
ggml_tensor * k_s = ggml_view_4d(ctx0, k, D, HKV, n_kv, 1,
@@ -453,15 +390,27 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
ggml_tensor * bs = ggml_pool_2d(ctx0, sc, GGML_OP_POOL_MAX, blk, 1, blk, 1, 0, 0);
cb(bs, "msa_bs", il);
// block-level 0/-inf keep mask on the CPU, tiny transfer
ggml_tensor * srcs[2] = { bs, bias_s };
ggml_tensor * bm = ggml_custom_4d(ctx0, GGML_TYPE_F16,
nblk, n_tps, Hd, 1,
srcs, 2, msa_block_mask_op, GGML_N_TASKS_MAX,
const_cast<msa_params *>(&mm.msa_p));
// bias the scores so locally-forced blocks always rank first
ggml_tensor * bsf = ggml_add(ctx0, bs, bias_s); // [nblk, Hd, n_tps]
cb(bsf, "msa_bsf", il);
ggml_tensor * idx = ggml_top_k(ctx0, bsf, K); // [K, Hd, n_tps] i32
ggml_tensor * ninf = ggml_cast(ctx0,
ggml_scale_bias(ctx0, bias_s, 0.0f, -1e30f),
GGML_TYPE_F16); // [nblk, 1, n_tps]
ninf = ggml_repeat_4d(ctx0, ninf, nblk, Hd, n_tps, 1);
ggml_tensor * zero = ggml_scale(ctx0,
ggml_cast(ctx0, idx, GGML_TYPE_F32), 0.0f);
ggml_tensor * bm = ggml_set_rows(ctx0,
ggml_reshape_3d(ctx0, ninf, 1, nblk, Hd*n_tps),
ggml_reshape_3d(ctx0, zero, 1, K, Hd*n_tps),
ggml_reshape_2d(ctx0, idx, K, Hd*n_tps));
bm = ggml_reshape_3d(ctx0, bm, nblk, Hd, n_tps);
bm = ggml_cont(ctx0, ggml_permute(ctx0, bm, 0, 2, 1, 3)); // [nblk, n_tps, Hd]
cb(bm, "msa_block_mask", il);
// expand block -> token granularity on the GPU (j = bk*blk + t),
// expand block -> token granularity (j = bk*blk + t),
// then combine with the causal mask in place
ggml_tensor * bmx = ggml_repeat_4d(ctx0,
ggml_reshape_3d(ctx0, bm, 1, nblk, n_tps*Hd),
src/models/models.h
+4
View File
@@ -2127,6 +2127,10 @@ struct llama_model_mimo2 : public llama_model_base {
graph(const llama_model & model, const llm_graph_params & params);
};
struct graph_mtp : public llm_graph_context {
graph_mtp(const llama_model & model, const llm_graph_params & params);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};