public_git
/
whisper.cpp
mirror of https://github.com/ggml-org/whisper.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'ggml-org:master' into master

This commit is contained in:
mits 2026-08-15 17:10:18 +09:00 committed by GitHub
parent 05c034f5f7 1fe009caed
commit 47154a1757
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
174 changed files with 10562 additions and 2410 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
examples/talk-llama/CMakeLists.txt
View File

@ -35,6 +35,7 @@ if (WHISPER_SDL2)
unicode-data.cpp
${SRC_MODELS})
target_include_directories(${TARGET} PRIVATE . ${SDL2_INCLUDE_DIRS})
target_compile_definitions(${TARGET} PRIVATE -DLLAMA_VERSION="0.0.0")
target_link_libraries(${TARGET} PRIVATE common common-sdl whisper ${SDL2_LIBRARIES} ${CMAKE_THREAD_LIBS_INIT})
install(TARGETS ${TARGET} RUNTIME)

12
examples/talk-llama/llama-arch.cpp
View File

@ -71,6 +71,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_OLMO, "olmo" },
{ LLM_ARCH_OLMO2, "olmo2" },
{ LLM_ARCH_OLMOE, "olmoe" },
{ LLM_ARCH_MUSE_GLIMMER, "muse-glimmer" },
{ LLM_ARCH_OPENELM, "openelm" },
{ LLM_ARCH_ARCTIC, "arctic" },
{ LLM_ARCH_DEEPSEEK, "deepseek" },
@ -100,6 +101,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_GRANITE, "granite" },
{ LLM_ARCH_GRANITE_MOE, "granitemoe" },
{ LLM_ARCH_GRANITE_HYBRID, "granitehybrid" },
{ LLM_ARCH_GRANITE_SWITCH, "graniteswitch" },
{ LLM_ARCH_CHAMELEON, "chameleon" },
{ LLM_ARCH_WAVTOKENIZER_DEC, "wavtokenizer-dec" },
{ LLM_ARCH_PLM, "plm" },
@ -144,6 +146,8 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_TALKIE, "talkie" },
{ LLM_ARCH_MELLUM, "mellum" },
{ LLM_ARCH_NANBEIGE, "nanbeige" },
{ LLM_ARCH_QWEN3TTS, "qwen3tts" },
{ LLM_ARCH_POCKETTTS, "pockettts" },
{ LLM_ARCH_UNKNOWN, "(unknown)" },
};
@ -219,6 +223,11 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_TIME_DECAY_EXTRA_DIM, "%s.time_decay_extra_dim" },
{ LLM_KV_RESIDUAL_SCALE, "%s.residual_scale" },
{ LLM_KV_EMBEDDING_SCALE, "%s.embedding_scale" },
{ LLM_KV_ADAPTER_COUNT, "%s.adapters.count" },
{ LLM_KV_ADAPTER_TOKEN_IDS_ACTIVATE, "%s.adapters.token_ids_activate" },
{ LLM_KV_ADAPTER_TOKEN_IDS_SUBSTITUTE, "%s.adapters.token_ids_substitute" },
{ LLM_KV_ADAPTER_LORA_RANK, "%s.adapters.lora_rank" },
{ LLM_KV_ADAPTER_ROUTER_GAIN, "%s.adapters.router_gain" },
{ LLM_KV_TOKEN_SHIFT_COUNT, "%s.token_shift_count" },
{ LLM_KV_INTERLEAVE_MOE_LAYER_STEP, "%s.interleave_moe_layer_step" },
{ LLM_KV_FULL_ATTENTION_INTERVAL, "%s.full_attention_interval" },
@ -992,6 +1001,8 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
case LLM_ARCH_QWEN35:
case LLM_ARCH_QWEN35MOE:
case LLM_ARCH_DEEPSEEK4:
case LLM_ARCH_NEMOTRON_H:
case LLM_ARCH_NEMOTRON_H_MOE:
return true;
default:
return false;
@ -1026,6 +1037,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
case LLM_ARCH_MINIMAX_M3:
case LLM_ARCH_MISTRAL4:
case LLM_ARCH_KIMI_LINEAR:
case LLM_ARCH_QWEN3TTS:
return false;
default:
return true;

9
examples/talk-llama/llama-arch.h
View File

@ -76,6 +76,7 @@ enum llm_arch {
LLM_ARCH_OLMO,
LLM_ARCH_OLMO2,
LLM_ARCH_OLMOE,
LLM_ARCH_MUSE_GLIMMER,
LLM_ARCH_OPENELM,
LLM_ARCH_ARCTIC,
LLM_ARCH_DEEPSEEK,
@ -105,6 +106,7 @@ enum llm_arch {
LLM_ARCH_GRANITE,
LLM_ARCH_GRANITE_MOE,
LLM_ARCH_GRANITE_HYBRID,
LLM_ARCH_GRANITE_SWITCH,
LLM_ARCH_CHAMELEON,
LLM_ARCH_WAVTOKENIZER_DEC,
LLM_ARCH_PLM,
@ -149,6 +151,8 @@ enum llm_arch {
LLM_ARCH_MINIMAX_M3,
LLM_ARCH_DFLASH,
LLM_ARCH_NANBEIGE,
LLM_ARCH_QWEN3TTS,
LLM_ARCH_POCKETTTS,
LLM_ARCH_UNKNOWN,
};
@ -224,6 +228,11 @@ enum llm_kv {
LLM_KV_TIME_DECAY_EXTRA_DIM,
LLM_KV_RESIDUAL_SCALE,
LLM_KV_EMBEDDING_SCALE,
LLM_KV_ADAPTER_COUNT,
LLM_KV_ADAPTER_TOKEN_IDS_ACTIVATE,
LLM_KV_ADAPTER_TOKEN_IDS_SUBSTITUTE,
LLM_KV_ADAPTER_LORA_RANK,
LLM_KV_ADAPTER_ROUTER_GAIN,
LLM_KV_TOKEN_SHIFT_COUNT,
LLM_KV_INTERLEAVE_MOE_LAYER_STEP,
LLM_KV_FULL_ATTENTION_INTERVAL,

325
examples/talk-llama/llama-context.cpp
View File

@ -10,6 +10,7 @@
#include "llama-mmap.h"
#include "llama-model.h"
#include "llama-ext.h"
#include "llama-sampler.h"
#include "llama.h"
#include <cinttypes>
@ -102,7 +103,7 @@ llama_context::llama_context(
cparams.n_rs_seq = params.n_rs_seq;
if (cparams.n_rs_seq > 0 && !llm_arch_supports_rs_rollback(model.arch)) {
LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model arch does not support recurrent partial rollback; clamping to 0\n",
LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model does not support recurrent partial rollback; clamping to 0\n",
__func__, cparams.n_rs_seq);
cparams.n_rs_seq = 0;
}
@ -159,25 +160,6 @@ llama_context::llama_context(
}
}
// Initialize backend samplers here so they are part of the sampling graph
// before the reserve passes run later in this function. This avoids a later
// re-reserve when graph nodes change.
if (params.samplers != nullptr && params.n_samplers > 0) {
for (size_t i = 0; i < params.n_samplers; ++i) {
const auto & config = params.samplers[i];
if (llama_sampler_chain_get(config.sampler, -1) == nullptr) {
throw std::runtime_error("the backend samplers must be of type llama_sampler_chain");
}
if (set_sampler(config.seq_id, config.sampler)) {
const int n_samplers = llama_sampler_chain_n(config.sampler);
LLAMA_LOG_INFO("%s: setting backend sampler for seq_id %d (n = %d)\n", __func__, config.seq_id, n_samplers);
}
}
}
auto rope_scaling_type = params.rope_scaling_type;
if (rope_scaling_type == LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED) {
rope_scaling_type = hparams.rope_scaling_type_train;
@ -265,6 +247,27 @@ llama_context::llama_context(
cparams.n_ubatch = std::min(cparams.n_batch, params.n_ubatch == 0 ? params.n_batch : params.n_ubatch);
cparams.n_outputs_max = params.n_outputs_max == 0 || llama_model_has_encoder(&model) ? cparams.n_batch : params.n_outputs_max;
cparams.n_outputs_max_per_seq = params.n_outputs_max_per_seq == 0 ?
cparams.n_outputs_max : std::min(params.n_outputs_max_per_seq, cparams.n_outputs_max);
// Initialize backend samplers here so they are part of the sampling graph
// before the reserve passes run later in this function. This avoids a later
// re-reserve when graph nodes change.
if (params.samplers != nullptr && params.n_samplers > 0) {
for (size_t i = 0; i < params.n_samplers; ++i) {
const auto & config = params.samplers[i];
if (llama_sampler_chain_get(config.sampler, -1) == nullptr) {
throw std::runtime_error("the backend samplers must be of type llama_sampler_chain");
}
if (set_sampler(config.seq_id, config.sampler)) {
const int n_samplers = llama_sampler_chain_n(config.sampler);
LLAMA_LOG_INFO("%s: setting backend sampler for seq_id %d (n = %d)\n", __func__, config.seq_id, n_samplers);
}
}
}
cparams.op_offload = params.op_offload;
cparams.kv_unified = params.kv_unified;
@ -300,18 +303,19 @@ llama_context::llama_context(
}
}
LLAMA_LOG_INFO("%s: n_seq_max = %u\n", __func__, cparams.n_seq_max);
LLAMA_LOG_INFO("%s: n_ctx = %u\n", __func__, cparams.n_ctx);
LLAMA_LOG_INFO("%s: n_ctx_seq = %u\n", __func__, cparams.n_ctx_seq);
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: causal_attn = %d\n", __func__, cparams.causal_attn);
LLAMA_LOG_INFO("%s: flash_attn = %s\n", __func__, llama_flash_attn_type_name(params.flash_attn_type));
LLAMA_LOG_INFO("%s: kv_unified = %s\n", __func__, cparams.kv_unified ? "true" : "false");
LLAMA_LOG_INFO("%s: freq_base = %.1f\n", __func__, cparams.rope_freq_base);
LLAMA_LOG_INFO("%s: freq_scale = %g\n", __func__, cparams.rope_freq_scale);
LLAMA_LOG_INFO("%s: n_rs_seq = %u\n", __func__, cparams.n_rs_seq);
LLAMA_LOG_INFO("%s: n_outputs_max = %u\n", __func__, cparams.n_outputs_max);
LLAMA_LOG_INFO("%s: n_seq_max = %u\n", __func__, cparams.n_seq_max);
LLAMA_LOG_INFO("%s: n_ctx = %u\n", __func__, cparams.n_ctx);
LLAMA_LOG_INFO("%s: n_ctx_seq = %u\n", __func__, cparams.n_ctx_seq);
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: causal_attn = %d\n", __func__, cparams.causal_attn);
LLAMA_LOG_INFO("%s: flash_attn = %s\n", __func__, llama_flash_attn_type_name(params.flash_attn_type));
LLAMA_LOG_INFO("%s: kv_unified = %s\n", __func__, cparams.kv_unified ? "true" : "false");
LLAMA_LOG_INFO("%s: freq_base = %.1f\n", __func__, cparams.rope_freq_base);
LLAMA_LOG_INFO("%s: freq_scale = %g\n", __func__, cparams.rope_freq_scale);
LLAMA_LOG_INFO("%s: n_rs_seq = %u\n", __func__, cparams.n_rs_seq);
LLAMA_LOG_INFO("%s: n_outputs_max = %u\n", __func__, cparams.n_outputs_max);
LLAMA_LOG_INFO("%s: n_outputs_max_per_seq = %u\n", __func__, cparams.n_outputs_max_per_seq);
if (cparams.n_ctx_seq < hparams.n_ctx_train) {
LLAMA_LOG_INFO("%s: n_ctx_seq (%u) < n_ctx_train (%u) -- the full capacity of the model will not be utilized\n",
@ -1231,7 +1235,7 @@ bool llama_context::set_sampler(llama_seq_id seq_id, llama_sampler * sampler) {
if (sampler && can_offload) {
auto * buft = ggml_backend_dev_buffer_type(model.dev_output());
sampler->iface->backend_init(sampler, buft);
sampler->iface->backend_init(sampler, buft, cparams.n_outputs_max_per_seq);
sampling.samplers[seq_id] = sampler;
@ -1576,108 +1580,38 @@ int llama_context::encode(const llama_batch & batch_inp) {
return 0;
}
static std::map<llama_seq_id, uint32_t> build_seq_to_output_row(const llama_ubatch & ubatch, uint32_t row_offset) {
std::map<llama_seq_id, uint32_t> seq_to_row;
// how many output tokens we have seen so far for this ubatch.
uint32_t local = 0;
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
// skip tokens that are not output.
if (!ubatch.output[i]) {
continue;
}
const llama_seq_id seq_id = ubatch.seq_id[i][0];
// row_offset is the number of output tokens before this ubatch.
seq_to_row[seq_id] = row_offset + local;
++local;
}
return seq_to_row;
}
static void copy_tensor_async_ints(
const std::map<llama_seq_id, ggml_tensor*> & tensor_map,
const buffer_view<llama_token> & sampled,
const std::map<llama_seq_id, uint32_t> & seq_to_row,
ggml_backend_sched_t sched) {
if (!sampled.has_data()) {
return;
}
for (const auto & [seq_id, tensor] : tensor_map) {
auto it = seq_to_row.find(seq_id);
if (it == seq_to_row.end()) {
continue;
}
const uint32_t row = it->second;
GGML_ASSERT(row < sampled.size);
GGML_ASSERT(ggml_is_contiguous(tensor) && "sampled tokens tensor must be contiguous for async copy");
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor);
ggml_backend_tensor_get_async(backend, tensor, sampled.data + row, 0, sizeof(sampled.data[row]));
}
}
static void copy_tensor_async_floats(
const std::map<llama_seq_id, ggml_tensor*> & tensor_map,
const buffer_view<float> & dst,
template<typename T>
static void copy_tensor_async_rows(
const std::vector<ggml_tensor *> & tensors,
const buffer_view<T> & dst,
size_t stride,
std::vector<uint32_t> & counts,
const std::map<llama_seq_id, uint32_t> & seq_to_row,
ggml_backend_sched_t sched) {
uint32_t row_offset,
ggml_backend_sched_t sched,
std::vector<uint32_t> * counts = nullptr) {
if (!dst.has_data()) {
return;
}
for (const auto & [seq_id, tensor] : tensor_map) {
auto it = seq_to_row.find(seq_id);
if (it == seq_to_row.end()) {
for (size_t i = 0; i < tensors.size(); ++i) {
auto * tensor = tensors[i];
if (tensor == nullptr) {
continue;
}
const uint32_t row = it->second;
GGML_ASSERT(row < counts.size());
GGML_ASSERT(ggml_is_contiguous(tensor) && "logits/probs tensor must be contiguous for async copy");
const uint32_t row = row_offset + i;
const size_t n_elements = ggml_nelements(tensor);
GGML_ASSERT(ggml_is_contiguous(tensor) && "sampling tensor must be contiguous for async copy");
GGML_ASSERT(n_elements <= stride);
GGML_ASSERT((size_t) row * stride + n_elements <= dst.size);
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor);
float * row_ptr = dst.data + (size_t) row * stride;
T * row_ptr = dst.data + (size_t) row * stride;
ggml_backend_tensor_get_async(backend, tensor, row_ptr, 0, ggml_nbytes(tensor));
// Update the actual number of logits/probabilities that were written for this row.
counts[row] = ggml_nelements(tensor);
}
}
static void copy_tensor_async_candidates(
const std::map<llama_seq_id, ggml_tensor*> & tensor_map,
const buffer_view<llama_token> & dst,
size_t stride,
std::vector<uint32_t> & counts,
const std::map<llama_seq_id, uint32_t> & seq_to_row,
ggml_backend_sched_t sched) {
if (!dst.has_data()) {
return;
}
for (const auto & [seq_id, tensor] : tensor_map) {
auto it = seq_to_row.find(seq_id);
if (it == seq_to_row.end()) {
continue;
if (counts) {
GGML_ASSERT(row < counts->size());
(*counts)[row] = n_elements;
}
const uint32_t row = it->second;
GGML_ASSERT(row < counts.size());
GGML_ASSERT(ggml_is_contiguous(tensor) && "candidates tensor must be contiguous for async copy");
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor);
llama_token * row_ptr = dst.data + (size_t) row * stride;
ggml_backend_tensor_get_async(backend, tensor, row_ptr, 0, ggml_nbytes(tensor));
// Update the actual number of candidates that were written.
counts[row] = ggml_nelements(tensor);
}
}
@ -1726,12 +1660,12 @@ int llama_context::decode(const llama_batch & batch_inp) {
const uint32_t n_seq_max = cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max;
// TODO: avoid this workaround in the future
if (has_samplers && batch_inp.logits) {
// embedding contexts output every token even when batch.logits is not set
if (has_samplers && (output_all || batch_inp.logits)) {
std::vector<int32_t> seq_output_count(n_seq_max, 0);
for (int32_t i = 0; i < batch_inp.n_tokens; ++i) {
if (batch_inp.logits[i] == 0) {
if (!output_all && batch_inp.logits[i] == 0) {
continue;
}
@ -1740,10 +1674,17 @@ int llama_context::decode(const llama_batch & batch_inp) {
for (int32_t s = 0; s < ns; ++s) {
const llama_seq_id seq_id = batch_inp.seq_id ? batch_inp.seq_id[i][s] : 0;
if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) {
continue;
}
seq_output_count[seq_id]++;
if (seq_output_count[seq_id] > 1) {
LLAMA_LOG_ERROR("%s: backend sampling requires at most one output token per sequence (seq_id %d had %d)\n",
__func__, seq_id, seq_output_count[seq_id]);
auto sampler = sampling.samplers.find(seq_id);
if (sampler != sampling.samplers.end() &&
seq_output_count[seq_id] > (int32_t) cparams.n_outputs_max_per_seq) {
LLAMA_LOG_ERROR("%s: backend sampling supports at most %u outputs per sequence "
"(seq_id %d had %d)\n", __func__, cparams.n_outputs_max_per_seq,
seq_id, seq_output_count[seq_id]);
return -1;
}
}
@ -1843,6 +1784,11 @@ int llama_context::decode(const llama_batch & batch_inp) {
return -2;
};
// start a new sampling transaction for this logical batch
for (const auto & entry : sampling.samplers) {
llama_sampler_backend_begin(entry.second);
}
int64_t n_outputs_prev = 0;
int64_t n_tokens_prev = 0;
@ -2009,17 +1955,14 @@ int llama_context::decode(const llama_batch & batch_inp) {
}
}
// Copy backend sampling output if this ubatch produced any sampling tensors.
if (has_samplers && (!res->t_sampled.empty() || !res->t_sampled_probs.empty() || !res->t_sampled_logits.empty())) {
const auto seq_to_output_row = build_seq_to_output_row(ubatch, n_outputs_prev);
if (has_samplers) {
const auto stride = n_vocab;
// async copy the sampling data from the backend to the host
copy_tensor_async_ints(res->t_sampled, sampling.sampled, seq_to_output_row, sched.get());
copy_tensor_async_floats (res->t_sampled_logits, sampling.logits, stride, sampling.logits_count, seq_to_output_row, sched.get());
copy_tensor_async_floats (res->t_sampled_probs, sampling.probs, stride, sampling.probs_count, seq_to_output_row, sched.get());
copy_tensor_async_candidates(res->t_candidates, sampling.candidates, stride, sampling.candidates_count, seq_to_output_row, sched.get());
copy_tensor_async_rows(res->t_sampled, sampling.sampled, 1, n_outputs_prev, sched.get());
copy_tensor_async_rows(res->t_sampled_logits, sampling.logits, stride, n_outputs_prev, sched.get(), &sampling.logits_count);
copy_tensor_async_rows(res->t_sampled_probs, sampling.probs, stride, n_outputs_prev, sched.get(), &sampling.probs_count);
copy_tensor_async_rows(res->t_candidates, sampling.candidates, stride, n_outputs_prev, sched.get(), &sampling.candidates_count);
}
n_outputs_prev += n_outputs;
@ -2349,6 +2292,7 @@ void llama_context::output_reorder() {
//
uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
uint32_t res;
if (model.arch == LLM_ARCH_QWEN3NEXT ||
model.arch == LLM_ARCH_KIMI_LINEAR ||
model.arch == LLM_ARCH_QWEN35 ||
@ -2357,11 +2301,31 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
(model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) ||
model.arch == LLM_ARCH_NANBEIGE ||
model.arch == LLM_ARCH_MINIMAX_M3) {
return std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
res = std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
} else {
res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
for (const auto & lora : model.loras) {
res += lora->get_n_nodes();
}
}
uint32_t res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
for (const auto & lora : model.loras) {
res += lora->get_n_nodes();
uint32_t n_sampling_nodes = 0;
uint32_t n_sampling_nodes_max = 0;
for (const auto & [seq_id, sampler] : sampling.samplers) {
const uint32_t n_nodes = llama_sampler_backend_n_nodes(sampler);
n_sampling_nodes += n_nodes;
if (cparams.n_outputs_max_per_seq > 1) {
n_sampling_nodes_max = std::max(n_sampling_nodes_max, n_nodes);
}
}
const uint32_t n_sampling_outputs_max = std::min<uint64_t>(
std::min(n_tokens, cparams.n_outputs_max),
(uint64_t) cparams.n_seq_max * cparams.n_outputs_max_per_seq);
res += n_sampling_nodes;
if (n_sampling_outputs_max > 1) {
res += (n_sampling_outputs_max - 1) * n_sampling_nodes_max;
}
return res;
}
@ -2370,6 +2334,63 @@ llm_graph_result * llama_context::get_gf_res_reserve() const {
return static_cast<llm_graph_result *>(gf_res_reserve.get());
}
// pack sampler outputs into as few sequences as possible before using sequences without samplers
static void ubatch_prepare_reserve(
llama_ubatch & ubatch,
uint32_t n_outputs,
const std::map<llama_seq_id, llama_sampler *> & samplers,
uint32_t n_outputs_max_per_seq) {
const uint32_t n_seqs = ubatch.n_seqs;
const uint32_t n_seq_tokens = ubatch.n_seq_tokens;
for (uint32_t s = 0; s < n_seqs; ++s) {
for (uint32_t t = 0; t < n_seq_tokens; ++t) {
const uint32_t i = s * n_seq_tokens + t;
ubatch.n_seq_id[i] = 1;
ubatch.seq_id[i] = &ubatch.seq_id_unq[s];
}
}
// sequences with a sampler that fit in this ubatch
std::vector<uint32_t> sampler_seqs;
std::vector<bool> has_sampler(n_seqs, false);
for (const auto & entry : samplers) {
const llama_seq_id seq_id = entry.first;
if (seq_id < 0 || (uint32_t) seq_id >= n_seqs) {
continue;
}
sampler_seqs.push_back(seq_id);
has_sampler[seq_id] = true;
}
uint32_t n_outputs_set = 0;
const uint32_t n_outputs_per_seq = std::min(n_seq_tokens, n_outputs_max_per_seq);
for (uint32_t s : sampler_seqs) {
if (n_outputs_set >= n_outputs) {
break;
}
for (uint32_t t = 0; t < n_outputs_per_seq && n_outputs_set < n_outputs; ++t) {
ubatch.output[s * n_seq_tokens + t] = true;
++n_outputs_set;
}
}
// use sequences without samplers for any remaining outputs
for (uint32_t t = 0; t < n_seq_tokens && n_outputs_set < n_outputs; ++t) {
for (uint32_t s = 0; s < n_seqs && n_outputs_set < n_outputs; ++s) {
if (has_sampler[s]) {
continue;
}
ubatch.output[s * n_seq_tokens + t] = true;
++n_outputs_set;
}
}
}
ggml_cgraph * llama_context::graph_reserve(
uint32_t n_tokens, uint32_t n_seqs, uint32_t n_outputs, const llama_memory_context_i * mctx, bool split_only, size_t * sizes) {
LLAMA_LOG_DEBUG("%s: reserving a graph for ubatch with n_tokens = %4u, n_seqs = %2u, n_outputs = %4u\n", __func__, n_tokens, n_seqs, n_outputs);
@ -2394,14 +2415,7 @@ ggml_cgraph * llama_context::graph_reserve(
llama_batch_allocr balloc(model.hparams.n_pos_per_embd());
llama_ubatch ubatch = balloc.ubatch_reserve(n_tokens/n_seqs, n_seqs);
// set one output token per sequence in order to activate all backend samplers
std::vector<llama_seq_id> seq_ids(n_seqs);
for (uint32_t i = 0; i < n_seqs; ++i) {
seq_ids[i] = i;
ubatch.n_seq_id[i] = 1;
ubatch.seq_id[i] = &seq_ids[i];
ubatch.output[i] = true;
}
ubatch_prepare_reserve(ubatch, n_outputs, sampling.samplers, cparams.n_outputs_max_per_seq);
auto * res = gf_res_reserve.get();
@ -3096,6 +3110,17 @@ size_t llama_context::state_seq_load_file(llama_seq_id seq_id, const char * file
{
const uint32_t n_token_count = file.read_u32();
if (tokens_out == nullptr) {
const size_t n_token_max = (file.size() - file.tell()) / sizeof(llama_token);
if (n_token_count > n_token_max) {
LLAMA_LOG_ERROR("%s: token count in sequence state file exceeds the file size! %u > %zu\n", __func__, n_token_count, n_token_max);
return 0;
}
*n_token_count_out = n_token_count;
return file.tell();
}
if (n_token_count > n_token_capacity) {
LLAMA_LOG_ERROR("%s: token count in sequence state file exceeded capacity! %u > %zu\n", __func__, n_token_count, n_token_capacity);
return 0;
@ -3488,6 +3513,7 @@ llama_context_params llama_context_default_params() {
/*.n_seq_max =*/ 1,
/*.n_rs_seq =*/ 0,
/*.n_outputs_max =*/ 0,
/*.n_outputs_max_per_seq =*/ 1,
/*.n_threads =*/ GGML_DEFAULT_N_THREADS, // TODO: better default
/*.n_threads_batch =*/ GGML_DEFAULT_N_THREADS,
/*.ctx_type =*/ LLAMA_CONTEXT_TYPE_DEFAULT,
@ -3602,8 +3628,9 @@ llama_context * llama_init_from_model(
model->hparams.pooling_type, params.pooling_type);
}
// router_layer >= 0 means n_layer_nextn is repurposed for a router layer, not real MTP
if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP &&
model->hparams.n_layer_nextn == 0) {
(model->hparams.n_layer_nextn == 0 || model->hparams.router_layer >= 0)) {
LLAMA_LOG_WARN("%s: context type MTP requested but model doesn't contain MTP layers\n", __func__);
return nullptr;
}

1
examples/talk-llama/llama-cparams.h
View File

@ -15,6 +15,7 @@ struct llama_cparams {
uint32_t n_seq_max;
uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback
uint32_t n_outputs_max; // max outputs supported by the context
uint32_t n_outputs_max_per_seq;
int32_t n_threads; // number of threads to use for generation
int32_t n_threads_batch; // number of threads to use for batch processing

6
examples/talk-llama/llama-ext.h
View File

@ -124,3 +124,9 @@ LLAMA_API llama_context * llama_get_ctx_other(struct llama_context * ctx);
LLAMA_API const int32_t * llama_model_target_layer_ids (const struct llama_model * model);
// returns the number of extracted layers from target model
LLAMA_API uint32_t llama_model_target_layer_ids_n(const struct llama_model * model);
// retrieves the whole token embedding matrix in F32 format (n_embd * n_vocab)
// returns total number of elements or 0 on error
// if out is nullptr, returns the number of tokens without writing to out
// caller must allocate enough memory for out before calling
LLAMA_API uint32_t llama_model_get_tok_embd(const struct llama_model * model, float * out);

6
examples/talk-llama/llama-grammar.cpp
View File

@ -648,10 +648,12 @@ const char * llama_grammar_parser::parse_sequence(
} else {
throw std::runtime_error(std::string("expecting ',' at ") + pos);
}
bool has_max = max_times != UINT64_MAX;
if (min_times > MAX_REPETITION_THRESHOLD || (has_max && max_times > MAX_REPETITION_THRESHOLD)) {
if (min_times > MAX_REPETITION_THRESHOLD) {
throw std::runtime_error(std::string("number of repetitions exceeds sane defaults, please reduce the number of repetitions"));
}
if (max_times != UINT64_MAX && max_times > MAX_REPETITION_THRESHOLD) {
max_times = UINT64_MAX;
}
handle_repetitions(min_times, max_times);
} else {
break;

164
examples/talk-llama/llama-graph.cpp
View File

@ -4,6 +4,7 @@
#include "llama-model.h"
#include "llama-batch.h"
#include "llama-cparams.h"
#include "llama-sampler.h"
#include "llama-kv-cache.h"
#include "llama-kv-cache-iswa.h"
@ -1353,24 +1354,24 @@ void llm_graph_result::set_outputs(const llm_graph_params & params) {
}
}
}
for (auto & [seq_id, t] : t_sampled) {
if (t != nullptr) {
ggml_set_output(t);
for (auto * tensor : t_sampled) {
if (tensor != nullptr) {
ggml_set_output(tensor);
}
}
for (auto & [seq_id, t] : t_sampled_probs) {
if (t != nullptr) {
ggml_set_output(t);
for (auto * tensor : t_sampled_probs) {
if (tensor != nullptr) {
ggml_set_output(tensor);
}
}
for (auto & [seq_id, t] : t_sampled_logits) {
if (t != nullptr) {
ggml_set_output(t);
for (auto * tensor : t_sampled_logits) {
if (tensor != nullptr) {
ggml_set_output(tensor);
}
}
for (auto & [seq_id, t] : t_candidates) {
if (t != nullptr) {
ggml_set_output(t);
for (auto * tensor : t_candidates) {
if (tensor != nullptr) {
ggml_set_output(tensor);
}
}
}
@ -3649,77 +3650,102 @@ void llm_graph_context::build_sampling() const {
auto inp_sampling = std::make_unique<llm_graph_input_sampling>(samplers);
res->add_input(std::move(inp_sampling));
std::map<llama_seq_id, int32_t> seq_to_logit_row;
int32_t logit_row_idx = 0;
for (uint32_t i = 0; i < ubatch.n_tokens; i++) {
std::map<llama_seq_id, std::vector<uint32_t>> sampling_rows;
uint32_t n_rows = 0;
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
if (ubatch.output[i]) {
llama_seq_id seq_id = ubatch.seq_id[i][0];
seq_to_logit_row[seq_id] = logit_row_idx;
logit_row_idx++;
sampling_rows[ubatch.seq_id[i][0]].push_back(n_rows++);
}
}
res->t_sampled.resize(n_rows, nullptr);
res->t_sampled_probs.resize(n_rows, nullptr);
res->t_sampled_logits.resize(n_rows, nullptr);
res->t_candidates.resize(n_rows, nullptr);
// res->t_logits will contain logits for all tokens that want the logits calculated (logits=1 or output=1)
GGML_ASSERT(res->t_logits != nullptr && "missing t_logits tensor");
// add a dummy row of logits
// this trick makes the graph static, regardless of which samplers are activated
// this is important in order to minimize graph reallocations
// add a dummy row to keep the single-output graph static regardless of active samplers
// multi-output graphs can still vary with the number of output rows
ggml_tensor * logits_t = ggml_pad(ctx0, res->t_logits, 0, 1, 0, 0);
for (const auto & [seq_id, sampler] : samplers) {
const auto it = seq_to_logit_row.find(seq_id);
// inactive samplers always work on the first row
const auto row_idx = it != seq_to_logit_row.end() ? it->second : 0;
const int i_out = it != seq_to_logit_row.end() ? 1 : 0;
ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], row_idx * logits_t->nb[1]);
ggml_format_name(logits_seq, "logits_seq_%d", seq_id);
struct llama_sampler_data data = {
/*.logits =*/ logits_seq,
/*.probs =*/ nullptr,
/*.sampled =*/ nullptr,
/*.candidates =*/ nullptr,
};
assert(sampler->iface->backend_apply);
sampler->iface->backend_apply(sampler, ctx0, gf, &data);
if (data.sampled != nullptr) {
res->t_sampled[seq_id] = data.sampled;
outs[1] = data.sampled;
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
if (data.probs != nullptr) {
res->t_sampled_probs[seq_id] = data.probs;
outs[1] = data.probs;
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
if (data.logits != nullptr) {
res->t_sampled_logits[seq_id] = data.logits;
outs[1] = data.logits;
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
if (data.candidates != nullptr) {
res->t_candidates[seq_id] = data.candidates;
outs[1] = data.candidates;
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
for (const auto & entry : samplers) {
if (entry.second->iface->backend_reset) {
entry.second->iface->backend_reset(entry.second);
}
}
// TODO: Call llama_sampler_accept_ggml after all samplers have been applied.
static const std::vector<uint32_t> dummy_row = { 0 };
for (const auto & [seq_id, sampler] : samplers) {
const auto it = sampling_rows.find(seq_id);
// inactive samplers always work on the first row
const bool active = it != sampling_rows.end();
const auto & rows = active ? it->second : dummy_row;
const int i_out = active ? 1 : 0;
for (uint32_t i = 0; i < rows.size(); ++i) {
ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], rows[i] * logits_t->nb[1]);
ggml_format_name(logits_seq, "logits_seq_%d_%u", seq_id, i);
struct llama_sampler_data data = {
/*.logits =*/ logits_seq,
/*.probs =*/ nullptr,
/*.sampled =*/ nullptr,
/*.candidates =*/ nullptr,
};
assert(sampler->iface->backend_apply);
sampler->iface->backend_apply(sampler, ctx0, gf, &data);
if (data.sampled != nullptr) {
if (active) {
res->t_sampled[rows[i]] = data.sampled;
}
outs[1] = data.sampled;
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
if (data.probs != nullptr) {
if (active) {
res->t_sampled_probs[rows[i]] = data.probs;
}
outs[1] = data.probs;
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
if (data.logits != nullptr) {
if (active) {
res->t_sampled_logits[rows[i]] = data.logits;
}
outs[1] = data.logits;
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
if (data.candidates != nullptr) {
if (active) {
res->t_candidates[rows[i]] = data.candidates;
}
outs[1] = data.candidates;
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
}
}
}
// TODO: Call backend_accept after all samplers have been applied.
/*
for (const auto & [seq_id, sampler] : samplers) {
if (auto it = res->t_sampled.find(seq_id); it != res->t_sampled.end()) {
ggml_tensor * selected_token = it->second;
if (selected_token != nullptr) {
llama_sampler_accept_ggml(sampler, ctx0, gf, selected_token);
const auto it = sampling_rows.find(seq_id);
if (it == sampling_rows.end()) {
continue;
}
for (uint32_t row : it->second) {
ggml_tensor * selected_token = res->t_sampled[row];
if (selected_token != nullptr && sampler->iface->backend_accept) {
sampler->iface->backend_accept(sampler, ctx0, gf, selected_token);
}
}
}

8
examples/talk-llama/llama-graph.h
View File

@ -904,10 +904,10 @@ public:
std::vector<ggml_tensor *> t_layer_inp;
std::map<llama_seq_id, ggml_tensor *> t_sampled_logits;
std::map<llama_seq_id, ggml_tensor *> t_candidates;
std::map<llama_seq_id, ggml_tensor *> t_sampled;
std::map<llama_seq_id, ggml_tensor *> t_sampled_probs;
std::vector<ggml_tensor *> t_sampled;
std::vector<ggml_tensor *> t_sampled_probs;
std::vector<ggml_tensor *> t_sampled_logits;
std::vector<ggml_tensor *> t_candidates;
std::vector<llm_graph_input_ptr> inputs;
std::vector<llm_graph_fused_node> fused_nodes;

10
examples/talk-llama/llama-hparams.cpp
View File

@ -277,6 +277,16 @@ bool llama_hparams::has_kv(uint32_t il) const {
return true;
}
bool llama_hparams::has_rope(uint32_t il) const {
// the router layer stores adapter routing signal, not positional info,
// so it must not be RoPE-shifted
if (router_layer >= 0 && (int32_t) il == router_layer) {
return false;
}
return true;
}
uint32_t llama_hparams::n_layer() const {
return n_layer_all - n_layer_nextn;
}

6
examples/talk-llama/llama-hparams.h
View File

@ -53,6 +53,10 @@ struct llama_hparams {
uint32_t n_embd;
uint32_t n_layer_all;
uint32_t n_layer_nextn = 0;
// granite-switch: index of the single-head "router" KV layer that encodes
// per-token adapter selection. -1 when the model has no such layer.
int32_t router_layer = -1;
uint32_t n_expert = 0;
uint32_t n_expert_used = 0;
uint32_t n_rel_attn_bkts = 0;
@ -371,6 +375,8 @@ struct llama_hparams {
bool has_kv(uint32_t il) const;
bool has_rope(uint32_t il) const;
// number of effective layers (excludes nextn layers)
uint32_t n_layer() const;

4
examples/talk-llama/llama-kv-cache.cpp
View File

@ -1931,6 +1931,10 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co
for (const auto & layer : layers) {
const uint32_t il = layer.il;
if (!hparams.has_rope(il)) {
continue;
}
const int64_t n_head_kv = hparams.n_head_kv(il);
const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);

46
examples/talk-llama/llama-model-loader.cpp
View File

@ -543,7 +543,7 @@ llama_model_loader::llama_model_loader(
tensor_buft_overrides = param_tensor_buft_overrides_p;
this->use_mmap = load_mode == LLAMA_LOAD_MODE_MMAP || load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK;
this->use_mmap = load_mode == LLAMA_LOAD_MODE_MMAP || load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK || load_mode == LLAMA_LOAD_MODE_AUTO;
this->use_direct_io = load_mode == LLAMA_LOAD_MODE_DIRECT_IO;
if (!fname.empty()) {
@ -937,10 +937,11 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w
} break;
case GGML_OP_MUL_MAT_ID:
{
const int n_expert_used = hparams.n_expert_used;
GGML_ASSERT(n_expert_used > 0);
ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_expert_used, 512);
ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, 512);
// Used for either MoE expert routing or embedded adapter routing
const int n_ids_used = hparams.router_layer >= 0 ? 1 : hparams.n_expert_used;
GGML_ASSERT(n_ids_used > 0);
ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_ids_used, 512);
ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_ids_used, 512);
op_tensor = ggml_mul_mat_id(ctx, w, b, ids);
} break;
case GGML_OP_ADD:
@ -1001,7 +1002,7 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w
ggml_tensor * B = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_state, n_group, n_seq_tokens, n_seqs);
ggml_tensor * C = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_state, n_group, n_seq_tokens, n_seqs);
ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seqs);
op_tensor = ggml_ssm_scan(ctx, s, x, dt, w, B, C, ids);
op_tensor = ggml_ssm_scan(ctx, s, x, dt, w, B, C, ids, /*K=*/1);
} break;
case GGML_OP_RWKV_WKV6:
{
@ -1123,15 +1124,14 @@ struct ggml_tensor * llama_model_loader::create_tensor(
return nullptr;
}
// tensors with "bias" suffix are always used with GGML_OP_ADD or GGML_OP_ADD_ID
// tensors with "bias" suffix are always used with GGML_OP_ADD or GGML_OP_ADD_ID;
// embedded-adapter ".lora_a"/".lora_b" tensors are always used with GGML_OP_MUL_MAT_ID
ggml_op op;
bool bias = tn.suffix != nullptr && strcmp(tn.suffix, "bias") == 0;
if (bias) {
if (info.op == GGML_OP_MUL_MAT_ID) {
op = GGML_OP_ADD_ID;
} else {
op = GGML_OP_ADD;
}
if (tn.suffix != nullptr && strcmp(tn.suffix, "bias") == 0) {
op = info.op == GGML_OP_MUL_MAT_ID ? GGML_OP_ADD_ID : GGML_OP_ADD;
} else if (hparams.router_layer >= 0 && tn.suffix != nullptr &&
(strcmp(tn.suffix, "lora_a") == 0 || strcmp(tn.suffix, "lora_b") == 0)) {
op = GGML_OP_MUL_MAT_ID;
} else {
op = info.op;
}
@ -1249,7 +1249,13 @@ struct ggml_tensor * llama_model_loader::create_tensor(
for (size_t dim = 0; dim < GGML_MAX_DIMS; dim++) {
t_meta.ne[dim] = dim < ne.size() ? ne.begin()[dim] : 1;
GGML_ASSERT(t_meta.ne[dim] >= 1);
t_meta.nb[dim] = dim == 0 ? ggml_type_size(type) : t_meta.ne[dim-1]*t_meta.nb[dim-1];
if (dim == 0) {
t_meta.nb[dim] = ggml_type_size(type);
} else if (dim == 1) {
t_meta.nb[dim] = ggml_row_size(type, t_meta.ne[dim-1]);
} else {
t_meta.nb[dim] = t_meta.nb[dim-1]*t_meta.ne[dim-1];
}
GGML_ASSERT(t_meta.nb[dim] >= 1);
}
ggml_set_name(&t_meta, tn.str().c_str());
@ -1272,10 +1278,18 @@ struct ggml_tensor * llama_model_loader::create_tensor(
if (flags & TENSOR_ALLOW_RESHAPE) {
for (size_t dim = 0; dim < GGML_MAX_DIMS; dim++) {
t_meta.ne[dim] = dim < ne.size() ? ne.begin()[dim] : 1;
t_meta.nb[dim] = dim == 0 ? ggml_type_size(t_meta.type) : t_meta.ne[dim-1]*t_meta.nb[dim-1];
if (dim == 0) {
t_meta.nb[dim] = ggml_type_size(t_meta.type);
} else if (dim == 1) {
t_meta.nb[dim] = ggml_row_size(t_meta.type, t_meta.ne[dim-1]);
} else {
t_meta.nb[dim] = t_meta.ne[dim-1]*t_meta.nb[dim-1];
}
}
}
GGML_ASSERT(ggml_nbytes(&t_meta) == ggml_nbytes(cur));
ggml_backend_buffer_type_t buft = buft_for_tensor(&t_meta);
if (buft == nullptr) {
return nullptr;

3
examples/talk-llama/llama-model-saver.cpp
View File

@ -27,6 +27,7 @@ bool llama_model_saver_supports_arch(llm_arch arch) {
case LLM_ARCH_APERTUS:
case LLM_ARCH_MIMO2:
case LLM_ARCH_STEP35:
case LLM_ARCH_MUSE_GLIMMER:
case LLM_ARCH_MELLUM:
case LLM_ARCH_LAGUNA:
return false;
@ -213,7 +214,7 @@ void llama_model_saver::add_kv_from_model() {
add_kv(LLM_KV_FEED_FORWARD_LENGTH, hparams.n_ff_arr, true);
add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp);
add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_chexp);
add_kv(LLM_KV_EXPERT_CHUNK_FEED_FORWARD_LENGTH, hparams.n_ff_chexp);
add_kv(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp);
add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp);
add_kv(LLM_KV_USE_PARALLEL_RESIDUAL, hparams.use_par_res);

94
examples/talk-llama/llama-model.cpp
View File

@ -40,6 +40,8 @@
static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params & params) {
switch (arch) {
case LLM_ARCH_CLIP:
return new llama_model_clip(params);
case LLM_ARCH_LLAMA:
return new llama_model_llama(params);
case LLM_ARCH_LLAMA4:
@ -112,6 +114,10 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
return new llama_model_qwen3vl(params);
case LLM_ARCH_QWEN3VLMOE:
return new llama_model_qwen3vlmoe(params);
case LLM_ARCH_QWEN3TTS:
return new llama_model_qwen3tts(params);
case LLM_ARCH_POCKETTTS:
return new llama_model_pockettts(params);
case LLM_ARCH_PHI2:
return new llama_model_phi2(params);
case LLM_ARCH_PHI3:
@ -172,6 +178,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
return new llama_model_olmo2(params);
case LLM_ARCH_OLMOE:
return new llama_model_olmoe(params);
case LLM_ARCH_MUSE_GLIMMER:
return new llama_model_muse_glimmer(params);
case LLM_ARCH_OPENELM:
return new llama_model_openelm(params);
case LLM_ARCH_GPTNEOX:
@ -232,6 +240,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
return new llama_model_granite(params);
case LLM_ARCH_GRANITE_MOE:
return new llama_model_granite_moe(params);
case LLM_ARCH_GRANITE_SWITCH:
return new llama_model_granite_switch(params);
case LLM_ARCH_MINICPM:
return new llama_model_minicpm(params);
case LLM_ARCH_GRANITE_HYBRID:
@ -1112,6 +1122,9 @@ void llama_model_base::load_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_CONVNEXT_EMBEDDING_LENGTH, hparams.convnext.n_embd);
ml.get_key(LLM_KV_CONVNEXT_BLOCK_COUNT, hparams.convnext.n_layer);
GGML_ASSERT(hparams.posnet.n_layer <= hparams.n_layer_all);
GGML_ASSERT(hparams.convnext.n_layer <= hparams.n_layer_all);
}
GGML_ASSERT(hparams.n_expert <= LLAMA_MAX_EXPERTS);
@ -1263,8 +1276,23 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
this->ml = &ml; // to be used by create_tensor() and load_arch_tensors()
if (ml.use_mmap && params.load_mode == LLAMA_LOAD_MODE_AUTO) {
for (const auto & dev : devices) {
ggml_backend_dev_props props;
ggml_backend_dev_get_props(dev.dev, &props);
if (!props.caps.mmap_support) {
ml.use_mmap = false;
break;
}
}
}
const char * load_mode_name = params.load_mode == LLAMA_LOAD_MODE_AUTO
? llama_load_mode_name(ml.use_mmap ? LLAMA_LOAD_MODE_MMAP : LLAMA_LOAD_MODE_NONE)
: llama_load_mode_name(params.load_mode);
LLAMA_LOG_INFO("%s: loading model tensors, this can take a while... (load_mode = %s)\n",
__func__, llama_load_mode_name(params.load_mode));
__func__, load_mode_name);
// build a list of buffer types for the CPU and GPU devices
pimpl->cpu_buft_list = make_cpu_buft_list(devices, params.use_extra_bufts, params.no_host);
@ -1910,6 +1938,7 @@ void llama_model::print_info() const {
arch == LLM_ARCH_GRANITE ||
arch == LLM_ARCH_GRANITE_MOE ||
arch == LLM_ARCH_GRANITE_HYBRID ||
arch == LLM_ARCH_GRANITE_SWITCH ||
arch == LLM_ARCH_NEMOTRON_H_MOE) {
LLAMA_LOG_INFO("%s: f_embedding_scale = %f\n", __func__, hparams.f_embedding_scale);
LLAMA_LOG_INFO("%s: f_residual_scale = %f\n", __func__, hparams.f_residual_scale);
@ -2226,6 +2255,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
params.ctx_type == LLAMA_CONTEXT_TYPE_MTP &&
(arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE);
const bool mtp_on_hybrid_nemotron =
params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && arch == LLM_ARCH_NEMOTRON_H_MOE;
if (llm_arch_is_recurrent(arch)) {
res = new llama_memory_recurrent(
*this,
@ -2236,7 +2268,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
cparams.n_seq_max,
cparams.n_rs_seq,
nullptr);
} else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen) {
} else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen && !mtp_on_hybrid_nemotron) {
// The main difference between hybrid architectures is the
// layer filters, so pick the right one here
llama_memory_hybrid::layer_filter_cb filter_attn = nullptr;
@ -2317,7 +2349,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
};
}
if (mtp_on_hybrid_qwen) {
if (mtp_on_hybrid_qwen || mtp_on_hybrid_nemotron) {
filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
}
@ -2440,7 +2472,7 @@ llama_model_params llama_model_default_params() {
/*.tensor_buft_overrides =*/ nullptr,
/*.n_gpu_layers =*/ -1,
/*.split_mode =*/ LLAMA_SPLIT_MODE_LAYER,
/*.load_mode =*/ LLAMA_LOAD_MODE_MMAP,
/*.load_mode =*/ LLAMA_LOAD_MODE_AUTO,
/*.main_gpu =*/ 0,
/*.tensor_split =*/ nullptr,
/*.progress_callback =*/ nullptr,
@ -2589,11 +2621,13 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
case LLM_ARCH_DEEPSEEK2OCR:
case LLM_ARCH_DEEPSEEK32:
case LLM_ARCH_DEEPSEEK4:
case LLM_ARCH_MUSE_GLIMMER:
case LLM_ARCH_PLM:
case LLM_ARCH_CHATGLM:
case LLM_ARCH_GRANITE:
case LLM_ARCH_GRANITE_MOE:
case LLM_ARCH_GRANITE_HYBRID:
case LLM_ARCH_GRANITE_SWITCH:
case LLM_ARCH_CHAMELEON:
case LLM_ARCH_BAILINGMOE:
case LLM_ARCH_NEO_BERT:
@ -2608,6 +2642,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
case LLM_ARCH_MAINCODER:
case LLM_ARCH_GLM_DSA:
case LLM_ARCH_NANBEIGE:
case LLM_ARCH_POCKETTTS:
return LLAMA_ROPE_TYPE_NORM;
// the pairs of head values are offset by n_rot/2
@ -2693,6 +2728,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
case LLM_ARCH_QWEN3VLMOE:
case LLM_ARCH_QWEN35:
case LLM_ARCH_QWEN35MOE:
case LLM_ARCH_QWEN3TTS:
return LLAMA_ROPE_TYPE_IMROPE;
case LLM_ARCH_GLM4:
@ -2887,6 +2923,21 @@ void llama_model_base::create_tensor_qkv(llama_layer & layer, int bid,
int64_t n_embd_, int64_t n_embd_q_, int64_t n_embd_k_, int64_t n_embd_v_,
int flags) {
const int64_t n_embd_qkv = n_embd_q_ + n_embd_k_ + n_embd_v_;
if (flags & TENSOR_SKIP) {
const int skip = TENSOR_NOT_REQUIRED | TENSOR_SKIP;
create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", bid), {n_embd_, n_embd_qkv}, skip | TENSOR_SKIP_IF_VIRTUAL);
create_tensor(tn(LLM_TENSOR_ATTN_QKV, "bias", bid), {n_embd_qkv}, skip | TENSOR_SKIP_IF_VIRTUAL);
create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", bid), {n_embd_, n_embd_q_}, skip);
create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", bid), {n_embd_, n_embd_k_}, skip);
create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", bid), {n_embd_, n_embd_v_}, skip);
create_tensor(tn(LLM_TENSOR_ATTN_Q, "bias", bid), {n_embd_q_}, skip);
create_tensor(tn(LLM_TENSOR_ATTN_K, "bias", bid), {n_embd_k_}, skip);
create_tensor(tn(LLM_TENSOR_ATTN_V, "bias", bid), {n_embd_v_}, skip);
return;
}
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", bid), {n_embd_, n_embd_qkv}, TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL);
if (layer.wqkv) {
layer.wqkv_b = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "bias", bid), {n_embd_qkv}, TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL);
@ -2908,3 +2959,38 @@ const int32_t * llama_model_target_layer_ids(const struct llama_model * model) {
uint32_t llama_model_target_layer_ids_n(const struct llama_model * model) {
return (uint32_t) model->target_layer_ids.size();
}
uint32_t llama_model_get_tok_embd(const struct llama_model * model, float * out) {
if (model->vocab.n_tokens() == 0 || model->tok_embd == nullptr) {
return 0;
}
const ggml_tensor * tensor = model->tok_embd;
const size_t nelements = ggml_nelements(tensor);
GGML_ASSERT(nelements <= UINT32_MAX); // for the return type
if (out == nullptr) {
return (uint32_t) nelements;
}
if (tensor->type == GGML_TYPE_F32) {
ggml_backend_tensor_get(tensor, out, 0, nelements * sizeof(float));
return (uint32_t) nelements;
}
std::vector<uint8_t> buf(ggml_nbytes(tensor));
ggml_backend_tensor_get(tensor, buf.data(), 0, buf.size());
const ggml_type_traits * traits = ggml_get_type_traits(tensor->type);
if (tensor->type == GGML_TYPE_F16) {
ggml_fp16_to_fp32_row((const ggml_fp16_t *) buf.data(), out, nelements);
} else if (tensor->type == GGML_TYPE_BF16) {
ggml_bf16_to_fp32_row((const ggml_bf16_t *) buf.data(), out, nelements);
} else if (ggml_is_quantized(tensor->type) && traits->to_float != nullptr) {
traits->to_float(buf.data(), out, nelements);
} else {
GGML_ABORT("unsupported tensor type for dequantization: %s", ggml_type_name(tensor->type));
}
return (uint32_t) nelements;
}

25
examples/talk-llama/llama-model.h
View File

@ -223,6 +223,24 @@ struct llama_layer_nextn {
struct ggml_tensor * shared_head_norm = nullptr;
};
struct llama_layer_switch_lora {
struct ggml_tensor * a_q = nullptr;
struct ggml_tensor * b_q = nullptr;
struct ggml_tensor * a_k = nullptr;
struct ggml_tensor * b_k = nullptr;
struct ggml_tensor * a_v = nullptr;
struct ggml_tensor * b_v = nullptr;
struct ggml_tensor * a_o = nullptr;
struct ggml_tensor * b_o = nullptr;
struct ggml_tensor * a_gate = nullptr;
struct ggml_tensor * b_gate = nullptr;
struct ggml_tensor * a_up = nullptr;
struct ggml_tensor * b_up = nullptr;
struct ggml_tensor * a_down = nullptr;
struct ggml_tensor * b_down = nullptr;
};
struct llama_layer {
// normalization
struct ggml_tensor * attn_norm = nullptr;
@ -533,6 +551,8 @@ struct llama_layer {
struct llama_layer_shortconv shortconv;
struct llama_layer_nextn nextn;
struct llama_layer_switch_lora switch_lora;
};
struct llama_device {
@ -603,8 +623,9 @@ struct llama_model {
struct ggml_tensor * per_layer_model_proj = nullptr;
struct ggml_tensor * per_layer_proj_norm = nullptr;
// eagle3
struct ggml_tensor * fc = nullptr; // feature fusion layer
// eagle3 / dflash feature fusion layer
struct ggml_tensor * fc = nullptr;
struct ggml_tensor * fc_s = nullptr;
struct ggml_tensor * d2t = nullptr; // draft to target vocabulary mapping
// dspark

489
examples/talk-llama/llama-sampler.cpp
View File

@ -467,9 +467,11 @@ static void llama_sampler_empty_free(struct llama_sampler * smpl) {
static bool llama_sampler_empty_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
GGML_UNUSED(smpl);
GGML_UNUSED(buft);
GGML_UNUSED(n_outputs_max_per_seq);
return true;
}
@ -511,6 +513,8 @@ static struct llama_sampler_i llama_sampler_empty_i = {
/* .backend_accept = */ llama_sampler_empty_backend_accept,
/* .backend_apply = */ llama_sampler_empty_backend_apply,
/* .backend_set_input = */ llama_sampler_empty_backend_set_input,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_empty(const char * name) {
@ -551,6 +555,12 @@ struct llama_sampler_backend {
this->support = support;
}
// copy the state that is not tied to the current sampling graph
// samplers that hold only immutable configuration can use this as is
void copy_state(const llama_sampler_backend & src) {
GGML_UNUSED(src);
}
private:
std::string name;
std::string name_ext;
@ -559,6 +569,71 @@ private:
bool support;
};
// .copy_state for samplers deriving from llama_sampler_backend
template<typename T>
static void llama_sampler_backend_copy_state(const struct llama_sampler * src, struct llama_sampler * dst) {
((T *) dst->ctx)->copy_state(*(const T *) src->ctx);
}
struct llama_sampler_backend_probe {
ggml_context_ptr ctx;
ggml_cgraph * gf;
};
static llama_sampler_backend_probe llama_sampler_backend_probe_graph(
llama_sampler * sampler,
int64_t n_candidates,
uint32_t max_nodes,
bool with_candidates) {
ggml_init_params params = {
/*.mem_size =*/ max_nodes * ggml_tensor_overhead() + ggml_graph_overhead_custom(max_nodes, false),
/*.mem_buffer =*/ nullptr,
/*.no_alloc =*/ true,
};
ggml_context_ptr ctx_ptr { ggml_init(params) };
if (!ctx_ptr) {
throw std::runtime_error(format("failed to create ggml context"));
}
auto * ctx = ctx_ptr.get();
auto * gf = ggml_new_graph_custom(ctx, max_nodes, false);
llama_sampler_data data = {
/*.logits =*/ ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_candidates),
/*.probs =*/ nullptr,
/*.sampled =*/ nullptr,
/*.candidates =*/ with_candidates ? ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_candidates) : nullptr,
};
if (sampler->iface->backend_reset) {
sampler->iface->backend_reset(sampler);
}
sampler->iface->backend_apply(sampler, ctx, gf, &data);
for (auto * output : { data.logits, data.probs, data.sampled, data.candidates }) {
if (output) {
ggml_build_forward_expand(gf, output);
}
}
if (sampler->iface->backend_reset) {
sampler->iface->backend_reset(sampler);
}
return { std::move(ctx_ptr), gf };
}
static uint32_t llama_sampler_backend_probe_n_nodes(const llama_sampler_backend_probe & probe) {
uint32_t n_tensors = 0;
for (auto * tensor = ggml_get_first_tensor(probe.ctx.get()); tensor;
tensor = ggml_get_next_tensor(probe.ctx.get(), tensor)) {
++n_tensors;
}
return std::max<uint32_t>(ggml_graph_n_nodes(probe.gf), n_tensors);
}
// check if all ggml ops used by the sampler are supported by the backend
static bool llama_sampler_backend_support(
llama_sampler * smpl,
@ -569,50 +644,10 @@ static bool llama_sampler_backend_support(
return true;
}
ggml_init_params params = {
/*.mem_size =*/ 128*ggml_tensor_overhead() + ggml_graph_overhead(),
/*.mem_buffer =*/ NULL,
/*.no_alloc =*/ true,
};
auto probe = llama_sampler_backend_probe_graph(smpl, 1024*1024, GGML_DEFAULT_GRAPH_SIZE, true);
ggml_context_ptr ctx_ptr { ggml_init(params) };
if (!ctx_ptr) {
throw std::runtime_error(format("failed to create ggml context"));
}
ggml_context * ctx = ctx_ptr.get();
const int64_t n = 1024*1024;
llama_sampler_data data = {
/*.logits = */ ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n),
/*.probs = */ nullptr,
/*.sampled = */ nullptr,
/*.candidates = */ ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n),
};
ggml_cgraph * gf = ggml_new_graph(ctx);
smpl->iface->backend_apply(smpl, ctx, gf, &data);
if (data.logits) {
ggml_build_forward_expand(gf, data.logits);
}
if (data.probs) {
ggml_build_forward_expand(gf, data.probs);
}
if (data.sampled) {
ggml_build_forward_expand(gf, data.sampled);
}
if (data.candidates) {
ggml_build_forward_expand(gf, data.candidates);
}
for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
struct ggml_tensor * op = ggml_graph_node(gf, i);
for (int i = 0; i < ggml_graph_n_nodes(probe.gf); i++) {
struct ggml_tensor * op = ggml_graph_node(probe.gf, i);
if (!ggml_backend_dev_supports_op(device, op)) {
LLAMA_LOG_WARN("%s: device '%s' does not have support for op %s needed for sampler '%s'\n",
@ -697,7 +732,8 @@ static void llama_sampler_chain_free(struct llama_sampler * smpl) {
static bool llama_sampler_chain_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * chain = (llama_sampler_chain *) smpl->ctx;
GGML_ASSERT(chain->is_init == false && "llama_sampler_chain_backend_init() called twice");
@ -705,26 +741,32 @@ static bool llama_sampler_chain_backend_init(
chain->is_init = true;
bool res = true;
bool backend_prefix = true;
for (auto & smpl : chain->samplers) {
bool res_cur = true;
bool cur_prefix = backend_prefix;
// to be able to run a sampler on the backend, it has to:
// - have the .backend_init() API implemented
// - return true during .backend_init()
if (smpl.ptr->iface->backend_init) {
if (!smpl.ptr->iface->backend_init(smpl.ptr, buft)) {
res_cur = false;
// - support the requested per-sequence output limit
if (cur_prefix && smpl.ptr->iface->backend_init) {
if (!smpl.ptr->iface->backend_init(smpl.ptr, buft, n_outputs_max_per_seq)) {
cur_prefix = false;
}
} else {
res_cur = false;
cur_prefix = false;
}
smpl.is_backend = res_cur;
smpl.is_backend = cur_prefix;
backend_prefix = cur_prefix;
res = res && res_cur;
res = res && cur_prefix;
}
auto probe = llama_sampler_backend_probe_graph(smpl, 1024*1024, GGML_DEFAULT_GRAPH_SIZE, false);
chain->n_nodes = llama_sampler_backend_probe_n_nodes(probe);
return res;
}
@ -780,6 +822,36 @@ static void llama_sampler_chain_backend_set_input(struct llama_sampler * smpl) {
}
}
static void llama_sampler_chain_backend_reset(struct llama_sampler * smpl) {
auto * chain = (llama_sampler_chain *) smpl->ctx;
for (auto & entry : chain->samplers) {
if (!entry.is_backend) {
break;
}
if (entry.ptr->iface->backend_reset) {
entry.ptr->iface->backend_reset(entry.ptr);
}
}
}
static void llama_sampler_chain_copy_state(const struct llama_sampler * src, struct llama_sampler * dst) {
const auto * src_chain = (const llama_sampler_chain *) src->ctx;
auto * dst_chain = (llama_sampler_chain *) dst->ctx;
GGML_ASSERT(src_chain->samplers.size() == dst_chain->samplers.size());
for (size_t i = 0; i < src_chain->samplers.size(); ++i) {
llama_sampler_copy(src_chain->samplers[i].ptr, dst_chain->samplers[i].ptr);
}
// note: is_init, n_nodes and is_backend belong to the current sampling graph
dst_chain->params = src_chain->params;
dst_chain->cur = src_chain->cur;
dst_chain->t_sample_us = src_chain->t_sample_us;
dst_chain->n_sample = src_chain->n_sample;
}
static struct llama_sampler_i llama_sampler_chain_i = {
/* .name = */ llama_sampler_chain_name,
/* .accept = */ llama_sampler_chain_accept,
@ -791,22 +863,35 @@ static struct llama_sampler_i llama_sampler_chain_i = {
/* .backend_accept = */ llama_sampler_chain_backend_accept,
/* .backend_apply = */ llama_sampler_chain_backend_apply,
/* .backend_set_input = */ llama_sampler_chain_backend_set_input,
/* .backend_reset = */ llama_sampler_chain_backend_reset,
/* .copy_state = */ llama_sampler_chain_copy_state,
};
struct llama_sampler * llama_sampler_chain_init(struct llama_sampler_chain_params params) {
return llama_sampler_init(
/* .iface = */ &llama_sampler_chain_i,
/* .ctx = */ new llama_sampler_chain {
/* .params = */ params,
/* .is_init = */ false,
/* .samplers = */ {},
/* .cur = */ {},
/* .t_sample_us = */ 0,
/* .n_sample = */ 0,
/* .params = */ params,
/* .is_init = */ false,
/* .n_nodes = */ 0,
/* .samplers = */ {},
/* .cur = */ {},
/* .t_sample_us = */ 0,
/* .n_sample = */ 0,
}
);
}
uint32_t llama_sampler_backend_n_nodes(const llama_sampler * sampler) {
GGML_ASSERT(sampler != nullptr);
GGML_ASSERT(sampler->iface == &llama_sampler_chain_i);
const auto * chain = (const llama_sampler_chain *) sampler->ctx;
GGML_ASSERT(chain->is_init);
return chain->n_nodes;
}
llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_context * ctx, int32_t idx) {
const llama_token sampled_token = llama_get_sampled_token_ith (ctx, idx);
const float * sampled_probs = llama_get_sampled_probs_ith (ctx, idx);
@ -816,6 +901,7 @@ llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_conte
// If a backend sampler has already sampled a token, return it.
if (sampled_token != LLAMA_TOKEN_NULL) {
LLAMA_LOG_DEBUG("%s: Backend sampler selected token for idx %d. Skipping CPU samplers\n", __func__, idx);
llama_sampler_accept(smpl, sampled_token);
return sampled_token;
}
@ -975,8 +1061,10 @@ static void llama_sampler_greedy_apply(struct llama_sampler * /*smpl*/, llama_to
static bool llama_sampler_greedy_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * sctx = (llama_sampler_greedy *) smpl->ctx;
GGML_UNUSED(n_outputs_max_per_seq);
const bool res = llama_sampler_backend_support(smpl, buft);
@ -1012,6 +1100,8 @@ static struct llama_sampler_i llama_sampler_greedy_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_greedy_backend_apply,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_greedy>,
};
struct llama_sampler * llama_sampler_init_greedy() {
@ -1031,7 +1121,25 @@ struct llama_sampler_dist : public llama_sampler_backend {
std::mt19937 rng;
ggml_tensor * inp_uniform;
// TODO: refactor + fix naming
// https://github.com/ggml-org/llama.cpp/pull/25532/changes#r3749906719
// use a temporary RNG for multi-output sampling so rejected tokens do not advance rng
bool backend_transactional;
std::mt19937 rng_backend;
size_t n_backend_draws_generated;
size_t n_backend_draws_committed;
// inputs for the current sampling graph
std::vector<ggml_tensor *> inp_uniforms;
void copy_state(const llama_sampler_dist & src) {
// note: inp_uniforms and backend_transactional belong to the current sampling graph
seed_cur = src.seed_cur;
rng = src.rng;
rng_backend = src.rng_backend;
n_backend_draws_generated = src.n_backend_draws_generated;
n_backend_draws_committed = src.n_backend_draws_committed;
}
};
static const char * llama_sampler_dist_name(const struct llama_sampler * smpl) {
@ -1050,7 +1158,11 @@ static void llama_sampler_dist_apply(struct llama_sampler * smpl, llama_token_da
cur_p->selected = 0;
std::uniform_real_distribution<double> dist(0.0f, 1.0f);
if (cur_p->size == 1) {
// keep the RNG state aligned with backend sampling, which draws once per output
dist(ctx->rng);
cur_p->data[0].p = 1.0f;
return;
}
@ -1075,7 +1187,6 @@ static void llama_sampler_dist_apply(struct llama_sampler * smpl, llama_token_da
// sample from the obtained probabilities and normalize the probs in a single pass
// this is ~3x faster on Mac with full gpt-oss vocab than the version below
//
std::uniform_real_distribution<double> dist(0.0f, 1.0f);
const double rnd = dist(ctx->rng);
double sum_run = 0.0f;
@ -1115,6 +1226,9 @@ static void llama_sampler_dist_reset(struct llama_sampler * smpl) {
auto * ctx = (llama_sampler_dist *) smpl->ctx;
ctx->seed_cur = get_rng_seed(ctx->seed);
ctx->rng.seed(ctx->seed_cur);
ctx->rng_backend = ctx->rng;
ctx->n_backend_draws_generated = 0;
ctx->n_backend_draws_committed = 0;
}
static struct llama_sampler * llama_sampler_dist_clone(const struct llama_sampler * smpl) {
@ -1125,7 +1239,12 @@ static struct llama_sampler * llama_sampler_dist_clone(const struct llama_sample
{
auto * result_ctx = (llama_sampler_dist *) result->ctx;
result_ctx->rng = ctx->rng;
result_ctx->seed_cur = ctx->seed_cur;
result_ctx->rng = ctx->rng;
result_ctx->backend_transactional = ctx->backend_transactional;
result_ctx->rng_backend = ctx->rng_backend;
result_ctx->n_backend_draws_generated = ctx->n_backend_draws_generated;
result_ctx->n_backend_draws_committed = ctx->n_backend_draws_committed;
}
return result;
@ -1137,12 +1256,17 @@ static void llama_sampler_dist_free(struct llama_sampler * smpl) {
static bool llama_sampler_dist_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * sctx = (llama_sampler_dist *) smpl->ctx;
const bool res = llama_sampler_backend_support(smpl, buft);
sctx->init(res);
sctx->backend_transactional = n_outputs_max_per_seq > 1;
sctx->rng_backend = sctx->rng;
sctx->n_backend_draws_generated = 0;
sctx->n_backend_draws_committed = 0;
return res;
}
@ -1156,9 +1280,10 @@ static void llama_sampler_dist_backend_apply(
auto * sctx = (llama_sampler_dist *) smpl->ctx;
sctx->inp_uniform = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
ggml_set_name (sctx->inp_uniform, "uniform");
ggml_set_input(sctx->inp_uniform);
ggml_tensor * inp_uniform = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
ggml_format_name(inp_uniform, "uniform_%zu", sctx->inp_uniforms.size());
ggml_set_input(inp_uniform);
sctx->inp_uniforms.push_back(inp_uniform);
// flatten
struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits));
@ -1174,7 +1299,7 @@ static void llama_sampler_dist_backend_apply(
// Recall that each entry in cumsum is the cumulative probability up to that
// index so values stay negative while the cumulative total is below the
// random value, and become zero/positive once the threshold is crossed.
struct ggml_tensor * diff = ggml_sub(ctx, cumsum, sctx->inp_uniform);
struct ggml_tensor * diff = ggml_sub(ctx, cumsum, inp_uniform);
ggml_set_name(diff, "dist_cumsum");
// The ggml_step function produces a tensor where entries are 1 if the
@ -1189,6 +1314,9 @@ static void llama_sampler_dist_backend_apply(
struct ggml_tensor * idxf = ggml_sum(ctx, mask);
ggml_set_name(idxf, "dist_index_f32");
// Clamp to prevent out-of-bounds access when computing the index.
idxf = ggml_clamp(ctx, idxf, 1.0f, mask->ne[0]);
// Use ggml_scale_bias to scale the index value by -1 and then add the size
// of the mask to that value so we get the correct index ((-1 * idxf) + n).
struct ggml_tensor * idx = ggml_cast(ctx, ggml_scale_bias(ctx, idxf, -1.0f, mask->ne[0]), GGML_TYPE_I32);
@ -1210,22 +1338,52 @@ static void llama_sampler_dist_backend_apply(
static void llama_sampler_dist_backend_set_input(struct llama_sampler * smpl) {
auto * sctx = (llama_sampler_dist *) smpl->ctx;
GGML_ASSERT(sctx->inp_uniform != nullptr);
GGML_ASSERT(!sctx->inp_uniforms.empty());
// We sample in double precision and cast to float to match rnd numbers of
// llama_dampler_dist which uses double precision (sampling from
// llama_sampler_dist which uses double precision (sampling from
// std::uniform_real_distribution<double> and
// std::uniform_real_distribution<float> with same rng will produce
// different sequences).
std::uniform_real_distribution<double> dist(0.0f, 1.0f);
const float rnd = dist(sctx->rng);
ggml_backend_tensor_set(sctx->inp_uniform, &rnd, 0, sizeof(float));
auto & rng = sctx->backend_transactional ? sctx->rng_backend : sctx->rng;
for (auto * inp_uniform : sctx->inp_uniforms) {
GGML_ASSERT(inp_uniform != nullptr);
const float rnd = dist(rng);
ggml_backend_tensor_set(inp_uniform, &rnd, 0, sizeof(float));
if (sctx->backend_transactional) {
++sctx->n_backend_draws_generated;
}
}
}
static void llama_sampler_dist_backend_reset(struct llama_sampler * smpl) {
auto * sctx = (llama_sampler_dist *) smpl->ctx;
sctx->inp_uniforms.clear();
}
static void llama_sampler_dist_accept(struct llama_sampler * smpl, llama_token token) {
GGML_UNUSED(token);
auto * sctx = (llama_sampler_dist *) smpl->ctx;
if (!sctx->backend_transactional ||
sctx->n_backend_draws_committed >= sctx->n_backend_draws_generated) {
return;
}
std::uniform_real_distribution<double> dist(0.0f, 1.0f);
dist(sctx->rng);
++sctx->n_backend_draws_committed;
}
static struct llama_sampler_i llama_sampler_dist_i = {
/* .name = */ llama_sampler_dist_name,
/* .accept = */ nullptr,
/* .accept = */ llama_sampler_dist_accept,
/* .apply = */ llama_sampler_dist_apply,
/* .reset = */ llama_sampler_dist_reset,
/* .clone = */ llama_sampler_dist_clone,
@ -1234,6 +1392,8 @@ static struct llama_sampler_i llama_sampler_dist_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_dist_backend_apply,
/* .backend_set_input = */ llama_sampler_dist_backend_set_input,
/* .backend_reset = */ llama_sampler_dist_backend_reset,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_dist>,
};
struct llama_sampler * llama_sampler_init_dist(uint32_t seed) {
@ -1242,14 +1402,39 @@ struct llama_sampler * llama_sampler_init_dist(uint32_t seed) {
/* .iface = */ &llama_sampler_dist_i,
/* .ctx = */ new llama_sampler_dist {
("dist"),
/* .seed = */ seed,
/* .seed_cur = */ seed_cur,
/* .rng = */ std::mt19937(seed_cur),
/* .inp_uniform = */ nullptr,
/* .seed = */ seed,
/* .seed_cur = */ seed_cur,
/* .rng = */ std::mt19937(seed_cur),
/* .backend_transactional = */ false,
/* .rng_backend = */ std::mt19937(seed_cur),
/* .n_backend_draws_generated = */ 0,
/* .n_backend_draws_committed = */ 0,
/* .inp_uniforms = */ {},
}
);
}
void llama_sampler_backend_begin(llama_sampler * sampler) {
GGML_ASSERT(sampler != nullptr);
if (sampler->iface == &llama_sampler_chain_i) {
auto * chain = (llama_sampler_chain *) sampler->ctx;
for (auto & entry : chain->samplers) {
if (!entry.is_backend) {
break;
}
llama_sampler_backend_begin(entry.ptr);
}
} else if (sampler->iface == &llama_sampler_dist_i) {
auto * ctx = (llama_sampler_dist *) sampler->ctx;
if (ctx->backend_transactional) {
ctx->rng_backend = ctx->rng;
ctx->n_backend_draws_generated = 0;
ctx->n_backend_draws_committed = 0;
}
}
}
// top-k
struct llama_sampler_top_k : public llama_sampler_backend {
@ -1277,8 +1462,10 @@ static void llama_sampler_top_k_free(struct llama_sampler * smpl) {
static bool llama_sampler_top_k_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * sctx = (llama_sampler_top_k *) smpl->ctx;
GGML_UNUSED(n_outputs_max_per_seq);
const bool res = llama_sampler_backend_support(smpl, buft);
@ -1325,6 +1512,8 @@ static struct llama_sampler_i llama_sampler_top_k_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_top_k_backend_apply,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_top_k>,
};
struct llama_sampler * llama_sampler_init_top_k(int32_t k) {
@ -1423,8 +1612,10 @@ static void llama_sampler_top_p_free(struct llama_sampler * smpl) {
static bool llama_sampler_top_p_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * sctx = (llama_sampler_top_p *) smpl->ctx;
GGML_UNUSED(n_outputs_max_per_seq);
const bool res = llama_sampler_backend_support(smpl, buft);
@ -1521,6 +1712,8 @@ static struct llama_sampler_i llama_sampler_top_p_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_top_p_backend_apply,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_top_p>,
};
struct llama_sampler * llama_sampler_init_top_p(float p, size_t min_keep) {
@ -1618,8 +1811,10 @@ static void llama_sampler_min_p_free(struct llama_sampler * smpl) {
static bool llama_sampler_min_p_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * sctx = (llama_sampler_min_p *) smpl->ctx;
GGML_UNUSED(n_outputs_max_per_seq);
const bool res = llama_sampler_backend_support(smpl, buft);
@ -1680,6 +1875,8 @@ static struct llama_sampler_i llama_sampler_min_p_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_min_p_backend_apply,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_min_p>,
};
struct llama_sampler * llama_sampler_init_min_p(float p, size_t min_keep) {
@ -1790,6 +1987,8 @@ static struct llama_sampler_i llama_sampler_typical_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_typical(float p, size_t min_keep) {
@ -1866,8 +2065,10 @@ static void llama_sampler_backend_temp_sampling(
static bool llama_sampler_temp_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * sctx = (llama_sampler_temp *) smpl->ctx;
GGML_UNUSED(n_outputs_max_per_seq);
const bool res = llama_sampler_backend_support(smpl, buft);
@ -1896,6 +2097,8 @@ static struct llama_sampler_i llama_sampler_temp_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_temp_backend_apply,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_temp>,
};
struct llama_sampler * llama_sampler_init_temp(float temp) {
@ -2009,8 +2212,10 @@ static void llama_sampler_temp_ext_free(struct llama_sampler * smpl) {
static bool llama_sampler_temp_ext_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * sctx = (llama_sampler_temp_ext *) smpl->ctx;
GGML_UNUSED(n_outputs_max_per_seq);
const bool res = llama_sampler_backend_support(smpl, buft);
@ -2095,6 +2300,8 @@ static struct llama_sampler_i llama_sampler_temp_ext_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_temp_ext_backend_apply,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_temp_ext>,
};
struct llama_sampler * llama_sampler_init_temp_ext(float temp, float delta, float exponent) {
@ -2202,6 +2409,8 @@ static struct llama_sampler_i llama_sampler_xtc_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_xtc(float p, float t, size_t min_keep, uint32_t seed) {
@ -2290,7 +2499,7 @@ static struct llama_sampler * llama_sampler_mirostat_clone(const struct llama_sa
// copy the state
{
auto * result_ctx = (llama_sampler_mirostat *) smpl->ctx;
auto * result_ctx = (llama_sampler_mirostat *) result->ctx;
result_ctx->mu = ctx->mu;
result_ctx->rng = ctx->rng;
@ -2321,6 +2530,8 @@ static struct llama_sampler_i llama_sampler_mirostat_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_mirostat(int32_t n_vocab, uint32_t seed, float tau, float eta, int32_t m) {
@ -2425,6 +2636,8 @@ static struct llama_sampler_i llama_sampler_mirostat_v2_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_mirostat_v2(uint32_t seed, float tau, float eta) {
@ -2546,6 +2759,8 @@ static struct llama_sampler_i llama_sampler_grammar_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
static struct llama_sampler * llama_sampler_init_grammar_impl(
@ -2661,6 +2876,12 @@ struct llama_sampler_penalties : public llama_sampler_backend {
std::vector<int32_t> host_token_ids;
std::vector<int32_t> host_counts;
void copy_state(const llama_sampler_penalties & src) {
// note: inp_token_ids/inp_counts belong to the current sampling graph
prev = src.prev;
token_count = src.token_count;
}
static bool is_disabled(
int32_t penalty_last_n,
float penalty_repeat,
@ -2790,9 +3011,15 @@ static void llama_sampler_penalties_free(struct llama_sampler * smpl) {
static bool llama_sampler_penalties_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
auto * sctx = (llama_sampler_penalties *) smpl->ctx;
if (n_outputs_max_per_seq > 1) {
sctx->init(false);
return false;
}
const bool res = llama_sampler_backend_support(smpl, buft);
sctx->init(res);
@ -2952,6 +3179,12 @@ static void llama_sampler_penalties_backend_set_input(struct llama_sampler * smp
ggml_backend_tensor_set(sctx->inp_counts, sctx->host_counts.data(), 0, sctx->n_max * sizeof(int32_t));
}
static void llama_sampler_penalties_backend_reset(struct llama_sampler * smpl) {
auto * sctx = (llama_sampler_penalties *) smpl->ctx;
sctx->inp_token_ids = nullptr;
sctx->inp_counts = nullptr;
}
static struct llama_sampler_i llama_sampler_penalties_i = {
/* .name = */ llama_sampler_penalties_name,
/* .accept = */ llama_sampler_penalties_accept,
@ -2963,6 +3196,8 @@ static struct llama_sampler_i llama_sampler_penalties_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_penalties_backend_apply,
/* .backend_set_input = */ llama_sampler_penalties_backend_set_input,
/* .backend_reset = */ llama_sampler_penalties_backend_reset,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_penalties>,
};
struct llama_sampler * llama_sampler_init_penalties(
@ -3058,6 +3293,8 @@ static struct llama_sampler_i llama_sampler_top_n_sigma_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_top_n_sigma(float n) {
@ -3078,8 +3315,6 @@ struct llama_sampler * llama_sampler_init_top_n_sigma(float n) {
// DRY
struct llama_sampler_dry {
int32_t total_context_size;
const float dry_multiplier;
const float dry_base;
const int32_t dry_allowed_length;
@ -3155,8 +3390,7 @@ static void llama_sampler_dry_apply(struct llama_sampler * smpl, llama_token_dat
return;
}
int32_t effective_dry_penalty_last_n = (ctx->dry_penalty_last_n == -1) ? ctx->total_context_size : std::max(ctx->dry_penalty_last_n, 0);
int last_n_repeat = std::min(std::min((int)ctx->last_tokens.size(), effective_dry_penalty_last_n), ctx->total_context_size);
int last_n_repeat = std::min((int) ctx->last_tokens.size(), ctx->dry_penalty_last_n);
if (last_n_repeat <= ctx->dry_allowed_length) {
return;
@ -3369,7 +3603,7 @@ static struct llama_sampler * llama_sampler_dry_clone(const struct llama_sampler
llama_vocab dummy_vocab;
// dummy vocab is passed because it is only needed for raw sequence breaker processing, which we have already done and will simply be copying
auto * result = llama_sampler_init_dry(&dummy_vocab, ctx->total_context_size, ctx->dry_multiplier, ctx->dry_base, ctx->dry_allowed_length, ctx->dry_penalty_last_n, NULL, 0);
auto * result = llama_sampler_init_dry(&dummy_vocab, ctx->dry_multiplier, ctx->dry_base, ctx->dry_allowed_length, ctx->dry_penalty_last_n, NULL, 0);
// Copy the state, including the processed breakers
{
@ -3398,10 +3632,12 @@ static struct llama_sampler_i llama_sampler_dry_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_dry(const struct llama_vocab * vocab, int32_t n_ctx_train, float dry_multiplier, float dry_base, int32_t dry_allowed_length, int32_t dry_penalty_last_n, const char** seq_breakers, size_t num_breakers) {
int32_t effective_dry_penalty_last_n = (dry_penalty_last_n == -1) ? n_ctx_train : std::max(dry_penalty_last_n, 0);
struct llama_sampler * llama_sampler_init_dry(const struct llama_vocab * vocab, float dry_multiplier, float dry_base, int32_t dry_allowed_length, int32_t dry_penalty_last_n, const char** seq_breakers, size_t num_breakers) {
dry_penalty_last_n = std::max(dry_penalty_last_n, 0);
std::unordered_multimap<llama_token, std::vector<llama_token>> processed_breakers;
const int MAX_CHAR_LEN = 40;
const int MAX_SEQ_LEN = 20;
@ -3438,23 +3674,22 @@ struct llama_sampler * llama_sampler_init_dry(const struct llama_vocab * vocab,
return llama_sampler_init(
/* .iface = */ &llama_sampler_dry_i,
/* .ctx = */ new llama_sampler_dry {
/* .total_context_size = */ n_ctx_train,
/* .dry_multiplier = */ dry_multiplier,
/* .dry_base = */ dry_base,
/* .dry_allowed_length = */ dry_allowed_length,
/* .dry_penalty_last_n = */ dry_penalty_last_n,
/* .dry_processed_breakers = */ std::move(processed_breakers),
/* .dry_repeat_count = */ dry_enabled ? std::vector<int>(effective_dry_penalty_last_n, 0) : std::vector<int>{},
/* .dry_repeat_count = */ dry_enabled ? std::vector<int>(dry_penalty_last_n, 0) : std::vector<int>{},
/* .dry_max_token_repeat = */ {},
/* .last_tokens = */ dry_enabled ? ring_buffer<llama_token>(effective_dry_penalty_last_n) : ring_buffer<llama_token>(0),
/* .last_tokens = */ dry_enabled ? ring_buffer<llama_token>(dry_penalty_last_n) : ring_buffer<llama_token>(0),
}
);
}
// wrapper for test-sampling.cpp
struct llama_sampler * llama_sampler_init_dry_testing(int32_t context_size, float dry_multiplier, float dry_base, int32_t dry_allowed_length, int32_t dry_penalty_last_n, const std::vector<std::vector<llama_token>>& seq_breakers) {
struct llama_sampler * llama_sampler_init_dry_testing(float dry_multiplier, float dry_base, int32_t dry_allowed_length, int32_t dry_penalty_last_n, const std::vector<std::vector<llama_token>>& seq_breakers) {
llama_vocab dummy_vocab;
auto * result = llama_sampler_init_dry(&dummy_vocab, context_size, dry_multiplier, dry_base, dry_allowed_length, dry_penalty_last_n, NULL, 0);
auto * result = llama_sampler_init_dry(&dummy_vocab, dry_multiplier, dry_base, dry_allowed_length, dry_penalty_last_n, NULL, 0);
auto * ctx = (llama_sampler_dry *) result->ctx;
// Process the token-based sequence breakers
@ -3618,6 +3853,8 @@ static struct llama_sampler_i llama_sampler_adaptive_p_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_adaptive_p(
@ -3719,13 +3956,17 @@ static void llama_sampler_logit_bias_backend_apply(
const size_t n = sctx->logit_bias.size();
sctx->inp_logit_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n);
ggml_set_name(sctx->inp_logit_bias, "logit_bias");
ggml_set_input(sctx->inp_logit_bias);
if (sctx->inp_logit_bias == nullptr) {
GGML_ASSERT(sctx->inp_logit_idxs == nullptr);
sctx->inp_logit_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n);
ggml_set_name(sctx->inp_logit_idxs, "logit_idxs");
ggml_set_input(sctx->inp_logit_idxs);
sctx->inp_logit_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n);
ggml_set_name(sctx->inp_logit_bias, "logit_bias");
ggml_set_input(sctx->inp_logit_bias);
sctx->inp_logit_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n);
ggml_set_name(sctx->inp_logit_idxs, "logit_idxs");
ggml_set_input(sctx->inp_logit_idxs);
}
ggml_tensor * cur = ggml_fill(ctx, data->logits, 0.0f);
@ -3760,10 +4001,18 @@ static void llama_sampler_logit_bias_backend_set_input(struct llama_sampler * sm
ggml_backend_tensor_set(sctx->inp_logit_idxs, data_logit_idxs.data(), 0, ggml_nbytes(sctx->inp_logit_idxs));
}
static void llama_sampler_logit_bias_backend_reset(struct llama_sampler * smpl) {
auto * sctx = (llama_sampler_logit_bias *) smpl->ctx;
sctx->inp_logit_bias = nullptr;
sctx->inp_logit_idxs = nullptr;
}
static bool llama_sampler_logit_bias_backend_init(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft) {
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq) {
GGML_UNUSED(buft);
GGML_UNUSED(n_outputs_max_per_seq);
auto * sctx = (llama_sampler_logit_bias *) smpl->ctx;
@ -3787,6 +4036,8 @@ static struct llama_sampler_i llama_sampler_logit_bias_i = {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ llama_sampler_logit_bias_backend_apply,
/* .backend_set_input = */ llama_sampler_logit_bias_backend_set_input,
/* .backend_reset = */ llama_sampler_logit_bias_backend_reset,
/* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_logit_bias>,
};
struct llama_sampler * llama_sampler_init_logit_bias(
@ -4026,10 +4277,12 @@ static struct llama_sampler_i llama_sampler_infill_i = {
/* .reset = */ nullptr,
/* .clone = */ llama_sampler_infill_clone,
/* .free = */ llama_sampler_infill_free,
/* .backend_apply = */ nullptr,
/* .backend_accept = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_init = */ nullptr,
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
/* .backend_reset = */ nullptr,
/* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * vocab) {
@ -4043,6 +4296,32 @@ struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * voca
);
}
void llama_sampler_copy(const struct llama_sampler * src, struct llama_sampler * dst) {
if (!src || !dst || src == dst) {
return;
}
GGML_ASSERT(src->iface == dst->iface && "llama_sampler_copy: cannot copy between different sampler types");
if (dst->iface->copy_state) {
dst->iface->copy_state(src, dst);
return;
}
// build a temporary sampler carrying src's current state
llama_sampler * tmp = llama_sampler_clone(src);
// free dst's old state (frees dst->ctx, including children for a chain)
if (dst->iface->free) {
dst->iface->free(dst);
}
// transplant tmp's state into dst, then destroy the (now empty) temp shell
dst->ctx = tmp->ctx;
tmp->ctx = nullptr;
delete tmp;
}
// utils
uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl) {

6
examples/talk-llama/llama-sampler.h
View File

@ -15,6 +15,8 @@ struct llama_sampler_chain {
// has .backend_init() been called?
bool is_init = false;
uint32_t n_nodes = 0;
struct info {
bool is_backend;
@ -33,8 +35,10 @@ struct llama_sampler_chain {
mutable int32_t n_sample;
};
uint32_t llama_sampler_backend_n_nodes(const llama_sampler * sampler);
void llama_sampler_backend_begin(llama_sampler * sampler);
struct llama_sampler * llama_sampler_init_dry_testing(
int32_t context_size,
float dry_multiplier,
float dry_base,
int32_t dry_allowed_length,

19
examples/talk-llama/llama-vocab.cpp
View File

@ -1373,8 +1373,10 @@ struct llm_tokenizer_plamo2 : llm_tokenizer {
if (vocab.is_byte(token_id)) {
if (entry.text.length() == 6 && entry.text.substr(0, 3) == "<0x" && entry.text.back() == '>') {
std::string hex_str = entry.text.substr(3, 2);
int byte_val = std::stoi(hex_str, nullptr, 16);
bytes_[byte_val] = static_cast<llama_token>(token_id);
if (std::isxdigit(static_cast<unsigned char>(hex_str[0])) && std::isxdigit(static_cast<unsigned char>(hex_str[1]))) {
int byte_val = std::stoi(hex_str, nullptr, 16);
bytes_[byte_val] = static_cast<llama_token>(token_id);
}
}
continue;
}
@ -3625,12 +3627,15 @@ int32_t llama_vocab::impl::token_to_piece(llama_token token, char * buf, int32_t
if (vocab.is_byte(token)) {
// Handle byte tokens like <0xXX>
if (token_text.length() == 6 && token_text.substr(0, 3) == "<0x" && token_text.back() == '>') {
int hex_val = std::stoi(token_text.substr(3, 2), nullptr, 16);
if (length < 1) {
return -1;
std::string hex_str = token_text.substr(3, 2);
if (std::isxdigit(static_cast<unsigned char>(hex_str[0])) && std::isxdigit(static_cast<unsigned char>(hex_str[1]))) {
int hex_val = std::stoi(hex_str, nullptr, 16);
if (length < 1) {
return -1;
}
buf[0] = static_cast<char>(hex_val);
return 1;
}
buf[0] = static_cast<char>(hex_val);
return 1;
}
}

21
examples/talk-llama/llama.cpp
View File

@ -48,6 +48,8 @@ const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_ty
const char * llama_load_mode_name(enum llama_load_mode load_mode) {
switch (load_mode) {
case LLAMA_LOAD_MODE_AUTO:
return "auto";
case LLAMA_LOAD_MODE_NONE:
return "none";
case LLAMA_LOAD_MODE_MMAP:
@ -63,11 +65,12 @@ const char * llama_load_mode_name(enum llama_load_mode load_mode) {
}
enum llama_load_mode llama_load_mode_from_str(const char * str) {
if (std::strcmp(str, "none") == 0) { return LLAMA_LOAD_MODE_NONE; }
if (std::strcmp(str, "mmap") == 0) { return LLAMA_LOAD_MODE_MMAP; }
if (std::strcmp(str, "mlock") == 0) { return LLAMA_LOAD_MODE_MLOCK; }
if (std::strcmp(str, "auto") == 0) { return LLAMA_LOAD_MODE_AUTO; }
if (std::strcmp(str, "none") == 0) { return LLAMA_LOAD_MODE_NONE; }
if (std::strcmp(str, "mmap") == 0) { return LLAMA_LOAD_MODE_MMAP; }
if (std::strcmp(str, "mlock") == 0) { return LLAMA_LOAD_MODE_MLOCK; }
if (std::strcmp(str, "mmap+mlock") == 0) { return LLAMA_LOAD_MODE_MMAP_MLOCK; }
if (std::strcmp(str, "dio") == 0) { return LLAMA_LOAD_MODE_DIRECT_IO; }
if (std::strcmp(str, "dio") == 0) { return LLAMA_LOAD_MODE_DIRECT_IO; }
throw std::invalid_argument(std::string("unknown load mode: ") + str);
}
@ -111,6 +114,10 @@ bool llama_supports_rpc(void) {
return ggml_backend_reg_by_name("RPC") != nullptr;
}
const char * llama_version(void) {
return LLAMA_VERSION;
}
void llama_backend_init(void) {
ggml_time_init();
@ -250,7 +257,11 @@ static bool llama_prepare_model_devices(const llama_model_params & params, llama
}
case GGML_BACKEND_DEVICE_TYPE_IGPU:
if (igpus.empty()) {
// igpus.empty() - workaround for integrated devices seen by multiple backends
// ref: https://github.com/ggml-org/llama.cpp/pull/23897
// ggml_backend_dev_backend_reg - allow devices of the same backend regardless if integrated
// ref: https://github.com/ggml-org/llama.cpp/pull/23897#issuecomment-5264222997
if (igpus.empty() || ggml_backend_dev_backend_reg(dev) == ggml_backend_dev_backend_reg(igpus.back().dev)) {
igpus.push_back({false, dev});
}
break;

55
examples/talk-llama/llama.h
View File

@ -203,11 +203,12 @@ extern "C" {
};
enum llama_load_mode {
LLAMA_LOAD_MODE_NONE = 0, // no special loading mode
LLAMA_LOAD_MODE_MMAP = 1, // memory map the model
LLAMA_LOAD_MODE_MLOCK = 2, // force system to keep model in RAM rather than swapping or compressing
LLAMA_LOAD_MODE_MMAP_MLOCK = 3, // mmap + force system to keep model in RAM rather than swapping or compressing
LLAMA_LOAD_MODE_DIRECT_IO = 4, // use direct I/O if available
LLAMA_LOAD_MODE_AUTO = -1, // auto-detect based on device capabilities
LLAMA_LOAD_MODE_NONE = 0, // no special loading mode
LLAMA_LOAD_MODE_MMAP = 1, // memory map the model
LLAMA_LOAD_MODE_MLOCK = 2, // force system to keep model in RAM rather than swapping or compressing
LLAMA_LOAD_MODE_MMAP_MLOCK = 3, // mmap + force system to keep model in RAM rather than swapping or compressing
LLAMA_LOAD_MODE_DIRECT_IO = 4, // use direct I/O if available
};
LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode);
@ -348,14 +349,15 @@ extern "C" {
// NOTE: changing the default values of parameters marked as [EXPERIMENTAL] may cause crashes or incorrect results in certain configurations
// https://github.com/ggml-org/llama.cpp/pull/7544
struct llama_context_params {
uint32_t n_ctx; // text context, 0 = from model
uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode
uint32_t n_ubatch; // physical maximum batch size
uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models)
uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL]
uint32_t n_outputs_max; // max outputs in a ubatch (0 = n_batch)
int32_t n_threads; // number of threads to use for generation
int32_t n_threads_batch; // number of threads to use for batch processing
uint32_t n_ctx; // text context, 0 = from model
uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode
uint32_t n_ubatch; // physical maximum batch size
uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models)
uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL]
uint32_t n_outputs_max; // max outputs in a ubatch (0 = n_batch)
uint32_t n_outputs_max_per_seq; // max outputs per sequence (0 = n_outputs_max)
int32_t n_threads; // number of threads to use for generation
int32_t n_threads_batch; // number of threads to use for batch processing
enum llama_context_type ctx_type; // set the context type (e.g. MTP)
enum llama_rope_scaling_type rope_scaling_type; // RoPE scaling type, from `enum llama_rope_scaling_type`
@ -455,6 +457,8 @@ extern "C" {
// lora adapter
struct llama_adapter_lora;
LLAMA_API const char * llama_version(void);
// Helpers for getting default parameters
// TODO: update API to start accepting pointers to params structs (https://github.com/ggml-org/llama.cpp/discussions/9172)
LLAMA_API struct llama_model_params llama_model_default_params(void);
@ -881,6 +885,7 @@ extern "C" {
const llama_token * tokens,
size_t n_token_count);
// If tokens_out is NULL, only the token count is reported through n_token_count_out and no state is loaded
LLAMA_API size_t llama_state_seq_load_file(
struct llama_context * ctx,
const char * filepath,
@ -1054,6 +1059,9 @@ extern "C" {
//
// Get the backend sampled token for the ith token.
// With multiple outputs, sampler state advances when the token is accepted,
// not when it is read through this function.
// When accepting multiple outputs, accept a contiguous prefix in output order.
// Returns LLAMA_TOKEN_NULL if no token was sampled.
LLAMA_API llama_token llama_get_sampled_token_ith(struct llama_context * ctx, int32_t i);
@ -1270,9 +1278,12 @@ extern "C" {
// [EXPERIMENTAL]
// backend sampling interface:
// return true if the backend supports all ops needed by the sampler
// return true if the backend supports all ops needed by the sampler and can handle up to n_outputs_max_per_seq outputs per sequence
// note: call once per sampler
bool (*backend_init)(struct llama_sampler * smpl, ggml_backend_buffer_type_t buft);
bool (*backend_init)(
struct llama_sampler * smpl,
ggml_backend_buffer_type_t buft,
uint32_t n_outputs_max_per_seq);
// call after .backend_apply()
void (*backend_accept)(
@ -1290,6 +1301,13 @@ extern "C" {
// called before graph execution to set inputs for the current ubatch
void (*backend_set_input)(struct llama_sampler * smpl);
// called before rebuilding a sampling graph to clear any internal sampler state
void (*backend_reset)(struct llama_sampler * smpl);
// copy mutable state from src into dst while keeping dst's references to the current sampling graph
// src and dst must have the same type and configuration
void (*copy_state)(const struct llama_sampler * src, struct llama_sampler * dst);
};
struct llama_sampler {
@ -1310,6 +1328,7 @@ extern "C" {
LLAMA_API void llama_sampler_apply ( struct llama_sampler * smpl, llama_token_data_array * cur_p);
LLAMA_API void llama_sampler_reset ( struct llama_sampler * smpl);
LLAMA_API struct llama_sampler * llama_sampler_clone (const struct llama_sampler * smpl);
LLAMA_API void llama_sampler_copy (const struct llama_sampler * src, struct llama_sampler * dst);
// important: do not free if the sampler has been added to a llama_sampler_chain (via llama_sampler_chain_add)
LLAMA_API void llama_sampler_free ( struct llama_sampler * smpl);
@ -1425,7 +1444,7 @@ extern "C" {
/// NOTE: Avoid using on the full vocabulary as searching for repeated tokens can become slow. For example, apply top-k or top-p sampling first.
LLAMA_API struct llama_sampler * llama_sampler_init_penalties(
int32_t n_vocab,
int32_t penalty_last_n, // last n tokens to penalize (0 = disable penalty, -1 = context size)
int32_t penalty_last_n, // last n tokens to penalize (0 = disable penalty)
float penalty_repeat, // must be > 0.0, 1.0 = disabled
float penalty_freq, // must be finite, 0.0 = disabled
float penalty_present); // must be finite, 0.0 = disabled
@ -1433,11 +1452,10 @@ extern "C" {
/// @details DRY sampler, designed by p-e-w, as described in: https://github.com/oobabooga/text-generation-webui/pull/5677, porting Koboldcpp implementation authored by pi6am: https://github.com/LostRuins/koboldcpp/pull/982
LLAMA_API struct llama_sampler * llama_sampler_init_dry(
const struct llama_vocab * vocab,
int32_t n_ctx_train,
float dry_multiplier,
float dry_base,
int32_t dry_allowed_length,
int32_t dry_penalty_last_n,
int32_t dry_penalty_last_n, // last n tokens to penalize (0 = disable penalty)
const char ** seq_breakers,
size_t num_breakers);
@ -1500,6 +1518,7 @@ extern "C" {
LLAMA_API uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl);
/// @details Sample and accept a token from the idx-th output of the last evaluation
// For multiple outputs from one sampler, call this function in output order without gaps.
//
// Shorthand for:
// const auto * logits = llama_get_logits_ith(ctx, idx);

18
examples/talk-llama/models/clip.cpp Normal file
View File

@ -0,0 +1,18 @@
#include "models.h"
// Stub to allow llama-quantize to open mmproj GGUFs
[[noreturn]]
void llama_model_clip::load_arch_hparams(llama_model_loader &) {
GGML_ABORT("CLIP is a quant-only stub; load_arch_hparams should not be called");
}
[[noreturn]]
void llama_model_clip::load_arch_tensors(llama_model_loader &) {
GGML_ABORT("CLIP is a quant-only stub; load_arch_tensors should not be called");
}
[[noreturn]]
std::unique_ptr<llm_graph_context> llama_model_clip::build_arch_graph(const llm_graph_params &) const {
GGML_ABORT("CLIP has no inference graph via llama_model dispatch; runtime lives in tools/mtmd/clip.cpp");
}

41
examples/talk-llama/models/dflash.cpp
View File

@ -14,11 +14,14 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
hparams.n_embd_inp_enc_impl = (uint32_t) target_layer_ids.size() * hparams.n_embd;
LLAMA_LOG_INFO("%s: DFlash extract_layers = [", __func__);
for (size_t i = 0; i < target_layer_ids.size(); ++i) {
LLAMA_LOG_INFO("%d%s", target_layer_ids[i], i + 1 < target_layer_ids.size() ? ", " : "");
std::string layers;
const char * sep = "";
for (const auto id : target_layer_ids) {
layers += sep;
layers += std::to_string(id);
sep = ", ";
}
LLAMA_LOG_INFO("]\n");
LLAMA_LOG_INFO("%s: DFlash extract_layers = [%s]\n", __func__, layers.c_str());
// DeepSeek-V4 DSpark backbone: stages are full DSV4 blocks, uniform sliding window (the draft KV ring)
ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult, false);
@ -66,7 +69,7 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
// DFlash has a single rope, so the SWA rope == main rope.
if (ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false) && hparams.n_swa > 0) {
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer());
ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl);
hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
}
@ -79,6 +82,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
const int64_t n_embd_inp = hparams.n_embd_inp_enc();
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
// DSpark = DFlash + a semi-autoregressive Markov head and Confidence head
//
// TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4)
@ -97,6 +101,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
}
fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0);
fc_s = create_tensor(tn(LLM_TENSOR_FC, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc)
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm
@ -125,7 +130,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head}, 0);
layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", i), {n_embd, n_embd_head}, 0);
layer.attn_kv_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM, "weight", i), {n_embd_head}, 0);
layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank * o_groups}, 0);
layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank, o_groups}, TENSOR_ALLOW_RESHAPE);
layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_groups * o_lora_rank, n_embd}, 0);
layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0);
@ -205,7 +210,7 @@ template <>
llama_model_dflash::graph<true>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
ggml_tensor * cur = build_inp_embd_enc();
cur = build_lora_mm(model.fc, cur);
cur = build_lora_mm(model.fc, cur, model.fc_s);
cb(cur, "fc_out", -1);
cur = build_norm(cur, model.output_norm_enc, NULL, LLM_NORM_RMS, -1);
@ -460,9 +465,9 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
cb(cur, "ffn_norm", il);
cur = build_ffn(cur,
layer.ffn_up, NULL, NULL,
layer.ffn_gate, NULL, NULL,
layer.ffn_down, NULL, NULL,
layer.ffn_up, NULL, layer.ffn_up_s,
layer.ffn_gate, NULL, layer.ffn_gate_s,
layer.ffn_down, NULL, layer.ffn_down_s,
NULL,
LLM_FFN_SILU, LLM_FFN_PAR, il);
cb(cur, "ffn_out", il);
@ -479,15 +484,17 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
res->t_embd = cur;
// lm_head from the target model (shared via ctx_other)
auto * output = model.output;
auto * output = model.output;
auto * output_s = model.output_s;
if (output == nullptr) {
GGML_ASSERT(cparams.ctx_other != nullptr);
const auto * model_other = llama_get_model(cparams.ctx_other);
GGML_ASSERT(model_other->output != nullptr && "DFlash decoder requires the target model's output projection");
output = model_other->output;
output = model_other->output;
output_s = model_other->output_s;
}
cur = build_lora_mm(output, cur);
cur = build_lora_mm(output, cur, output_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
@ -655,15 +662,17 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
cb(cur, "result_norm", -1);
// lm_head from the target model (shared via ctx_other)
auto * output = model.output;
auto * output = model.output;
auto * output_s = model.output_s;
if (output == nullptr) {
GGML_ASSERT(cparams.ctx_other != nullptr);
const auto * model_other = llama_get_model(cparams.ctx_other);
GGML_ASSERT(model_other->output != nullptr && "DSpark decoder requires the target model's output projection");
output = model_other->output;
output = model_other->output;
output_s = model_other->output_s;
}
cur = build_lora_mm(output, cur);
cur = build_lora_mm(output, cur, output_s);
cb(cur, "result_output", -1);
res->t_logits = cur;

6
examples/talk-llama/models/exaone4.cpp
View File

@ -1,6 +1,9 @@
#include "models.h"
void llama_model_exaone4::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer");
if (hparams.n_layer() == 64) { // 32B
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
hparams.n_swa = 4096;
@ -15,9 +18,6 @@ void llama_model_exaone4::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer");
switch (hparams.n_layer()) {
case 30: type = LLM_TYPE_1_2B; break;

426
examples/talk-llama/models/granite-switch.cpp Normal file
View File

@ -0,0 +1,426 @@
#include "models.h"
#include <cmath>
void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale);
ml.get_key(LLM_KV_RESIDUAL_SCALE, hparams.f_residual_scale, false);
ml.get_key(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale, false);
ml.get_key(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale, false);
bool rope_finetuned = true;
ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false);
hparams.rope_finetuned = rope_finetuned;
switch (hparams.n_layer()) {
case 40: type = hparams.n_embd == 4096 ? LLM_TYPE_8B : LLM_TYPE_3B; break;
case 64: type = LLM_TYPE_30B; break;
default: type = LLM_TYPE_UNKNOWN;
}
ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, /* required */ false);
ml.get_key(LLM_KV_ADAPTER_COUNT, n_adapters);
ml.get_key(LLM_KV_ADAPTER_LORA_RANK, max_lora_rank);
ml.get_key(LLM_KV_ADAPTER_ROUTER_GAIN, router_gain, /* required */ false);
// bound counts that size tensors
if (n_adapters > 4096) {
throw std::runtime_error(format("graniteswitch: invalid adapter count %u", n_adapters));
}
if (max_lora_rank > 4096) {
throw std::runtime_error(format("graniteswitch: invalid lora rank %u", max_lora_rank));
}
std::vector<llama_token> token_ids;
std::vector<llama_token> substitute_ids;
ml.get_arr(LLM_KV_ADAPTER_TOKEN_IDS_ACTIVATE, token_ids);
ml.get_arr(LLM_KV_ADAPTER_TOKEN_IDS_SUBSTITUTE, substitute_ids);
if (token_ids.size() != n_adapters || substitute_ids.size() != n_adapters) {
throw std::runtime_error(format(
"graniteswitch: adapter token id arrays (%zu activate, %zu substitute) do not match adapter count %u",
token_ids.size(), substitute_ids.size(), n_adapters));
}
adapter_token_to_slot.clear();
adapter_token_to_substitute.clear();
for (uint32_t i = 0; i < n_adapters; ++i) {
// adapter i -> stacked slot i+1 (slot 0 is the base/zero delta)
adapter_token_to_slot[token_ids[i]] = (int32_t) (i + 1);
adapter_token_to_substitute[token_ids[i]] = substitute_ids[i];
}
// extra single-head attention layer at the END (index n_real) holds the router
// K/V. reusing n_layer_nextn keeps n_layer() == n_real, so the regular layers
// keep their indices and the KV cache shift/defrag skips the router layer.
// n_layer_nextn is repurposed here (no MTP): it leaks as 1 into the
// llama_model_n_layer_nextn() getter and a re-saved nextn_predict_layers
const uint32_t n_real = hparams.n_layer();
if (n_real >= LLAMA_MAX_LAYERS) {
throw std::runtime_error(format("graniteswitch: block count %u exceeds LLAMA_MAX_LAYERS", n_real));
}
hparams.router_layer = (int32_t) n_real;
hparams.n_layer_all = n_real + 1;
hparams.n_layer_nextn = 1;
hparams.n_head_arr[n_real] = 1;
hparams.n_head_kv_arr[n_real] = 1;
hparams.n_ff_arr[n_real] = 0;
}
void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) {
LLAMA_LOAD_LOCALS;
const int64_t n_slots = (int64_t) n_adapters + 1; // slot 0 = base/zero delta
const int64_t n_rank = (int64_t) max_lora_rank;
const int64_t n_embd_q = n_embd_head_k * n_head;
const int64_t n_embd_kv = n_embd_k_gqa;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
// substitute ids index tok_embd rows directly; range-check against n_vocab
for (const auto & kv : adapter_token_to_substitute) {
const llama_token sub = kv.second;
if (sub < 0 || (int64_t) sub >= n_vocab) {
throw std::runtime_error(format(
"graniteswitch: substitute token id %d out of range [0, %d)", sub, (int) n_vocab));
}
}
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
if (output == NULL) {
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
}
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, n_embd_q + 2*n_embd_kv}, 0);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_q, n_embd}, 0);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
auto & sl = layer.switch_lora;
sl.a_q = create_tensor(tn(LLM_TENSOR_ATTN_Q, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
sl.b_q = create_tensor(tn(LLM_TENSOR_ATTN_Q, "lora_b", i), {n_rank, n_embd_q, n_slots}, 0);
sl.a_k = create_tensor(tn(LLM_TENSOR_ATTN_K, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
sl.b_k = create_tensor(tn(LLM_TENSOR_ATTN_K, "lora_b", i), {n_rank, n_embd_kv, n_slots}, 0);
sl.a_v = create_tensor(tn(LLM_TENSOR_ATTN_V, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
sl.b_v = create_tensor(tn(LLM_TENSOR_ATTN_V, "lora_b", i), {n_rank, n_embd_kv, n_slots}, 0);
sl.a_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "lora_a", i), {n_embd_q, n_rank, n_slots}, 0);
sl.b_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "lora_b", i), {n_rank, n_embd, n_slots}, 0);
sl.a_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
sl.b_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "lora_b", i), {n_rank, n_ff, n_slots}, 0);
sl.a_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
sl.b_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "lora_b", i), {n_rank, n_ff, n_slots}, 0);
sl.a_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "lora_a", i), { n_ff, n_rank, n_slots}, 0);
sl.b_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "lora_b", i), {n_rank, n_embd, n_slots}, 0);
}
}
class llm_graph_input_switch : public llm_graph_input_i {
public:
llm_graph_input_switch(const llama_model_granite_switch & smodel) : smodel(smodel) {}
virtual ~llm_graph_input_switch() = default;
void set_input(const llama_ubatch * ubatch) override;
ggml_tensor * sub_tokens = nullptr; // I32 [n_tokens] adapter-substituted token ids
ggml_tensor * router_ksig = nullptr; // F32 [n_tokens] router K signal (+/-gain)
ggml_tensor * router_vval = nullptr; // F32 [n_tokens] router V value (adapter slot / 0)
ggml_tensor * router_q = nullptr; // F32 [n_tokens] router Q value (constant 1.0)
const llama_model_granite_switch & smodel;
};
// K dim-0 is +gain for an adapter token, -gain otherwise; the causal softmax then
// lets a single visible adapter token dominate so the readback recovers its slot.
void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) {
if (!ubatch->token) {
return;
}
const int64_t n_tokens = ubatch->n_tokens;
std::vector<int32_t> sub (n_tokens);
std::vector<float> ksig(n_tokens);
std::vector<float> vval(n_tokens);
std::vector<float> q (n_tokens, 1.0f);
for (int64_t i = 0; i < n_tokens; ++i) {
const llama_token tok = ubatch->token[i];
const auto it = smodel.adapter_token_to_slot.find(tok);
if (it != smodel.adapter_token_to_slot.end()) {
ksig[i] = +smodel.router_gain;
vval[i] = (float) it->second;
} else {
ksig[i] = -smodel.router_gain;
vval[i] = 0.0f;
}
const auto sit = smodel.adapter_token_to_substitute.find(tok);
sub[i] = (sit != smodel.adapter_token_to_substitute.end())
? (int32_t) sit->second
: (int32_t) tok;
}
ggml_backend_tensor_set(sub_tokens, sub.data(), 0, n_tokens*ggml_element_size(sub_tokens));
ggml_backend_tensor_set(router_ksig, ksig.data(), 0, n_tokens*ggml_element_size(router_ksig));
ggml_backend_tensor_set(router_vval, vval.data(), 0, n_tokens*ggml_element_size(router_vval));
ggml_backend_tensor_set(router_q, q.data(), 0, n_tokens*ggml_element_size(router_q));
}
std::unique_ptr<llm_graph_context> llama_model_granite_switch::build_arch_graph(const llm_graph_params & params) const {
return std::make_unique<graph>(*this, params);
}
// per-token switched LoRA delta: B_a*(A_a*x), adapter selected per token via ids.
// cur: {n_in, n_tokens}, ids: {n_tokens} -> {n_out, n_tokens}
ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_delta(
ggml_tensor * lora_a,
ggml_tensor * lora_b,
ggml_tensor * cur,
ggml_tensor * ids) {
const int64_t n_in = cur->ne[0];
const int64_t n_tokens = cur->ne[1];
ggml_tensor * x = ggml_reshape_3d(ctx0, cur, n_in, 1, n_tokens);
ggml_tensor * ids2 = ggml_reshape_2d(ctx0, ids, 1, n_tokens);
ggml_tensor * a = ggml_mul_mat_id(ctx0, lora_a, x, ids2); // {max_rank, 1, n_tokens}
ggml_tensor * d = ggml_mul_mat_id(ctx0, lora_b, a, ids2); // {n_out, 1, n_tokens}
return ggml_reshape_2d(ctx0, d, d->ne[0], n_tokens);
}
ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_mm(
ggml_tensor * w,
ggml_tensor * lora_a,
ggml_tensor * lora_b,
ggml_tensor * cur,
ggml_tensor * ids) {
ggml_tensor * base = ggml_mul_mat(ctx0, w, cur);
ggml_tensor * delta = build_switched_lora_delta(lora_a, lora_b, cur, ids);
return ggml_add(ctx0, base, delta);
}
llama_model_granite_switch::graph::graph(
const llama_model & model,
const llm_graph_params & params)
: llm_graph_context(params) {
const auto & smodel = static_cast<const llama_model_granite_switch &>(model);
// TODO: support raw embedding input (multimodal / pre-embedded tokens) when needed
GGML_ASSERT(ubatch.token && "granite-switch requires token input");
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
GGML_ASSERT(n_embd_head == n_rot);
auto inp_switch = std::make_unique<llm_graph_input_switch>(smodel);
inp_switch->sub_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
inp_switch->router_ksig = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens);
inp_switch->router_vval = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens);
inp_switch->router_q = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens);
ggml_set_input(inp_switch->sub_tokens);
ggml_set_input(inp_switch->router_ksig);
ggml_set_input(inp_switch->router_vval);
ggml_set_input(inp_switch->router_q);
ggml_tensor * sub_tokens = inp_switch->sub_tokens;
ggml_tensor * router_ksig = inp_switch->router_ksig;
ggml_tensor * router_vval = inp_switch->router_vval;
ggml_tensor * router_q = inp_switch->router_q;
res->add_input(std::move(inp_switch));
// embed the substituted ids directly; build_inp_embd would embed the raw tokens
ggml_tensor * inpL = ggml_get_rows(ctx0, model.tok_embd, sub_tokens);
if (hparams.f_embedding_scale != 0.0f) {
inpL = ggml_scale(ctx0, inpL, hparams.f_embedding_scale);
}
cb(inpL, "inp_embd", -1);
ggml_tensor * inp_pos = nullptr;
if (hparams.rope_finetuned) {
inp_pos = build_inp_pos();
}
auto * inp_attn = build_attn_inp_kv();
// single causal head at layer R recovers the adapter index in-graph: only dim 0
// carries signal (Q[0]=1, K[0]=+/-gain, V[0]=slot/0), the rest is zero-padded.
const int R = hparams.router_layer;
GGML_ASSERT(R >= 0);
auto router_lane = [&](ggml_tensor * sig1d) {
ggml_tensor * t = ggml_reshape_3d(ctx0, sig1d, 1, 1, n_tokens);
return ggml_pad(ctx0, t, (int) n_embd_head - 1, 0, 0, 0);
};
ggml_tensor * Qr = router_lane(router_q);
ggml_tensor * Kr = router_lane(router_ksig);
ggml_tensor * Vr = router_lane(router_vval);
ggml_tensor * router_out = build_attn(inp_attn,
nullptr, nullptr, nullptr,
Qr, Kr, Vr, nullptr, nullptr, nullptr, /*kq_scale=*/1.0f, /*il=*/R);
cb(router_out, "router_out", R);
// row 0 of router_out is the attended slot; clamp+round to an I32 index
ggml_tensor * slot_f = ggml_cont(ctx0,
ggml_view_2d(ctx0, router_out, 1, n_tokens, router_out->nb[1], 0));
slot_f = ggml_reshape_1d(ctx0, slot_f, n_tokens);
slot_f = ggml_clamp(ctx0, slot_f, 0.0f, (float) smodel.n_adapters);
slot_f = ggml_round(ctx0, slot_f);
ggml_tensor * adapter_ids = ggml_cast(ctx0, slot_f, GGML_TYPE_I32);
cb(adapter_ids, "adapter_ids", -1);
ggml_tensor * inp_out_ids = build_inp_out_ids();
ggml_tensor * cur;
for (int il = 0; il < n_layer; ++il) {
ggml_tensor * inpSA = inpL;
cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "attn_norm", il);
cur = build_attention_layer(cur, inp_pos, adapter_ids, inp_attn, model, n_embd_head, il);
if (il == n_layer - 1 && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
// keep adapter_ids aligned to the kept rows (2D round-trip for get_rows)
const int64_t n_out = inp_out_ids->ne[0];
adapter_ids = ggml_get_rows(ctx0,
ggml_reshape_2d(ctx0, adapter_ids, 1, adapter_ids->ne[0]), inp_out_ids);
adapter_ids = ggml_reshape_1d(ctx0, adapter_ids, n_out);
}
cur = build_layer_ffn(cur, inpSA, adapter_ids, model, il);
inpL = cur;
}
cur = inpL;
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
cb(cur, "result_norm", -1);
res->t_embd = cur;
cur = build_lora_mm(model.output, cur, model.output_s);
cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
ggml_tensor * llama_model_granite_switch::graph::build_attention_layer(
ggml_tensor * cur,
ggml_tensor * inp_pos,
ggml_tensor * adapter_ids,
llm_graph_input_attn_kv * inp_attn,
const llama_model & model,
const int64_t n_embd_head,
const int il) {
const auto & layer = model.layers[il];
const auto & sl = layer.switch_lora;
const int64_t n_head = hparams.n_head(il);
const int64_t n_head_kv = hparams.n_head_kv(il);
ggml_tensor * qkv = ggml_mul_mat(ctx0, layer.wqkv, cur);
cb(qkv, "wqkv", il);
const int64_t n_embd_q = n_embd_head * n_head;
const int64_t n_embd_kv = n_embd_head * n_head_kv;
// slice fused qkv into Q/K/V, made contiguous so LoRA deltas can be added
ggml_tensor * Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_q, qkv->ne[1], qkv->nb[1], 0));
ggml_tensor * Kcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_kv, qkv->ne[1], qkv->nb[1], n_embd_q*ggml_element_size(qkv)));
ggml_tensor * Vcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_kv, qkv->ne[1], qkv->nb[1], (n_embd_q + n_embd_kv)*ggml_element_size(qkv)));
Qcur = ggml_add(ctx0, Qcur, build_switched_lora_delta(sl.a_q, sl.b_q, cur, adapter_ids));
Kcur = ggml_add(ctx0, Kcur, build_switched_lora_delta(sl.a_k, sl.b_k, cur, adapter_ids));
Vcur = ggml_add(ctx0, Vcur, build_switched_lora_delta(sl.a_v, sl.b_v, cur, adapter_ids));
Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
if (hparams.rope_finetuned) {
ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
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(Kcur, "Kcur", il);
cb(Vcur, "Vcur", il);
const float kq_scale = hparams.f_attention_scale == 0.0f
? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
// wo = nullptr so build_attn returns concatenated heads; o-proj is switched below
ggml_tensor * attn = build_attn(inp_attn,
nullptr, nullptr, nullptr,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
cb(attn, "attn_pre_o", il);
cur = build_switched_lora_mm(layer.wo, sl.a_o, sl.b_o, attn, adapter_ids);
cb(cur, "attn_out", il);
return cur;
}
ggml_tensor * llama_model_granite_switch::graph::build_layer_ffn(
ggml_tensor * cur,
ggml_tensor * inpSA,
ggml_tensor * adapter_ids,
const llama_model & model,
const int il) {
const auto & layer = model.layers[il];
const auto & sl = layer.switch_lora;
if (hparams.f_residual_scale) {
cur = ggml_scale(ctx0, cur, hparams.f_residual_scale);
}
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
cb(ffn_inp, "ffn_inp", il);
cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "ffn_norm", il);
ggml_tensor * g = build_switched_lora_mm(layer.ffn_gate, sl.a_gate, sl.b_gate, cur, adapter_ids);
ggml_tensor * u = build_switched_lora_mm(layer.ffn_up, sl.a_up, sl.b_up, cur, adapter_ids);
g = ggml_silu(ctx0, g);
ggml_tensor * gu = ggml_mul(ctx0, g, u);
cur = build_switched_lora_mm(layer.ffn_down, sl.a_down, sl.b_down, gu, adapter_ids);
cb(cur, "ffn_out", il);
if (hparams.f_residual_scale) {
cur = ggml_scale(ctx0, cur, hparams.f_residual_scale);
}
cur = ggml_add(ctx0, cur, ffn_inp);
cur = build_cvec(cur, il);
cb(cur, "l_out", il);
return cur;
}

48
examples/talk-llama/models/mamba-base.cpp
View File

@ -2,6 +2,8 @@
#include "llama-memory-recurrent.h"
#include <algorithm>
llm_build_mamba_base::llm_build_mamba_base(const llm_graph_params & params) : llm_graph_context(params) {}
ggml_tensor * llm_build_mamba_base::build_mamba_layer(llm_graph_input_rs * inp,
@ -118,7 +120,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba_layer(llm_graph_input_rs * inp,
// Custom operator to optimize the parallel associative scan
// as described in the Annex D of the Mamba paper.
// => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids);
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, /*K=*/1);
};
ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);
@ -153,7 +155,8 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
int il) const {
const auto * mctx_cur = inp->mctx;
const auto kv_head = mctx_cur->get_head();
const auto kv_head = mctx_cur->get_head();
const auto mem_size = mctx_cur->get_size();
const int64_t d_conv = hparams.ssm_d_conv;
const int64_t d_inner = hparams.ssm_d_inner;
@ -164,6 +167,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
const int64_t n_seqs = ubatch.n_seqs;
const int64_t n_seq_tokens = ubatch.n_seq_tokens;
const int64_t K = cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1;
GGML_ASSERT(n_seqs != 0);
GGML_ASSERT(ubatch.equal_seqs());
@ -173,6 +177,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il);
const int64_t state_slots = ssm_states_all->ne[1];
ggml_tensor * conv = build_rs(inp, conv_states_all, hparams.n_embd_r(), n_seqs);
conv = ggml_reshape_3d(ctx0, conv, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs);
@ -198,15 +203,19 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
// => {d_conv - 1 + n_seq_tokens, d_inner + 2*n_group*d_state, n_seqs}
ggml_tensor * conv_x = ggml_concat(ctx0, conv, ggml_transpose(ctx0, xBC), 0);
// copy last (d_conv - 1) columns back into the state cache
ggml_tensor * last_conv = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs,
conv_x->nb[1], conv_x->nb[2], n_seq_tokens * (conv_x->nb[0]));
const int64_t row_count = (d_conv - 1) * (d_inner + 2 * n_group * d_state);
const size_t row_size = ggml_row_size(conv_states_all->type, row_count);
const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv,
ggml_view_1d(ctx0, conv_states_all,
(d_conv - 1) * (d_inner + 2 * n_group * d_state) * (n_seqs),
kv_head * (d_conv - 1) * (d_inner + 2 * n_group * d_state) *
ggml_element_size(conv_states_all))));
for (int64_t slot = 0; slot < n_written; ++slot) {
ggml_tensor * last_conv = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs,
conv_x->nb[1], conv_x->nb[2], (n_seq_tokens - slot) * conv_x->nb[0]);
ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv,
ggml_view_2d(ctx0, conv_states_all, row_count, n_seqs,
conv_states_all->nb[1],
((size_t) slot * mem_size + kv_head) * row_size)));
}
// 1D convolution
// The equivalent is to make a self-overlapping view of conv_x
@ -244,20 +253,27 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
// (this is necessary in order to properly use the states before they are overwritten,
// while avoiding to make unnecessary copies of the states)
auto get_ssm_rows = [&](ggml_context * ctx, ggml_tensor * states, ggml_tensor * ids) {
ggml_tensor * ssm = ggml_reshape_4d(ctx, states, d_state, head_dim, n_head, mctx_cur->get_size());
ggml_tensor * ssm = ggml_reshape_4d(ctx, states, d_state, head_dim, n_head, state_slots);
// TODO: use semistructured matrices to implement state-space duality
// => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids);
// K > 1 asks the backend to return rollback snapshots in addition to the final state.
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, K);
};
ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);
const int64_t D = d_state * d_inner;
const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
const size_t row_size = ggml_row_size(ssm_states_all->type, D);
const size_t y_row_size = ggml_row_size(y_ssm->type, D);
const size_t state_offset = ggml_nelements(x) * ggml_element_size(x);
// store last states
ggml_build_forward_expand(
gf, ggml_cpy(ctx0, ggml_view_1d(ctx0, y_ssm, d_state * d_inner * n_seqs, ggml_nelements(x) * x->nb[0]),
ggml_view_1d(ctx0, ssm_states_all, d_state * d_inner * n_seqs,
kv_head * d_state * d_inner * ggml_element_size(ssm_states_all))));
gf, ggml_cpy(ctx0,
ggml_view_3d(ctx0, y_ssm, D, n_seqs, n_written,
y_row_size, y_row_size * n_seqs, state_offset),
ggml_view_3d(ctx0, ssm_states_all, D, n_seqs, n_written,
ssm_states_all->nb[1], (size_t) mem_size * row_size, kv_head * row_size)));
ggml_tensor * y = ggml_view_4d(ctx0, y_ssm, head_dim, n_head, n_seq_tokens, n_seqs, x->nb[1], n_head * x->nb[1],
n_seq_tokens * n_head * x->nb[1], 0);

101
examples/talk-llama/models/models.h
View File

@ -386,6 +386,22 @@ struct llama_model_bloom : public llama_model_base {
};
// Quant-only stub for mmproj GGUFs
// none of these are ever called, they only exist to satisfy the llama_model_base interface
struct llama_model_clip : public llama_model_base {
llama_model_clip(const struct llama_model_params & params) : llama_model_base(params) {}
[[noreturn]]
void load_arch_hparams(llama_model_loader & ml) override;
[[noreturn]]
void load_arch_tensors(llama_model_loader & ml) override;
[[noreturn]]
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
struct llama_model_mpt : public llama_model_base {
llama_model_mpt(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
@ -596,6 +612,11 @@ struct llama_model_qwen3vlmoe : public llama_model_base {
};
struct llama_model_qwen3tts : public llama_model_qwen3vl {
llama_model_qwen3tts(const struct llama_model_params & params) : llama_model_qwen3vl(params) {}
};
struct llama_model_phi2 : public llama_model_base {
llama_model_phi2(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
@ -692,6 +713,19 @@ struct llama_model_gpt2 : public llama_model_base {
};
struct llama_model_pockettts : public llama_model_base {
llama_model_pockettts(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
struct graph : public llm_graph_context {
graph(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;
};
struct llama_model_codeshell : public llama_model_base {
llama_model_codeshell(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
@ -1023,6 +1057,19 @@ struct llama_model_olmoe : public llama_model_base {
};
struct llama_model_muse_glimmer : public llama_model_base {
llama_model_muse_glimmer(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
struct graph : public llm_graph_context {
graph(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;
};
struct llama_model_openelm : public llama_model_base {
llama_model_openelm(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
@ -1456,6 +1503,10 @@ struct llama_model_nemotron_h_moe : public llama_model_nemotron_h {
using graph = llama_model_nemotron_h::graph;
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;
};
@ -1591,6 +1642,56 @@ struct llama_model_granite_moe : public llama_model_base {
};
struct llama_model_granite_switch : public llama_model_base {
llama_model_granite_switch(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
uint32_t n_adapters = 0;
uint32_t max_lora_rank = 0;
float router_gain = 15.0f;
std::unordered_map<llama_token, int32_t> adapter_token_to_slot;
std::unordered_map<llama_token, llama_token> adapter_token_to_substitute;
struct graph : public llm_graph_context {
graph(const llama_model & model, const llm_graph_params & params);
private:
ggml_tensor * build_switched_lora_delta(
ggml_tensor * lora_a,
ggml_tensor * lora_b,
ggml_tensor * cur,
ggml_tensor * ids);
ggml_tensor * build_switched_lora_mm(
ggml_tensor * w,
ggml_tensor * lora_a,
ggml_tensor * lora_b,
ggml_tensor * cur,
ggml_tensor * ids);
ggml_tensor * build_attention_layer(
ggml_tensor * cur,
ggml_tensor * inp_pos,
ggml_tensor * adapter_ids,
llm_graph_input_attn_kv * inp_attn,
const llama_model & model,
const int64_t n_embd_head,
const int il);
ggml_tensor * build_layer_ffn(
ggml_tensor * cur,
ggml_tensor * inpSA,
ggml_tensor * adapter_ids,
const llama_model & model,
const int il);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
struct llama_model_minicpm : public llama_model_base {
llama_model_minicpm(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;

208
examples/talk-llama/models/muse-glimmer.cpp Normal file
View File

@ -0,0 +1,208 @@
#include "models.h"
void llama_model_muse_glimmer::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false);
ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale);
hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
uint32_t swa_period = 4;
if (ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, swa_period, false)) {
hparams.set_swa_pattern(swa_period);
} else {
ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer());
}
switch (hparams.n_layer()) {
case 52: type = LLM_TYPE_30B; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
void llama_model_muse_glimmer::load_arch_tensors(llama_model_loader &) {
LLAMA_LOAD_LOCALS;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
// Pre/post-attention norms (Muse Glimmer's `weight + 1` applied at conversion time).
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 0);
// Q/K/V/O projections.
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);
// QK-norm. Weights are synthesized at conversion time to absorb `qk_scale_factor`.
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);
// Attention output gate: sigmoid(gate) * attn_out before o_proj (same as afmoe).
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_embd_head_k * n_head}, 0);
// Pre/post-FFN norms (FFN_PRE_NORM is aliased to LLM_TENSOR_FFN_NORM).
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0);
// Dense FFN (unlike afmoe, no MoE branches).
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
}
}
llama_model_muse_glimmer::graph::graph(const llama_model & model, const llm_graph_params & params)
: llm_graph_context(params) {
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
// Different to f_norm_rms_eps for post-attn / post-FFN norms
const float post_norm_eps = 1e-8f;
ggml_tensor * cur;
ggml_tensor * inpL;
inpL = build_inp_embd(model.tok_embd);
inpL = build_norm(inpL, nullptr, nullptr, LLM_NORM_RMS, -1);
cb(inpL, "embd_norm", -1);
ggml_tensor * inp_pos = build_inp_pos();
auto * inp_attn = build_attn_inp_kv_iswa();
ggml_tensor * inp_out_ids = build_inp_out_ids();
const float kq_scale = 1.0f / sqrtf(float(n_embd_head));
for (int il = 0; il < n_layer; ++il) {
// expose per-layer residual for speculative drafts (see LLM_KV_TARGET_LAYERS).
res->t_layer_inp[il] = inpL;
const float freq_base_l = model.get_rope_freq_base (cparams, il);
const float freq_scale_l = model.get_rope_freq_scale(cparams, il);
ggml_tensor * inpSA = inpL;
// RoPE runs on the SWA layers, NoPE on full ones.
const bool use_rope = hparams.is_swa(il);
// pre-attention norm (weight+1 folded at conversion time)
cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "attn_norm", il);
// self-attention: attention output gate around SDPA (afmoe.cpp:147-191)
{
ggml_tensor * attn_inp = cur; // save input for gate computation
auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
n_embd_head, n_head, n_head_kv, il);
// gate = wqkv_gate @ attn_inp (from pre-attn hidden state)
ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, attn_inp);
cb(gate, "attn_gate_proj", il);
// QK-norm. attn_q_norm weight was synthesized at conversion to broadcast
// qk_scale_factor across head_dim; attn_k_norm is identity (ones).
Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);
Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);
cb(Qcur, "Qcur_normed", il);
cb(Kcur, "Kcur_normed", il);
if (use_rope) {
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);
cb(Qcur, "Qcur_rope", il);
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(Kcur, "Kcur_rope", il);
}
// SDPA. wo is deferred; the gate goes between attn_out and o_proj.
cur = build_attn(inp_attn,
NULL, NULL, NULL,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
cb(cur, "attn_out", il);
gate = ggml_sigmoid(ctx0, gate);
cb(gate, "attn_gate_sig", il);
cur = ggml_mul(ctx0, cur, gate);
cb(cur, "attn_gated", il);
cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);
cb(cur, "attn_o_proj", il);
}
cur = ggml_rms_norm(ctx0, cur, post_norm_eps);
cur = ggml_mul(ctx0, cur, model.layers[il].attn_post_norm);
cb(cur, "attn_post_norm", il);
if (il == n_layer - 1 && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
cb(ffn_inp, "ffn_inp", il);
// pre-FFN norm
cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "ffn_norm", il);
// SwiGLU dense FFN
cur = build_ffn(cur,
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, il);
cb(cur, "ffn_out", il);
cur = ggml_rms_norm(ctx0, cur, post_norm_eps);
cur = ggml_mul(ctx0, cur, model.layers[il].ffn_post_norm);
cb(cur, "ffn_post_norm", il);
cur = ggml_add(ctx0, cur, ffn_inp);
cur = build_cvec(cur, il);
cb(cur, "l_out", il);
inpL = cur;
}
cur = inpL;
// final norm
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
cb(cur, "result_norm", -1);
res->t_embd = cur;
// lm_head, followed by output multiplier
cur = build_lora_mm(model.output, cur, model.output_s);
cur = ggml_scale(ctx0, cur, hparams.f_logit_scale);
// Final logit tanh softcap (from gemma3.cpp).
if (hparams.f_final_logit_softcapping) {
cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping);
cur = ggml_tanh(ctx0, cur);
cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping);
}
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
std::unique_ptr<llm_graph_context> llama_model_muse_glimmer::build_arch_graph(const llm_graph_params & params) const {
return std::make_unique<graph>(*this, params);
}

150
examples/talk-llama/models/nemotron-h-moe.cpp
View File

@ -1,6 +1,156 @@
#include "models.h"
std::unique_ptr<llm_graph_context> llama_model_nemotron_h_moe::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);
}
// MTP draft head for Nemotron-H MoE
llama_model_nemotron_h_moe::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
: llm_graph_context(params) {
GGML_ASSERT(hparams.n_layer_nextn == 1 && "NEMOTRON_H_MOE MTP currently supports a single MTP block");
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
const int il = hparams.n_layer();
const auto & layer = model.layers[il];
GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm);
GGML_ASSERT(layer.ffn_gate_inp);
// token embedding weights
ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
GGML_ASSERT(tok_embd_w != nullptr && "NEMOTRON_H_MOE MTP requires token embeddings");
auto inp = std::make_unique<llm_graph_input_embd_h>(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_inp(), n_tokens);
ggml_set_input(inp->embd);
ggml_tensor * tok_embd;
if (ubatch.token) {
tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
} else {
tok_embd = inp->embd;
}
cb(tok_embd, "mtp_tok_embd", il);
inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
ggml_set_input(inp->h);
ggml_set_name(inp->h, "mtp_h_input");
ggml_tensor * h_embd = inp->h;
res->add_input(std::move(inp));
ggml_tensor * inp_out_ids = build_inp_out_ids();
// attention fills KV over all tokens, but the MoE is position-wise: gather output rows before
// it to save FFN compute (unless unmasked embeddings_nextn needs the full-length hidden state)
const bool emit_h_nextn = cparams.embeddings_nextn;
const bool crop_before_ffn = inp_out_ids && (!emit_h_nextn || cparams.embeddings_nextn_masked);
auto * inp_attn = build_attn_inp_kv();
ggml_tensor * h_norm = build_norm(h_embd, 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);
// dense NoPE attention sub-layer (mtp.layers.0)
ggml_tensor * inpSA = cur;
cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "mtp_attn_norm", il);
{
auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il);
const float kq_scale = hparams.f_attention_scale == 0.0f
? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
cur = build_attn(inp_attn, layer.wo, layer.wo_b, layer.wo_s,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
cb(cur, "mtp_attn_out", il);
}
cur = ggml_add(ctx0, cur, inpSA);
cb(cur, "mtp_attn_residual", il);
// gather the output rows here so the MoE FFN below only runs on the positions we keep
if (crop_before_ffn) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
// MoE FFN sub-layer (mtp.layers.1)
ggml_tensor * ffn_residual = cur;
cur = build_norm(cur, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "mtp_attn_post_norm", il);
{
ggml_tensor * router_logits = build_lora_mm(layer.ffn_gate_inp, cur);
cb(router_logits, "mtp_ffn_moe_logits", il);
ggml_tensor * moe_out =
build_moe_ffn(cur,
layer.ffn_gate_inp,
layer.ffn_up_exps,
nullptr, // no gate
layer.ffn_down_exps,
layer.ffn_exp_probs_b,
n_expert, n_expert_used,
LLM_FFN_RELU_SQR, hparams.expert_weights_norm,
hparams.expert_weights_scale,
LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID,
il,
router_logits, nullptr,
layer.ffn_up_exps_s,
nullptr, // no gate
layer.ffn_down_exps_s);
cb(moe_out, "mtp_ffn_moe_out", il);
ggml_tensor * ffn_shexp = build_ffn(cur,
layer.ffn_up_shexp, NULL, layer.ffn_up_shexp_s,
NULL, NULL, NULL,
layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s,
NULL,
LLM_FFN_RELU_SQR, LLM_FFN_PAR, il);
cb(ffn_shexp, "mtp_ffn_shexp", il);
cur = ggml_add(ctx0, moe_out, ffn_shexp);
cb(cur, "mtp_ffn_out", il);
}
cur = ggml_add(ctx0, cur, ffn_residual);
cb(cur, "mtp_post_ffn", il);
// final head norm: the MTP head has its own LayerNorm
GGML_ASSERT(layer.nextn.shared_head_norm && "NEMOTRON_H_MOE MTP: missing final head norm");
cur = build_norm(cur, layer.nextn.shared_head_norm, nullptr, LLM_NORM, -1);
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
if (!crop_before_ffn && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
// LM head
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 != nullptr && "NEMOTRON_H_MOE MTP requires an output projection");
cur = build_lora_mm(head_w, cur, head_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}

106
examples/talk-llama/models/nemotron-h.cpp
View File

@ -7,13 +7,18 @@ void llama_model_nemotron_h::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);
ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);
// NextN/MTP: optional draft head appended as extra trailing block(s)
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer_all");
// A layer is recurrent IFF the n_head_kv value is set to 0 and
// the n_ff value is set to 0
for (uint32_t i = 0; i < hparams.n_layer(); ++i) {
hparams.is_recr_impl[i] = (hparams.n_head_kv(i) == 0 && hparams.n_ff(i) == 0);
// the n_ff value is set to 0. Appended MTP blocks are dense (non-recurrent)
for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
hparams.is_recr_impl[i] = i < hparams.n_layer() && hparams.n_head_kv(i) == 0 && hparams.n_ff(i) == 0;
}
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); // MTP head final_layernorm
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false);
ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
@ -30,9 +35,13 @@ void llama_model_nemotron_h::load_arch_hparams(llama_model_loader & ml) {
}
}
void llama_model_nemotron_h::load_arch_tensors(llama_model_loader &) {
void llama_model_nemotron_h::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const bool mtp_only = hparams.n_layer_nextn > 0 && ml.get_weight("blk.0.attn_norm.weight") == nullptr;
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
const int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
// mamba2 Mixer SSM params
// NOTE: int64_t for tensor dimensions
const int64_t d_conv = hparams.ssm_d_conv;
@ -60,61 +69,94 @@ void llama_model_nemotron_h::load_arch_tensors(llama_model_loader &) {
auto & layer = layers[i];
// all blocks use the attn norm
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, trunk_flags);
if (hparams.is_recr(i)) {
// ssm layers
layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, d_in_proj}, 0);
layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, d_in_proj}, trunk_flags);
layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner + 2*n_group*d_state}, 0);
layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner + 2*n_group*d_state}, trunk_flags);
layer.ssm_conv1d_b = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "bias", i), {d_inner + 2*n_group*d_state}, TENSOR_NOT_REQUIRED);
layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {n_ssm_head}, 0);
layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {n_ssm_head}, trunk_flags);
// no "weight" suffix for these
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_ssm_head}, 0);
layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, i), {1, n_ssm_head}, 0);
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_ssm_head}, trunk_flags);
layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, i), {1, n_ssm_head}, trunk_flags);
layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {d_inner / n_group, n_group}, 0);
layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {d_inner / n_group, n_group}, trunk_flags);
// out_proj
layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), {d_inner, n_embd}, 0);
layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), {d_inner, n_embd}, trunk_flags);
} else if (hparams.n_ff(i) == 0) {
// attention layers (with optional bias)
const int64_t n_head_i = hparams.n_head(i);
const int64_t n_embd_k_gqa_i = hparams.n_embd_k_gqa(i);
const int64_t n_embd_v_gqa_i = hparams.n_embd_v_gqa(i);
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, 0);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, 0);
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, trunk_flags);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, trunk_flags);
layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED);
} else {
if (n_expert != 0) {
const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used;
const int64_t n_ff_shexp = hparams.n_ff_shexp;
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert}, 0);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert }, 0);
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert}, trunk_flags);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert }, trunk_flags);
// MoE branch
layer.ffn_latent_down = create_tensor(tn(LLM_TENSOR_FFN_LATENT_DOWN, "weight", i), {n_embd, moe_n_embd}, TENSOR_NOT_REQUIRED);
layer.ffn_latent_up = create_tensor(tn(LLM_TENSOR_FFN_LATENT_UP, "weight", i), {moe_n_embd, n_embd}, TENSOR_NOT_REQUIRED);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, 0);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, 0);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, trunk_flags);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, trunk_flags);
// Shared expert branch
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, 0);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, trunk_flags);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, trunk_flags);
} else {
// mlp layers
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { hparams.n_ff(i), n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, hparams.n_ff(i)}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { hparams.n_ff(i), n_embd}, trunk_flags);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, hparams.n_ff(i)}, trunk_flags);
layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED);
layer.ffn_up_b = create_tensor(tn(LLM_TENSOR_FFN_UP, "bias", i), {hparams.n_ff(i)}, TENSOR_NOT_REQUIRED);
}
}
}
// NextN/MTP draft head: each predict layer folds an attention sub-layer and a MoE
// sub-layer into a single trailing block
for (int i = n_layer; i < n_layer_all; ++i) {
auto & layer = layers[i];
const int64_t n_head_i = hparams.n_head(i);
const int64_t n_embd_k_gqa_i = hparams.n_embd_k_gqa(i);
const int64_t n_embd_v_gqa_i = hparams.n_embd_v_gqa(i);
const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used;
const int64_t n_ff_shexp = hparams.n_ff_shexp;
// NextN input-fusion tensors
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, mtp_flags);
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, mtp_flags);
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2*n_embd, n_embd}, mtp_flags);
layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, mtp_flags);
// attention sub-layer
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, mtp_flags);
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, mtp_flags);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, mtp_flags);
layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, mtp_flags | TENSOR_NOT_REQUIRED);
// MoE sub-layer
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, mtp_flags);
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, mtp_flags);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, mtp_flags);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, mtp_flags);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, mtp_flags);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, mtp_flags);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, mtp_flags);
}
}
std::unique_ptr<llm_graph_context> llama_model_nemotron_h::build_arch_graph(const llm_graph_params & params) const {
@ -135,8 +177,11 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
auto * inp = build_inp_mem_hybrid();
ggml_tensor * inp_out_ids = build_inp_out_ids();
const bool extract_final_inp = (size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer];
for (int il = 0; il < n_layer; ++il) {
res->t_layer_inp[il] = inpL;
struct ggml_tensor * inpSA = inpL;
// norm
@ -153,7 +198,7 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
cur = build_ffn_layer(cur, model, il);
}
if (il == n_layer - 1 && inp_out_ids) {
if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked && !extract_final_inp) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@ -167,9 +212,24 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
}
cur = inpL;
if (extract_final_inp) {
res->t_layer_inp[n_layer] = cur;
if (inp_out_ids && cparams.embeddings_nextn_masked) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
}
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
// seed for the MTP/NextN draft head
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);
}
cb(cur, "result_norm", -1);
res->t_embd = cur;

2
examples/talk-llama/models/plamo2.cpp
View File

@ -382,7 +382,7 @@ ggml_tensor * llama_model_plamo2::graph::build_plamo2_mamba_layer(llm_graph_inpu
// Custom operator to optimize the parallel associative scan
// as described in the Annex D of the Mamba paper.
// => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids);
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, /*K=*/1);
};
ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);

146
examples/talk-llama/models/pockettts.cpp Normal file
View File

@ -0,0 +1,146 @@
#include "models.h"
// backbone of the pocket-tts CALM pipeline: the "text" side of a flow language model.
// it has no lm_head, the audio latents are produced by the flow net inside the mmproj
void llama_model_pockettts::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
switch (hparams.n_layer()) {
case 6: type = LLM_TYPE_109M; break;
case 24: type = LLM_TYPE_335M; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
void llama_model_pockettts::load_arch_tensors(llama_model_loader &) {
LLAMA_LOAD_LOCALS;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output_norm_b = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "bias"), {n_embd}, 0);
// no output head, the logits are unused; reuse the embedding table so a sampler can still run
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.attn_norm_b = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "bias", i), {n_embd}, 0);
create_tensor_qkv(layer, i, n_embd, n_embd, n_embd_gqa, n_embd_gqa, TENSOR_NOT_REQUIRED);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd, n_embd}, 0);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_norm_b = create_tensor(tn(LLM_TENSOR_FFN_NORM, "bias", i), {n_embd}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
}
}
std::unique_ptr<llm_graph_context> llama_model_pockettts::build_arch_graph(const llm_graph_params & params) const {
return std::make_unique<graph>(*this, params);
}
llama_model_pockettts::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
GGML_ASSERT(n_embd_head == n_rot);
ggml_tensor * cur;
ggml_tensor * inpL;
inpL = build_inp_embd(model.tok_embd);
ggml_tensor * inp_pos = build_inp_pos();
auto * inp_attn = build_attn_inp_kv();
ggml_tensor * inp_out_ids = build_inp_out_ids();
for (int il = 0; il < n_layer; ++il) {
cur = build_norm(inpL,
model.layers[il].attn_norm,
model.layers[il].attn_norm_b,
LLM_NORM, il);
cb(cur, "attn_norm", il);
// self-attention
{
auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
n_embd_head, n_head, n_head_kv, il);
Qcur = ggml_rope_ext(
ctx0, Qcur, inp_pos, nullptr,
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, nullptr,
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);
cb(Vcur, "Vcur", il);
cur = build_attn(inp_attn,
model.layers[il].wo, NULL, model.layers[il].wo_s,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);
}
if (il == n_layer - 1 && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
}
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL);
cb(ffn_inp, "ffn_inp", il);
// FF
{
cur = build_norm(ffn_inp,
model.layers[il].ffn_norm,
model.layers[il].ffn_norm_b,
LLM_NORM, il);
cb(cur, "ffn_norm", il);
cur = build_ffn(cur,
model.layers[il].ffn_up, NULL, NULL,
NULL, NULL, NULL,
model.layers[il].ffn_down, NULL, NULL,
NULL,
LLM_FFN_GELU, LLM_FFN_SEQ, il);
cb(cur, "ffn_out", il);
}
cur = ggml_add(ctx0, cur, ffn_inp);
cur = build_cvec(cur, il);
cb(cur, "l_out", il);
// input for next layer
inpL = cur;
}
cur = build_norm(inpL,
model.output_norm,
model.output_norm_b,
LLM_NORM, -1);
cb(cur, "result_norm", -1);
res->t_embd = cur;
cur = build_lora_mm(model.output, cur, model.output_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}

3
examples/talk-llama/models/qwen3tts.cpp Normal file
View File

@ -0,0 +1,3 @@
#include "models.h"
// llama_model_qwen3tts reuses llama_model_qwen3vl's hparams/tensors/graph logic

25
examples/talk-llama/models/qwen3vl.cpp
View File

@ -16,11 +16,16 @@ void llama_model_qwen3vl::load_arch_hparams(llama_model_loader & ml) {
void llama_model_qwen3vl::load_arch_tensors(llama_model_loader &) {
LLAMA_LOAD_LOCALS;
int64_t n_vocab_out = n_vocab;
if (arch == LLM_ARCH_QWEN3TTS) {
n_vocab_out = 3072;
}
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
// output
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab_out}, TENSOR_NOT_REQUIRED);
// if output is NULL, init from the input tok embed
if (output == NULL) {
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
@ -166,6 +171,24 @@ llama_model_qwen3vl::graph::graph(const llama_model & model, const llm_graph_par
// lm_head
cur = build_lora_mm(model.output, cur, model.output_s);
int64_t n_vocab_in = model.tok_embd->ne[1];
int64_t n_vocab_out = model.output->ne[1];
if (n_vocab_in > n_vocab_out) {
// case: Qwen3TTS model with codec_head as output
GGML_ASSERT(model.output_norm);
int64_t pad = n_vocab_in - n_vocab_out;
// using this trick to get a scalar -inf tensor to pad the output
ggml_tensor * neg_inf = ggml_scale_bias(ctx0,
ggml_view_1d(ctx0, model.output_norm, 1, 0),
0.0f, -INFINITY);
neg_inf = ggml_repeat_4d(ctx0, neg_inf, pad, cur->ne[1], 1, 1);
cur = ggml_concat(ctx0, neg_inf, cur, 0); // [padded .. n_vocab_out, n_stream]
} else if (n_vocab_in < n_vocab_out) {
GGML_ABORT("invalid case");
}
cb(cur, "result_output", -1);
res->t_logits = cur;

8
ggml/CMakeLists.txt
View File

@ -4,8 +4,8 @@ project("ggml" C CXX ASM)
### GGML Version
set(GGML_VERSION_MAJOR 0)
set(GGML_VERSION_MINOR 18)
set(GGML_VERSION_PATCH 1)
set(GGML_VERSION_MINOR 20)
set(GGML_VERSION_PATCH 0)
set(GGML_VERSION_BASE "${GGML_VERSION_MAJOR}.${GGML_VERSION_MINOR}.${GGML_VERSION_PATCH}")
list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/")
@ -402,7 +402,7 @@ configure_package_config_file(
GGML_BIN_INSTALL_DIR)
write_basic_package_version_file(
${CMAKE_CURRENT_BINARY_DIR}/ggml-version.cmake
${CMAKE_CURRENT_BINARY_DIR}/ggml-config-version.cmake
VERSION ${GGML_INSTALL_VERSION}
COMPATIBILITY SameMajorVersion)
@ -414,7 +414,7 @@ message(STATUS "ggml version: ${GGML_INSTALL_VERSION}")
message(STATUS "ggml commit: ${GGML_BUILD_COMMIT}")
install(FILES ${CMAKE_CURRENT_BINARY_DIR}/ggml-config.cmake
${CMAKE_CURRENT_BINARY_DIR}/ggml-version.cmake
${CMAKE_CURRENT_BINARY_DIR}/ggml-config-version.cmake
DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/ggml)
if (MSVC)

7
ggml/cmake/ggml-config.cmake.in
View File

@ -113,6 +113,7 @@ set_and_check(GGML_LIB_DIR "@PACKAGE_GGML_LIB_INSTALL_DIR@")
if(NOT TARGET ggml::ggml)
find_package(Threads REQUIRED)
unset(GGML_LIBRARY CACHE)
find_library(GGML_LIBRARY ggml
REQUIRED
HINTS ${GGML_LIB_DIR}
@ -121,8 +122,10 @@ if(NOT TARGET ggml::ggml)
add_library(ggml::ggml UNKNOWN IMPORTED)
set_target_properties(ggml::ggml
PROPERTIES
IMPORTED_LOCATION "${GGML_LIBRARY}")
IMPORTED_LOCATION "${GGML_LIBRARY}"
INTERFACE_INCLUDE_DIRECTORIES "${GGML_INCLUDE_DIR}")
unset(GGML_BASE_LIBRARY CACHE)
find_library(GGML_BASE_LIBRARY ggml-base
REQUIRED
HINTS ${GGML_LIB_DIR}
@ -132,6 +135,7 @@ if(NOT TARGET ggml::ggml)
set_target_properties(ggml::ggml-base
PROPERTIES
IMPORTED_LOCATION "${GGML_BASE_LIBRARY}"
INTERFACE_INCLUDE_DIRECTORIES "${GGML_INCLUDE_DIR}"
INTERFACE_LINK_LIBRARIES "${GGML_BASE_INTERFACE_LINK_LIBRARIES}")
set(_ggml_all_targets "")
@ -140,6 +144,7 @@ if(NOT TARGET ggml::ggml)
string(REPLACE "-" "_" _ggml_backend_pfx "${_ggml_backend}")
string(TOUPPER "${_ggml_backend_pfx}" _ggml_backend_pfx)
unset(${_ggml_backend_pfx}_LIBRARY CACHE)
find_library(${_ggml_backend_pfx}_LIBRARY ${_ggml_backend}
REQUIRED
HINTS ${GGML_LIB_DIR}

2
ggml/include/ggml-backend.h
View File

@ -154,6 +154,8 @@ extern "C" {
bool buffer_from_host_ptr;
// event synchronization
bool events;
// mmap is supported for loading
bool mmap_support;
};
// all the device properties

9
ggml/include/ggml.h
View File

@ -2459,7 +2459,8 @@ extern "C" {
struct ggml_tensor * A,
struct ggml_tensor * B,
struct ggml_tensor * C,
struct ggml_tensor * ids);
struct ggml_tensor * ids,
int64_t K);
// partition into non-overlapping windows with padding if needed
// example:
@ -2788,6 +2789,12 @@ extern "C" {
struct ggml_cgraph * cgraph,
struct ggml_tensor * tensor);
// add the tensor and its parents to the graph without marking them for compute
// the flag is set later, when the tensor is reached from a node that computes
GGML_API void ggml_build_forward_order(
struct ggml_cgraph * cgraph,
struct ggml_tensor * tensor);
GGML_API void ggml_build_backward_expand(
struct ggml_context * ctx, // context for gradient computation
struct ggml_cgraph * cgraph,

2
ggml/src/ggml-backend-meta.cpp
View File

@ -132,6 +132,7 @@ static void ggml_backend_meta_device_get_props(ggml_backend_dev_t dev, ggml_back
/* .host_buffer = */ false, // Not implemented.
/* .buffer_from_host_ptr = */ false, // Not implemented.
/* .events = */ false, // Not implemented.
/* .mmap_support = */ true,
};
for (ggml_backend_dev_t simple_dev : meta_dev_ctx->simple_devs) {
ggml_backend_dev_props tmp_props;
@ -140,6 +141,7 @@ static void ggml_backend_meta_device_get_props(ggml_backend_dev_t dev, ggml_back
props->caps.host_buffer = props->caps.host_buffer && tmp_props.caps.host_buffer;
props->caps.buffer_from_host_ptr = props->caps.buffer_from_host_ptr && tmp_props.caps.buffer_from_host_ptr;
props->caps.events = props->caps.events && tmp_props.caps.events;
props->caps.mmap_support = props->caps.mmap_support && tmp_props.caps.mmap_support;
}
}

1
ggml/src/ggml-blas/ggml-blas.cpp
View File

@ -367,6 +367,7 @@ static void ggml_backend_blas_device_get_props(ggml_backend_dev_t dev, struct gg
/* .host_buffer = */ false,
/* .buffer_from_host_ptr = */ true,
/* .events = */ false,
/* .mmap_support = */ true,
};
}

1
ggml/src/ggml-cann/ggml-cann.cpp
View File

@ -2815,6 +2815,7 @@ static void ggml_backend_cann_device_get_props(ggml_backend_dev_t dev, ggml_back
/* .host_buffer = */ host_buffer,
/* .buffer_from_host_ptr = */ false,
/* .events = */ true,
/* .mmap_support = */ true,
};
}

74
ggml/src/ggml-cpu/arch/arm/cpu-feats.cpp
View File

@ -1,81 +1,19 @@
#include "ggml-backend-impl.h"
#include "ggml-feats.h"
#if defined(__aarch64__)
#if defined(__linux__)
#include <sys/auxv.h>
#elif defined(__APPLE__)
#include <sys/sysctl.h>
#endif
#if !defined(HWCAP2_SVE2)
#define HWCAP2_SVE2 (1 << 1)
#endif
#if !defined(HWCAP2_I8MM)
#define HWCAP2_I8MM (1 << 13)
#endif
#if !defined(HWCAP2_SME)
#define HWCAP2_SME (1 << 23)
#endif
struct aarch64_features {
// has_neon not needed, aarch64 has NEON guaranteed
bool has_dotprod = false;
bool has_fp16_va = false;
bool has_sve = false;
bool has_sve2 = false;
bool has_i8mm = false;
bool has_sme = false;
bool has_sme2 = false;
aarch64_features() {
#if defined(__linux__)
uint32_t hwcap = getauxval(AT_HWCAP);
uint32_t hwcap2 = getauxval(AT_HWCAP2);
has_dotprod = !!(hwcap & HWCAP_ASIMDDP);
has_fp16_va = !!(hwcap & HWCAP_FPHP);
has_sve = !!(hwcap & HWCAP_SVE);
has_sve2 = !!(hwcap2 & HWCAP2_SVE2);
has_i8mm = !!(hwcap2 & HWCAP2_I8MM);
has_sme = !!(hwcap2 & HWCAP2_SME);
#elif defined(__APPLE__)
int oldp = 0;
size_t size = sizeof(oldp);
if (sysctlbyname("hw.optional.arm.FEAT_DotProd", &oldp, &size, NULL, 0) == 0) {
has_dotprod = static_cast<bool>(oldp);
}
if (sysctlbyname("hw.optional.arm.FEAT_I8MM", &oldp, &size, NULL, 0) == 0) {
has_i8mm = static_cast<bool>(oldp);
}
if (sysctlbyname("hw.optional.arm.FEAT_SME", &oldp, &size, NULL, 0) == 0) {
has_sme = static_cast<bool>(oldp);
}
if (sysctlbyname("hw.optional.arm.FEAT_SME2", &oldp, &size, NULL, 0) == 0) {
has_sme2 = static_cast<bool>(oldp);
}
// Apple apparently does not implement SVE yet
#endif
}
};
#if defined(__aarch64__) || defined(_M_ARM64)
static int ggml_backend_cpu_aarch64_score() {
int score = 1;
aarch64_features af;
const ggml_feats_arch64_runtime_t af = ggml_feats_get_arch64_runtime();
GGML_UNUSED(af);
#ifdef GGML_USE_DOTPROD
if (!af.has_dotprod) { return 0; }
score += 1<<1;
#endif
#ifdef GGML_USE_FP16_VECTOR_ARITHMETIC
if (!af.has_fp16_va) { return 0; }
if (!af.has_fp16) { return 0; }
score += 1<<2;
#endif
#ifdef GGML_USE_SVE
@ -100,4 +38,4 @@ static int ggml_backend_cpu_aarch64_score() {
GGML_BACKEND_DL_SCORE_IMPL(ggml_backend_cpu_aarch64_score)
# endif // defined(__aarch64__)
# endif // defined(__aarch64__) || defined(_M_ARM64)

7
ggml/src/ggml-cpu/ggml-cpu.c
View File

@ -2608,7 +2608,7 @@ static bool ggml_thread_apply_priority(int32_t prio) {
return true;
}
#elif defined(__gnu_linux__)
#elif defined(__linux__)
// TODO: this may not work on BSD, to be verified
static bool ggml_thread_apply_affinity(const bool * mask) {
@ -2795,6 +2795,11 @@ struct ggml_cplan ggml_graph_plan(
n_threads = 1;
#endif
#if defined(__wasi__)
// WASI doesn't support parallelism yet
n_threads = 1;
#endif
size_t work_size = 0;
struct ggml_cplan cplan;

3
ggml/src/ggml-cpu/ggml-cpu.cpp
View File

@ -397,6 +397,7 @@ static void ggml_backend_cpu_device_get_props(ggml_backend_dev_t dev, struct ggm
/* .host_buffer = */ false,
/* .buffer_from_host_ptr = */ true,
/* .events = */ false,
/* .mmap_support = */ true,
};
}
@ -471,6 +472,8 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
case GGML_OP_CONV_2D:
return ggml_is_contiguous(op->src[0]);
case GGML_OP_SSM_SCAN:
return ggml_get_op_params_i32(op, 0) == 1 || op->src[3]->ne[0] == 1;
default:
return true;
}

312
ggml/src/ggml-cpu/kleidiai/kleidiai.cpp
View File

@ -2,10 +2,12 @@
// SPDX-License-Identifier: MIT
//
#include <arm_neon.h>
#include <assert.h>
#include <stdio.h>
#include <cassert>
#include <cstdio>
#include <cstdlib>
#include <atomic>
#include <cfloat>
#include <cctype>
#include <algorithm>
#include <cmath>
#include <stdexcept>
@ -17,25 +19,21 @@
#include <cstddef>
#include <cstdint>
#include <fstream>
#include <set>
#include <map>
#include <iostream>
#include <climits>
#include <charconv>
#include <system_error>
#if defined(__linux__)
#include <asm/hwcap.h>
#include <dirent.h>
#include <sys/auxv.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#ifndef HWCAP2_SME2
#define HWCAP2_SME2 (1UL << 37)
#endif
#elif defined(__APPLE__)
#include <string_view>
#include <sys/sysctl.h>
#include <sys/types.h>
#elif defined(_WIN32)
#include <windows.h>
#include <excpt.h>
#endif
#include "kleidiai.h"
@ -43,6 +41,7 @@
#include "ggml-cpu.h"
#include "ggml-cpu-impl.h"
#include "ggml-impl.h"
#include "ggml-feats.h"
#include "ggml-backend-impl.h"
#include "ggml-threading.h"
#include "traits.h"
@ -64,8 +63,8 @@ struct ggml_kleidiai_context {
ggml_kleidiai_kernels * kernels_q4;
ggml_kleidiai_kernels * kernels_q8;
ggml_kleidiai_kernels * kernels_f32;
int sme_thread_cap; // <= 0 means “SME disabled/unknown”;
int thread_hint; // <= 0 means “no hint”
int sme_thread_cap; // <= 0 means "SME disabled/unknown"
int thread_hint; // <= 0 means "no hint"
int chunk_multiplier;
} static ctx = { CPU_FEATURE_NONE, nullptr, nullptr, nullptr, 0, -1, 4 };
@ -93,24 +92,117 @@ static const char* cpu_feature_to_string(cpu_feature f) {
}
}
#if defined(__linux__) && defined(__aarch64__)
static bool parse_cpu_dir_name(const char* name, size_t* cpu) {
if (strncmp(name, "cpu", 3) != 0 ||
name[3] < '0' || name[3] > '9') {
return false;
}
const char* first = name + 3;
const char* last = name + strlen(name);
size_t value = 0;
const auto [end, ec] = std::from_chars(first, last, value, 10);
if (ec != std::errc{} || end != last) {
return false;
}
*cpu = value;
return true;
}
static std::vector<size_t> detect_cpu_ids() {
std::vector<size_t> cpus;
DIR * dir = opendir("/sys/devices/system/cpu");
if (dir == nullptr) {
return cpus;
}
while (dirent * entry = readdir(dir)) {
size_t cpu = 0;
if (parse_cpu_dir_name(entry->d_name, &cpu)) {
cpus.push_back(cpu);
}
}
closedir(dir);
std::sort(cpus.begin(), cpus.end());
cpus.erase(std::unique(cpus.begin(), cpus.end()), cpus.end());
return cpus;
}
#endif
#if defined(__APPLE__) && defined(__aarch64__)
static bool apple_sme_counted_perf_level(std::string name) {
for (std::string::size_type i = 0; i < name.size(); ++i) {
name[i] = (char) std::tolower((unsigned char) name[i]);
}
// Conservative ceiling: only count perf-level names observed to provide full SME throughput.
// Future names should be calibrated here before they raise the automatic SME thread cap.
return name.find("super") != std::string::npos ||
name.find("performance") != std::string::npos;
}
#endif
static void add_smcus_from_smidr(uint64_t smidr, size_t & num_private, std::map<uint32_t, size_t> & shared_counts) {
// Arm ARM: SMIDR_EL1. SH==0 is implementation-defined; keep the existing
// conservative policy and only treat zero affinity as private.
const uint32_t sh = (uint32_t)((smidr >> 13) & 0x3);
const uint32_t nsmc = (uint32_t)((smidr >> 56) & 0xF);
const size_t shared_count = nsmc == 0xF ? 1 : (size_t)nsmc + 1;
const uint32_t affinity = (uint32_t)(smidr & 0xFFFu);
const uint32_t affinity2 = (uint32_t)((smidr >> 32) & 0xFFFFFu);
const uint32_t id = (affinity2 << 12) | affinity;
if (nsmc == 0xF) {
GGML_LOG_WARN("kleidiai: NSMC detected as 0xF indicating reseved value, setting min safe shared SMCU count to 1");
}
switch (sh) {
case 2: // private SMCU
++num_private;
break;
case 3: // shared SMCU
if (shared_counts[id] < shared_count) {
shared_counts[id] = shared_count;
}
break;
case 0:
if (id == 0) {
++num_private;
} else if (shared_counts[id] < shared_count) {
shared_counts[id] = shared_count;
}
break;
default:
break;
}
}
static size_t detect_num_smcus() {
if (!ggml_cpu_has_sme()) {
const auto runtime_feat = ggml_feats_get_arch64_runtime();
if (!runtime_feat.has_sme) {
return 0;
}
#if defined(__linux__) && defined(__aarch64__)
// Linux/aarch64: Best-effort count of Streaming Mode Compute Units (SMCUs) via SMIDR_EL1 sysfs.
size_t num_private = 0;
std::set<uint32_t> shared_ids;
std::map<uint32_t, size_t> shared_counts;
for (size_t cpu = 0;; ++cpu) {
const std::vector<size_t> cpus = detect_cpu_ids();
for (const size_t cpu : cpus) {
const std::string path =
"/sys/devices/system/cpu/cpu" + std::to_string(cpu) +
"/regs/identification/smidr_el1";
std::ifstream file(path);
if (!file.is_open()) {
break;
continue;
}
uint64_t smidr = 0;
@ -118,54 +210,69 @@ static size_t detect_num_smcus() {
continue;
}
// Arm ARM: SMIDR_EL1
const uint32_t sh = (uint32_t)((smidr >> 13) & 0x3);
// Build an "affinity-like" identifier for shared SMCUs.
// Keep the original packing logic, but isolate it here.
const uint32_t id = (uint32_t)((smidr & 0xFFFu) | ((smidr >> 20) & 0xFFFFF000u));
switch (sh) {
case 0b10: // private SMCU
++num_private;
break;
case 0b11: // shared SMCU
shared_ids.emplace(id);
break;
case 0b00:
// Ambiguous / implementation-defined. Be conservative:
// treat id==0 as private, otherwise as shared.
if (id == 0) ++num_private;
else shared_ids.emplace(id);
break;
default:
break;
}
add_smcus_from_smidr(smidr, num_private, shared_counts);
}
return num_private + shared_ids.size();
size_t total = num_private;
for (const auto & entry : shared_counts) {
total += entry.second;
}
return total;
#elif defined(__APPLE__) && defined(__aarch64__)
// table for known M4 variants. Users can override via GGML_KLEIDIAI_SME=<n>.
char chip_name[256] = {};
size_t size = sizeof(chip_name);
int perf_levels = 0;
size_t size = sizeof(perf_levels);
if (sysctlbyname("hw.nperflevels", &perf_levels, &size, nullptr, 0) != 0 ||
size != sizeof(perf_levels) || perf_levels <= 0) {
return 0;
}
if (sysctlbyname("machdep.cpu.brand_string", chip_name, &size, nullptr, 0) == 0) {
const std::string brand(chip_name);
size_t units = 0;
for (int i = 0; i < perf_levels; ++i) {
char key[64] = {};
int physical_cpus = 0;
int cpus_per_l2 = 0;
struct ModelSMCU { const char *match; size_t smcus; };
static const ModelSMCU table[] = {
{ "M4 Ultra", 2 },
{ "M4 Max", 2 },
{ "M4 Pro", 2 },
{ "M4", 1 },
};
snprintf(key, sizeof(key), "hw.perflevel%d.physicalcpu", i);
size = sizeof(physical_cpus);
if (sysctlbyname(key, &physical_cpus, &size, nullptr, 0) != 0 ||
size != sizeof(physical_cpus) || physical_cpus <= 0) {
continue;
}
for (const auto &e : table) {
if (brand.find(e.match) != std::string::npos) {
return e.smcus;
}
snprintf(key, sizeof(key), "hw.perflevel%d.cpusperl2", i);
size = sizeof(cpus_per_l2);
if (sysctlbyname(key, &cpus_per_l2, &size, nullptr, 0) != 0 ||
size != sizeof(cpus_per_l2) || cpus_per_l2 <= 0) {
continue;
}
snprintf(key, sizeof(key), "hw.perflevel%d.name", i);
size = 0;
if (sysctlbyname(key, nullptr, &size, nullptr, 0) != 0 || size == 0) {
continue;
}
std::string name(size, '\0');
if (sysctlbyname(key, &name[0], &size, nullptr, 0) != 0) {
continue;
}
name.resize(size);
while (!name.empty() && name.back() == '\0') {
name.pop_back();
}
if (apple_sme_counted_perf_level(name)) {
units += (size_t) ((physical_cpus + cpus_per_l2 - 1) / cpus_per_l2);
}
}
return units;
#elif defined(_WIN32) && (defined(_M_ARM64) || defined(__aarch64__))
// No verified Windows arm64 SMCU detection path yet. Return unknown and use
// GGML_KLEIDIAI_SME=N as a diagnostics/debug override for SME thread cap
// calibration until a detection mechanism is verified on real hardware.
return 0;
#else
@ -198,15 +305,18 @@ static void init_kleidiai_context(void) {
if (!initialized) {
initialized = true;
// Optional diagnostics/debug overrides; production defaults come from runtime detection.
const char *env_sme = getenv("GGML_KLEIDIAI_SME");
const char *env_threads = getenv("GGML_TOTAL_THREADS");
const char *env_chunk_mult = getenv("GGML_KLEIDIAI_CHUNK_MULTIPLIER");
const auto runtime_feat = ggml_feats_get_arch64_runtime();
size_t detected_smcus = 0;
ctx.features = (ggml_cpu_has_dotprod() ? CPU_FEATURE_DOTPROD : CPU_FEATURE_NONE) |
(ggml_cpu_has_matmul_int8() ? CPU_FEATURE_I8MM : CPU_FEATURE_NONE) |
((ggml_cpu_has_sve() && ggml_cpu_get_sve_cnt() == QK8_0) ? CPU_FEATURE_SVE : CPU_FEATURE_NONE);
ctx.features = (runtime_feat.has_dotprod ? CPU_FEATURE_DOTPROD : CPU_FEATURE_NONE) |
(runtime_feat.has_i8mm ? CPU_FEATURE_I8MM : CPU_FEATURE_NONE) |
(runtime_feat.sve_cnt == QK8_0 ? CPU_FEATURE_SVE : CPU_FEATURE_NONE);
if (env_threads) {
bool ok = false;
@ -224,54 +334,54 @@ static void init_kleidiai_context(void) {
}
}
// SME policy:
// - env unset => auto-detect SMCUs; enable SME only if detected > 0.
// - env=0 => force off.
// - env>0 => force N cores, if the binary was built with SME.
int sme_cores = 0;
bool sme_env_ok = false;
bool sme_env_set = (env_sme != nullptr);
const bool has_supported_sme_family = runtime_feat.has_sme;
bool sme_cap_detected = false;
if (has_supported_sme_family) {
detected_smcus = detect_num_smcus();
sme_cap_detected = detected_smcus > 0;
// Some platforms expose SME without exposing a calibrated SMCU count.
// Use one SME thread as the conservative default; add platform SMCU detection to raise it.
sme_cores = sme_cap_detected ? (int)detected_smcus : 1;
if (!sme_env_set && !sme_cap_detected) {
GGML_LOG_INFO("kleidiai: SME detected; SMCU count unavailable, using conservative SME thread cap=1\n");
}
}
// Runtime-detect SME support and available SMCUs first. The detected SMCU
// count is used as the SME thread cap, and GGML_KLEIDIAI_SME can debug-override that:
// - unset: use runtime detection.
// - 0: disable SME-family kernels.
// - N > 0: use N as the SME thread cap, if an SME-family kernel is selectable.
if (sme_env_set) {
bool ok = false;
int v = parse_uint_env(env_sme, "GGML_KLEIDIAI_SME", &ok);
sme_env_ok = ok;
if (!ok) {
GGML_LOG_WARN("kleidiai: GGML_KLEIDIAI_SME set but parsing failed; falling back to runtime SME-core detection\n");
detected_smcus = detect_num_smcus();
sme_cores = detected_smcus > 0 ? (int)detected_smcus : 0;
} else if (v == 0) {
sme_cores = 0;
} else if (!ggml_cpu_has_sme()) {
GGML_LOG_WARN("kleidiai: GGML_KLEIDIAI_SME=%d but the binary was not built with SME; disabling SME\n", v);
sme_cores = 0;
if (ok) {
if (has_supported_sme_family) {
sme_cores = v;
} else {
if (v > 0) {
GGML_LOG_WARN("kleidiai: GGML_KLEIDIAI_SME=%d but SME is not supported on this CPU; disabling SME-family kernels\n", v);
}
sme_cores = 0;
}
} else {
sme_cores = v;
GGML_LOG_WARN("kleidiai: GGML_KLEIDIAI_SME set but parsing failed; using automatic SME thread cap\n");
}
} else {
detected_smcus = detect_num_smcus();
sme_cores = detected_smcus > 0 ? (int)detected_smcus : 0;
}
if (!sme_env_set && ggml_cpu_has_sme() && sme_cores == 0) {
GGML_LOG_WARN("kleidiai: runtime SME-core detection returned 0; falling back to NEON\n");
}
if (sme_cores > 0) {
if (sme_cores > 0 && has_supported_sme_family) {
ctx.features |= CPU_FEATURE_SME;
#if defined(__aarch64__) && defined(__linux__)
// ARM guarantees SME2 implies SME, so only check SME2 when SME is enabled.
if (getauxval(AT_HWCAP2) & HWCAP2_SME2) {
if (runtime_feat.has_sme2) {
ctx.features |= CPU_FEATURE_SME2;
}
#elif defined(__aarch64__) && defined(__APPLE__)
int feat_sme2 = 0;
size_t size = sizeof(feat_sme2);
if (sysctlbyname("hw.optional.arm.FEAT_SME2", &feat_sme2, &size, NULL, 0) == 0 && feat_sme2) {
ctx.features |= CPU_FEATURE_SME2;
}
#endif
}
// Kernel selection
@ -297,16 +407,19 @@ static void init_kleidiai_context(void) {
GGML_LOG_INFO("kleidiai: primary f32 kernel feature %s\n", cpu_feature_to_string(ctx.kernels_f32->required_cpu));
}
ctx.sme_thread_cap = (ctx.features & CPU_FEATURE_SME) ? sme_cores : 0;
const bool has_selected_sme_family_kernel =
(ctx.kernels_q4 && is_sme_family(ctx.kernels_q4->required_cpu)) ||
(ctx.kernels_q8 && is_sme_family(ctx.kernels_q8->required_cpu)) ||
(ctx.kernels_f32 && is_sme_family(ctx.kernels_f32->required_cpu));
ctx.sme_thread_cap = has_selected_sme_family_kernel ? sme_cores : 0;
if (ctx.features & CPU_FEATURE_SME) {
const bool has_sme2 = (ctx.features & CPU_FEATURE_SME2) != CPU_FEATURE_NONE;
if (has_selected_sme_family_kernel) {
if (sme_env_set && sme_env_ok && sme_cores > 0) {
GGML_LOG_INFO("kleidiai: SME%s enabled (GGML_KLEIDIAI_SME=%d override)\n",
has_sme2 ? "2" : "", sme_cores);
GGML_LOG_INFO("kleidiai: SME enabled (GGML_KLEIDIAI_SME=%d debug override)\n", sme_cores);
} else if (sme_cap_detected) {
GGML_LOG_INFO("kleidiai: SME enabled (runtime-detected SME thread cap=%d)\n", sme_cores);
} else {
GGML_LOG_INFO("kleidiai: SME%s enabled (runtime-detected SME cores=%d)\n",
has_sme2 ? "2" : "", sme_cores);
GGML_LOG_INFO("kleidiai: SME enabled (runtime SME detected, conservative thread cap=%d)\n", sme_cores);
}
} else {
GGML_LOG_INFO("kleidiai: SME disabled\n");
@ -467,7 +580,7 @@ static int kleidiai_collect_kernel_chain_common(
}
if (is_sme_family(primary->required_cpu)) {
const cpu_feature fallback_mask = static_cast<cpu_feature>(features & ~CPU_FEATURE_SME & ~CPU_FEATURE_SME2);
const cpu_feature fallback_mask = static_cast<cpu_feature>(features & ~(CPU_FEATURE_SME | CPU_FEATURE_SME2));
if (fallback_mask != CPU_FEATURE_NONE) {
ggml_kleidiai_kernels * fallback = select_fallback(fallback_mask);
if (fallback && fallback != primary &&
@ -1077,13 +1190,14 @@ class tensor_traits : public ggml::cpu::tensor_traits {
const int ith_total = params->ith;
int sme_slot = -1;
int non_sme_slot = -1;
for (int i = 0; i < runtime_count; ++i) {
if (is_sme_family(runtime[i].kernels->required_cpu)) {
sme_slot = i;
break;
}
}
int non_sme_slot = -1;
for (int i = 0; i < runtime_count; ++i) {
if (!is_sme_family(runtime[i].kernels->required_cpu)) {
non_sme_slot = i;

14
ggml/src/ggml-cpu/ops.cpp
View File

@ -8941,7 +8941,7 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
for (int tk = 0; tk < kv_tile; tk++) {
const char * v_data = (const char *)v->data + (ic + tk)*nbv1 + iv2*nbv2 + iv3*nbv3;
if (kv_type == GGML_TYPE_F16) {
ggml_fp16_to_fp32_row((const ggml_fp16_t *)v_data, V32 + tk * DV, DV);
ggml_cpu_fp16_to_fp32((const ggml_fp16_t *)v_data, V32 + tk * DV, DV);
} else {
memcpy(V32 + tk * DV, v_data, DV * sizeof(float));
}
@ -9644,11 +9644,13 @@ static void ggml_compute_forward_ssm_scan_f32(
const int64_t ng = src4->ne[1];
const int64_t nt = src1->ne[2]; // number of tokens per sequence
const int64_t ns = src1->ne[3]; // number of sequences in the batch
const int64_t K = ggml_get_op_params_i32(dst, 0);
// can't use ggml_nbytes because src1 is not necessarily contiguous
const int64_t s_off = ggml_nelements(src1) * ggml_element_size(src1);
GGML_ASSERT(ggml_nelements(src1) + nc*nr*nh*ns == ggml_nelements(dst));
GGML_ASSERT(K >= 1);
GGML_ASSERT(ggml_nelements(src1) + K*nc*nr*nh*ns == ggml_nelements(dst));
GGML_ASSERT(src0->nb[0] == sizeof(float));
GGML_ASSERT(src1->nb[0] == sizeof(float));
GGML_ASSERT(src2->nb[0] == sizeof(float));
@ -9657,6 +9659,7 @@ static void ggml_compute_forward_ssm_scan_f32(
GGML_ASSERT(src5->nb[0] == sizeof(float));
GGML_ASSERT(src6->nb[0] == sizeof(int32_t));
GGML_ASSERT(nh % ng == 0);
GGML_ASSERT(src3->ne[0] == 1 || K == 1);
// heads per thread
const int dh = (nh + nth - 1)/nth;
@ -9831,6 +9834,13 @@ static void ggml_compute_forward_ssm_scan_f32(
}
}
}
const int64_t slot = nt - 1 - i2;
if (K > 1 && slot > 0 && slot < K) {
float * s_snapshot = (float *) ((char *) dst->data + s_off + (slot*ns + i3)*(src0->nb[3]));
for (int h = ih0; h < ih1; ++h) {
memcpy((char *) s_snapshot + h*src0->nb[2], (char *) s + h*src0->nb[2], src0->nb[2]);
}
}
// use the output as the source when it's not the first token-wise iteration
s0 = s;
}

2
ggml/src/ggml-cpu/spacemit/ime.cpp
View File

@ -195,6 +195,7 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_
case GGML_TYPE_Q4_K:
case GGML_TYPE_Q6_K:
case GGML_TYPE_Q8_0:
case GGML_TYPE_Q5_0:
case GGML_TYPE_Q5_1:
case GGML_TYPE_Q5_K:
//case GGML_TYPE_MXFP4:
@ -214,6 +215,7 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_
case GGML_TYPE_Q4_K:
case GGML_TYPE_Q6_K:
case GGML_TYPE_Q8_0:
case GGML_TYPE_Q5_0:
case GGML_TYPE_Q5_1:
case GGML_TYPE_Q5_K:
//case GGML_TYPE_MXFP4:

44
ggml/src/ggml-cuda/cpy.cu
View File

@ -253,9 +253,9 @@ static void ggml_cpy_f32_q8_0_cuda(
const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {
GGML_ASSERT(ne % QK8_0 == 0);
const int64_t num_blocks = ne / QK8_0;
const int64_t num_blocks = (ne/QK8_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
GGML_ASSERT(num_blocks <= INT_MAX);
cpy_f32_q<cpy_blck_f32_q8_0, QK8_0><<<num_blocks, 1, 0, stream>>>
cpy_f32_q<cpy_blck_f32_q8_0, QK8_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
}
@ -264,9 +264,9 @@ static void ggml_cpy_q8_0_f32_cuda(
const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t nb00, const int64_t nb01, const int64_t nb02,
const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {
const int64_t num_blocks = ne;
const int64_t num_blocks = (ne/QK8_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
GGML_ASSERT(num_blocks <= INT_MAX);
cpy_q_f32<cpy_blck_q8_0_f32, QK8_0><<<num_blocks, 1, 0, stream>>>
cpy_q_f32<cpy_blck_q8_0_f32, QK8_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
}
@ -276,9 +276,9 @@ static void ggml_cpy_f32_q4_0_cuda(
const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {
GGML_ASSERT(ne % QK4_0 == 0);
const int64_t num_blocks = ne / QK4_0;
const int64_t num_blocks = (ne/QK4_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
GGML_ASSERT(num_blocks <= INT_MAX);
cpy_f32_q<cpy_blck_f32_q4_0, QK4_0><<<num_blocks, 1, 0, stream>>>
cpy_f32_q<cpy_blck_f32_q4_0, QK4_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
}
@ -289,9 +289,9 @@ static void ggml_cpy_q4_0_f32_cuda(
const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12,
const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13,
cudaStream_t stream) {
const int64_t num_blocks = ne;
const int64_t num_blocks = (ne/QK4_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
GGML_ASSERT(num_blocks <= INT_MAX);
cpy_q_f32<cpy_blck_q_f32<dequantize_q4_0, QK4_0>, QK4_0><<<num_blocks, 1, 0, stream>>>(
cpy_q_f32<cpy_blck_q_f32<dequantize_q4_0, QK4_0>, QK4_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>(
cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
ne10, ne11, ne12, nb10, nb11, nb12, nb13);
}
@ -302,9 +302,9 @@ static void ggml_cpy_f32_q4_1_cuda(
const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {
GGML_ASSERT(ne % QK4_1 == 0);
const int64_t num_blocks = ne / QK4_1;
const int64_t num_blocks = (ne/QK4_1 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
GGML_ASSERT(num_blocks <= INT_MAX);
cpy_f32_q<cpy_blck_f32_q4_1, QK4_1><<<num_blocks, 1, 0, stream>>>
cpy_f32_q<cpy_blck_f32_q4_1, QK4_1><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
}
@ -315,9 +315,9 @@ static void ggml_cpy_q4_1_f32_cuda(
const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12,
const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13,
cudaStream_t stream) {
const int64_t num_blocks = ne;
const int64_t num_blocks = (ne/QK4_1 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
GGML_ASSERT(num_blocks <= INT_MAX);
cpy_q_f32<cpy_blck_q_f32<dequantize_q4_1, QK4_1>, QK4_1><<<num_blocks, 1, 0, stream>>>(
cpy_q_f32<cpy_blck_q_f32<dequantize_q4_1, QK4_1>, QK4_1><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>(
cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
ne10, ne11, ne12, nb10, nb11, nb12, nb13);
}
@ -328,9 +328,9 @@ static void ggml_cpy_f32_q5_0_cuda(
const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {
GGML_ASSERT(ne % QK5_0 == 0);
const int64_t num_blocks = ne / QK5_0;
const int64_t num_blocks = (ne/QK5_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
GGML_ASSERT(num_blocks <= INT_MAX);
cpy_f32_q<cpy_blck_f32_q5_0, QK5_0><<<num_blocks, 1, 0, stream>>>
cpy_f32_q<cpy_blck_f32_q5_0, QK5_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
}
@ -341,9 +341,9 @@ static void ggml_cpy_q5_0_f32_cuda(
const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12,
const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13,
cudaStream_t stream) {
const int64_t num_blocks = ne;
const int64_t num_blocks = (ne/QK5_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
GGML_ASSERT(num_blocks <= INT_MAX);
cpy_q_f32<cpy_blck_q_f32<dequantize_q5_0, QK5_0>, QK5_0><<<num_blocks, 1, 0, stream>>>(
cpy_q_f32<cpy_blck_q_f32<dequantize_q5_0, QK5_0>, QK5_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>(
cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
ne10, ne11, ne12, nb10, nb11, nb12, nb13);
}
@ -354,9 +354,9 @@ static void ggml_cpy_f32_q5_1_cuda(
const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {
GGML_ASSERT(ne % QK5_1 == 0);
const int64_t num_blocks = ne / QK5_1;
const int64_t num_blocks = (ne/QK5_1 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
GGML_ASSERT(num_blocks <= INT_MAX);
cpy_f32_q<cpy_blck_f32_q5_1, QK5_1><<<num_blocks, 1, 0, stream>>>
cpy_f32_q<cpy_blck_f32_q5_1, QK5_1><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
}
@ -367,9 +367,9 @@ static void ggml_cpy_q5_1_f32_cuda(
const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12,
const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13,
cudaStream_t stream) {
const int64_t num_blocks = ne;
const int64_t num_blocks = (ne/QK5_1 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
GGML_ASSERT(num_blocks <= INT_MAX);
cpy_q_f32<cpy_blck_q_f32<dequantize_q5_1, QK5_1>, QK5_1><<<num_blocks, 1, 0, stream>>>(
cpy_q_f32<cpy_blck_q_f32<dequantize_q5_1, QK5_1>, QK5_1><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>(
cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
ne10, ne11, ne12, nb10, nb11, nb12, nb13);
}
@ -380,9 +380,9 @@ static void ggml_cpy_f32_iq4_nl_cuda(
const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {
GGML_ASSERT(ne % QK4_NL == 0);
const int64_t num_blocks = ne / QK4_NL;
const int64_t num_blocks = (ne/QK4_NL + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
GGML_ASSERT(num_blocks <= INT_MAX);
cpy_f32_q<cpy_blck_f32_iq4_nl, QK4_NL><<<num_blocks, 1, 0, stream>>>
cpy_f32_q<cpy_blck_f32_iq4_nl, QK4_NL><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
}

144
ggml/src/ggml-cuda/ggml-cuda.cu
View File

@ -1865,6 +1865,37 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor
ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst);
}
// returns true when ggml_cuda_mul_mat_id takes the fallback path that requires stream synchronization
// [TAG_MUL_MAT_ID_CUDA_GRAPHS]
static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int cc) {
const ggml_tensor * src0 = dst->src[0];
const ggml_tensor * src1 = dst->src[1];
if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
return true;
}
if (dst->ne[2] <= MMVQ_MAX_BATCH_SIZE) {
if (ggml_is_quantized(src0->type)) {
if (dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc)) {
return false;
}
} else if (GGML_CUDA_CC_IS_AMD(cc)) {
return false;
}
}
if (ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[2], /*n_experts=*/src0->ne[2])) {
return false;
}
if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
return false;
}
return true;
}
static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0];
const ggml_tensor * src1 = dst->src[1];
@ -1907,7 +1938,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor *
}
// note: this path should not be reached when recording CUDA graphs, because it requires stream synchronization
// TODO: add asserts to verify this. should work with CUDA, HIP, etc.
GGML_ASSERT(ggml_cuda_mul_mat_id_needs_sync(dst, cc));
cudaStream_t stream = ctx.stream();
GGML_ASSERT(nb12 % nb11 == 0);
@ -2522,10 +2553,8 @@ static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) {
// [TAG_MUL_MAT_ID_CUDA_GRAPHS]
if (node->op == GGML_OP_MUL_MAT_ID) {
const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
const int mmvq_mmid_max = get_mmvq_mmid_max_batch(node->src[0]->type, cc);
if (!ggml_is_quantized(node->src[0]->type) || node->ne[2] > mmvq_mmid_max) {
// under these conditions, the mul_mat_id operation will need to synchronize the stream, so we cannot use CUDA graphs
// TODO: figure out a way to enable for larger batch sizes, without hurting performance
if (ggml_cuda_mul_mat_id_needs_sync(node, cc)) {
// the mul_mat_id fallback path synchronizes the stream, so we cannot use CUDA graphs
// ref: https://github.com/ggml-org/llama.cpp/pull/18958
use_cuda_graph = false;
#ifndef NDEBUG
@ -2651,6 +2680,52 @@ static bool ggml_cuda_should_fuse_rope_set_rows(const ggml_tensor * rope,
return true;
}
static bool ggml_cuda_should_fuse_rms_norm_mul_rope(const ggml_tensor * rms_norm,
const ggml_tensor * mul,
const ggml_tensor * rope) {
if (rms_norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || rope->op != GGML_OP_ROPE) {
return false;
}
if (rms_norm->src[0]->type != GGML_TYPE_F32 || rms_norm->type != GGML_TYPE_F32 ||
mul->src[0]->type != GGML_TYPE_F32 || mul->src[1]->type != GGML_TYPE_F32 ||
mul->type != GGML_TYPE_F32 || rope->type != GGML_TYPE_F32) {
return false;
}
if (rope->src[0] != mul) {
return false;
}
//if rms norm is the B operand, then we don't handle broadcast
if (rms_norm == mul->src[1] && !ggml_are_same_shape(mul->src[0], rms_norm)) {
return false;
}
if (!ggml_are_same_shape(rms_norm, mul)) {
return false;
}
//rms_norm kernel assumes contiguous rows
if (!ggml_is_contiguous_rows(rms_norm->src[0]) ||
!ggml_is_contiguous_rows(mul->src[0]) || !ggml_is_contiguous_rows(mul->src[1])) {
return false;
}
// the fused kernel handles the norm/neox rope modes only
const int mode = ((const int32_t *) rope->op_params)[2];
if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX) {
return false;
}
const int n_dims = ((const int32_t *) rope->op_params)[1];
if (n_dims % 2 != 0 || rope->src[0]->ne[0] % 2 != 0) {
return false;
}
return true;
}
// match gated_delta_net + the strided cpy that scatters its state snapshots into the cache
// (slot i -> rollback group i, slot 0 newest), so the kernel can write them and skip the cpy.
static int ggml_cuda_try_gdn_cache_fusion(
@ -2980,6 +3055,36 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph,
}
}
std::initializer_list<enum ggml_op> rms_norm_mul_rope_ops = { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE };
std::initializer_list<enum ggml_op> rms_norm_mul_rope_set_rows_ops = { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS };
if (is_equal(rms_norm_mul_rope_set_rows_ops, ops) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 4 })) {
const ggml_tensor * rms_norm = cgraph->nodes[node_idx];
const ggml_tensor * mul = cgraph->nodes[node_idx + 1];
const ggml_tensor * rope = cgraph->nodes[node_idx + 2];
const ggml_tensor * view = cgraph->nodes[node_idx + 3];
const ggml_tensor * set_rows = cgraph->nodes[node_idx + 4];
if (ggml_check_edges(cgraph, node_idx, {{1, 0, 0}, {2, 0, 1}, {3, 0, 2}, {4, 0, 3}}) &&
ggml_cuda_should_fuse_rms_norm_mul_rope(rms_norm, mul, rope) &&
ggml_cuda_should_fuse_rope_set_rows(rope, view, set_rows)) {
int out_nodes[] = { node_idx + 4 };
return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1);
}
}
if (is_equal(rms_norm_mul_rope_ops, ops) && ggml_can_fuse(cgraph, node_idx, ops)) {
const ggml_tensor * rms_norm = cgraph->nodes[node_idx];
const ggml_tensor * mul = cgraph->nodes[node_idx + 1];
const ggml_tensor * rope = cgraph->nodes[node_idx + 2];
if (ggml_cuda_should_fuse_rms_norm_mul_rope(rms_norm, mul, rope)) {
int out_nodes[] = { node_idx + 2 };
return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1);
}
return false;
}
std::initializer_list<enum ggml_op> rope_set_rows_ops = { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS };
if (is_equal(rope_set_rows_ops, ops) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 2 })) {
@ -2988,7 +3093,8 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph,
const ggml_tensor * set_rows = cgraph->nodes[node_idx + 2];
if (ggml_cuda_should_fuse_rope_set_rows(rope, view, set_rows)) {
return true;
int out_nodes[] = { node_idx + 2 };
return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1);
}
}
@ -3840,6 +3946,16 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
return fused_node_count - 1;
}
if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, {})) {
ggml_cuda_op_rms_norm_mul_rope_fused(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2], cgraph->nodes[i + 4]);
return 4;
}
if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE }, {})) {
ggml_cuda_op_rms_norm_mul_rope_fused(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2], nullptr);
return 2;
}
if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }, {})) {
ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]);
return 2;
@ -4033,7 +4149,11 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud
continue;
}
#ifndef NDEBUG
assert(node->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device));
// On integrated GPUs (APUs, e.g. RDNA3.5) the scheduler may place a
// node's output on the host-visible buffer, which the compute path
// handles. Allow that here, mirroring the src-tensor check below.
assert(node->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) ||
(integrated && ggml_backend_buft_is_cuda_host(node->buffer->buft)));
for (int j = 0; j < GGML_MAX_SRC; j++) {
if (node->src[j] != nullptr) {
assert(node->src[j]->buffer);
@ -4710,6 +4830,7 @@ static void ggml_backend_cuda_device_get_props(ggml_backend_dev_t dev, ggml_back
/* .host_buffer = */ host_buffer,
/* .buffer_from_host_ptr = */ false,
/* .events = */ events,
/* .mmap_support = */ props->type != GGML_BACKEND_DEVICE_TYPE_IGPU,
};
}
@ -5068,11 +5189,17 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
(op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) &&
(op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16);
case GGML_OP_SSM_SCAN: {
const int32_t K = ggml_get_op_params_i32(op, 0);
if (op->src[3]->ne[0] == 1) {
// Mamba2
// (kernel only supports (d_state == 128 || d_state == 256) && d_head % 16 == 0)
return (op->src[0]->ne[0] == 128 || op->src[0]->ne[0] == 256) && op->src[0]->ne[1] % 16 == 0;
} else {
if (K > 1) {
return false;
}
// Mamba
// (kernel only supports d_state == 16, d_head == 1, n_head % 128 == 0, n_group == 1)
return op->src[0]->ne[0] == 16 && op->src[0]->ne[1] == 1 && op->src[0]->ne[2] % 128 == 0 && op->src[4]->ne[1] == 1;
@ -5094,7 +5221,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
return max_bias == 0.0f;
}
case GGML_OP_ROLL:
if(op->src[0]->type == GGML_TYPE_F32) {
if(op->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(op->src[0])) {
return true;
}
return false;
@ -5205,6 +5332,7 @@ static bool ggml_backend_cuda_device_offload_op(ggml_backend_dev_t dev, const gg
static ggml_backend_event_t ggml_backend_cuda_device_event_new(ggml_backend_dev_t dev) {
#ifdef GGML_CUDA_NO_PEER_COPY
GGML_UNUSED(dev);
return nullptr;
#else
ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *)dev->context;

2
ggml/src/ggml-cuda/quantize.cu
View File

@ -8,7 +8,6 @@ struct __builtin_align__(32) float8 {
float x; float y; float z; float w;
float p; float q; float r; float s;
};
#endif
#if CUDART_VERSION >= 12080
static __device__ __forceinline__ float nvfp4_native_scale_error(
@ -49,6 +48,7 @@ static __device__ __forceinline__ float nvfp4_native_scale_error(
return err;
}
#endif // CUDART_VERSION >= 12080
#endif // defined(BLACKWELL_MMA_AVAILABLE)
__launch_bounds__(CUDA_QUANTIZE_BLOCK_SIZE, 1)
static __global__ void quantize_q8_1(

235
ggml/src/ggml-cuda/rope.cu
View File

@ -670,3 +670,238 @@ void ggml_cuda_op_rope_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
void ggml_cuda_op_rope_fused(ggml_backend_cuda_context & ctx, ggml_tensor * rope, ggml_tensor * set_rows) {
ggml_cuda_op_rope_impl<true>(ctx, rope, set_rows);
}
// fused RMS_NORM + MUL + ROPE (+ VIEW + SET_ROWS)
// one block per row: block_reduce gives the norm scale, then each thread applies mul and rope to the elements it owns
template <int block_size, bool has_ff, typename D>
static __global__ void rms_norm_mul_rope_f32(
const float * x, D * dst, const int ncols,
const int64_t s01, const int64_t s02, const int64_t s03,
const int64_t s1, const int64_t s2, const int64_t s3,
const float eps,
const float * mul,
const int64_t mul_s01, const int64_t mul_s02, const int64_t mul_s03,
const uint3 mul_ncols_packed, const uint3 mul_nrows_packed,
const uint3 mul_nchannels_packed, const uint3 mul_nsamples_packed,
const int n_dims, const int32_t * pos,
const float freq_scale, const float ext_factor, const float attn_factor,
const rope_corr_dims corr_dims, const float theta_scale,
const float * freq_factors,
const int64_t * row_indices, const int set_rows_stride,
const bool is_neox) {
ggml_cuda_pdl_lc();
const int row = blockIdx.x;
const int channel = blockIdx.y;
const int sample = blockIdx.z;
const int tid = threadIdx.x;
x += sample*s03 + channel*s02 + row*s01;
const uint32_t mul_row = fastmodulo(row, mul_nrows_packed);
const uint32_t mul_channel = fastmodulo(channel, mul_nchannels_packed);
const uint32_t mul_sample = fastmodulo(sample, mul_nsamples_packed);
mul += mul_sample*mul_s03 + mul_channel*mul_s02 + mul_row*mul_s01;
float tmp = 0.0f;
ggml_cuda_pdl_sync();
for (int col = tid; col < ncols; col += block_size) {
const float xi = x[col];
tmp += xi * xi;
}
extern __shared__ float s_sum[];
tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum);
const float scale = rsqrtf(tmp/ncols + eps);
int64_t idst = sample*s3 + channel*s2 + row*s1;
if (set_rows_stride != 0) {
idst = row*s1 + row_indices[channel]*set_rows_stride;
}
dst += idst;
for (int i0 = 2*tid; i0 < ncols; i0 += 2*block_size) {
int ix0;
int ix1;
if (is_neox && i0 < n_dims) {
ix0 = i0/2;
ix1 = i0/2 + n_dims/2;
} else {
ix0 = i0 + 0;
ix1 = i0 + 1;
}
const float x0 = scale * x[ix0] * mul[fastmodulo(ix0, mul_ncols_packed)];
const float x1 = scale * x[ix1] * mul[fastmodulo(ix1, mul_ncols_packed)];
if (i0 >= n_dims) {
dst[ix0] = ggml_cuda_cast<D>(x0);
dst[ix1] = ggml_cuda_cast<D>(x1);
continue;
}
const float theta_base = pos[channel]*powf(theta_scale, i0/2.0f);
const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f;
float cos_theta;
float sin_theta;
rope_yarn<true>(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta);
dst[ix0] = ggml_cuda_cast<D>(x0*cos_theta - x1*sin_theta);
dst[ix1] = ggml_cuda_cast<D>(x0*sin_theta + x1*cos_theta);
}
}
template <typename D>
static void rms_norm_mul_rope_cuda(
const float * x, D * dst,
const int ncols, const int nrows, const int nchannels, const int nsamples,
const int64_t s01, const int64_t s02, const int64_t s03,
const int64_t s1, const int64_t s2, const int64_t s3,
const float eps,
const float * mul,
const int64_t mul_s01, const int64_t mul_s02, const int64_t mul_s03,
const uint32_t mul_ncols, const uint32_t mul_nrows,
const uint32_t mul_nchannels, const uint32_t mul_nsamples,
const int n_dims, const int32_t * pos,
const float freq_scale, const float freq_base, const float ext_factor, const float attn_factor,
const rope_corr_dims corr_dims,
const float * freq_factors,
const int64_t * row_indices, const int set_rows_stride,
const bool is_neox, cudaStream_t stream) {
GGML_ASSERT(ncols % 2 == 0);
const dim3 blocks_num(nrows, nchannels, nsamples);
const float theta_scale = powf(freq_base, -2.0f/n_dims);
const uint3 mul_ncols_packed = init_fastdiv_values(mul_ncols);
const uint3 mul_nrows_packed = init_fastdiv_values(mul_nrows);
const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels);
const uint3 mul_nsamples_packed = init_fastdiv_values(mul_nsamples);
if (ncols < 1024) {
const dim3 block_dims(256, 1, 1);
const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, 32*sizeof(float), stream};
if (freq_factors == nullptr) {
ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<256, false, D>, launch_params,
x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale,
freq_factors, row_indices, set_rows_stride, is_neox);
} else {
ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<256, true, D>, launch_params,
x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale,
freq_factors, row_indices, set_rows_stride, is_neox);
}
} else {
const dim3 block_dims(1024, 1, 1);
const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, 32*sizeof(float), stream};
if (freq_factors == nullptr) {
ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<1024, false, D>, launch_params,
x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale,
freq_factors, row_indices, set_rows_stride, is_neox);
} else {
ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<1024, true, D>, launch_params,
x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale,
freq_factors, row_indices, set_rows_stride, is_neox);
}
}
}
void ggml_cuda_op_rms_norm_mul_rope_fused(ggml_backend_cuda_context & ctx,
ggml_tensor * rms_norm, ggml_tensor * mul, ggml_tensor * rope, ggml_tensor * set_rows) {
const ggml_tensor * x = rms_norm->src[0];
const ggml_tensor * mul_src = mul->src[0] == rms_norm ? mul->src[1] : mul->src[0];
float eps = 0.0f;
memcpy(&eps, rms_norm->op_params, sizeof(float));
GGML_ASSERT(eps >= 0.0f);
GGML_ASSERT(x->type == GGML_TYPE_F32);
GGML_ASSERT(mul_src->type == GGML_TYPE_F32);
GGML_ASSERT(rope->type == GGML_TYPE_F32);
void * dst_d = rope->data;
ggml_type dst_type = rope->type;
const int64_t * row_indices = nullptr;
int set_rows_stride = 0;
if (set_rows != nullptr) {
dst_d = set_rows->data;
dst_type = set_rows->type;
row_indices = (const int64_t *) set_rows->src[1]->data;
set_rows_stride = set_rows->nb[1] / ggml_type_size(set_rows->type);
}
const int n_dims = ((const int32_t *) rope->op_params)[1];
const int mode = ((const int32_t *) rope->op_params)[2];
const int n_ctx_orig = ((const int32_t *) rope->op_params)[4];
float freq_base;
float freq_scale;
float ext_factor;
float attn_factor;
float beta_fast;
float beta_slow;
memcpy(&freq_base, (const int32_t *) rope->op_params + 5, sizeof(float));
memcpy(&freq_scale, (const int32_t *) rope->op_params + 6, sizeof(float));
memcpy(&ext_factor, (const int32_t *) rope->op_params + 7, sizeof(float));
memcpy(&attn_factor, (const int32_t *) rope->op_params + 8, sizeof(float));
memcpy(&beta_fast, (const int32_t *) rope->op_params + 9, sizeof(float));
memcpy(&beta_slow, (const int32_t *) rope->op_params + 10, sizeof(float));
const bool is_neox = mode & GGML_ROPE_TYPE_NEOX;
const int32_t * pos = (const int32_t *) rope->src[1]->data;
const float * freq_factors = rope->src[2] != nullptr ? (const float *) rope->src[2]->data : nullptr;
rope_corr_dims corr_dims;
ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims.v);
const size_t ts0 = ggml_type_size(x->type);
GGML_ASSERT(x->nb[0] == ts0);
const int64_t s01 = x->nb[1] / ts0;
const int64_t s02 = x->nb[2] / ts0;
const int64_t s03 = x->nb[3] / ts0;
const size_t ts_mul = ggml_type_size(mul_src->type);
GGML_ASSERT(mul_src->nb[0] == ts_mul);
const int64_t mul_s01 = mul_src->nb[1] / ts_mul;
const int64_t mul_s02 = mul_src->nb[2] / ts_mul;
const int64_t mul_s03 = mul_src->nb[3] / ts_mul;
const size_t ts_dst = ggml_type_size(rope->type);
const int64_t s1 = rope->nb[1] / ts_dst;
const int64_t s2 = rope->nb[2] / ts_dst;
const int64_t s3 = rope->nb[3] / ts_dst;
cudaStream_t stream = ctx.stream();
if (dst_type == GGML_TYPE_F32) {
rms_norm_mul_rope_cuda((const float *) x->data, (float *) dst_d,
x->ne[0], x->ne[1], x->ne[2], x->ne[3], s01, s02, s03, s1, s2, s3, eps,
(const float *) mul_src->data, mul_s01, mul_s02, mul_s03,
mul_src->ne[0], mul_src->ne[1], mul_src->ne[2], mul_src->ne[3],
n_dims, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims,
freq_factors, row_indices, set_rows_stride, is_neox, stream);
} else if (dst_type == GGML_TYPE_F16) {
rms_norm_mul_rope_cuda((const float *) x->data, (half *) dst_d,
x->ne[0], x->ne[1], x->ne[2], x->ne[3], s01, s02, s03, s1, s2, s3, eps,
(const float *) mul_src->data, mul_s01, mul_s02, mul_s03,
mul_src->ne[0], mul_src->ne[1], mul_src->ne[2], mul_src->ne[3],
n_dims, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims,
freq_factors, row_indices, set_rows_stride, is_neox, stream);
} else {
GGML_ABORT("fatal error");
}
}

2
ggml/src/ggml-cuda/rope.cuh
View File

@ -7,3 +7,5 @@ void ggml_cuda_op_rope(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_rope_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_rope_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * set_rows);
void ggml_cuda_op_rms_norm_mul_rope_fused(ggml_backend_cuda_context & ctx, ggml_tensor * rms_norm, ggml_tensor * mul, ggml_tensor * rope, ggml_tensor * set_rows);

27
ggml/src/ggml-cuda/ssm-scan.cu
View File

@ -149,7 +149,7 @@ __global__ void __launch_bounds__(d_state, 1)
const int src0_nb2, const int src0_nb3, const int src1_nb2, const int src1_nb3,
const int src2_nb1, const int src2_nb2, const int src3_nb1,
const int src4_nb2, const int src4_nb3, const int src5_nb2, const int src5_nb3,
const int64_t s_off, const int64_t n_head, const int64_t d_head, const int64_t n_group, const int64_t n_tok) {
const int64_t s_off, const int64_t n_head, const int64_t d_head, const int64_t n_group, const int64_t n_tok, const int64_t K) {
const float * GGML_CUDA_RESTRICT src0 = src0_ptr;
const float * GGML_CUDA_RESTRICT src1 = src1_ptr;
const float * GGML_CUDA_RESTRICT src2 = src2_ptr;
@ -217,6 +217,16 @@ __global__ void __launch_bounds__(d_state, 1)
if (lane == 0) {
y_warp[i * stride_y] = state_sum;
}
// Slot 0 is the final state written below; slots 1..K-1 are rollback snapshots.
const int64_t slot = n_tok - 1 - i;
if (K > 1 && slot > 0 && slot < K) {
float * s_snapshot_warp = (float *) ((char *) dst + s_off + (slot * gridDim.y + seq_idx) * src0_nb3 + head_idx * src0_nb2 + head_off * d_state);
#pragma unroll
for (int j = 0; j < c_factor; j++) {
s_snapshot_warp[WARP_SIZE * j + lane] = state[j];
}
}
}
// write back the state
@ -232,7 +242,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
const int src2_nb2, const int src3_nb1, const int src4_nb2, const int src4_nb3, const int src5_nb2,
const int src5_nb3, const int64_t s_off, const int64_t d_state, const int64_t head_dim,
const int64_t n_head, const int64_t n_group, const int64_t n_tok, const int64_t n_seq,
cudaStream_t stream) {
const int64_t K, cudaStream_t stream) {
// NOTE: if you change conditions here, be sure to update the corresponding supports_op condition!
if (src3_nb1 == sizeof(float)) {
// Mamba-2
@ -245,7 +255,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
ggml_cuda_kernel_launch(ssm_scan_f32_group<128/WARP_SIZE, 128>, launch_params,
src0, src1, src2, src3, src4, src5, src6, dst,
src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok);
src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K);
} else if (d_state == 256) { // Falcon-H1
constexpr int threads = 256;
constexpr int num_warps = threads/WARP_SIZE;
@ -255,12 +265,13 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
ggml_cuda_kernel_launch(ssm_scan_f32_group<256/WARP_SIZE, 256>, launch_params,
src0, src1, src2, src3, src4, src5, src6, dst,
src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok);
src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K);
} else {
GGML_ABORT("doesn't support d_state!=(128 or 256).");
}
} else {
// Mamba-1
GGML_ASSERT(K == 1);
constexpr int threads = 128;
GGML_ASSERT(n_head % threads == 0);
GGML_ASSERT(head_dim == 1);
@ -769,10 +780,12 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const int64_t ng = src4->ne[1]; // n_group
const int64_t n_t = src1->ne[2]; // number of tokens per sequence
const int64_t n_s = src1->ne[3]; // number of sequences in the batch
const int32_t K_param = ggml_get_op_params_i32(dst, 0);
const int64_t K = K_param > 0 ? K_param : 1;
const int64_t s_off = ggml_nelements(src1) * sizeof(float);
GGML_ASSERT(ggml_nelements(src1) + nc*nr*nh*n_s == ggml_nelements(dst));
GGML_ASSERT(ggml_nelements(src1) + K*nc*nr*nh*n_s == ggml_nelements(dst));
GGML_ASSERT(src0->nb[0] == sizeof(float));
GGML_ASSERT(src1->nb[0] == sizeof(float));
GGML_ASSERT(src2->nb[0] == sizeof(float));
@ -780,6 +793,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
GGML_ASSERT(src4->nb[0] == sizeof(float));
GGML_ASSERT(src5->nb[0] == sizeof(float));
GGML_ASSERT(src6->nb[0] == sizeof(int32_t));
GGML_ASSERT(src3->ne[0] == 1 || K == 1);
const float * src0_d = (const float *) src0->data;
const float * src1_d = (const float *) src1->data;
@ -814,6 +828,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const bool is_mamba2 = (src3->nb[1] == sizeof(float));
const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
const bool use_ssd = is_mamba2 && n_t > SSM_SSD_MIN_TOKENS
&& K == 1
&& n_t <= SSM_SSD_MAX_TOKENS
&& GGML_CUDA_CC_IS_NVIDIA(cc)
&& cc >= GGML_CUDA_CC_TURING
@ -841,5 +856,5 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d,
src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2],
src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3],
s_off, nc, nr, nh, ng, n_t, n_s, stream);
s_off, nc, nr, nh, ng, n_t, n_s, K, stream);
}

56
ggml/src/ggml-cuda/wkv.cu
View File

@ -141,6 +141,57 @@ static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, cons
}
}
template <int rows_per_block>
static __global__ void __launch_bounds__(WARP_SIZE * rows_per_block, 2)
rwkv_wkv7_f32_t1_warp_row(const int T, const int C, const int H, const float * r, const float * w, const float * k, const float * v, const float * a, const float * b, const float * s, float * dst) {
constexpr int head_size = CUDA_WKV_BLOCK_SIZE;
constexpr int half_head = head_size / 2;
const int lane = threadIdx.x;
const int row = blockIdx.y * rows_per_block + threadIdx.y;
const int bid = blockIdx.x;
const int batch_i = bid / H;
const int head_i = bid % H;
const int state_size = C * head_size;
const int head_off = head_i * head_size;
const int t = batch_i * C + head_off + row;
__shared__ float _r[head_size], _w[head_size], _k[head_size], _a[head_size], _b[head_size];
if (threadIdx.y == 0) {
_r[lane] = r[batch_i * C + head_off + lane];
_w[lane] = w[batch_i * C + head_off + lane];
_k[lane] = k[batch_i * C + head_off + lane];
_a[lane] = a[batch_i * C + head_off + lane];
_b[lane] = b[batch_i * C + head_off + lane];
_r[lane + half_head] = r[batch_i * C + head_off + lane + half_head];
_w[lane + half_head] = w[batch_i * C + head_off + lane + half_head];
_k[lane + half_head] = k[batch_i * C + head_off + lane + half_head];
_a[lane + half_head] = a[batch_i * C + head_off + lane + half_head];
_b[lane + half_head] = b[batch_i * C + head_off + lane + half_head];
}
__syncthreads();
const int64_t state_base = batch_i * state_size + head_i * head_size * head_size + row * head_size;
const float s0 = s[state_base + lane];
const float s1 = s[state_base + lane + half_head];
const float sa = warp_reduce_sum(_a[lane] * s0 + _a[lane + half_head] * s1);
const float vt = v[t];
const float st0 = s0 * _w[lane] + _k[lane] * vt + sa * _b[lane];
const float st1 = s1 * _w[lane + half_head] + _k[lane + half_head] * vt + sa * _b[lane + half_head];
const float y = warp_reduce_sum(st0 * _r[lane] + st1 * _r[lane + half_head]);
dst[T * C + state_base + lane] = st0;
dst[T * C + state_base + lane + half_head] = st1;
if (lane == 0) {
dst[t] = y;
}
}
void ggml_cuda_op_rwkv_wkv6(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const float * k_d = (const float *)dst->src[0]->data;
const float * v_d = (const float *)dst->src[1]->data;
@ -191,7 +242,10 @@ void ggml_cuda_op_rwkv_wkv7(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
GGML_ASSERT(C % H == 0);
GGML_ASSERT(C / H == CUDA_WKV_BLOCK_SIZE || C / H == CUDA_WKV_BLOCK_SIZE * 2);
if (C / H == CUDA_WKV_BLOCK_SIZE) {
if (T / B == 1 && C / H == CUDA_WKV_BLOCK_SIZE) {
constexpr int rows_per_block = 4;
rwkv_wkv7_f32_t1_warp_row<rows_per_block><<<dim3(B * H, CUDA_WKV_BLOCK_SIZE / rows_per_block), dim3(WARP_SIZE, rows_per_block), 0, stream>>>(T, C, H, r_d, w_d, k_d, v_d, a_d, b_d, s_d, dst_d);
} else if (C / H == CUDA_WKV_BLOCK_SIZE) {
rwkv_wkv7_f32<CUDA_WKV_BLOCK_SIZE><<<B * H, C / H, 0, stream>>>(B, T, C, H, r_d, w_d, k_d, v_d, a_d, b_d, s_d, dst_d);
} else {
rwkv_wkv7_f32<CUDA_WKV_BLOCK_SIZE * 2><<<B * H, C / H, 0, stream>>>(B, T, C, H, r_d, w_d, k_d, v_d, a_d, b_d, s_d, dst_d);

15
ggml/src/ggml-et/et-kernels/src/ssm_scan_f32.c
View File

@ -12,7 +12,8 @@ struct ggml_et_ssm_scan_params {
struct ggml_tensor src4; // B: [d_state, n_group, n_seq_tokens, n_seqs]
struct ggml_tensor src5; // C: [d_state, n_group, n_seq_tokens, n_seqs]
struct ggml_tensor src6; // ids: [n_seqs] i32
struct ggml_tensor dst; // packed [y, final_state]
struct ggml_tensor dst; // packed [y, states]
int32_t K;
};
static inline float softplus_f32(float x) {
@ -72,6 +73,7 @@ int entry_point(struct ggml_et_ssm_scan_params * params, void * env) {
const int64_t n_seq_tokens = src1->ne[2];
const int64_t n_seqs = src1->ne[3];
const int64_t y_elems = src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3];
const int64_t K = params->K;
if (src0->nb[0] != sizeof(float) || src1->nb[0] != sizeof(float) || src2->nb[0] != sizeof(float) ||
src3->nb[0] != sizeof(float) || src4->nb[0] != sizeof(float) || src5->nb[0] != sizeof(float) ||
@ -79,7 +81,7 @@ int entry_point(struct ggml_et_ssm_scan_params * params, void * env) {
return -1;
}
if (n_group <= 0 || n_head % n_group != 0) {
if (K < 1 || n_group <= 0 || n_head % n_group != 0) {
return -1;
}
@ -260,6 +262,15 @@ int entry_point(struct ggml_et_ssm_scan_params * params, void * env) {
sumf += st * C_row[state_idx];
}
const int64_t slot = n_seq_tokens - 1 - token_idx;
if (slot > 0 && slot < K) {
float * state_snapshot =
(float *) ((char *) state_dst + (size_t) slot * n_seqs * src0->nb[3]);
for (int64_t i = 0; i < d_state; ++i) {
state_snapshot[i] = state_dst[i];
}
}
dst_data[seq_idx * (n_seq_tokens * n_head * head_dim) + token_idx * (n_head * head_dim) +
head_idx * head_dim + dim_idx] = sumf;
}

1
ggml/src/ggml-et/ggml-et-ops.cpp
View File

@ -2064,6 +2064,7 @@ bool ggml_et_op_ssm_scan(ggml_backend_et_device_context * dev_ctx, const ggml_te
params.src5 = *node->src[5];
params.src6 = *node->src[6];
params.dst = *node;
params.K = ggml_get_op_params_i32(node, 0);
bool kernel_result = ggml_et_launch_kernel(dev_ctx, "ssm_scan_f32", &params, sizeof(params), 0xFFFFFFFF);

3
ggml/src/ggml-et/ggml-et-ops.h
View File

@ -218,7 +218,8 @@ struct ggml_et_ssm_scan_params {
ggml_tensor src4; // B: [d_state, n_group, n_seq_tokens, n_seqs]
ggml_tensor src5; // C: [d_state, n_group, n_seq_tokens, n_seqs]
ggml_tensor src6; // ids: [n_seqs] i32
ggml_tensor dst; // [y, final_state] packed output from ggml_ssm_scan()
ggml_tensor dst; // [y, states] packed output from ggml_ssm_scan()
int32_t K;
};
struct ggml_et_rwkv_wkv6_params {

1
ggml/src/ggml-et/ggml-et.cpp
View File

@ -1646,6 +1646,7 @@ static void ggml_backend_et_device_get_props(ggml_backend_dev_t dev, struct ggml
/* .host_buffer = */ false,
/* .buffer_from_host_ptr = */ false,
/* .events = */ false,
/* .mmap_support = */ true,
};
}

166
ggml/src/ggml-feats.h Normal file
View File

@ -0,0 +1,166 @@
#pragma once
#if defined(__aarch64__) || defined(_M_ARM64)
#if defined(__linux__)
#include <sys/auxv.h>
#include <sys/prctl.h>
#if !defined(HWCAP2_SVE2)
#define HWCAP2_SVE2 (1ULL << 1)
#endif
#if !defined(HWCAP_FPHP)
#define HWCAP_FPHP (1 << 9)
#endif
#if !defined(HWCAP_ASIMDHP)
#define HWCAP_ASIMDHP (1 << 10)
#endif
#if !defined(HWCAP2_I8MM)
#define HWCAP2_I8MM (1ULL << 13)
#endif
#if !defined(HWCAP_ASIMDDP)
#define HWCAP_ASIMDDP (1 << 20)
#endif
#if !defined(HWCAP_SVE)
#define HWCAP_SVE (1 << 22)
#endif
#if !defined(HWCAP2_SME)
#define HWCAP2_SME (1ULL << 23)
#endif
#if !defined(HWCAP2_SME2)
#define HWCAP2_SME2 (1ULL << 37)
#endif
#if !defined(PR_SVE_GET_VL)
#define PR_SVE_GET_VL 51
#endif
#if !defined(PR_SVE_VL_LEN_MASK)
#define PR_SVE_VL_LEN_MASK 0xffff
#endif
#elif defined(__APPLE__)
#include <sys/sysctl.h>
#elif defined(_WIN32)
#include <windows.h>
#if !defined(PF_ARM_V82_DP_INSTRUCTIONS_AVAILABLE)
#define PF_ARM_V82_DP_INSTRUCTIONS_AVAILABLE 43
#endif
#if !defined(PF_ARM_SVE_INSTRUCTIONS_AVAILABLE)
#define PF_ARM_SVE_INSTRUCTIONS_AVAILABLE 46
#endif
#if !defined(PF_ARM_SVE2_INSTRUCTIONS_AVAILABLE)
#define PF_ARM_SVE2_INSTRUCTIONS_AVAILABLE 47
#endif
#if !defined(PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE)
#define PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE 66
#endif
#if !defined(PF_ARM_V82_FP16_INSTRUCTIONS_AVAILABLE)
#define PF_ARM_V82_FP16_INSTRUCTIONS_AVAILABLE 67
#endif
#if !defined(PF_ARM_SME_INSTRUCTIONS_AVAILABLE)
#define PF_ARM_SME_INSTRUCTIONS_AVAILABLE 70
#endif
#if !defined(PF_ARM_SME2_INSTRUCTIONS_AVAILABLE)
#define PF_ARM_SME2_INSTRUCTIONS_AVAILABLE 71
#endif
#endif
typedef struct ggml_feats_arch64_runtime {
bool has_dotprod;
bool has_fp16;
bool has_sve;
bool has_sve2;
bool has_i8mm;
bool has_sme;
bool has_sme2;
int sve_cnt;
} ggml_feats_arch64_runtime_t;
static inline ggml_feats_arch64_runtime_t ggml_feats_get_arch64_runtime(void) {
ggml_feats_arch64_runtime_t runtime_feat = {};
#if defined(__linux__)
const unsigned long hwcap = getauxval(AT_HWCAP);
const unsigned long hwcap2 = getauxval(AT_HWCAP2);
runtime_feat.has_dotprod = !!(hwcap & HWCAP_ASIMDDP);
runtime_feat.has_fp16 = !!(hwcap & HWCAP_FPHP) && !!(hwcap & HWCAP_ASIMDHP);;
runtime_feat.has_sve = !!(hwcap & HWCAP_SVE);
runtime_feat.has_sve2 = !!(hwcap2 & HWCAP2_SVE2);
runtime_feat.has_i8mm = !!(hwcap2 & HWCAP2_I8MM);
runtime_feat.has_sme = !!(hwcap2 & HWCAP2_SME);
runtime_feat.has_sme2 = !!(hwcap2 & HWCAP2_SME2);
if (runtime_feat.has_sve) {
const int vl = prctl(PR_SVE_GET_VL);
if (vl >= 0) {
runtime_feat.sve_cnt = vl & PR_SVE_VL_LEN_MASK;
}
}
#elif defined(__APPLE__)
int oldp = 0;
size_t size = sizeof(oldp);
if (sysctlbyname("hw.optional.arm.FEAT_DotProd", &oldp, &size, nullptr, 0) == 0) {
runtime_feat.has_dotprod = static_cast<bool>(oldp);
}
if (sysctlbyname("hw.optional.arm.FEAT_FP16", &oldp, &size, nullptr, 0) == 0) {
runtime_feat.has_fp16 = static_cast<bool>(oldp);
}
if (sysctlbyname("hw.optional.arm.FEAT_SVE", &oldp, &size, nullptr, 0) == 0) {
runtime_feat.has_sve = static_cast<bool>(oldp);
}
if (sysctlbyname("hw.optional.arm.FEAT_SVE2", &oldp, &size, nullptr, 0) == 0) {
runtime_feat.has_sve2 = static_cast<bool>(oldp);
}
if (sysctlbyname("hw.optional.arm.FEAT_I8MM", &oldp, &size, nullptr, 0) == 0) {
runtime_feat.has_i8mm = static_cast<bool>(oldp);
}
if (sysctlbyname("hw.optional.arm.FEAT_SME", &oldp, &size, nullptr, 0) == 0) {
runtime_feat.has_sme = static_cast<bool>(oldp);
}
if (sysctlbyname("hw.optional.arm.FEAT_SME2", &oldp, &size, nullptr, 0) == 0) {
runtime_feat.has_sme2 = static_cast<bool>(oldp);
}
// Apple does not support userspace non-streaming SVE; keep SVE vector length unknown.
runtime_feat.sve_cnt = 0;
#elif defined (_WIN32)
runtime_feat.has_dotprod = IsProcessorFeaturePresent(PF_ARM_V82_DP_INSTRUCTIONS_AVAILABLE) != 0;
runtime_feat.has_fp16 = IsProcessorFeaturePresent(PF_ARM_V82_FP16_INSTRUCTIONS_AVAILABLE) != 0;
runtime_feat.has_sve = IsProcessorFeaturePresent(PF_ARM_SVE_INSTRUCTIONS_AVAILABLE) != 0;
runtime_feat.has_sve2 = IsProcessorFeaturePresent(PF_ARM_SVE2_INSTRUCTIONS_AVAILABLE) != 0;
runtime_feat.has_i8mm = IsProcessorFeaturePresent(PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE) != 0;
runtime_feat.has_sme = IsProcessorFeaturePresent(PF_ARM_SME_INSTRUCTIONS_AVAILABLE) != 0;
runtime_feat.has_sme2 = IsProcessorFeaturePresent(PF_ARM_SME2_INSTRUCTIONS_AVAILABLE) != 0;
// Windows exposes SVE feature presence, but not the runtime SVE vector length here.
runtime_feat.sve_cnt = 0;
#endif
return runtime_feat;
}
#endif // defined(__aarch64__) || defined(_M_ARM64)

1
ggml/src/ggml-hexagon/ggml-hexagon.cpp
View File

@ -3930,6 +3930,7 @@ static void ggml_backend_hexagon_device_get_props(ggml_backend_dev_t dev, struct
/* .host_buffer = */ (bool) opt_hostbuf,
/* .buffer_from_host_ptr = */ false,
/* .events = */ false,
/* .mmap_support = */ false,
};
}

3
ggml/src/ggml-hip/CMakeLists.txt
View File

@ -126,9 +126,6 @@ if (GGML_HIP_EXPORT_METRICS)
set(CMAKE_HIP_FLAGS "${CMAKE_HIP_FLAGS} -Rpass-analysis=kernel-resource-usage --save-temps")
endif()
# Fast math for HIP, like CUDA's -use_fast_math. Not -ffast-math: that implies -ffinite-math-only, which breaks ggml's INFINITY masking and produces NaNs.
set(CMAKE_HIP_FLAGS "${CMAKE_HIP_FLAGS} -funsafe-math-optimizations")
if (NOT GGML_CUDA_FA)
add_compile_definitions(GGML_CUDA_NO_FA)
endif()

10
ggml/src/ggml-metal/ggml-metal-device.cpp
View File

@ -953,6 +953,11 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta
nr0 = N_R0_IQ4_XS;
smem = 32*sizeof(float);
} break;
case GGML_TYPE_TQ2_0:
{
nsg = N_SG_TQ2_0;
nr0 = N_R0_TQ2_0;
} break;
default:
{
GGML_LOG_ERROR("Asserting on type %d\n", (int) tsrc0);
@ -1182,6 +1187,11 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id(ggml_m
nr0 = N_R0_IQ4_XS;
smem = 32*sizeof(float);
} break;
case GGML_TYPE_TQ2_0:
{
nsg = N_SG_TQ2_0;
nr0 = N_R0_TQ2_0;
} break;
default:
{
GGML_LOG_ERROR("Asserting on type %d\n", (int)op->src[2]->type);

9
ggml/src/ggml-metal/ggml-metal-device.m
View File

@ -1268,8 +1268,9 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
case GGML_OP_ARGSORT:
case GGML_OP_TOP_K:
case GGML_OP_ARANGE:
case GGML_OP_ROLL:
return true;
case GGML_OP_ROLL:
return ggml_is_contiguous(op->src[0]);
case GGML_OP_FLASH_ATTN_EXT:
// for new head sizes, add checks here
if (op->src[0]->ne[0] != 32 &&
@ -1375,9 +1376,10 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
ggml_is_contiguous_rows(op->src[1]) &&
ggml_is_contiguous_rows(op->src[2]) &&
ggml_is_contiguous_rows(op->src[3]);
case GGML_OP_SSM_CONV:
case GGML_OP_SSM_SCAN:
return has_simdgroup_reduction;
case GGML_OP_SSM_CONV:
return has_simdgroup_reduction;
case GGML_OP_RWKV_WKV6:
case GGML_OP_RWKV_WKV7:
return true;
@ -1406,6 +1408,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
case GGML_TYPE_Q5_0:
case GGML_TYPE_Q5_1:
case GGML_TYPE_IQ4_NL:
case GGML_TYPE_TQ2_0:
case GGML_TYPE_I32:
return true;
default:
@ -1434,6 +1437,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
case GGML_TYPE_Q5_0:
case GGML_TYPE_Q5_1:
case GGML_TYPE_Q8_0:
case GGML_TYPE_TQ2_0:
switch (op->type) {
case GGML_TYPE_F32:
case GGML_TYPE_F16:
@ -1469,6 +1473,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
case GGML_TYPE_Q5_0:
case GGML_TYPE_Q5_1:
case GGML_TYPE_IQ4_NL:
case GGML_TYPE_TQ2_0:
return true;
default:
return false;

4
ggml/src/ggml-metal/ggml-metal-impl.h
View File

@ -87,6 +87,9 @@
#define N_R0_IQ4_XS 2
#define N_SG_IQ4_XS 2
#define N_R0_TQ2_0 4
#define N_SG_TQ2_0 2
// function constants offsets
#define FC_FLASH_ATTN_EXT_PAD 100
#define FC_FLASH_ATTN_EXT_BLK 200
@ -877,6 +880,7 @@ typedef struct {
int64_t n_group;
int64_t n_seq_tokens;
int64_t n_seqs;
int64_t K;
uint64_t s_off;
uint64_t nb00;
uint64_t nb01;

7
ggml/src/ggml-metal/ggml-metal-ops.cpp
View File

@ -1710,6 +1710,10 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) {
const int64_t n_group = ne41;
const int64_t n_seq_tokens = ne12;
const int64_t n_seqs = ne13;
const int64_t K = ggml_get_op_params_i32(op, 0);
GGML_ASSERT(K >= 1);
GGML_ASSERT(ggml_nelements(op->src[1]) + K*d_state*d_inner*n_head*n_seqs == ggml_nelements(op));
ggml_metal_kargs_ssm_scan args = {
/*.d_state =*/ d_state,
@ -1718,6 +1722,7 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) {
/*.n_group =*/ n_group,
/*.n_seq_tokens =*/ n_seq_tokens,
/*.n_seqs =*/ n_seqs,
/*.K =*/ K,
/*.s_off =*/ ggml_nelements(op->src[1]) * sizeof(float),
/*.nb00 =*/ nb00,
/*.nb01 =*/ nb01,
@ -3816,7 +3821,7 @@ int ggml_metal_op_norm(ggml_metal_op_t ctx, int idx) {
}
nth = std::min(nth, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
nth = std::min(nth, args.ne00_t);
nth = std::min(nth, (args.ne00_t + 31)/32*32);
const size_t smem = pipeline.smem;

1
ggml/src/ggml-metal/ggml-metal.cpp
View File

@ -681,6 +681,7 @@ static void ggml_backend_metal_device_get_props(ggml_backend_dev_t dev, ggml_bac
/* .host_buffer = */ false,
/* .buffer_from_host_ptr = */ true,
/* .events = */ true,
/* .mmap_support = */ true,
};
}

221
ggml/src/ggml-metal/ggml-metal.metal
View File

@ -468,6 +468,34 @@ void quantize_iq4_nl(device const float * src, device block_iq4_nl & dst) {
dst.d = sumq2 > 0 ? sumqx/sumq2 : d;
}
void quantize_tq2_0(device const float * src, device block_tq2_0 & dst) {
#pragma METAL fp math_mode(safe)
float amax = 0.0f; // absolute max
for (int j = 0; j < QK_K; j++) {
const float v = src[j];
amax = MAX(amax, fabs(v));
}
const float d = amax;
const float id = d ? 1.0f/d : 0.0f;
dst.d = (half) d;
for (int j = 0; j < QK_K/4; j += 32) {
for (int m = 0; m < 32; ++m) {
uint8_t q = 0;
for (int n = 0; n < 4; ++n) {
// -1, 0, 1 -> 0, 1, 2
int xi = (int)round(src[m + n*32] * id) + 1;
q += (uint8_t)((xi & 3) << (2*n));
}
dst.qs[j + m] = q;
}
src += 4*32;
}
}
template <typename type4x4>
void dequantize_q4_1(device const block_q4_1 * xb, short il, thread type4x4 & reg) {
device const uint16_t * qs = ((device const uint16_t *)xb + 2);
@ -1021,6 +1049,25 @@ void dequantize_iq4_xs(device const block_iq4_xs * xb, short il, thread type4x4
}
}
template <typename type4x4>
void dequantize_tq2_0(device const block_tq2_0 * xb, short il, thread type4x4 & reg) {
device const uint8_t * qs = xb->qs;
const float d = xb->d;
float4x4 reg_f;
// 2 bits per element, 4 elements per byte, 128 elements per 32-byte group
const short base = il * 16;
for (int k = 0; k < 16; k++) {
const int i = base + k;
const int byte = ((i >> 7) & 1) * 32 + (i & 31);
const int l = (i >> 5) & 3;
reg_f[k/4][k%4] = d * (float)(((qs[byte] >> (2*l)) & 3) - 1);
}
reg = (type4x4) reg_f;
}
enum ggml_sort_order {
GGML_SORT_ORDER_ASC,
GGML_SORT_ORDER_DESC,
@ -2382,6 +2429,8 @@ kernel void kernel_ssm_scan_f32(
const int32_t nh = args.n_head;
const int32_t ng = args.n_group;
const int32_t n_t = args.n_seq_tokens;
const int32_t n_s = args.n_seqs;
const int32_t K = args.K;
const int32_t s_off = args.s_off;
@ -2440,6 +2489,12 @@ kernel void kernel_ssm_scan_f32(
// recurse
s0 = s;
const int32_t slot = n_t - 1 - (i2 + t);
if (slot > 0 && slot < K) {
device float * s_snapshot = (device float *) ((device char *) s_buff + (int64_t) slot*n_s*args.nb03);
s_snapshot[i] = s;
}
B += args.ns42;
C += args.ns52;
}
@ -8001,6 +8056,7 @@ template [[host_name("kernel_cpy_f32_q4_1")]] kernel cpy_f_q_t kernel_cpy_f32_
template [[host_name("kernel_cpy_f32_q5_0")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK5_0, block_q5_0, quantize_q5_0>;
template [[host_name("kernel_cpy_f32_q5_1")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK5_1, block_q5_1, quantize_q5_1>;
template [[host_name("kernel_cpy_f32_iq4_nl")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK4_NL, block_iq4_nl, quantize_iq4_nl>;
template [[host_name("kernel_cpy_f32_tq2_0")]] kernel cpy_f_q_t kernel_cpy_f32_q<QK_K, block_tq2_0, quantize_tq2_0>;
template<typename T4x4, typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread T4x4 &)>
kernel void kernel_cpy_q_f32(
@ -8048,6 +8104,8 @@ template [[host_name("kernel_cpy_q5_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<
template [[host_name("kernel_cpy_q5_1_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q5_1, 2, dequantize_q5_1>;
template [[host_name("kernel_cpy_q8_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q8_0, 2, dequantize_q8_0>;
template [[host_name("kernel_cpy_tq2_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_tq2_0, QK_NL, dequantize_tq2_0>;
template [[host_name("kernel_cpy_q1_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q1_0, 8, dequantize_q1_0>;
template [[host_name("kernel_cpy_q2_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q2_0, 4, dequantize_q2_0>;
template [[host_name("kernel_cpy_q4_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q4_0, 2, dequantize_q4_0>;
@ -8056,6 +8114,8 @@ template [[host_name("kernel_cpy_q5_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<
template [[host_name("kernel_cpy_q5_1_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q5_1, 2, dequantize_q5_1>;
template [[host_name("kernel_cpy_q8_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q8_0, 2, dequantize_q8_0>;
template [[host_name("kernel_cpy_tq2_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_tq2_0, QK_NL, dequantize_tq2_0>;
template<typename T>
kernel void kernel_concat(
constant ggml_metal_kargs_concat & args,
@ -9822,6 +9882,121 @@ kernel void kernel_mul_mv_mxfp4_f32(
kernel_mul_mv_mxfp4_f32_impl<N_R0_MXFP4, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
}
template<int nr0, typename args_t>
void kernel_mul_mv_tq2_0_f32_impl(
args_t args,
device const char * src0,
device const char * src1,
device char * dst,
threadgroup char * shmem,
uint3 tgpig,
ushort tiisg,
ushort sgitg) {
const short NSG = FC_mul_mv_nsg;
const int nb = args.ne00/QK_K;
const int r0 = tgpig.x;
const int r1 = tgpig.y;
const int im = tgpig.z;
const int first_row = (r0 * NSG + sgitg) * nr0;
const uint i12 = im%FC_mul_mv_ne12;
const uint i13 = im/FC_mul_mv_ne12;
const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13;
device const float * y = (device const float *) (src1 + offset1);
device const block_tq2_0 * ax[nr0];
for (int row = 0; row < nr0; ++row) {
const uint64_t offset0 = (first_row + row)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
ax[row] = (device const block_tq2_0 *) ((device char *) src0 + offset0);
}
float sumf[nr0] = {0.f};
// 8 threads per block, NBLOCK blocks per pass, 2 halves per block per pass
constexpr short NBLOCK = 4;
constexpr short NB = N_SIMDWIDTH/NBLOCK; // threads per block
const short blk = tiisg / NB; // 0..NBLOCK-1, block handled by this thread
const short htg = tiisg % NB; // 0..NB-1, thread within block (0..7)
// byte and y base offsets within the block (32 elements per thread, 4 per byte)
device const float4 * yb4 = (device const float4 *)(y + 4*htg + blk*QK_K);
// hoisted per-byte coefficients (from y) and total y-sum, shared across rows
// ref: https://github.com/ggml-org/llama.cpp/pull/26980
float4 coef[4];
for (int ib = blk; ib < nb; ib += NBLOCK) {
FOR_UNROLL (short h0 = 0; h0 < 2; ++h0) {
const float4 y0 = yb4[ 0 + 32*h0];
const float4 y1 = yb4[ 8 + 32*h0];
const float4 y2 = yb4[16 + 32*h0];
const float4 y3 = yb4[24 + 32*h0];
float sumy = 0.f;
FOR_UNROLL (short j = 0; j < 4; ++j) {
coef[j] = float4(
y0[j],
y1[j] - 4.0f*y0[j],
y2[j] - 4.0f*y1[j],
y3[j] - 4.0f*y2[j]);
sumy += (y0[j] + y1[j]) + (y2[j] + y3[j]);
}
FOR_UNROLL (short row = 0; row < nr0; ++row) {
device const block_tq2_0 & xb = ax[row][ib];
device const uchar * qs = xb.qs + 4*htg + 32*h0;
float sum = -sumy;
FOR_UNROLL (short j = 0; j < 4; ++j) {
// express the 2-bit field shifts (v>>2, v>>4, v>>6) as float floor ops
const float v = (float)qs[j];
const float f0 = v;
const float f1 = floor(v*0.25f); // v>>2
const float f2 = floor(v*0.0625); // v>>4
const float f3 = floor(v*0.015625); // v>>6
sum += coef[j][0]*f0 + coef[j][1]*f1 + coef[j][2]*f2 + coef[j][3]*f3;
}
sumf[row] += xb.d * sum;
}
}
yb4 += QK_K * NBLOCK / 4;
}
device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
for (int row = 0; row < nr0; ++row) {
const float tot = simd_sum(sumf[row]);
if (tiisg == 0 && first_row + row < args.ne01) {
dst_f32[first_row + row] = tot;
}
}
}
[[host_name("kernel_mul_mv_tq2_0_f32")]]
kernel void kernel_mul_mv_tq2_0_f32(
constant ggml_metal_kargs_mul_mv & args,
device const char * src0,
device const char * src1,
device char * dst,
uint3 tgpig[[threadgroup_position_in_grid]],
ushort tiisg[[thread_index_in_simdgroup]],
ushort sgitg[[simdgroup_index_in_threadgroup]]) {
kernel_mul_mv_tq2_0_f32_impl<N_R0_TQ2_0, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
}
template<typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread float4x4 &)>
kernel void kernel_get_rows_q(
constant ggml_metal_kargs_get_rows & args,
@ -9915,6 +10090,38 @@ template [[host_name("kernel_get_rows_iq1_s")]] kernel get_rows_q_t kernel_get
template [[host_name("kernel_get_rows_iq1_m")]] kernel get_rows_q_t kernel_get_rows_q<block_iq1_m, QK_NL, dequantize_iq1_m>;
template [[host_name("kernel_get_rows_iq4_nl")]] kernel get_rows_q_t kernel_get_rows_q<block_iq4_nl, 2, dequantize_iq4_nl>;
template [[host_name("kernel_get_rows_iq4_xs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq4_xs, QK_NL, dequantize_iq4_xs>;
template [[host_name("kernel_get_rows_tq2_0")]] kernel get_rows_q_t kernel_get_rows_q<block_tq2_0, QK_NL, dequantize_tq2_0>;
template<typename TS, typename TI, short QK, typename block_q, void (*quantize_func)(device const float *, device block_q &)>
kernel void kernel_set_rows_q(
constant ggml_metal_kargs_set_rows & args,
device const void * src0,
device const void * src1,
device float * dst,
uint3 tgpig[[threadgroup_position_in_grid]],
uint tiitg[[thread_index_in_threadgroup]],
uint3 tptg [[threads_per_threadgroup]]) {
const int32_t i03 = tgpig.z;
const int32_t i02 = tgpig.y;
const int32_t i12 = i03%args.ne12;
const int32_t i11 = i02%args.ne11;
const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x;
if (i01 >= args.ne01) {
return;
}
const int32_t i10 = i01;
const TI i1 = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
device block_q * dst_row = ( device block_q *) (( device char *) dst + i1*args.nb1 + i02*args.nb2 + i03*args.nb3);
const device TS * src_row = (const device TS *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
quantize_func(src_row + QK*ind, dst_row[ind]);
}
}
template<typename TS, typename TI, typename block_q, void (*quantize_func)(device const float *, device block_q &)>
kernel void kernel_set_rows_q32(
@ -10011,6 +10218,11 @@ template [[host_name("kernel_set_rows_f32_i32_q5_1")]] kernel set_rows_q32_t k
template [[host_name("kernel_set_rows_f32_i64_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_iq4_nl, quantize_iq4_nl>;
template [[host_name("kernel_set_rows_f32_i32_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_iq4_nl, quantize_iq4_nl>;
typedef decltype(kernel_set_rows_q<float, int64_t, QK_K, block_tq2_0, quantize_tq2_0>) set_rows_qK_t;
template [[host_name("kernel_set_rows_f32_i64_tq2_0")]] kernel set_rows_qK_t kernel_set_rows_q<float, int64_t, QK_K, block_tq2_0, quantize_tq2_0>;
template [[host_name("kernel_set_rows_f32_i32_tq2_0")]] kernel set_rows_qK_t kernel_set_rows_q<float, int32_t, QK_K, block_tq2_0, quantize_tq2_0>;
kernel void kernel_diag_f32(
constant ggml_metal_kargs_diag & args,
device const char * src0,
@ -10786,6 +10998,7 @@ template [[host_name("kernel_mul_mm_iq1_s_f32")]] kernel mul_mm_t kernel_mul_m
template [[host_name("kernel_mul_mm_iq1_m_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq1_m, QK_NL, dequantize_iq1_m, float, float4x4, float, float2x4>;
template [[host_name("kernel_mul_mm_iq4_nl_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq4_nl, 2, dequantize_iq4_nl, float, float4x4, float, float2x4>;
template [[host_name("kernel_mul_mm_iq4_xs_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq4_xs, QK_NL, dequantize_iq4_xs, float, float4x4, float, float2x4>;
template [[host_name("kernel_mul_mm_tq2_0_f32")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_tq2_0, QK_NL, dequantize_tq2_0, float, float4x4, float, float2x4>;
template [[host_name("kernel_mul_mm_f32_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, float4x4, 1, dequantize_f32, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_f16_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, half4x4, 1, dequantize_f16, half, half4x4, half, half2x4>;
@ -10811,6 +11024,7 @@ template [[host_name("kernel_mul_mm_iq1_s_f16")]] kernel mul_mm_t kernel_mul_m
template [[host_name("kernel_mul_mm_iq1_m_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq1_m, QK_NL, dequantize_iq1_m, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_iq4_nl_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq4_nl, 2, dequantize_iq4_nl, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_iq4_xs_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq4_xs, QK_NL, dequantize_iq4_xs, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_tq2_0_f16")]] kernel mul_mm_t kernel_mul_mm<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_tq2_0, QK_NL, dequantize_tq2_0, float, float4x4, half, half2x4>;
//
// indirect matrix-matrix multiplication
@ -10845,6 +11059,7 @@ template [[host_name("kernel_mul_mm_id_iq1_s_f32")]] kernel mul_mm_id kernel_m
template [[host_name("kernel_mul_mm_id_iq1_m_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq1_m, QK_NL, dequantize_iq1_m, float, float4x4, float, float2x4>;
template [[host_name("kernel_mul_mm_id_iq4_nl_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq4_nl, 2, dequantize_iq4_nl, float, float4x4, float, float2x4>;
template [[host_name("kernel_mul_mm_id_iq4_xs_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq4_xs, QK_NL, dequantize_iq4_xs, float, float4x4, float, float2x4>;
template [[host_name("kernel_mul_mm_id_tq2_0_f32")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_tq2_0, QK_NL, dequantize_tq2_0, float, float4x4, float, float2x4>;
template [[host_name("kernel_mul_mm_id_f32_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, float4x4, 1, dequantize_f32, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_id_f16_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, half4x4, 1, dequantize_f16, half, half4x4, half, half2x4>;
@ -10870,6 +11085,7 @@ template [[host_name("kernel_mul_mm_id_iq1_s_f16")]] kernel mul_mm_id kernel_m
template [[host_name("kernel_mul_mm_id_iq1_m_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq1_m, QK_NL, dequantize_iq1_m, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_id_iq4_nl_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq4_nl, 2, dequantize_iq4_nl, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_id_iq4_xs_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq4_xs, QK_NL, dequantize_iq4_xs, float, float4x4, half, half2x4>;
template [[host_name("kernel_mul_mm_id_tq2_0_f16")]] kernel mul_mm_id kernel_mul_mm_id<half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_tq2_0, QK_NL, dequantize_tq2_0, float, float4x4, half, half2x4>;
//
// matrix-vector multiplication
@ -11027,6 +11243,7 @@ template [[host_name("kernel_mul_mv_id_iq3_s_f32")]] kernel kernel_mul_mv_id_t
template [[host_name("kernel_mul_mv_id_iq2_s_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_s_f32_impl <N_R0_IQ2_S>>>;
template [[host_name("kernel_mul_mv_id_iq4_nl_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq4_nl_f32_impl <N_R0_IQ4_NL>>>;
template [[host_name("kernel_mul_mv_id_iq4_xs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq4_xs_f32_impl <N_R0_IQ4_XS>>>;
template [[host_name("kernel_mul_mv_id_tq2_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_tq2_0_f32_impl <N_R0_TQ2_0>>>;
kernel void kernel_pool_2d_max_f32(
constant ggml_metal_kargs_pool_2d & args,
@ -11328,8 +11545,8 @@ kernel void kernel_lightning_indexer(
const int i_kv_0 = tgpig.x*NK; // first key of this threadgroup
const int i_kv = i_kv_0 + sgitg*NKPSG; // first key of this simdgroup
threadgroup half4x4 sk4x4[NK*DK16];
threadgroup half * sk = (threadgroup half *) sk4x4;
threadgroup half sk[NK * DK16 * 16];
threadgroup half4x4 * sk4x4 = (threadgroup half4x4 *) sk;
for (short i = tiitg; i < NK*DK16; i += NTG) {
const short ik = i/DK16;

48
ggml/src/ggml-opencl/ggml-opencl.cpp
View File

@ -73,6 +73,7 @@ typedef const void * (*get_adreno_bin_kernel_func_t)(
//------------------------------------------------------------------------------
bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor);
static bool ggml_cl_is_q4_0_soa(const ggml_tensor * tensor);
static bool ggml_cl_is_q8_0_soa(const ggml_tensor * tensor);
static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst);
@ -4629,6 +4630,23 @@ static std::string ggml_opencl_fa_compile_opts(ggml_backend_opencl_context * bac
if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) {
opts += " -D FA_C8_NO_SG_PIN";
}
// Transposed K tile in local memory: the KV rows the QK loop walks together become
// adjacent, so a group of them is ONE 128-bit local read instead of several narrow
// ones. The QK loop is LDS-read-issue-bound (a wrong-math probe that kept every FMA/dp4a
// but removed the LDS reads ran the kernel ~40% faster), so this is worth up to +26% on
// fa=1 prefill. Output is bit-identical -- only the layout moves.
//
// DK <= 128 only. At DK=256 (gemma-3-4b) it measures 1-2% NEGATIVE and reproduces across
// rounds; padding the row stride does not recover it, so the cause is not a simple bank
// conflict and the wider tile does not want this layout.
//
// Default on within that gate; GGML_OPENCL_FA_K_LDS_T=0 restores the row-major tile.
{
const char * e = getenv("GGML_OPENCL_FA_K_LDS_T");
if ((e == nullptr || e[0] != '0') && cfg->dk <= 128) {
opts += " -D FA_K_LDS_T";
}
}
return opts;
}
@ -4911,8 +4929,13 @@ static bool ggml_opencl_ensure_fa_variant(ggml_backend_opencl_context * backend_
const int x = (e && e[0]) ? atoi(e) : 0;
return (x == 8 || x == 16 || x == 32) ? x : 0; // 0 = per-gen default
}();
// X2E needs 16 to keep per-lane o_acc at 128B (the compiler spills the
// kernel-default width); X1E does not spill, but C=16 is still a measured
// +28-30% DK128-GQA4 decode win there (X1-85, kv 4096/8192), neutral on
// DK64 / GQA1 / quant-KV.
const int fa_cl_c_gqa4 = fa_cl_c_env ? fa_cl_c_env
: (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E ? 16 : 0);
: (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E ||
backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E ? 16 : 0);
const std::string opts_cl_c_gqa4 = fa_cl_c_gqa4
? " -D FA_CL_C=" + std::to_string(fa_cl_c_gqa4) : std::string();
const std::string fa_cl_c_g8_val = std::to_string(fa_cl_c_gqa4 ? fa_cl_c_gqa4 * 2 : 16);
@ -7058,6 +7081,19 @@ inline bool enable_adreno_trans_weight(const ggml_backend_opencl_context *backen
return ((elem_num < 128 * 1024 * 1024) && adreno_kernel && shape_ok); // max element num: 2**27
}
inline bool enable_adreno_trans_weight_q5_K(const ggml_backend_opencl_context *backend_ctx, const ggml_tensor *tensor) {
if (!use_adreno_kernels(backend_ctx, tensor)) {
return false;
}
const size_t elem_num = ggml_nelements(tensor);
const size_t q_img_width = elem_num / 8;
const size_t qh_img_width = elem_num / 16;
return q_img_width <= backend_ctx->image_max_buffer_size &&
qh_img_width <= backend_ctx->image_max_buffer_size;
}
static inline bool use_flat_gemv_for_large_m_q4_K(const ggml_tensor *tensor) {
// gemv_noshuffle variant perf drops for large M, use flat variant for large M.
// threshold is well above typical hidden/FFN dims, but below typical vocab sizes.
@ -9237,7 +9273,7 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer,
#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
cl_kernel kernel = backend_ctx->kernel_convert_block_q5_K;
if (use_adreno_kernels(backend_ctx, tensor)) {
if (enable_adreno_trans_weight_q5_K(backend_ctx, tensor)) {
kernel = backend_ctx->kernel_convert_block_q5_K_noshuffle;
}
#else
@ -9272,7 +9308,7 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer,
tensor->extra = extra;
#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
if (use_adreno_kernels(backend_ctx, tensor)) {
if (enable_adreno_trans_weight_q5_K(backend_ctx, tensor)) {
int M = tensor->ne[1];
int K = tensor->ne[0];
@ -10370,7 +10406,7 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
CL_CHECK(clReleaseMemObject(data_device));
return;
}
if (use_adreno_kernels(backend_ctx, tensor)) {
if (enable_adreno_trans_weight_q5_K(backend_ctx, tensor)) {
int M = tensor->ne[1];
int K = tensor->ne[0];
@ -10777,6 +10813,7 @@ static void ggml_backend_opencl_device_get_props(ggml_backend_dev_t dev, struct
/* .host_buffer = */ false,
/* .buffer_from_host_ptr = */ false,
/* .events = */ false,
/* .mmap_support = */ false,
};
}
@ -18909,7 +18946,8 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co
}
// q5_K x fp32
if (src0t == GGML_TYPE_Q5_K && src1t == GGML_TYPE_F32) {
if (src0t == GGML_TYPE_Q5_K && src1t == GGML_TYPE_F32 &&
enable_adreno_trans_weight_q5_K(backend_ctx, src0)) {
ggml_cl_mul_mat_q5_K_f32_adreno(backend, src0, src1, dst);
return;
}

51
ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl
View File

@ -211,7 +211,30 @@ __kernel void FA_TILE_NAME(
float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
#ifdef FA_K_LDS_T
// K tile transposed: [dk vec][kv row] instead of [kv row][dk vec].
//
// The QK loop walks 2 or 4 KV rows at a time against the same dk element. Row-major
// those are DK_VEC half4s apart, so each is its own 64-bit local read. Transposed they
// are adjacent, so a pair is one 128-bit read -- half the LDS issues for the same bytes,
// no extra registers, arithmetic untouched.
//
// This kernel looked like it should be FMA-bound (a half4 mad does ~4 ALU ops per LDS
// read, unlike the 1:1 of the dp4a loop), but it is NOT: a wrong-math probe that kept
// every FMA and removed the LDS reads ran it 38.6% faster (18.92 -> 11.62 ms/op).
// Explicitly 16-byte aligned: FA_LK_PAIR below reads two adjacent half4 as one float4,
// and the element type only obliges the compiler to align this array to 8. The indices
// are even so the offset is a multiple of 16, but the base has to be too, and relying
// on the compiler to over-align it is relying on luck.
__local KV_DATA_TYPE4 l_k[DK_VEC][BLOCK_N] __attribute__((aligned(16)));
#define FA_LK(ROW, C) l_k[C][ROW]
// Two adjacent KV rows as one 128-bit local read (half4 pair == 16 B). j is even and
// BLOCK_N is even, so &l_k[c][j] is 16 B past a 16 B-aligned base.
#define FA_LK_PAIR(C, J) as_half8(*(__local const float4 *)(&l_k[C][J]))
#else
__local KV_DATA_TYPE4 l_k[BLOCK_N][DK_VEC];
#define FA_LK(ROW, C) l_k[ROW][C]
#endif
__local KV_DATA_TYPE4 l_v[BLOCK_N][DV_VEC];
#if N_SPLIT > 1 && !defined(HAS_SUBGROUP_SHUFFLE)
@ -254,17 +277,17 @@ __kernel void FA_TILE_NAME(
#ifdef FA_K_IMG
if (use_kv_pad) {
const ulong k_row_offset = batch_idx * k_tile_nb3 + head_kv_idx * k_tile_nb2 + k_row_idx * k_nb1;
l_k[row][col] = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col];
FA_LK(row, col) = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col];
} else {
const int k_row_px = batch_idx * k_pitch_px_batch + head_kv_idx * k_pitch_px_head + k_row_idx * k_pitch_px_row;
l_k[row][col] = read_imageh(k_img, k_row_px + col);
FA_LK(row, col) = read_imageh(k_img, k_row_px + col);
}
#else
const ulong k_row_offset = batch_idx * k_tile_nb3 + head_kv_idx * k_tile_nb2 + k_row_idx * k_nb1;
l_k[row][col] = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col];
FA_LK(row, col) = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col];
#endif
} else {
l_k[row][col] = (KV_DATA_TYPE4)(0.0h);
FA_LK(row, col) = (KV_DATA_TYPE4)(0.0h);
}
}
for (int i = tid; i < BLOCK_N * DV_VEC; i += WG_SIZE) {
@ -292,8 +315,15 @@ __kernel void FA_TILE_NAME(
FA_UNROLL
for (int k = 0; k < SPLIT_DK_VEC; k++) {
const ACC_TYPE4 qk = q_priv[k];
#if defined(FA_K_LDS_T)
// 2 KV rows adjacent in the transposed tile: one 128-bit local read.
const half8 kk = FA_LK_PAIR(dk_off + k, j);
ACC_TYPE4 dot0 = qk * CONVERT_KV_ACC4(kk.lo);
ACC_TYPE4 dot1 = qk * CONVERT_KV_ACC4(kk.hi);
#else
ACC_TYPE4 dot0 = qk * CONVERT_KV_ACC4(l_k[j ][dk_off + k]);
ACC_TYPE4 dot1 = qk * CONVERT_KV_ACC4(l_k[j+1][dk_off + k]);
#endif
partial0 += dot0.s0 + dot0.s1 + dot0.s2 + dot0.s3;
partial1 += dot1.s0 + dot1.s1 + dot1.s2 + dot1.s3;
}
@ -359,7 +389,7 @@ __kernel void FA_TILE_NAME(
ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f);
FA_UNROLL
for (int k = 0; k < SPLIT_DK_VEC; k++) {
dot_acc = mad(q_priv[k], CONVERT_KV_ACC4(l_k[j][dk_off + k]), dot_acc);
dot_acc = mad(q_priv[k], CONVERT_KV_ACC4(FA_LK(j, dk_off + k)), dot_acc);
}
local_partial[j][tid] =
dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3;
@ -452,10 +482,21 @@ __kernel void FA_TILE_NAME(
FA_UNROLL
for (int k = 0; k < DK_VEC; k++) {
const ACC_TYPE4 qk = q_priv[k];
#if defined(FA_K_LDS_T)
// 4 KV rows adjacent in the transposed tile: two 128-bit local reads
// instead of four 64-bit ones.
const half8 kk01 = FA_LK_PAIR(k, j);
const half8 kk23 = FA_LK_PAIR(k, j + 2);
dot_acc0 = mad(qk, CONVERT_KV_ACC4(kk01.lo), dot_acc0);
dot_acc1 = mad(qk, CONVERT_KV_ACC4(kk01.hi), dot_acc1);
dot_acc2 = mad(qk, CONVERT_KV_ACC4(kk23.lo), dot_acc2);
dot_acc3 = mad(qk, CONVERT_KV_ACC4(kk23.hi), dot_acc3);
#else
dot_acc0 = mad(qk, CONVERT_KV_ACC4(l_k[j][k]), dot_acc0);
dot_acc1 = mad(qk, CONVERT_KV_ACC4(l_k[j+1][k]), dot_acc1);
dot_acc2 = mad(qk, CONVERT_KV_ACC4(l_k[j+2][k]), dot_acc2);
dot_acc3 = mad(qk, CONVERT_KV_ACC4(l_k[j+3][k]), dot_acc3);
#endif
}
ACC_TYPE s0 = (dot_acc0.s0 + dot_acc0.s1 + dot_acc0.s2 + dot_acc0.s3) * scale;
ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale;

47
ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl
View File

@ -1631,8 +1631,25 @@ __kernel void flash_attn_f32_q4_0(
float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
#ifdef FA_HAVE_INT_DOT
// Accessors so the staging code is layout-agnostic.
#ifdef FA_K_LDS_T
#define FA_K_PACKED(ROW, IDX) l_k_packed[IDX][ROW]
#define FA_K_SCALE(ROW, BLK) l_k_scale[BLK][ROW]
#else
#define FA_K_PACKED(ROW, IDX) l_k_packed[ROW][IDX]
#define FA_K_SCALE(ROW, BLK) l_k_scale[ROW][BLK]
#endif
#ifdef FA_K_LDS_T
// K tile transposed: the 4 KV rows the QK loop walks together become adjacent, so each
// (block, group) step is ONE 128-bit local read instead of four 32-bit ones. The QK
// loop is LDS-read-issue-bound.
__local uint l_k_packed[DK_Q4_BLOCKS_PREFILL * 8][BLOCK_N];
__local float l_k_scale [DK_Q4_BLOCKS_PREFILL][BLOCK_N];
#else
__local uint l_k_packed[BLOCK_N][DK_Q4_BLOCKS_PREFILL * 8];
__local float l_k_scale [BLOCK_N][DK_Q4_BLOCKS_PREFILL];
#endif
#else
__local half4 l_k[BLOCK_N][DK_VEC];
#endif
@ -1660,17 +1677,17 @@ __kernel void flash_attn_f32_q4_0(
const global char * blk_ptr = k_base + k_row_off + blk * Q4_0_BLOCK_SIZE;
const float df = (float) vload_half(0, (const global half *) blk_ptr);
const global uchar * qs = (const global uchar *)(blk_ptr + 2);
l_k_scale[row][blk] = df;
FA_K_SCALE(row, blk) = df;
uint k_packed[8];
pack_q4_0_nibbles(qs, k_packed);
#pragma unroll
for (int j = 0; j < 8; ++j) {
l_k_packed[row][blk * 8 + j] = k_packed[j];
FA_K_PACKED(row, blk * 8 + j) = k_packed[j];
}
} else {
l_k_scale[row][blk] = 0.0f;
FA_K_SCALE(row, blk) = 0.0f;
#pragma unroll
for (int j = 0; j < 8; ++j) l_k_packed[row][blk * 8 + j] = 0u;
for (int j = 0; j < 8; ++j) FA_K_PACKED(row, blk * 8 + j) = 0u;
}
}
#else
@ -1760,6 +1777,19 @@ __kernel void flash_attn_f32_q4_0(
for (int b_local = 0; b_local < SPLIT_DK_Q4_BLOCKS; ++b_local) {
const int b = k_blk_base + b_local;
int sum0 = 0, sum1 = 0, sum2 = 0, sum3 = 0;
#ifdef FA_K_LDS_T
// 4 KV rows are adjacent in the transposed tile: one 128-bit local
// read per (block, group) instead of four 32-bit ones.
#pragma unroll
for (int g = 0; g < 8; ++g) {
const uint qp = q_packed_pf[b_local * 8 + g];
const uint4 kq4 = vload4(0, &l_k_packed[b * 8 + g][j]);
sum0 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s0, sum0);
sum1 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s1, sum1);
sum2 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s2, sum2);
sum3 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s3, sum3);
}
#else
#pragma unroll
for (int g = 0; g < 8; ++g) {
const uint qp = q_packed_pf[b_local * 8 + g];
@ -1768,12 +1798,21 @@ __kernel void flash_attn_f32_q4_0(
sum2 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+2][b * 8 + g], sum2);
sum3 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+3][b * 8 + g], sum3);
}
#endif
const float qd = q_d_pf[b_local];
const int q_sum = q_sum_pf[b_local];
#ifdef FA_K_LDS_T
const float4 ks4 = vload4(0, &l_k_scale[b][j]);
s0 += (float)(sum0 - 8 * q_sum) * qd * ks4.s0;
s1 += (float)(sum1 - 8 * q_sum) * qd * ks4.s1;
s2 += (float)(sum2 - 8 * q_sum) * qd * ks4.s2;
s3 += (float)(sum3 - 8 * q_sum) * qd * ks4.s3;
#else
s0 += (float)(sum0 - 8 * q_sum) * qd * l_k_scale[j ][b];
s1 += (float)(sum1 - 8 * q_sum) * qd * l_k_scale[j+1][b];
s2 += (float)(sum2 - 8 * q_sum) * qd * l_k_scale[j+2][b];
s3 += (float)(sum3 - 8 * q_sum) * qd * l_k_scale[j+3][b];
#endif
}
#else
ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f);

53
ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl
View File

@ -1393,8 +1393,31 @@ __kernel void flash_attn_f32_q8_0(
float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
#ifdef FA_HAVE_INT_DOT
// Accessors so the staging code is layout-agnostic.
#ifdef FA_K_LDS_T
#define FA_K_PACKED(ROW, IDX) l_k_packed[IDX][ROW]
#define FA_K_SCALE(ROW, BLK) l_k_scale[BLK][ROW]
#else
#define FA_K_PACKED(ROW, IDX) l_k_packed[ROW][IDX]
#define FA_K_SCALE(ROW, BLK) l_k_scale[ROW][BLK]
#endif
#ifdef FA_K_LDS_T
// K tile transposed: [block*8 + g][kv row] instead of [kv row][block*8 + g].
//
// The QK loop walks 4 KV rows at a time against the same (b, g), so in the original
// layout those 4 values are BLOCK_N*8 uints apart and cost 4 separate 32-bit local
// reads. Transposed they are adjacent, so they are one 128-bit read -- 4x fewer LDS
// issues for the same bytes and no extra registers. That matters because the QK loop
// is LDS-read-issue-bound: a wrong-math probe that kept every dp4a but cut the LDS
// reads ran the whole kernel 41% faster (18.51 -> 10.91 ms/op), and deleting QK
// outright only reached 10.88 -- i.e. essentially ALL of QK's cost is these reads.
__local uint l_k_packed[DK_Q8_BLOCKS_PREFILL * 8][BLOCK_N];
__local float l_k_scale [DK_Q8_BLOCKS_PREFILL][BLOCK_N];
#else
__local uint l_k_packed[BLOCK_N][DK_Q8_BLOCKS_PREFILL * 8];
__local float l_k_scale [BLOCK_N][DK_Q8_BLOCKS_PREFILL];
#endif
#else
__local half4 l_k[BLOCK_N][DK_VEC];
#endif
@ -1427,7 +1450,7 @@ __kernel void flash_attn_f32_q8_0(
const global char * blk_ptr = k_base + k_row_off + blk * Q8_0_BLOCK_SIZE;
const float df = (float) vload_half(0, (const global half *) blk_ptr);
const global uchar * qs = (const global uchar *)(blk_ptr + 2);
l_k_scale[row][blk] = df;
FA_K_SCALE(row, blk) = df;
#pragma unroll
for (int j = 0; j < 8; ++j) {
uint k_packed =
@ -1435,12 +1458,12 @@ __kernel void flash_attn_f32_q8_0(
((uint) qs[j*4 + 1]) << 8 |
((uint) qs[j*4 + 2]) << 16 |
((uint) qs[j*4 + 3]) << 24;
l_k_packed[row][blk * 8 + j] = k_packed;
FA_K_PACKED(row, blk * 8 + j) = k_packed;
}
} else {
l_k_scale[row][blk] = 0.0f;
FA_K_SCALE(row, blk) = 0.0f;
#pragma unroll
for (int j = 0; j < 8; ++j) l_k_packed[row][blk * 8 + j] = 0u;
for (int j = 0; j < 8; ++j) FA_K_PACKED(row, blk * 8 + j) = 0u;
}
}
#else
@ -1556,6 +1579,19 @@ __kernel void flash_attn_f32_q8_0(
for (int b_local = 0; b_local < SPLIT_DK_Q8_BLOCKS; ++b_local) {
const int b = k_blk_base + b_local;
int sum0 = 0, sum1 = 0, sum2 = 0, sum3 = 0;
#if defined(FA_K_LDS_T)
// The 4 KV rows are adjacent in the transposed tile, so each (b, g)
// step is ONE 128-bit local read instead of four 32-bit ones.
#pragma unroll
for (int g = 0; g < 8; ++g) {
const uint qp = q_packed_pf[b_local * 8 + g];
const uint4 kq4 = vload4(0, &l_k_packed[b * 8 + g][j]);
sum0 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s0, sum0);
sum1 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s1, sum1);
sum2 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s2, sum2);
sum3 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s3, sum3);
}
#else
#pragma unroll
for (int g = 0; g < 8; ++g) {
const uint qp = q_packed_pf[b_local * 8 + g];
@ -1564,11 +1600,20 @@ __kernel void flash_attn_f32_q8_0(
sum2 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+2][b * 8 + g], sum2);
sum3 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+3][b * 8 + g], sum3);
}
#endif
const float qd = q_d_pf[b_local];
#ifdef FA_K_LDS_T
const float4 ks4 = vload4(0, &l_k_scale[b][j]);
s0 += (float)sum0 * qd * ks4.s0;
s1 += (float)sum1 * qd * ks4.s1;
s2 += (float)sum2 * qd * ks4.s2;
s3 += (float)sum3 * qd * ks4.s3;
#else
s0 += (float)sum0 * qd * l_k_scale[j ][b];
s1 += (float)sum1 * qd * l_k_scale[j+1][b];
s2 += (float)sum2 * qd * l_k_scale[j+2][b];
s3 += (float)sum3 * qd * l_k_scale[j+3][b];
#endif
}
#else
ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f);

496
ggml/src/ggml-openvino/ggml-decoder.cpp
View File

@ -16,6 +16,7 @@
#include <iomanip>
#include <map>
#include <memory>
#include <mutex>
#include <openvino/core/dimension.hpp>
#include <openvino/core/except.hpp>
#include <openvino/core/node.hpp>
@ -25,12 +26,13 @@
#include <openvino/core/type/float16.hpp>
#include <openvino/op/constant.hpp>
#include <openvino/op/convert.hpp>
#include <openvino/op/parameter.hpp>
#include <openvino/runtime/tensor.hpp>
#include <ostream>
#include <set>
#include <stdexcept>
#include <string>
#include <cstring>
#include <unordered_map>
#include <vector>
GgmlOvDecoder::GgmlOvDecoder(ggml_cgraph * cgraph,
@ -98,27 +100,119 @@ GgmlOvDecoder::GgmlOvDecoder(ggml_cgraph * cgraph, std::map<std::string, std::sh
}
}
namespace {
bool is_inplace_op(const ggml_tensor * node) {
return node->op == GGML_OP_SET_ROWS || node->op == GGML_OP_CPY || (node->op == GGML_OP_SCALE && node->view_src);
}
bool is_same_shape(const ggml_tensor * a, const ggml_tensor * b) {
for (int i = 0; i < GGML_MAX_DIMS; i++) {
if (a->ne[i] != b->ne[i]) {
return false;
}
}
return true;
}
bool is_conv_states_all_tensor(const ggml_tensor * tensor) {
return tensor != nullptr && strncmp(tensor->name, "conv_states_all", strlen("conv_states_all")) == 0;
}
// CPY writing the tail of conv_input (the concat of the previous conv state and the new tokens)
// back into a slot block of the recurrent state cache. Detected structurally because the rollback
// variant (cparams.n_rs_seq > 0) emits one such CPY per snapshot slot without naming them.
bool is_conv_state_writeback(const ggml_tensor * node) {
return node->op == GGML_OP_CPY && node->view_src != nullptr && GgmlOvDecoder::is_kvcache(node->view_src, nullptr) &&
node->src[0] != nullptr && node->src[0]->op == GGML_OP_VIEW && node->src[0]->src[0] != nullptr &&
node->src[0]->src[0]->op == GGML_OP_CONCAT && node->src[1] != nullptr && node->src[1]->op == GGML_OP_VIEW &&
node->src[1]->view_src == node->view_src;
}
// MoE expert aggregation (build_moe_ffn in llama-graph.cpp): each expert plane is
// `ggml_view_2d(experts, n_embd, n_tokens, experts->nb[2], i*experts->nb[1])` and the planes
// are summed with a chain of ADDs: moe_out = ((view_0 + view_1) + view_2) + ... + view_{n-1}.
// Detected structurally by walking the ADD chain and checking every leaf is a same-shape,
// same-stride VIEW of one common base tensor, indexed by a distinct expert-plane offset, and
// that the chain covers every plane of that base (leaf count == base->ne[1]). Only the
// outermost ADD of the chain satisfies this (inner ADDs see fewer leaves than base->ne[1]).
bool is_moe_expert_sum_add(const ggml_tensor * node) {
std::vector<const ggml_tensor *> leaves;
const ggml_tensor * cur = node;
while (cur->op == GGML_OP_ADD) {
if (cur->src[0] == nullptr || cur->src[1] == nullptr) {
return false;
}
leaves.push_back(cur->src[1]);
cur = cur->src[0];
}
leaves.push_back(cur);
const ggml_tensor * base = nullptr;
std::set<int64_t> plane_indices;
for (const ggml_tensor * leaf : leaves) {
if (leaf->op != GGML_OP_VIEW || leaf->src[0] == nullptr) {
return false;
}
const ggml_tensor * leaf_base = leaf->src[0];
if (base == nullptr) {
base = leaf_base;
} else if (leaf_base != base) {
return false;
}
if (leaf->ne[0] != base->ne[0] || leaf->ne[1] != base->ne[2] || leaf->ne[2] != 1 || leaf->ne[3] != 1 ||
leaf->nb[1] != base->nb[2]) {
return false;
}
if (base->nb[1] == 0 || leaf->view_offs % base->nb[1] != 0) {
return false;
}
int64_t plane = static_cast<int64_t>(leaf->view_offs / base->nb[1]);
if (plane < 0 || plane >= base->ne[1] || !plane_indices.insert(plane).second) {
return false;
}
}
return base != nullptr && base->ne[1] > 1 && plane_indices.size() == static_cast<size_t>(base->ne[1]);
}
} // namespace
static std::string get_tensor_ov_name(const ggml_cgraph * cgraph, const ggml_tensor * tensor) {
if (tensor == nullptr) {
return "";
}
const size_t hash_pos = ggml_hash_find(&cgraph->visited_hash_set, tensor);
if (((tensor->flags & GGML_TENSOR_FLAG_COMPUTE) || GgmlOvDecoder::is_kvcache(tensor, nullptr)) &&
hash_pos != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, hash_pos)) {
return std::string(tensor->name) + "#" + std::to_string(hash_pos);
}
return tensor->name;
}
static std::string get_tensor_graph_input_ov_name(const GgmlOvDecoder * decoder,
const ggml_cgraph * cgraph,
const ggml_tensor * tensor,
const ggml_tensor * op) {
if (GgmlOvDecoder::is_inp_pos(tensor, op)) {
return "inp_pos";
}
if (GgmlOvDecoder::is_inp_emb(tensor, op)) {
return "embd";
}
if (decoder->is_stateful() && GgmlOvDecoder::is_inp_mask(tensor, op)) {
return std::string(tensor->name).find("swa") == std::string::npos ? "self_kq_mask" : "self_kq_mask_swa";
}
return get_tensor_ov_name(cgraph, tensor);
}
void GgmlOvDecoder::set_input_output() {
for (int node_n = 0; node_n < m_cgraph->n_nodes; node_n++) {
auto node = m_cgraph->nodes[node_n];
auto * node = m_cgraph->nodes[node_n];
NodeInfo current_node_info;
auto node_name = std::string(node->name);
auto node_output_name = node_name;
auto * node_output = node;
if (node->op == GGML_OP_SET_ROWS) {
// SET_ROWS updates the tensor in place. For later ov op that uses the
// the view_src of SET_ROWS, we need to make sure they get the updated tensor
// by putting the view_src name in the tensor_map in
// <openvino>/src/frontends/ggml/src/translate_session.cpp
node_output_name = std::string(node->view_src->name);
node_output = node->view_src;
}
auto node_name = get_tensor_ov_name(m_cgraph, node);
current_node_info.node = node;
current_node_info.node_name = node_name;
current_node_info.node_output = node_output;
current_node_info.node_output_name = node_output_name;
current_node_info.node_op_case = 0;
current_node_info.data_addr = node->data;
@ -127,9 +221,9 @@ void GgmlOvDecoder::set_input_output() {
if (src == nullptr) {
continue;
}
auto src_name = std::string(src->name);
auto src_name = get_tensor_ov_name(m_cgraph, src);
if (src->flags & GGML_TENSOR_FLAG_INPUT) {
src_name = get_graph_input_ov_name(src, node);
src_name = get_tensor_graph_input_ov_name(this, m_cgraph, src, node);
}
current_node_info.node_inputs[src_name] = src;
current_node_info.node_inputs_names.push_back(src_name);
@ -140,9 +234,9 @@ void GgmlOvDecoder::set_input_output() {
auto current = src;
while (current != nullptr) {
auto current_name = std::string(current->name);
auto current_name = get_tensor_ov_name(m_cgraph, current);
if (current->flags & GGML_TENSOR_FLAG_INPUT) {
current_name = get_graph_input_ov_name(current, node);
current_name = get_tensor_graph_input_ov_name(this, m_cgraph, current, node);
}
view_chain.emplace_back(current_name, current);
// If current src is also a VIEW, continue traversing
@ -166,6 +260,7 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const {
int op_case = 0;
switch (node->op) {
case GGML_OP_RESHAPE: {
auto name = std::string(node->name);
auto * src = node->src[0];
if (src->op == GGML_OP_RESHAPE && src->src[0]->ne[0] == node->ne[0] && src->src[0]->ne[1] == node->ne[1]) {
op_case = 4;
@ -178,11 +273,12 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const {
}
} else if (src->ne[0] * src->ne[1] * src->ne[2] == node->ne[1]) {
op_case = 3;
} else if (src->ne[1] * src->ne[2] == node->ne[1]) {
op_case = 6;
}
if (op_case == 0 && ggml_nelements(node) == ggml_nelements(src)) {
} else if (name.find("linear_attn_qkv_mixed") == 0 || name.find("alpha") == 0) {
op_case = 6;
} else if (name.find("linear_attn_out") == 0) {
op_case = 7;
} else if (name.find("state_predelta") == 0) {
op_case = 8;
}
break;
}
@ -232,7 +328,14 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const {
}
case GGML_OP_GET_ROWS: {
if (node->src[1]->op == GGML_OP_VIEW) {
op_case = 2;
// GET_ROWS gathering recurrent state cache rows via the inp->s_copy index list:
// src[0] is a reshape of cache_r/cache_s, src[1] is a view of the s_copy leaf.
// op_case 3: main view (active sequences, view offset 0)
// op_case 4: extra view (defrag remainder, nonzero view offset)
if (node->src[0]->op == GGML_OP_RESHAPE && node->src[0]->src[0] != nullptr &&
is_kvcache(node->src[0]->src[0], nullptr)) {
op_case = node->src[1]->view_offs == 0 ? 1 : 2;
}
}
break;
}
@ -260,7 +363,7 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const {
// throw std::runtime_error("Unsupported VIEW case");
}
op_case = 0;
if (m_model_is_splitted && m_model_inputs.find(std::string(src->name)) != m_model_inputs.end()) {
if (m_model_is_splitted && m_model_inputs.find(get_tensor_ov_name(m_cgraph, src)) != m_model_inputs.end()) {
op_case = 0;
}
}
@ -295,6 +398,56 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const {
}
break;
}
case GGML_OP_RMS_NORM: {
if (node->src[0]->op == GGML_OP_VIEW) {
if (is_same_shape(node->src[0]->src[0], node->src[0])) {
op_case = 1;
} else if (node->src[0]->src[0]->op == GGML_OP_GATED_DELTA_NET) {
op_case = 2;
}
}
break;
}
case GGML_OP_CPY: {
if (node->src[0]->op == GGML_OP_VIEW) {
if (node->src[0]->src[0]->op == GGML_OP_GATED_DELTA_NET) {
op_case = 1;
} else if (is_conv_state_writeback(node)) {
op_case = 2;
break;
} else if (is_conv_states_all_tensor(node->view_src) && node->src[1] != nullptr &&
node->src[1]->op == GGML_OP_VIEW && node->src[1]->view_src == node->view_src) {
op_case = 4;
break;
}
} else if (node->src[0]->op == GGML_OP_GET_ROWS && node->src[1] != nullptr &&
node->src[1]->op == GGML_OP_VIEW && node->src[1]->view_src != nullptr &&
is_kvcache(node->src[1]->view_src, nullptr)) {
// s_copy defrag remainder writeback: gathered extra state rows copied back into the cache
op_case = 3;
}
break;
}
case GGML_OP_ADD: {
if (is_moe_expert_sum_add(node)) {
// Outermost ADD of a MoE expert-plane sum chain: translated as a single
// ReduceSum over the base tensor instead of N-1 chained Adds over N Slices.
op_case = 1;
}
break;
}
case GGML_OP_SCALE: {
if (node->view_src && node->buffer->usage == GGML_BACKEND_BUFFER_USAGE_ANY) {
op_case = 1;
}
break;
}
case GGML_OP_L2_NORM: {
if (std::string(node->name).find("predelta") != std::string::npos) {
op_case = 1;
}
break;
}
default:
break;
}
@ -476,6 +629,43 @@ std::pair<ModelParams, ComputeParams> GgmlOvDecoder::compute_llm_params(ggml_cgr
model_params.mixed_rope_params = true;
}
}
if (node->op == GGML_OP_GATED_DELTA_NET) {
model_params.state_size = node->src[0]->ne[0];
}
if (node->op == GGML_OP_SCALE && node->view_src != nullptr && is_kvcache(node->view_src, nullptr)) {
compute_params.cache_rs_reset_len = ggml_nelements(node) / node->view_src->ne[0];
compute_params.cache_rs_reset_idx = node->src[0]->view_offs / node->view_src->ne[0];
}
// Capture the destination slot block of every recurrent state cache writeback, plus the
// conv_input window the conv state writeback copies. The active sequences occupy a
// contiguous slot block [begin, begin + n_seqs) of the cache; the block and the window move
// with the batch, so they are fed to the cached model as runtime inputs.
if (node->op == GGML_OP_CPY && node->view_src != nullptr && is_kvcache(node->view_src, nullptr) &&
node->src[1] != nullptr && node->src[1]->op == GGML_OP_VIEW && node->src[1]->view_src == node->view_src) {
const bool is_conv = is_conv_state_writeback(node);
const bool is_gdn = node->src[0]->op == GGML_OP_VIEW && node->src[0]->src[0] != nullptr &&
node->src[0]->src[0]->op == GGML_OP_GATED_DELTA_NET;
const bool is_extra = node->src[0]->op == GGML_OP_GET_ROWS;
const ggml_tensor * dest_view = node->src[1];
const ggml_tensor * cache = node->view_src;
const size_t row_bytes = cache->ne[0] * ggml_type_size(cache->type);
if (row_bytes > 0 && (is_conv || is_gdn || is_extra)) {
ComputeParams::RsWriteback writeback;
writeback.slot_begin = (int) (dest_view->view_offs / row_bytes);
if (is_conv) {
// conv_input column the copied window starts at
writeback.src_begin = (int) (node->src[0]->view_offs / node->src[0]->view_src->nb[0]);
} else if (is_gdn) {
// first row of the state part of the gated-delta-net output
writeback.src_begin = (int) (node->src[0]->view_offs / node->src[0]->view_src->nb[1]);
}
compute_params.rs_writebacks[get_tensor_ov_name(cgraph, node)] = writeback;
}
if (is_conv || is_gdn) {
compute_params.s_copy_active_slot_len = (int) dest_view->ne[1];
}
}
}
auto * output_tensor = cgraph->nodes[cgraph->n_nodes - 1];
compute_params.output_len = output_tensor->ne[1];
@ -505,6 +695,10 @@ ov::PartialShape GgmlOvDecoder::get_graph_input_shape(const ggml_tensor * op,
if (is_inp_tok(input, op) || is_inp_pos(input, op)) {
// tokens or positions
int len = m_is_static ? (m_is_prefill ? m_prefill_chunk_size : 1) : -1;
if (m_is_static && is_inp_pos(input, op)) {
// IMROPE stacks n_planes (t/h/w/e) position planes back to back
len *= get_inp_pos_n_planes(op);
}
input_shape = ov::PartialShape{1, 1, 1, len};
} else if (is_output_idx(input, op)) {
@ -543,6 +737,9 @@ ov::PartialShape GgmlOvDecoder::get_graph_input_shape(const ggml_tensor * op,
int len = m_is_static ? (m_is_prefill ? m_prefill_chunk_size : 1) : -1;
input_shape = ov::PartialShape{1, 1, 1, len};
} else if (is_inp_s_copy(input, op) || is_s_copy_leaf(input)) {
input_shape = ov::PartialShape{1, 1, 1, -1};
} else {
input_shape = ov::PartialShape{get_shape(input)};
}
@ -558,6 +755,35 @@ ov::PartialShape GgmlOvDecoder::get_graph_input_shape(const ggml_tensor * op,
return input_shape;
}
bool GgmlOvDecoder::is_s_copy_leaf(const ggml_tensor * tensor) const {
if (tensor == nullptr || tensor->op != GGML_OP_NONE || m_cgraph == nullptr) {
return false;
}
for (int i = 0; i < m_cgraph->n_nodes; i++) {
const ggml_tensor * node = m_cgraph->nodes[i];
if (node->op != GGML_OP_GET_ROWS || node->src[0] == nullptr || node->src[1] == nullptr) {
continue;
}
// The index list may reach the s_copy leaf through one or more VIEWs.
const ggml_tensor * idx = node->src[1];
while (idx != nullptr && idx->op == GGML_OP_VIEW) {
idx = idx->src[0];
}
if (idx != tensor) {
continue;
}
// The gathered data must be a recurrent state cache (cache_r/cache_s).
const ggml_tensor * data = node->src[0];
while (data != nullptr && (data->op == GGML_OP_VIEW || data->op == GGML_OP_RESHAPE)) {
data = data->src[0];
}
if (data != nullptr && is_kvcache(data, nullptr)) {
return true;
}
}
return false;
}
void GgmlOvDecoder::add_extra_inputs() {
// Extra inputs:
// 1. `attention_size`, used in FLASH_ATTN where the shape of the matmul's are 256 aligned,
@ -565,21 +791,7 @@ void GgmlOvDecoder::add_extra_inputs() {
// 2. `n_seq_active` and `seq_active_start`, used in FLASH_ATTN_EXT to indicate the active sequences in the batch
auto create_1d_input = [this](const std::string & name, int64_t value) {
if (m_is_static) {
auto constant =
std::make_shared<ov::op::v0::Constant>(ov::element::i64, ov::Shape{1}, std::vector<int64_t>{value});
constant->set_friendly_name(name);
m_model_extra_inputs[name] = constant;
} else {
auto param_node = std::make_shared<ov::op::v0::Parameter>(ov::element::i64, ov::Shape{1});
param_node->set_friendly_name(name);
param_node->output(0).get_tensor().set_names({name});
m_model_extra_inputs[name] = param_node;
auto tensor = std::make_shared<ov::Tensor>(ov::element::i64, ov::Shape{1});
*tensor->data<int64_t>() = value;
m_model_extra_input_values[name] = tensor;
}
m_model_extra_inputs[name] = {ov::element::i64, ov::Shape{1}, value, !m_is_static};
};
if (m_compute_params.attention_size != -1) {
@ -595,6 +807,20 @@ void GgmlOvDecoder::add_extra_inputs() {
create_1d_input("token_len_per_seq", m_compute_params.token_len_per_seq);
}
// create_1d_input("token_len", m_compute_params.token_len_per_seq * m_compute_params.n_seq_active);
if (m_compute_params.cache_rs_reset_idx != -1) {
create_1d_input("cache_rs_reset_idx", m_compute_params.cache_rs_reset_idx);
create_1d_input("cache_rs_reset_len", m_compute_params.cache_rs_reset_len);
}
if (m_compute_params.s_copy_active_slot_len != -1) {
create_1d_input("s_copy_active_slot_len", m_compute_params.s_copy_active_slot_len);
}
for (const auto & [node_name, writeback] : m_compute_params.rs_writebacks) {
create_1d_input("rs_slot_begin_" + node_name, writeback.slot_begin);
create_1d_input("rs_src_begin_" + node_name, writeback.src_begin);
}
}
bool GgmlOvDecoder::node_is_used_as_src(const int node_idx) {
@ -617,14 +843,11 @@ void GgmlOvDecoder::compute_model_inputs() {
ggml_tensor * node = m_cgraph->nodes[i];
// the node op is NONE means this node maybe as input of later nodes, we should add it to model inputs for this node.
if (node->op == GGML_OP_NONE && node_is_used_as_src(i)) {
std::string node_name(node->name);
std::string node_name = get_tensor_ov_name(m_cgraph, node);
if (m_model_weights.find(node_name) == m_model_weights.end()) {
m_inputs[node_name] = node;
auto param_node = std::make_shared<ov::op::v0::Parameter>(
get_ov_type(node), get_graph_input_shape(node, nullptr, m_node_dynamic_dims[node]));
param_node->set_friendly_name(node_name);
param_node->output(0).get_tensor().set_names({node_name});
m_model_inputs[node_name] = param_node;
m_model_inputs[node_name] = {get_ov_type(node),
get_graph_input_shape(node, nullptr, m_node_dynamic_dims[node])};
}
continue;
}
@ -633,9 +856,9 @@ void GgmlOvDecoder::compute_model_inputs() {
if (src == nullptr) {
continue;
}
std::string src_name = std::string(src->name);
std::string src_name = get_tensor_ov_name(m_cgraph, src);
if (src->flags & GGML_TENSOR_FLAG_INPUT) {
src_name = get_graph_input_ov_name(src, node);
src_name = get_tensor_graph_input_ov_name(this, m_cgraph, src, node);
}
if (m_model_weights.find(src_name) != m_model_weights.end()) {
continue;
@ -668,14 +891,11 @@ void GgmlOvDecoder::compute_model_inputs() {
// Resolve nested VIEW nodes by following src[0] until the first non-VIEW tensor.
while (src->op == GGML_OP_VIEW && src->src[0] != nullptr) {
src = src->src[0];
src_name = std::string(src->name);
src_name = get_tensor_ov_name(m_cgraph, src);
}
m_inputs[src_name] = src;
ov::PartialShape param_shape = get_graph_input_shape(node, src, m_node_dynamic_dims[src]);
auto param_node = std::make_shared<ov::op::v0::Parameter>(get_ov_type(src), param_shape);
param_node->set_friendly_name(src_name);
param_node->output(0).get_tensor().set_names({src_name});
m_model_inputs[src_name] = param_node;
m_model_inputs[src_name] = {get_ov_type(src),
get_graph_input_shape(node, src, m_node_dynamic_dims[src])};
}
}
}
@ -691,8 +911,8 @@ void GgmlOvDecoder::compute_model_outputs() {
}
auto cur_node_use_count = m_cgraph->use_counts[ggml_hash_find(&m_cgraph->visited_hash_set, cur_node)];
if (cur_node_use_count == 0) {
// The output of SET_ROWS is the view_src tensor, which is updated in place. We should use the view_src name as the output name to make sure it can be correctly matched with the later ops that use the view_src.
if (cur_node != nullptr && cur_node->op == GGML_OP_SET_ROWS) {
// The output of in-place ops is the view_src tensor, which is updated in place. We should use the view_src name as the output name to make sure it can be correctly matched with the later ops that use the view_src.
if (cur_node != nullptr && ::is_inplace_op(cur_node) && ggml_nbytes(cur_node) > 0) {
cur_node = cur_node->view_src;
}
} else {
@ -710,9 +930,9 @@ void GgmlOvDecoder::compute_model_outputs() {
}
}
if (cur_node != nullptr) {
std::string node_output_name(cur_node->name);
m_model_outputs[node_output_name] = cur_node;
m_model_output_names.push_back(node_output_name);
std::string cur_node_name = get_tensor_ov_name(m_cgraph, cur_node);
m_model_outputs[cur_node_name] = cur_node;
m_model_output_names.insert(cur_node_name);
}
}
}
@ -740,7 +960,7 @@ const ggml_tensor * GgmlOvDecoder::get_tensor_from_name(const std::string & name
if (src == nullptr) {
break;
}
if (std::string(src->name) == name) {
if (get_tensor_ov_name(m_cgraph, src) == name) {
return src;
}
}
@ -756,6 +976,16 @@ std::map<std::string, std::string> GgmlOvDecoder::get_kv_param_res_names() const
return kv_param_res_names;
}
// MUL_MAT_ID's src[0] is the [k, m, n_expert] expert-weight tensor. It is always a constant per-expert
// weight table -- never a computed activation -- regardless of whether the backend happened to mark its
// buffer as GGML_BACKEND_BUFFER_USAGE_WEIGHTS (test-backend-ops, for example, never sets that usage
// flag, unlike real inference). Without this, non-quantized (F16/F32/BF16) expert weights would fall
// through the check below as "not a weight", get decoded as a Parameter/activation instead of a
// Constant, and crash GatherMatmul's "only constant weights are supported" check.
static bool is_mul_mat_id_expert_weight(const ggml_tensor * node, int src_index) {
return node->op == GGML_OP_MUL_MAT_ID && src_index == 0;
}
std::map<std::string, std::shared_ptr<ov::Node>> GgmlOvDecoder::create_weight_nodes(ggml_cgraph * cgraph, bool naive) {
std::map<std::string, std::shared_ptr<ov::Node>> model_weights;
auto * nodes = cgraph->nodes;
@ -768,13 +998,14 @@ std::map<std::string, std::shared_ptr<ov::Node>> GgmlOvDecoder::create_weight_no
continue;
}
std::string src_name(src->name);
std::string src_name = get_tensor_ov_name(cgraph, src);
if (is_rope_freqs_weight(src, node)) {
src_name = "rope_freqs.weight";
}
if (!src->view_src) {
ggml_backend_buffer * buffer = src->buffer;
if (buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS || ggml_is_quantized(src->type)) {
if (buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS || ggml_is_quantized(src->type) ||
is_mul_mat_id_expert_weight(node, i)) {
if (model_weights.find(src_name) == model_weights.end()) {
auto weight_node = create_weight_node(src, naive);
weight_node->set_friendly_name(src_name);
@ -787,6 +1018,42 @@ std::map<std::string, std::shared_ptr<ov::Node>> GgmlOvDecoder::create_weight_no
return model_weights;
}
// Process-lifetime cache for weight nodes built from NON-OpenVINO buffers (e.g. the
// token_embd.weight copy that lives in a CPU/mmap buffer and feeds GET_ROWS). Such
// tensors have no OV buffer context to own a cached extra, so without this they are
// re-extracted/re-requantized on every (re)compile — for token_embd that is a ~1-2 GB
// F32 dequant each time. Keyed by tensor->data, which is stable for the process and
// uniquely identifies the immutable weight bytes. OV-buffer weights keep using the
// per-tensor extra cache and never reach here.
static std::mutex g_nonov_weight_cache_mutex;
static std::unordered_map<const void *, std::shared_ptr<ov::Node>> g_nonov_weight_cache;
std::set<std::string> GgmlOvDecoder::collect_weight_names(ggml_cgraph * cgraph) {
// Mirrors the name-selection logic of create_weight_nodes() but builds no nodes,
// so topology checks don't trigger weight extraction/requantization.
std::set<std::string> names;
for (int node_i = 0; node_i < cgraph->n_nodes; node_i++) {
auto * node = cgraph->nodes[node_i];
for (int i = 0; i < GGML_MAX_SRC; i++) {
auto * src = node->src[i];
if (src == nullptr) {
continue;
}
std::string src_name(src->name);
if (is_rope_freqs_weight(src, node)) {
src_name = "rope_freqs.weight";
}
if (!src->view_src) {
ggml_backend_buffer * buffer = src->buffer;
if (buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS || ggml_is_quantized(src->type)) {
names.insert(src_name);
}
}
}
}
return names;
}
std::shared_ptr<ov::Node> GgmlOvDecoder::create_weight_node(ggml_tensor * tensor, bool naive) {
const bool is_ov_buffer = ggml_backend_buffer_is_openvino(tensor->buffer);
@ -826,6 +1093,21 @@ std::shared_ptr<ov::Node> GgmlOvDecoder::create_weight_node(ggml_tensor * tensor
return weight_node;
}
// Non-OV-buffer weights (CPU/mmap, e.g. the GET_ROWS token_embd copy) have no buffer
// context to cache an extra in, so memoize them here keyed by their (stable) data
// pointer to avoid re-extracting on every recompile. Opt-in via
// GGML_OPENVINO_REDUCE_COMPILE_MEM or GGML_OPENVINO_MEMORY_OPTIMIZE. Skip
// for `naive` (test/naive path) since use_bias changes the produced node.
const bool cacheable_nonov = ggml_openvino_reduce_compile_mem_enabled() && !is_ov_buffer &&
!naive && tensor->data != nullptr;
if (cacheable_nonov) {
std::lock_guard<std::mutex> lock(g_nonov_weight_cache_mutex);
auto it = g_nonov_weight_cache.find(tensor->data);
if (it != g_nonov_weight_cache.end()) {
return it->second;
}
}
// There are three cases where we need to create a new weight node:
// 1. weights are in openvino_host_buffer. Weight loading to host buffer will not trigger backend_buffer_set_tensor
// 2. weights are in cpu/cpu_mapped buffer. On token_embd.weight goes to case 1 or 2, depending on whether mmap or direct_io is used
@ -834,7 +1116,7 @@ std::shared_ptr<ov::Node> GgmlOvDecoder::create_weight_node(ggml_tensor * tensor
// GGML_LOG_DEBUG("%s: creating new weight node for %s\n", __func__, tensor->name);
static const std::set<ggml_type> weight_types = {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0,
GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_1, GGML_TYPE_Q4_K,
GGML_TYPE_Q5_K, GGML_TYPE_Q6_K};
GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_MXFP4};
if (weight_types.find(tensor->type) == weight_types.end()) {
throw std::runtime_error("Unexpected weight tensor type: " + std::string(tensor->name) + " with type " +
ggml_type_name(tensor->type));
@ -863,6 +1145,12 @@ std::shared_ptr<ov::Node> GgmlOvDecoder::create_weight_node(ggml_tensor * tensor
ov_weight.weight_node->set_friendly_name(tensor->name);
if (!is_ov_buffer) {
if (cacheable_nonov) {
std::lock_guard<std::mutex> lock(g_nonov_weight_cache_mutex);
// Another thread may have inserted concurrently; keep the first.
auto [it, inserted] = g_nonov_weight_cache.emplace(tensor->data, ov_weight.weight_node);
return it->second;
}
return ov_weight.weight_node;
}
@ -1178,7 +1466,7 @@ std::string GgmlOvDecoder::get_view_input_name(int node_idx, const std::string &
auto it = m_node_info_list[node_idx].node_inputs_views.find(name);
if (it != m_node_info_list[node_idx].node_inputs_views.end()) {
if (view_index < it->second.size()) {
return it->second[view_index].second->name;
return it->second[view_index].first;
}
}
return "";
@ -1190,7 +1478,7 @@ std::string GgmlOvDecoder::get_view_input_src_name(int node_idx, const std::stri
if (view_index < it->second.size()) {
auto * view_tensor = it->second[view_index].second;
if (view_tensor && view_tensor->src[0]) {
return view_tensor->src[0]->name;
return get_tensor_ov_name(m_cgraph, view_tensor->src[0]);
}
}
}
@ -1214,7 +1502,7 @@ std::vector<std::string> GgmlOvDecoder::get_input_names(int node_idx) const {
}
ov::PartialShape GgmlOvDecoder::get_output_shape(int node_idx) const {
auto * ggml_tensor = m_node_info_list[node_idx].node_output;
auto * ggml_tensor = m_node_info_list[node_idx].node;
return ov::PartialShape(get_shape(ggml_tensor));
}
@ -1228,7 +1516,28 @@ std::vector<size_t> GgmlOvDecoder::get_output_stride(int node_idx) const {
}
std::vector<std::string> GgmlOvDecoder::get_output_names(int node_idx) const {
return {m_node_info_list[node_idx].node_output_name};
return {m_node_info_list[node_idx].node_name};
}
std::string GgmlOvDecoder::get_inplace_op_src(int node_idx) const {
auto * node = m_node_info_list[node_idx].node;
if (!::is_inplace_op(node) || node->view_src == nullptr || ggml_nbytes(node) == 0) {
return "";
}
const int op_case = m_node_info_list[node_idx].node_op_case;
if (node->op == GGML_OP_CPY && (op_case == 1 || op_case == 2 || op_case == 3) &&
m_compute_params.s_copy_active_slot_len == -1) {
return "";
}
return get_tensor_ov_name(m_cgraph, node->view_src);
}
bool GgmlOvDecoder::is_view_like_alias_of(int node_idx, const std::string & view_src_name) const {
auto * node = m_node_info_list[node_idx].node;
if (node->view_src == nullptr || get_tensor_ov_name(m_cgraph, node->view_src) != view_src_name) {
return false;
}
return node->op == GGML_OP_RESHAPE || node->op == GGML_OP_VIEW;
}
const std::string & GgmlOvDecoder::get_op_name() const {
@ -1404,14 +1713,18 @@ void GgmlOvDecoder::compute_node_dynamic_dims() {
}
if (m_node_dynamic_dims[node] != -1 && dynamic_dim_value != node->ne[m_node_dynamic_dims[node]]) {
m_node_dynamic_dims[node] = -1;
// std::cout << "Warning: Dynamic dim value mismatch for node: " << node->name
// << " and its src[0]: " << node->src[0]->name << std::endl;
GGML_LOG_WARN("ggml-openvino: dynamic dim value mismatch for VIEW node '%s', src[0]: '%s'\n",
node->name, node->src[0]->name);
}
}
break;
}
case GGML_OP_TRANSPOSE:
case GGML_OP_RESHAPE: {
if (is_same_shape(node->src[0], node)) {
m_node_dynamic_dims[node] = m_node_dynamic_dims[node->src[0]];
break;
}
// RESHAPE requires src[0] to be contiguous, so both src and result
// have standard compact strides: nb[i] = type_size * prod(ne[0..i-1]).
// Match src->nb[dynamic_dim] against result->nb[i] to find the output
@ -1429,7 +1742,7 @@ void GgmlOvDecoder::compute_node_dynamic_dims() {
}
}
if (m_node_dynamic_dims[node] == -1) {
// std::cout << "Cannot determine dynamic dim for RESHAPE node: " << node->name << std::endl;
GGML_LOG_WARN("ggml-openvino: cannot determine dynamic dim for RESHAPE node '%s'\n", node->name);
}
}
break;
@ -1480,15 +1793,29 @@ void GgmlOvDecoder::compute_node_dynamic_dims() {
}
if (matched_dim_count != 1) {
m_node_dynamic_dims[node] = -1;
// std::cout << "Warning: Cannot determine dynamic dim for CONT node: " << node->name
// << " and its src[0]: " << node->src[0]->name << std::endl;
GGML_LOG_WARN("ggml-openvino: cannot determine dynamic dim for CONT node '%s', src[0]: '%s'\n",
node->name, node->src[0]->name);
}
}
}
break;
case GGML_OP_CONCAT:
for (int i = 0; i < GGML_MAX_DIMS; i++) {
if (node->src[0]->ne[i] != node->ne[i]) {
m_node_dynamic_dims[node] = i;
break;
}
}
break;
case GGML_OP_SSM_CONV:
case GGML_OP_GATED_DELTA_NET:
m_node_dynamic_dims[node] = 1;
break;
case GGML_OP_RMS_NORM:
case GGML_OP_L2_NORM:
case GGML_OP_NORM:
case GGML_OP_ADD:
case GGML_OP_SUB:
case GGML_OP_GLU:
case GGML_OP_ROPE:
case GGML_OP_SCALE:
@ -1496,9 +1823,31 @@ void GgmlOvDecoder::compute_node_dynamic_dims() {
case GGML_OP_ARGSORT:
case GGML_OP_ADD_ID:
case GGML_OP_UNARY:
case GGML_OP_CUMSUM:
case GGML_OP_FILL:
case GGML_OP_SET:
case GGML_OP_DIAG:
case GGML_OP_TRI:
case GGML_OP_REPEAT:
// Shape-preserving elementwise ops: the dynamic dim is unchanged from src[0].
// DIV/CLAMP are used in the MoE routing-weight normalization
// (sum_rows -> clamp -> div). If they are left untracked here the dynamic
// (token) dim is lost there, the captured prefill token count gets baked into
// the downstream reshapes, and every decoder layer after layer 0 turns static
// (which then triggers the GPU in-place-concat KV-cache corruption).
case GGML_OP_DIV:
case GGML_OP_CLAMP:
case GGML_OP_PAD:
m_node_dynamic_dims[node] = m_node_dynamic_dims[node->src[0]];
break;
case GGML_OP_SUM_ROWS:
// SUM_ROWS reduces ggml axis 0 to size 1 and preserves all other axes, so the
// dynamic dim is preserved unless it was axis 0 (then it is summed away).
m_node_dynamic_dims[node] =
(m_node_dynamic_dims[node->src[0]] == 0) ? -1 : m_node_dynamic_dims[node->src[0]];
break;
case GGML_OP_MUL_MAT_ID:
case GGML_OP_SOLVE_TRI:
m_node_dynamic_dims[node] = m_node_dynamic_dims[node->src[1]];
break;
case GGML_OP_CPY:
@ -1534,7 +1883,8 @@ void GgmlOvDecoder::compute_node_dynamic_dims() {
break;
}
default:
// std::cout << "Doesn't handle node name: " << node->name << " op: " << ggml_op_name(node->op) << std::endl;
GGML_LOG_DEBUG("ggml-openvino: compute_node_dynamic_dims: unhandled op %s for node '%s'\n",
ggml_op_name(node->op), node->name);
break;
}
};

98
ggml/src/ggml-openvino/ggml-decoder.h
View File

@ -11,6 +11,8 @@
#include <memory>
#include <openvino/core/partial_shape.hpp>
#include <optional>
#include <set>
#include <string>
#include <vector>
struct ModelParams {
@ -20,6 +22,7 @@ struct ModelParams {
int n_seq = 1;
int n_heads_kv = -1;
int head_size = -1;
int state_size = -1; // for SSM molels, eg qwen35
int32_t rope_params[15];
bool mixed_rope_params = false;
std::vector<int> swa_layers;
@ -48,6 +51,47 @@ struct ComputeParams {
int token_len_per_seq = -1;
int past_kv_len = -1;
int output_len = 1;
int cache_rs_reset_idx = -1;
int cache_rs_reset_len = -1;
// SSM/DeltaNet models otionally clear cache_r and cache_s of certain slots in the cgraph
// 3: [ 18432, 4, 1, 1] RESHAPE cache_r_l0 (reshaped)
// [ 18432, 4, 1, 1] 0: NONE cache_r_l0
// 4: [ 18432, 1, 1, 1] VIEW cache_r_l0 (reshaped) (view)
// [ 18432, 4, 1, 1] 0: RESHAPE cache_r_l0 (reshaped)
// 5: [ 18432, 1, 1, 1] SCALE cache_r_l0 (reshaped) (view) (view)
// [ 18432, 1, 1, 1] 0: VIEW cache_r_l0 (reshaped) (view)
int s_copy_active_slot_len = -1;
// SSM/DeltaNet models otionally reorder slots of state cache, to make the active slots contiguous
// leaf_5 is the inp->s_copy in llama-graph.cpp, eg if there are 8 slots in total and slot 3 and 7
// are active in the current batch, leaf_5 will be [3, 7, 5, 6, 4]
// 6: [ 2, 1, 1, 1] VIEW (view)
// [ 2, 1, 1, 1] 0: NONE leaf_5
// 7: [ 18432, 2, 1, 1] GET_ROWS conv_states-0
// [ 18432, 4, 1, 1] 0: RESHAPE cache_r_l0 (reshaped)
// [ 2, 1, 1, 1] 1: VIEW (view)
// 8: [ 0, 1, 1, 1] VIEW (view)
// [ 2, 1, 1, 1] 0: NONE leaf_5
// 9: [ 18432, 0, 1, 1] GET_ROWS node_9
// [ 18432, 4, 1, 1] 0: RESHAPE cache_r_l0 (reshaped)
// [ 0, 1, 1, 1] 1: VIEW (view)
// 10: [ 18432, 0, 1, 1] VIEW cache_r_l0 (view)
// [ 18432, 4, 1, 1] 0: NONE cache_r_l0
// 11: [ 18432, 0, 1, 1] CPY cache_r_l0 (view) (copy of )
// [ 18432, 0, 1, 1] 0: GET_ROWS node_9
// [ 18432, 0, 1, 1] 1: VIEW cache_r_l0 (view)
struct RsWriteback {
int slot_begin = 0; // first cache slot written by the CPY
int src_begin = 0; // where the copied data starts in the source tensor (in rows of it)
};
std::map<std::string, RsWriteback> rs_writebacks;
// Offsets of the state cache writeback CPY nodes, keyed by node name. They change with the
// batch (kv head, active sequence count, token count) and, with rollback enabled
// (cparams.n_rs_seq > 0), the conv state is written back once per snapshot slot, each snapshot
// taking a different conv_input window. Passed to the cached model as runtime inputs.
};
class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder {
@ -59,8 +103,6 @@ public:
std::map<std::string, ggml_tensor *> node_inputs;
std::map<std::string, std::vector<std::pair<std::string, ggml_tensor *>>> node_inputs_views;
std::vector<std::string> node_inputs_names;
ggml_tensor * node_output;
std::string node_output_name;
int node_op_case = 0;
void * data_addr;
};
@ -156,6 +198,10 @@ public:
virtual std::vector<std::string> get_output_names(int node_idx) const override;
virtual std::string get_inplace_op_src(int node_idx) const override;
virtual bool is_view_like_alias_of(int node_idx, const std::string & view_src_name) const override;
virtual const std::string & get_op_type() const override;
virtual const std::string & get_op_type(int node_idx) const override;
@ -173,23 +219,19 @@ public:
virtual int get_op_case(int node_idx) const override { return m_node_info_list[node_idx].node_op_case; }
virtual const std::map<std::string, std::shared_ptr<ov::Node>> & get_model_inputs() const override {
virtual const std::map<std::string, ov::frontend::ggml::ModelInputInfo> & get_model_inputs() const override {
return m_model_inputs;
}
virtual const std::map<std::string, std::shared_ptr<ov::Node>> & get_model_extra_inputs() const override {
virtual const std::map<std::string, ov::frontend::ggml::ModelExtraInputInfo> & get_model_extra_inputs() const override {
return m_model_extra_inputs;
}
virtual const std::map<std::string, std::shared_ptr<ov::Tensor>> & get_model_extra_input_values() const {
return m_model_extra_input_values;
}
virtual const std::map<std::string, std::shared_ptr<ov::Node>> & get_model_weights() const override {
return m_model_weights;
}
virtual std::vector<std::string> get_model_output_names() const override { return m_model_output_names; }
virtual std::set<std::string> get_model_output_names() const override { return m_model_output_names; }
const std::map<std::string, ggml_tensor *> & get_model_outputs() const { return m_model_outputs; }
@ -214,6 +256,8 @@ public:
virtual bool has_mixed_rope_params() const override { return m_model_params.mixed_rope_params; }
virtual int get_ssm_state_size() const override { return m_model_params.state_size; }
virtual std::map<std::string, std::string> get_kv_param_res_names() const override;
virtual bool is_static() const override { return m_is_static; }
@ -235,6 +279,11 @@ public:
static std::map<std::string, std::shared_ptr<ov::Node>> create_weight_nodes(ggml_cgraph * cgraph,
bool naive = false);
// Collect just the set of weight-tensor names referenced by the graph, without
// building (or requantizing) any OV weight nodes. Used by topology checks like
// is_model_splitted that only need name membership.
static std::set<std::string> collect_weight_names(ggml_cgraph * cgraph);
const ggml_tensor * get_tensor_used_op(const ggml_tensor * tensor) const;
const ggml_tensor * get_tensor_from_name(const std::string & name) const;
@ -274,6 +323,12 @@ public:
return op->op == GGML_OP_ROPE && tensor == op->src[1];
}
// IMROPE packs 4 stacked position planes (t/h/w/e) into inp_pos, each of length
// n_tokens; other modes carry a single position per token.
inline static int get_inp_pos_n_planes(const ggml_tensor * op) {
return op->op_params[2] == GGML_ROPE_TYPE_IMROPE ? 4 : 1;
}
inline static bool is_inp_emb(const ggml_tensor * tensor, const ggml_tensor * op) {
return tensor->op == GGML_OP_GET_ROWS && op->op == GGML_OP_RMS_NORM;
}
@ -287,8 +342,12 @@ public:
return op->op == GGML_OP_ROPE && tensor == op->src[2];
}
// also returns true for cache_s and cache_r in SSM/DeltaNet models
inline static bool is_kvcache(const ggml_tensor * tensor, const ggml_tensor * op) {
return tensor->buffer->usage == GGML_BACKEND_BUFFER_USAGE_ANY ||
if (tensor == nullptr) {
return false;
}
return (tensor->buffer != nullptr && tensor->buffer->usage == GGML_BACKEND_BUFFER_USAGE_ANY) ||
(op != nullptr && op->op == GGML_OP_SET_ROWS && op->src[2] == tensor);
}
@ -301,7 +360,13 @@ public:
op->src[1]->op == GGML_OP_NONE;
}
std::string get_graph_input_ov_name(const ggml_tensor * tensor, const ggml_tensor * op) {
// the state permutation index input used in SSM/DeltaNet models (inp->s_copy in llama-graph.cpp)
inline static bool is_inp_s_copy(const ggml_tensor * tensor, const ggml_tensor * op) {
return op->op == GGML_OP_GET_ROWS && tensor == op->src[1] &&
op->src[0]->buffer->usage == GGML_BACKEND_BUFFER_USAGE_ANY;
}
std::string get_graph_input_ov_name(const ggml_tensor * tensor, const ggml_tensor * op) const {
if (is_inp_pos(tensor, op)) {
return "inp_pos";
}
@ -321,6 +386,10 @@ private:
void compute_model_inputs();
void compute_model_outputs();
// True if tensor is the inp->s_copy index leaf gathered by a recurrent state cache GET_ROWS
// (possibly through a VIEW), so it gets a dynamic [1,1,1,-1] graph-input shape.
bool is_s_copy_leaf(const ggml_tensor * tensor) const;
// Infer and propagate dynamic-dimension indices for all tensors in the GGML graph.
void compute_node_dynamic_dims();
@ -329,12 +398,11 @@ private:
ggml_cgraph * m_cgraph = nullptr;
std::map<std::string, ggml_tensor *> m_inputs;
std::map<std::string, std::shared_ptr<ov::Node>> m_model_inputs;
std::map<std::string, std::shared_ptr<ov::Node>> m_model_extra_inputs;
std::map<std::string, std::shared_ptr<ov::Tensor>> m_model_extra_input_values;
std::map<std::string, ov::frontend::ggml::ModelInputInfo> m_model_inputs;
std::map<std::string, ov::frontend::ggml::ModelExtraInputInfo> m_model_extra_inputs;
std::map<std::string, std::shared_ptr<ov::Node>> m_model_weights;
std::map<std::string, ggml_tensor *> m_model_outputs;
std::vector<std::string> m_model_output_names;
std::set<std::string> m_model_output_names;
std::vector<NodeInfo> m_node_info_list;
std::map<ggml_tensor *, int> m_node_dynamic_dims;

59
ggml/src/ggml-openvino/ggml-openvino-extra.cpp
View File

@ -31,6 +31,7 @@ void ggml_openvino_device_config::init() {
// String values (use ggml_openvino_getenv_str)
"GGML_OPENVINO_DEVICE",
"GGML_OPENVINO_CACHE_DIR",
"GGML_OPENVINO_DEBUG_NODE",
// Integer values (use ggml_openvino_getenv_int)
"GGML_OPENVINO_PREFILL_CHUNK_SIZE",
// Boolean toggles (treated as int flags via ggml_openvino_getenv_int)
@ -44,7 +45,12 @@ void ggml_openvino_device_config::init() {
"GGML_OPENVINO_ENABLE_CACHE",
"GGML_OPENVINO_DISABLE_CACHE",
"GGML_OPENVINO_DISABLE_KV_SLICE",
"GGML_OPENVINO_ENABLE_FALLBACK",
"GGML_OPENVINO_MANUAL_GQA_ATTN",
"GGML_OPENVINO_MEMORY_OPTIMIZE",
"GGML_OPENVINO_RELEASE_WEIGHTS",
"GGML_OPENVINO_REDUCE_COMPILE_MEM",
"GGML_OPENVINO_COMPILED_MODEL_CACHE_DIR",
};
for (const char * const & env_var : env_var_names) {
@ -168,6 +174,22 @@ int ggml_openvino_getenv_int(const char * var, int default_value) {
return v ? std::atoi(v) : default_value;
}
bool ggml_openvino_reduce_compile_mem_enabled() {
const char * reduce_compile_mem = ggml_openvino_getenv_str("GGML_OPENVINO_REDUCE_COMPILE_MEM");
if (reduce_compile_mem != nullptr) {
return ggml_openvino_getenv_int("GGML_OPENVINO_REDUCE_COMPILE_MEM") != 0;
}
return ggml_openvino_getenv_int("GGML_OPENVINO_MEMORY_OPTIMIZE") != 0;
}
bool ggml_openvino_release_weights_enabled(const std::string & device) {
const char * release_weights = ggml_openvino_getenv_str("GGML_OPENVINO_RELEASE_WEIGHTS");
if (release_weights != nullptr) {
return device == "GPU" && ggml_openvino_getenv_int("GGML_OPENVINO_RELEASE_WEIGHTS") != 0;
}
return device == "GPU" && ggml_openvino_getenv_int("GGML_OPENVINO_MEMORY_OPTIMIZE") != 0;
}
// Check if running on NPU
bool ggml_openvino_is_npu() {
return ggml_openvino_get_device_config().is_npu;
@ -252,14 +274,31 @@ ggml_openvino_extracted_layout ggml_openvino_get_extracted_layout(const ggml_ten
return layout;
}
// Only handle 2D weight tensors
if (tensor->ne[2] != 1 || tensor->ne[3] != 1) {
// Most quantized weights use the existing 2D extraction path. 3D expert weights for
// MUL_MAT_ID (MoE) are also supported, either as MXFP4 (packed, dedicated branch below) or via the
// generic sizing math below, which is shape-agnostic (based on total element count). Only reject 4D.
if (tensor->ne[3] != 1) {
return layout;
}
// 3D MoE expert weights that are not requantized (see below) always use the exact f16
// zero-point extraction (see extract_quantized_weights), which needs a wider zp slot than
// the packed integer zero point -- must be kept in sync with that function so the buffer
// sizing here matches what process_weight_tensor actually writes.
const bool for_gather_matmul = tensor->ne[2] > 1;
int64_t n_elements = ggml_nelements(tensor);
const size_t alignment = 64; // Good for SIMD
if (tensor->type == GGML_TYPE_MXFP4 && (tensor->ne[2] > 1 || tensor->ne[3] > 1)) {
layout.weights_per_block = 32;
layout.is_symmetric = true;
layout.weights_size = ggml_nbytes(tensor);
layout.weights_offset = 0;
layout.total_size = layout.weights_size;
return layout;
}
// Check if requantization is needed (NPU-specific)
auto requant_type = ggml_openvino_get_requant_type(tensor, use_bias);
if (requant_type.has_value()) {
@ -334,6 +373,11 @@ ggml_openvino_extracted_layout ggml_openvino_get_extracted_layout(const ggml_ten
layout.is_symmetric = false;
switch (tensor->type) {
case GGML_TYPE_MXFP4:
layout.is_u4 = true;
layout.is_symmetric = true;
break;
case GGML_TYPE_Q4_0:
layout.is_u4 = true;
layout.is_symmetric = true;
@ -369,12 +413,17 @@ ggml_openvino_extracted_layout ggml_openvino_get_extracted_layout(const ggml_ten
// Weights: U4 = n_elements/2 bytes, U8 = n_elements bytes
layout.weights_size = layout.is_u4 ? (n_elements / 2) : n_elements;
// Scales: F16 per block
// Scales: F16 per block, except MXFP4 which stores one E8M0 byte per block.
int64_t n_blocks = n_elements / layout.weights_per_block;
layout.scales_size = n_blocks * sizeof(uint16_t); // F16 = 2 bytes
// For symmetric quantization, no zp needed (weights stored as signed)
layout.scales_size = n_blocks * (tensor->type == GGML_TYPE_MXFP4 ? sizeof(uint8_t) : sizeof(uint16_t));
// For symmetric quantization, no zp needed (weights stored as signed). Asymmetric
// for_gather_matmul (3D MoE expert) weights use an exact f16 zero point (see
// extract_quantized_weights/make_int8_weights/make_int4_weights), which needs one f16 per
// block instead of a packed u4/u8 integer zero point.
if (layout.is_symmetric) {
layout.zp_size = 0;
} else if (use_bias || for_gather_matmul) {
layout.zp_size = n_blocks * sizeof(uint16_t);
} else {
layout.zp_size = layout.is_u4 ? ((n_blocks + 1) / 2) : n_blocks;
}

13
ggml/src/ggml-openvino/ggml-openvino-extra.h
View File

@ -96,9 +96,22 @@ const std::string & ggml_openvino_get_device_name();
const char * ggml_openvino_getenv_str(const char * var, const char * default_value = nullptr);
int ggml_openvino_getenv_int(const char * var, int default_value = 0);
// Memory optimization toggles. GGML_OPENVINO_MEMORY_OPTIMIZE is an umbrella
// switch; the fine-grained env vars still override it when explicitly set.
bool ggml_openvino_reduce_compile_mem_enabled();
bool ggml_openvino_release_weights_enabled(const std::string & device);
// Check if running on NPU
bool ggml_openvino_is_npu();
// Host weight-buffer release (GGML_OPENVINO_RELEASE_WEIGHTS, GPU only).
// register: record a host weight buffer (idempotent per data pointer).
// release: madvise(MADV_DONTNEED) all registered buffers, dropping their RSS.
// released: true once release has run (used to fail-fast on post-release recompile).
void ggml_openvino_register_weight_buffer(void * data, size_t size);
void ggml_openvino_release_weight_buffers();
bool ggml_openvino_weight_buffers_released();
// Get requantization type for a tensor type (returns nullopt if no requant needed)
std::optional<ExtraQuantType> ggml_openvino_get_requant_type(const ggml_tensor * tensor, bool no_requant = false);

349
ggml/src/ggml-openvino/ggml-openvino.cpp
View File

@ -32,6 +32,7 @@
# endif
# include <windows.h>
#else
# include <sys/mman.h>
# include <unistd.h>
#endif
@ -135,6 +136,81 @@ struct ggml_backend_openvino_buffer_type_context {
std::string name;
};
// =====================================================
// Host weight-buffer release (GGML_OPENVINO_RELEASE_WEIGHTS)
// =====================================================
// The OpenVINO weight Constants are zero-copy views into the host buffers
// allocated here (ggml_aligned_malloc, anonymous memory). On GPU the plugin
// holds its own device copy after compile_model, so the host pages are dead
// weight for inference and can be dropped to reclaim RSS (~weights size).
//
// We do NOT free the buffer (ggml owns its lifetime and tensors still point
// into it); instead madvise(MADV_DONTNEED) drops the resident pages while
// keeping the mapping valid. A later recompile would re-read these Constants
// from now-zeroed memory and produce garbage, so once released we fail fast
// if the cache-miss compile branch is reached again (see utils.cpp).
namespace {
struct ov_weight_buffer_registry {
std::mutex mutex;
// (data, size) of every non-remote weight buffer, for madvise.
std::vector<std::pair<void *, size_t>> buffers;
bool released = false;
};
ov_weight_buffer_registry & ov_weight_registry() {
static ov_weight_buffer_registry reg;
return reg;
}
} // namespace
void ggml_openvino_register_weight_buffer(void * data, size_t size) {
if (data == nullptr || size == 0) {
return;
}
auto & reg = ov_weight_registry();
std::lock_guard<std::mutex> lock(reg.mutex);
for (const auto & b : reg.buffers) {
if (b.first == data) {
return; // already registered
}
}
reg.buffers.emplace_back(data, size);
}
bool ggml_openvino_weight_buffers_released() {
auto & reg = ov_weight_registry();
std::lock_guard<std::mutex> lock(reg.mutex);
return reg.released;
}
void ggml_openvino_release_weight_buffers() {
auto & reg = ov_weight_registry();
std::lock_guard<std::mutex> lock(reg.mutex);
if (reg.released) {
return;
}
size_t total = 0;
#if !defined(_WIN32)
for (const auto & b : reg.buffers) {
// Align down/up to page boundaries so madvise only drops whole pages
// fully owned by this buffer.
const long page = sysconf(_SC_PAGESIZE);
uintptr_t start = reinterpret_cast<uintptr_t>(b.first);
uintptr_t end = start + b.second;
uintptr_t astart = (start + page - 1) & ~(uintptr_t) (page - 1);
uintptr_t aend = end & ~(uintptr_t) (page - 1);
if (aend > astart) {
if (madvise(reinterpret_cast<void *>(astart), aend - astart, MADV_DONTNEED) == 0) {
total += aend - astart;
}
}
}
#endif
reg.released = true;
GGML_LOG_INFO("%s: released %zu MB of host weight buffers (%zu buffers)\n", __func__, total / 1024 / 1024,
reg.buffers.size());
}
// Buffer interface functions
static void ggml_backend_openvino_buffer_free_buffer(ggml_backend_buffer_t buffer) {
ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context;
@ -235,10 +311,12 @@ static void ggml_backend_openvino_buffer_set_tensor(ggml_backend_buffer_t buffer
bool is_weight_buffer = (buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
// Full tensor set: offset=0, full size, not a view
bool is_full_tensor_set = (offset == 0 && size == ggml_nbytes(tensor) && tensor->view_src == nullptr);
// 2D tensor (typical weight shape)
// 2D tensor (typical weight shape), or a 3D quantized MoE expert weight (MUL_MAT_ID). Dense 3D
// expert weights are handled later in create_weight_node instead.
bool is_2d = (tensor->ne[2] == 1 && tensor->ne[3] == 1);
bool is_supported_weight_shape = is_2d || (tensor->ne[3] == 1 && ggml_is_quantized(tensor->type));
if (is_weight_buffer && is_full_tensor_set && is_2d) {
if (is_weight_buffer && is_full_tensor_set && is_supported_weight_shape) {
try {
auto result = process_weight_tensor(tensor, data, tensor->data);
result.weight_node->set_friendly_name(tensor->name);
@ -274,6 +352,22 @@ static void ggml_backend_openvino_buffer_set_tensor(ggml_backend_buffer_t buffer
ctx->tensor_extras[tensor] = extra;
tensor->extra = extra;
// Register the host buffer so its pages can be dropped after the GPU
// plugin has its own device copy (GGML_OPENVINO_RELEASE_WEIGHTS).
if (!ctx->is_remote) {
// Weights are set once at model load. Setting a weight after a release
// means a second model is loading while the first's compiled graph is
// pinned — that graph would be wrongly reused with this model's key.
// Fail loud rather than return silently-wrong results.
if (ggml_openvino_weight_buffers_released()) {
GGML_ABORT(
"ggml-openvino: loading a new model while GGML_OPENVINO_RELEASE_WEIGHTS pinned a previous "
"model's compiled graph. This mode supports a single model per process; unset it for "
"multi-model runs.");
}
ggml_openvino_register_weight_buffer(ctx->data, ctx->size);
}
} catch (const std::exception & e) {
GGML_LOG_ERROR("%s: failed to process weight tensor for %s: %s\n", __func__, tensor->name, e.what());
memcpy((char *) tensor->data + offset, data, size);
@ -458,8 +552,8 @@ static size_t ggml_backend_openvino_buffer_type_get_alloc_size(ggml_backend_buff
const ggml_tensor * tensor) {
GGML_UNUSED(buft);
// For quantized 2D tensors (weights), we need extra space for extracted data
if (ggml_is_quantized(tensor->type) && tensor->ne[2] == 1 && tensor->ne[3] == 1) {
// For quantized weight tensors, we need extra space for extracted data.
if (ggml_is_quantized(tensor->type) && tensor->ne[3] == 1) {
ggml_openvino_extracted_layout layout = ggml_openvino_get_extracted_layout(tensor);
if (layout.total_size > 0) {
// GGML_LOG_DEBUG("%s: tensor %s needs %zu bytes (original %zu, extracted: weights=%zu scales=%zu zp=%zu)\n",
@ -618,7 +712,13 @@ static void ggml_backend_openvino_free(ggml_backend_t backend) {
if (ctx->runtime_context) {
auto r_ctx = std::static_pointer_cast<ov_runtime_context>(ctx->runtime_context);
if (--r_ctx->backend_count == 0) {
r_ctx->clear_caches();
// If host weight buffers were released (GGML_OPENVINO_RELEASE_WEIGHTS), the
// dropped pages can never be repopulated, so a recompile is impossible. Keep
// the compiled-model cache alive across backend teardown so the next context
// reuses it instead of recompiling against zeroed weights.
if (!ggml_openvino_weight_buffers_released()) {
r_ctx->clear_caches();
}
}
}
@ -763,6 +863,7 @@ static void ggml_backend_openvino_device_get_props(ggml_backend_dev_t dev, ggml_
/* .host_buffer = */ false,
/* .buffer_from_host_ptr = */ false,
/* .events = */ false,
/* .mmap_support = */ true,
};
}
@ -855,6 +956,32 @@ static bool checked_mul_size(size_t a, size_t b, size_t & out) {
return true;
}
static bool tensor_view_fits_src_buffer(const ggml_tensor * tensor) {
if (tensor->view_src == nullptr) {
return true;
}
const size_t src_nbytes = ggml_nbytes(tensor->view_src);
if (tensor->view_offs > src_nbytes) {
return false;
}
const size_t tensor_nbytes = ggml_nbytes(tensor);
return tensor_nbytes <= src_nbytes - tensor->view_offs;
}
static bool cpy_output_view_is_supported(const ggml_tensor * op) {
if (op->view_src == nullptr) {
return true;
}
if (!tensor_view_fits_src_buffer(op)) {
return false;
}
return ggml_nbytes(op) == 0 || ggml_is_contiguous(op);
}
static bool mul_mat_id_requires_large_tmp(const ggml_tensor * op) {
const ggml_tensor * as = op->src[0];
const ggml_tensor * ids = op->src[2];
@ -862,9 +989,10 @@ static bool mul_mat_id_requires_large_tmp(const ggml_tensor * op) {
return true;
}
// The current OpenVINO translation materializes selected expert weights with
// shape [n_tokens, n_used, rows, k]. Skip cases that would create a very
// large temporary on GPU and let the scheduler fall back instead.
// The MXFP4 MUL_MAT_ID translation (translate_mul_mat_id_mxfp4_packed in mul_mat_id.cpp)
// materializes selected expert weights with shape [n_tokens, n_used, rows, k]. Skip cases that
// would create a very large temporary and let the scheduler fall back instead. Every other weight
// type goes through GatherMatmul, which never materializes this temporary.
size_t tmp_elems = 1;
if (!checked_mul_size(tmp_elems, static_cast<size_t>(ids->ne[1]), tmp_elems) ||
!checked_mul_size(tmp_elems, static_cast<size_t>(ids->ne[0]), tmp_elems) ||
@ -882,12 +1010,56 @@ static bool mul_mat_id_requires_large_tmp(const ggml_tensor * op) {
return tmp_bytes > mul_mat_id_tmp_limit;
}
static bool tensor_name_starts_with(const ggml_tensor * tensor, const char * prefix) {
return tensor != nullptr && strncmp(tensor->name, prefix, strlen(prefix)) == 0;
}
static bool is_msa_block_mask_expansion(const ggml_tensor * op) {
if (tensor_name_starts_with(op, "msa_")) {
return true;
}
const ggml_tensor * src = op->src[0];
while (src != nullptr && (src->op == GGML_OP_RESHAPE || src->op == GGML_OP_REPEAT)) {
if (tensor_name_starts_with(src, "msa_block_mask")) {
return true;
}
src = src->src[0];
}
return tensor_name_starts_with(src, "msa_block_mask");
}
static bool is_op_unsupported_case(const ggml_tensor * op) {
if (is_msa_block_mask_expansion(op)) {
return true;
}
switch (op->op) {
case GGML_OP_CONCAT: {
if (op->type == GGML_TYPE_I64) {
return true;
}
if (ggml_openvino_get_device_name() == "GPU" && op->type == GGML_TYPE_BF16 && has_view_op_input(op)) {
return true;
}
break;
}
case GGML_OP_SET: {
const auto nb1 = static_cast<size_t>(op->op_params[0]);
const auto nb2 = static_cast<size_t>(op->op_params[1]);
const auto nb3 = static_cast<size_t>(op->op_params[2]);
// OpenVINO SET translation currently supports dst layouts that match src0 strides.
if (op->src[0] == nullptr || nb1 != op->src[0]->nb[1] || nb2 != op->src[0]->nb[2] || nb3 != op->src[0]->nb[3]) {
// std::cout << "Unsupported SET op with dst nb1=" << nb1 << ", nb2=" << nb2 << ", nb3=" << nb3
// << " that does not match src0 strides nb[1]="
// << (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[1]) : "null")
// << ", nb[2]=" << (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[2]) : "null")
// << ", nb[3]=" << (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[3]) : "null")
// << std::endl;
return true;
}
break;
}
case GGML_OP_GET_ROWS:
@ -895,23 +1067,24 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
if (op->ne[3] != 1) {
return true;
}
if (op->ne[0] == 256 && (op->src[0]->type == GGML_TYPE_Q4_K || op->src[0]->type == GGML_TYPE_Q5_K)) {
// ERR = 0.000000306 > 0.000000100 GET_ROWS(type=q4_K,n=256,m=5,r=4,be1=1,be2=1,v=0)
// ERR = 0.000000197 > 0.000000100 GET_ROWS(type=q5_K,n=256,m=5,r=4,be1=1,be2=1,v=0)
if (op->op == GGML_OP_GET_ROWS && ggml_openvino_get_device_name() == "GPU" &&
op->src[0]->type == GGML_TYPE_BF16) {
return true;
}
if (op->ne[0] == 256 && (op->src[0]->type == GGML_TYPE_Q4_K || op->src[0]->type == GGML_TYPE_Q5_K ||
op->src[0]->type == GGML_TYPE_Q4_1 || op->src[0]->type == GGML_TYPE_Q5_1)) {
// These are all f16-arithmetic dequant rounding errors that intermittently exceed the
// tight 1e-7 NMSE threshold depending on the random test data (see ggml-quants.cpp
// make_int8_weights/make_int4_weights: dequant is done in f16, not f32, to keep the
// Convert/Subtract/Multiply chain fusable into GatherMatmulCompressed/FullyConnectedCompressed
// for the shared non-test code paths).
return true;
}
// Keep the MoE routing weights gather on CPU for GPU runs. Splitting
// only at the later SUM/CLAMP/DIV nodes still leaves this routing path
// numerically unstable for arctic-style MoE graphs.
if (strncmp(op->name, "ffn_moe_weights", sizeof("ffn_moe_weights") - 1) == 0) {
return true;
}
break;
}
case GGML_OP_RESHAPE: {
if (strncmp(op->name, "ffn_moe_weights", sizeof("ffn_moe_weights") - 1) == 0 ||
strncmp(op->name, "ffn_norm_exps", sizeof("ffn_norm_exps") - 1) == 0) {
if (strncmp(op->name, "ffn_norm_exps", sizeof("ffn_norm_exps") - 1) == 0) {
return true;
}
break;
@ -938,69 +1111,22 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
break;
}
case GGML_OP_DIV: {
bool requires_broadcast = false;
for (int i = 0; i < 4; i++) {
if (op->src[0]->ne[i] == op->src[1]->ne[i]) {
continue;
}
if (op->src[0]->ne[i] != 1 && op->src[1]->ne[i] != 1) {
return true;
}
requires_broadcast = true;
}
// The GPU plugin can fuse broadcast DIV into the preceding FFN GEMM path
// and produce infs for per-channel scale vectors. Keep those DIVs on CPU
// until the fused GPU kernel is reliable. (falied case llama-arch-test mpt)
if (requires_broadcast && ggml_openvino_get_device_name() == "GPU") {
return true;
}
// qwen3next MoE weight normalization is numerically sensitive on the GPU
// path. Keep the normalization divide on CPU to match the reference.
if (strncmp(op->name, "ffn_moe_weights_norm", sizeof("ffn_moe_weights_norm") - 1) == 0) {
return true;
}
break;
}
case GGML_OP_SOFT_MAX: {
if (op->src[2] != nullptr) {
// GGML_LOG_WARN("OpenVINO backend does not support SOFT_MAX with sinks\n");
return true;
}
if (strncmp(op->name, "ffn_moe_probs", sizeof("ffn_moe_probs") - 1) == 0) {
return true;
}
// GPU execution of the MoE routing weights softmax is numerically unstable
// when fused with the surrounding GET_ROWS/reshape path. Keep this softmax
// on CPU so the scheduler splits at the same boundary that restores parity.
if (op->src[0] != nullptr && op->src[0]->op == GGML_OP_RESHAPE && op->src[0]->src[0] != nullptr &&
strncmp(op->src[0]->src[0]->name, "ffn_moe_weights", sizeof("ffn_moe_weights") - 1) == 0) {
if (ggml_openvino_get_device_name() == "GPU" && op->src[1]->ne[0] == op->ne[0] &&
op->src[1]->ne[1] == 1 && op->src[1]->ne[2] == 1 && op->src[1]->ne[3] == 1) {
return true;
}
break;
}
case GGML_OP_SUM_ROWS: {
if (strncmp(op->name, "ffn_moe_weights_sum", sizeof("ffn_moe_weights_sum") - 1) == 0) {
return true;
}
// if the input is PERMUTE skip
if (op->src[0]->op == GGML_OP_PERMUTE) {
return true;
}
break;
}
case GGML_OP_CLAMP: {
if (strncmp(op->name, "ffn_moe_weights_sum_clamped", sizeof("ffn_moe_weights_sum_clamped") - 1) == 0) {
return true;
}
break;
}
case GGML_OP_FLASH_ATTN_EXT: {
float scale = 1.0f;
float max_bias = 0.0f;
@ -1047,23 +1173,29 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
// GGML_LOG_WARN("OpenVINO backend does not support CPY with non-contiguous data or bf16 types\n");
return true;
}
// CPY to a quantized destination (e.g. f32 -> q4_0) is numerically unstable with OpenVINO backend.
if (ggml_is_quantized(op->type)) {
return true;
}
if (ggml_nelements(op->src[0]) != ggml_nelements(op->src[1])) {
return true;
}
// op test case with non-contiguous src or dst
if ((op->ne[0] == 3 && op->ne[1] == 4 && op->ne[2] == 3 && op->ne[3] == 2) ||
(op->ne[0] == 1 && op->ne[1] == 4 && op->ne[2] == 3 && op->ne[3] == 2) ||
(op->ne[0] == 2 && op->ne[1] == 4 && op->ne[2] == 3 && op->ne[3] == 2)) {
return true;
}
// CPY into a strided view of a larger buffer (recurrent-state snapshots) not supported
if (op->view_src && ggml_nbytes(op) != ggml_nbytes(op->view_src)) {
if (!cpy_output_view_is_supported(op)) {
return true;
}
break;
}
case GGML_OP_MUL_MAT: {
if (ggml_openvino_get_device_name() == "GPU" && op->src[1]->op == GGML_OP_SOFT_MAX &&
op->src[0]->op == GGML_OP_CONT && op->src[0]->src[0] != nullptr &&
op->src[0]->src[0]->op == GGML_OP_TRANSPOSE && op->src[0]->src[0]->src[0] != nullptr &&
op->src[0]->src[0]->src[0]->op == GGML_OP_PERMUTE) {
if (ggml_openvino_get_device_name() == "GPU" && op->src[0] != nullptr && op->src[1] != nullptr &&
ggml_is_quantized(op->src[0]->type) && strcmp(op->src[0]->name, "a") == 0 &&
strcmp(op->src[1]->name, "b") == 0 && op->src[0]->ne[1] == 1 && op->src[1]->ne[1] == 64 &&
op->src[0]->ne[0] == 256 && op->src[1]->ne[0] == 256) {
return true;
}
if (op->src[0]->ne[3] != op->src[1]->ne[3] && op->src[0]->ne[3] != 1 && op->src[1]->ne[3] != 1) {
@ -1075,12 +1207,18 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
break;
}
case GGML_OP_MUL_MAT_ID: {
if (strncmp(op->name, "ffn_moe_gate_up", sizeof("ffn_moe_gate_up") - 1) == 0 ||
strncmp(op->name, "ffn_moe_down", sizeof("ffn_moe_down") - 1) == 0) {
// Single-expert (or empty) MUL_MAT_ID is a degenerate shape that stresses GatherMatmul edge
// cases and never occurs in real MoE; let it fall back to CPU.
if (op->src[0] != nullptr && op->src[0]->ne[2] <= 1) {
return true;
}
if (mul_mat_id_requires_large_tmp(op)) {
if (ggml_openvino_get_device_name() == "GPU" && op->src[0] != nullptr && op->src[0]->type == GGML_TYPE_BF16) {
return true;
}
// GPU MUL_MAT_ID uses a Gather+MatMul fallback because the GPU plugin rejects internal
// GatherMatmul for these test shapes. Skip cases that would materialize a large selected
// expert-weight temporary.
if (ggml_openvino_get_device_name() == "GPU" && mul_mat_id_requires_large_tmp(op)) {
return true;
}
break;
@ -1093,8 +1231,10 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
// GGML_LOG_WARN("OpenVINO backend does not support ROPE with mode %d\n", mode);
return true;
}
if (n_dims != 0.0f && n_dims != op->src[0]->ne[0]) {
// GGML_LOG_WARN("OpenVINO backend does not support ROPE with n_dims %d != src[0]->ne[0] %ld\n", n_dims,
const int64_t head_dim = op->src[0]->ne[0];
const int64_t rope_dims = n_dims == 0 ? head_dim : n_dims;
if (rope_dims <= 0 || rope_dims > head_dim || (rope_dims % 2) != 0) {
// GGML_LOG_WARN("OpenVINO backend does not support ROPE with n_dims %d and src[0]->ne[0] %ld\n", n_dims,
// op->src[0]->ne[0]);
return true;
}
@ -1127,9 +1267,15 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
}
break;
}
case GGML_OP_REPEAT: {
if (ggml_openvino_get_device_name() == "GPU" && op->type == GGML_TYPE_BF16) {
return true;
}
break;
}
case GGML_OP_GATED_DELTA_NET: {
// enable after https://github.com/openvinotoolkit/openvino/pull/35917 is included in OV release
return true;
// return true;
// if (ggml_openvino_get_device_name() == "GPU" && op->src[0]->ne[2] > 1) {
// // CVS-186471
// return true;
@ -1141,13 +1287,8 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
if (op->src[3]->ne[0] != 1) {
return true;
}
// v_repeat > 1 (GQA): ggml uses modulo head mapping (h_q = h_v % H_k)
// but the fused op uses consecutive mapping (h_q = h_v / group_size)
if (op->src[2]->ne[1] != op->src[0]->ne[1]) {
return true;
}
// K > 1 (multiple state snapshots) not supported by fused op
if (op->src[5]->ne[1] > 1) {
if (((const int32_t *) op->op_params)[0] > 1) {
return true;
}
break;
@ -1155,11 +1296,12 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
case GGML_OP_SSM_CONV: {
// qwen3next is numerically unstable with OpenVINO SSM_CONV.
// Keep this op on CPU until the OpenVINO implementation is fixed.
return true;
// return true;
break;
}
case GGML_OP_VIEW: {
// Skip TOPK_MOE fused tests until it is fully supported
// the argsort_top_k VIEW wrapping ARGSORT is named "selected_experts" in test_topk_moe
// Skip TOPK_MOE fused tests until it is fully supported.
// The argsort_top_k VIEW wrapping ARGSORT is named "selected_experts" in test_topk_moe.
if (strcmp(op->name, "selected_experts") == 0) {
return true;
}
@ -1176,7 +1318,8 @@ static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, con
static std::unordered_set<ggml_type> supported_types{
GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_I64, GGML_TYPE_I32, GGML_TYPE_Q4_0,
GGML_TYPE_Q4_1, GGML_TYPE_Q4_K, GGML_TYPE_Q5_1, GGML_TYPE_Q5_K, GGML_TYPE_Q8_0, GGML_TYPE_Q6_K};
GGML_TYPE_Q4_1, GGML_TYPE_Q4_K, GGML_TYPE_Q5_1, GGML_TYPE_Q5_K, GGML_TYPE_Q8_0, GGML_TYPE_Q6_K,
GGML_TYPE_MXFP4};
// derive supported op sets from the op_table map, keys in
// the map use the full macro name (e.g. "GGML_OP_ADD"), while
@ -1223,6 +1366,9 @@ static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, con
// GGML_LOG_WARN("OpenVINO backend does not support unary op %s\n", ggml_unary_op_name(ggml_get_unary_op(op)));
return false;
}
if (ggml_get_unary_op(op) == GGML_UNARY_OP_EXP && op->type == GGML_TYPE_F32) {
return false;
}
break;
}
case GGML_OP_GLU: {
@ -1231,11 +1377,11 @@ static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, con
// GGML_LOG_WARN("OpenVINO backend does not support GLU op %s\n", ggml_glu_op_name(ggml_get_glu_op(op)));
return false;
}
if (has_view_op_input(op)) {
// GGML_LOG_WARN("OpenVINO backend does not support unary op %s with view input\n",
// ggml_glu_op_name(ggml_get_glu_op(op)));
return false;
}
// if (has_view_op_input(op)) {
// // GGML_LOG_WARN("OpenVINO backend does not support unary op %s with view input\n",
// // ggml_glu_op_name(ggml_get_glu_op(op)));
// return false;
// }
if (op->src[1] == nullptr && op->src[0]->ne[0] % 2 != 0) {
// triggers bug in ov gpu
return false;
@ -1248,16 +1394,11 @@ static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, con
// GGML_LOG_WARN("OpenVINO backend does not support op %s\n", ggml_op_name(op->op));
return false;
}
static std::set<ggml_op> ops_not_support_view_input{
GGML_OP_L2_NORM,
};
static std::set<ggml_op> ops_not_support_view_input{};
if (ops_not_support_view_input.find(op->op) != ops_not_support_view_input.end() && has_view_op_input(op)) {
// GGML_LOG_WARN("OpenVINO backend does not support op %s with view input\n", ggml_op_name(op->op));
return false;
}
if (op->op == GGML_OP_RMS_NORM && has_non_contiguous_view_input(op)) {
return false;
}
}
}
@ -1274,7 +1415,9 @@ static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, con
// GGML_LOG_WARN("OpenVINO backend does not support tensor type %s\n", ggml_type_name(src->type));
return false;
}
if (ggml_is_quantized(src->type) && src->ne[2] != 1) {
const bool is_supported_3d_moe_expert =
op->op == GGML_OP_MUL_MAT_ID && i == 0 && (src->type == GGML_TYPE_MXFP4 || src->ne[3] == 1);
if (ggml_is_quantized(src->type) && src->ne[2] != 1 && !is_supported_3d_moe_expert) {
// GGML_LOG_WARN("OpenVINO backend does not support 3D quantized tensors\n");
return false;
}

381
ggml/src/ggml-openvino/ggml-quants.cpp
View File

@ -2,6 +2,7 @@
#include "ggml-common.h"
#include "ggml-impl.h"
#include "ggml-openvino-extra.h"
#include "ggml.h"
#include <algorithm>
@ -19,6 +20,8 @@
#include <openvino/core/type/element_type.hpp>
#include <openvino/core/type/element_type_traits.hpp>
#include <openvino/core/type/float16.hpp>
#include <openvino/core/type/float4_e2m1.hpp>
#include <openvino/core/type/float8_e8m0.hpp>
#include <openvino/op/add.hpp>
#include <openvino/op/constant.hpp>
#include <openvino/op/convert.hpp>
@ -26,6 +29,7 @@
#include <openvino/op/reshape.hpp>
#include <openvino/op/subtract.hpp>
#include <openvino/op/util/attr_types.hpp>
#include <openvino/pass/constant_folding.hpp>
#include <openvino/runtime/tensor.hpp>
#include <string>
#include <vector>
@ -44,6 +48,38 @@ void unpack_32_4(const uint8_t * data, uint8_t * dst) {
}
}
static constexpr size_t MXFP4_BLOCK_SIZE = 32;
static constexpr size_t MXFP4_BLOCK_QS_SIZE = MXFP4_BLOCK_SIZE / 2;
static constexpr size_t MXFP4_BLOCK_BYTES = sizeof(uint8_t) + MXFP4_BLOCK_QS_SIZE;
static void pack_32_mxfp4_for_openvino(const uint8_t * data, uint8_t * dst) {
for (int j = 0; j < static_cast<int>(MXFP4_BLOCK_QS_SIZE); j += 2) {
const uint8_t v0 = data[j] & 0x0F;
const uint8_t v1 = (data[j + 1] & 0x0F) << 4;
const uint8_t v16 = data[j] >> 4;
const uint8_t v17 = data[j + 1] & 0xF0;
dst[j / 2] = v0 | v1;
dst[MXFP4_BLOCK_SIZE / 4 + j / 2] = v16 | v17;
}
}
void extract_mxfp4_data(const ggml_tensor * tensor, ov::Tensor & weights_arr, ov::Tensor & scales_arr) {
GGML_ASSERT(tensor->type == GGML_TYPE_MXFP4);
GGML_ASSERT(weights_arr.get_element_type() == ov::element::f4e2m1);
GGML_ASSERT(scales_arr.get_element_type() == ov::element::f8e8m0);
const auto * data = static_cast<const uint8_t *>(tensor->data);
auto * weights = static_cast<uint8_t *>(weights_arr.data());
auto * scales = scales_arr.data<ov::element_type_traits<ov::element::f8e8m0>::value_type>();
const size_t n_blocks = scales_arr.get_size();
ov::parallel_for(n_blocks, [&](size_t i) {
const uint8_t * block = data + i * MXFP4_BLOCK_BYTES;
pack_32_mxfp4_for_openvino(block + sizeof(uint8_t), weights + i * MXFP4_BLOCK_QS_SIZE);
scales[i] = ov::float8_e8m0::from_bits(block[0]);
});
}
// Extracts (weight, scales, zp) from Q4_0 tensors.
// Data layout is: |16 bit scale|32 x 4bit weights|.
// When zp_arr is empty (symmetric), weights are stored as signed i4 (value - 8).
@ -470,22 +506,34 @@ void extract_q5_k_data(const ggml_tensor * tensor,
// TODO Reorder for make_intX_weights
// If for_gather_matmul is true, weight may be N-D (e.g. 3D MoE expert weights [n_expert, rows, cols]).
// The dequantization chain below is built as usual but left in f16 (no final Convert to f32) --
// ov::pass::MarkDequantization (registered in translate_session.cpp) marks the chain so it survives
// model-build-time ConstantFolding. mul_mat_id.cpp constructs ov::op::internal::GatherMatmul directly
// on top of the resulting f16 chain.
ov::Output<ov::Node> make_int8_weights(ov::Tensor & weight,
ov::Tensor & scales,
ov::Tensor & zp,
size_t group_size,
bool use_bias) {
bool use_bias,
bool for_gather_matmul) {
ov::Shape orig_shape = weight.get_shape();
bool is_signed = (weight.get_element_type() == ov::element::i8); // Symmetric: signed weights, no ZP
// Expand dimensions for scales and zp/bias
auto scale_shape = scales.get_shape();
ov::Shape packed_shape = {orig_shape[0], orig_shape[1] / group_size, group_size};
// Group the innermost (last) dimension. For 2D weights [rows, cols] this yields
// [rows, cols/group_size, group_size]; for 3D MoE experts [n_expert, rows, cols] this yields
// [n_expert, rows, cols/group_size, group_size].
ov::Shape packed_shape = orig_shape;
packed_shape.back() /= group_size;
packed_shape.push_back(group_size);
const size_t group_dim = packed_shape.size() - 2;
if (packed_shape[1] == 1) {
if (packed_shape[group_dim] == 1) {
// Requantized channel-wise case
packed_shape.erase(packed_shape.begin() + 1);
packed_shape.erase(packed_shape.begin() + group_dim);
} else {
scale_shape.push_back(1);
scales.set_shape(scale_shape);
@ -505,7 +553,8 @@ ov::Output<ov::Node> make_int8_weights(ov::Tensor & weight,
static_cast<uint8_t *>(weight.data()), nullptr);
weights_node->get_rt_info()["__gguf_tensor_holder"] = weight;
auto weights_f16 = std::make_shared<ov::op::v0::Convert>(weights_node, ov::element::f16);
result = std::make_shared<ov::op::v1::Multiply>(weights_f16, scales_f16, ov::op::AutoBroadcastType::NUMPY);
auto mul = std::make_shared<ov::op::v1::Multiply>(weights_f16, scales_f16, ov::op::AutoBroadcastType::NUMPY);
result = mul;
} else {
// Unsigned path
auto weights_node = std::make_shared<ov::op::v0::Constant>(ov::element::u8, packed_shape,
@ -514,11 +563,25 @@ ov::Output<ov::Node> make_int8_weights(ov::Tensor & weight,
auto weights_f16 = std::make_shared<ov::op::v0::Convert>(weights_node, ov::element::f16);
if (use_bias && zp.get_size() > 0) {
// Bias path: w * s + b (zp tensor holds f16 bias values)
auto bias_f16 = std::make_shared<ov::op::v0::Constant>(zp);
auto w_s =
std::make_shared<ov::op::v1::Multiply>(weights_f16, scales_f16, ov::op::AutoBroadcastType::NUMPY);
result = std::make_shared<ov::op::v1::Add>(w_s, bias_f16, ov::op::AutoBroadcastType::NUMPY);
// Accurate dequant in the FUSABLE zero-point form: (w - zp) * s, where the zero
// point is an exact f16 value zp = -bias/scale (the zp tensor holds bias values
// coming in). Algebraically equal to w*s + bias, but unlike an Add(bias) graph this
// matches CompressedWeightsBlock's pattern (Constant->Convert->Subtract->Multiply),
// so for_gather_matmul weights still fuse into GatherMatmulCompressed. Also avoids
// the round(min/scale) error of an integer zero point. Convert bias -> zero-point IN
// PLACE in the (possibly buffer-backed) zp tensor to avoid a duplicate allocation.
auto * bias_zp_data = zp.data<ov::float16>();
const auto * scale_data = scales.data<ov::float16>();
const size_t n = zp.get_size();
for (size_t i = 0; i < n; i++) {
float s = static_cast<float>(scale_data[i]);
float b = static_cast<float>(bias_zp_data[i]);
bias_zp_data[i] = ov::float16(s != 0.0f ? -b / s : 0.0f);
}
auto zero_point_f16 = std::make_shared<ov::op::v0::Constant>(zp);
auto w_zp =
std::make_shared<ov::op::v1::Subtract>(weights_f16, zero_point_f16, ov::op::AutoBroadcastType::NUMPY);
result = std::make_shared<ov::op::v1::Multiply>(w_zp, scales_f16, ov::op::AutoBroadcastType::NUMPY);
} else {
// Zero point path: (w - zp) * s
auto zero_point = std::make_shared<ov::op::v0::Constant>(zp);
@ -529,37 +592,49 @@ ov::Output<ov::Node> make_int8_weights(ov::Tensor & weight,
auto zero_point_f16 = std::make_shared<ov::op::v0::Convert>(zero_point, ov::element::f16);
auto w_zp =
std::make_shared<ov::op::v1::Subtract>(weights_f16, zero_point_f16, ov::op::AutoBroadcastType::NUMPY);
result = std::make_shared<ov::op::v1::Multiply>(w_zp, scales_f16, ov::op::AutoBroadcastType::NUMPY);
auto mul = std::make_shared<ov::op::v1::Multiply>(w_zp, scales_f16, ov::op::AutoBroadcastType::NUMPY);
result = mul;
}
}
if (packed_shape.size() != 2) {
if (packed_shape.size() != orig_shape.size()) {
// If not requantized channel-wise case, reshape back to original shape
auto final_shape =
std::make_shared<ov::op::v0::Constant>(ov::element::i64, ov::Shape{orig_shape.size()}, orig_shape);
result = std::make_shared<ov::op::v1::Reshape>(result, final_shape, false);
auto reshaped = std::make_shared<ov::op::v1::Reshape>(result, final_shape, false);
result = reshaped;
}
if (for_gather_matmul) {
return result;
}
return std::make_shared<ov::op::v0::Convert>(result, ov::element::f32);
}
// See make_int8_weights for the meaning of for_gather_matmul.
ov::Output<ov::Node> make_int4_weights(ov::Tensor & weight,
ov::Tensor & scales,
ov::Tensor & zp,
size_t group_size,
bool use_bias) {
bool use_bias,
bool for_gather_matmul) {
ov::Shape orig_weight_shape = weight.get_shape();
bool is_signed = (weight.get_element_type() == ov::element::i4); // Symmetric: signed weights, no ZP
// Expand dimensions for scales and zp/bias
ov::Shape scale_shape = scales.get_shape();
// Create INT4 weight tensor
ov::Shape packed_shape = {orig_weight_shape[0], orig_weight_shape[1] / group_size, group_size};
// Create INT4 weight tensor. Group the innermost (last) dimension: for 2D weights
// [rows, cols] this yields [rows, cols/group_size, group_size]; for 3D MoE experts
// [n_expert, rows, cols] this yields [n_expert, rows, cols/group_size, group_size].
ov::Shape packed_shape = orig_weight_shape;
packed_shape.back() /= group_size;
packed_shape.push_back(group_size);
const size_t group_dim = packed_shape.size() - 2;
if (packed_shape[1] == 1) {
if (packed_shape[group_dim] == 1) {
// Requantized channel-wise case
packed_shape.erase(packed_shape.begin() + 1);
packed_shape.erase(packed_shape.begin() + group_dim);
} else {
scale_shape.push_back(1);
scales.set_shape(scale_shape);
@ -579,7 +654,8 @@ ov::Output<ov::Node> make_int4_weights(ov::Tensor & weight,
static_cast<uint8_t *>(weight.data()), nullptr);
weights_node->get_rt_info()["__gguf_tensor_holder"] = weight;
auto weights_f16 = std::make_shared<ov::op::v0::Convert>(weights_node, ov::element::f16);
result = std::make_shared<ov::op::v1::Multiply>(weights_f16, scales_f16, ov::op::AutoBroadcastType::NUMPY);
auto mul = std::make_shared<ov::op::v1::Multiply>(weights_f16, scales_f16, ov::op::AutoBroadcastType::NUMPY);
result = mul;
} else {
// Unsigned path
auto weights_node = std::make_shared<ov::op::v0::Constant>(ov::element::u4, packed_shape,
@ -588,11 +664,23 @@ ov::Output<ov::Node> make_int4_weights(ov::Tensor & weight,
auto weights_f16 = std::make_shared<ov::op::v0::Convert>(weights_node, ov::element::f16);
if (use_bias && zp.get_size() > 0) {
// Bias path: w * s + b (zp tensor holds f16 bias values)
auto bias_f16 = std::make_shared<ov::op::v0::Constant>(zp);
auto w_s =
std::make_shared<ov::op::v1::Multiply>(weights_f16, scales_f16, ov::op::AutoBroadcastType::NUMPY);
result = std::make_shared<ov::op::v1::Add>(w_s, bias_f16, ov::op::AutoBroadcastType::NUMPY);
// Accurate dequant in the FUSABLE zero-point form: (w - zp) * s with an exact f16
// zp = -bias/scale. Equivalent to w*s + bias but matches CompressedWeightsBlock's
// pattern so for_gather_matmul weights still fuse into GatherMatmulCompressed, and
// avoids the round(min/scale) error of an integer zp. Convert bias -> zero-point IN
// PLACE in the (possibly buffer-backed) zp tensor to avoid a duplicate allocation.
auto * bias_zp_data = zp.data<ov::float16>();
const auto * scale_data = scales.data<ov::float16>();
const size_t n = zp.get_size();
for (size_t i = 0; i < n; i++) {
float s = static_cast<float>(scale_data[i]);
float b = static_cast<float>(bias_zp_data[i]);
bias_zp_data[i] = ov::float16(s != 0.0f ? -b / s : 0.0f);
}
auto zero_points_f16 = std::make_shared<ov::op::v0::Constant>(zp);
auto w_zp =
std::make_shared<ov::op::v1::Subtract>(weights_f16, zero_points_f16, ov::op::AutoBroadcastType::NUMPY);
result = std::make_shared<ov::op::v1::Multiply>(w_zp, scales_f16, ov::op::AutoBroadcastType::NUMPY);
} else {
// Zero point path: (w - zp) * s
auto zero_points_node = std::make_shared<ov::op::v0::Constant>(zp);
@ -603,20 +691,61 @@ ov::Output<ov::Node> make_int4_weights(ov::Tensor & weight,
auto zero_points_f16 = std::make_shared<ov::op::v0::Convert>(zero_points_node, ov::element::f16);
auto w_zp =
std::make_shared<ov::op::v1::Subtract>(weights_f16, zero_points_f16, ov::op::AutoBroadcastType::NUMPY);
result = std::make_shared<ov::op::v1::Multiply>(w_zp, scales_f16, ov::op::AutoBroadcastType::NUMPY);
auto mul = std::make_shared<ov::op::v1::Multiply>(w_zp, scales_f16, ov::op::AutoBroadcastType::NUMPY);
result = mul;
}
}
if (packed_shape.size() != 2) {
if (packed_shape.size() != orig_weight_shape.size()) {
// If not requantized channel-wise case, reshape back to original shape
auto final_shape = std::make_shared<ov::op::v0::Constant>(ov::element::i64, ov::Shape{orig_weight_shape.size()},
orig_weight_shape);
result = std::make_shared<ov::op::v1::Reshape>(result, final_shape, false);
auto reshaped = std::make_shared<ov::op::v1::Reshape>(result, final_shape, false);
result = reshaped;
}
if (for_gather_matmul) {
return result;
}
return std::make_shared<ov::op::v0::Convert>(result, ov::element::f32);
}
ov::Output<ov::Node> make_mxfp4_weights(ov::Tensor & weight, ov::Tensor & scales) {
const ov::Shape final_shape = weight.get_shape();
GGML_ASSERT(!final_shape.empty());
GGML_ASSERT(final_shape.back() % MXFP4_BLOCK_SIZE == 0);
ov::Shape packed_shape = final_shape;
packed_shape.back() /= MXFP4_BLOCK_SIZE;
packed_shape.push_back(MXFP4_BLOCK_SIZE);
ov::Shape scale_shape = packed_shape;
scale_shape.back() = 1;
scales.set_shape(scale_shape);
auto weights_node = std::make_shared<ov::op::v0::Constant>(ov::element::f4e2m1, packed_shape,
static_cast<uint8_t *>(weight.data()), nullptr);
weights_node->get_rt_info()["__gguf_tensor_holder"] = weight;
auto weights_f32 = std::make_shared<ov::op::v0::Convert>(weights_node, ov::element::f32);
auto scales_node = std::make_shared<ov::op::v0::Constant>(scales);
auto scales_f32 = std::make_shared<ov::op::v0::Convert>(scales_node, ov::element::f32);
ov::Output<ov::Node> result =
std::make_shared<ov::op::v1::Multiply>(weights_f32, scales_f32, ov::op::AutoBroadcastType::NUMPY);
auto final_shape_node =
std::make_shared<ov::op::v0::Constant>(ov::element::i64, ov::Shape{final_shape.size()}, final_shape);
return std::make_shared<ov::op::v1::Reshape>(result, final_shape_node, false);
}
ov::Output<ov::Node> make_mxfp4_moe_packed_weights(ov::Tensor & weight) {
auto weights_node = std::make_shared<ov::op::v0::Constant>(ov::element::u8, weight.get_shape(),
static_cast<uint8_t *>(weight.data()), nullptr);
weights_node->get_rt_info()["__gguf_tensor_holder"] = weight;
weights_node->get_rt_info()["__ggml_openvino_mxfp4_moe_packed"] = true;
return weights_node;
}
// Extract quantized weights from tensor and create weight subgraph
std::shared_ptr<ov::Node> extract_quantized_weights(const ggml_tensor * tensor,
const void * data,
@ -628,6 +757,13 @@ std::shared_ptr<ov::Node> extract_quantized_weights(const ggml_tensor * tensor,
ggml_tensor temp_tensor = *tensor;
temp_tensor.data = const_cast<void *>(data);
if (tensor->type == GGML_TYPE_MXFP4) {
extract_mxfp4_data(&temp_tensor, weights, scales);
auto result = make_mxfp4_weights(weights, scales).get_node_shared_ptr();
result->set_friendly_name(tensor->name);
return result;
}
// Determine block size based on tensor type
int64_t weights_per_block;
bool is_u4;
@ -653,6 +789,13 @@ std::shared_ptr<ov::Node> extract_quantized_weights(const ggml_tensor * tensor,
std::string(ggml_type_name(tensor->type)));
}
// 3D MoE expert weights (for_gather_matmul) always use the exact f16 zero-point extraction
// (see make_int8_weights/make_int4_weights) rather than the rounded integer zero point --
// round(min/scale) error is what corrupts Q4_K/Q5_1 experts, and the f16-zp form still fuses
// into GatherMatmulCompressed since it stays a Subtract, not an Add.
const bool for_gather_matmul = tensor->ne[2] > 1;
use_bias = use_bias || for_gather_matmul;
// Extract quantized data
switch (tensor->type) {
case GGML_TYPE_Q4_0:
@ -680,12 +823,13 @@ std::shared_ptr<ov::Node> extract_quantized_weights(const ggml_tensor * tensor,
throw std::runtime_error("Unsupported quantized type: " + std::string(ggml_type_name(tensor->type)));
}
// Create the OpenVINO weight subgraph
// Create the OpenVINO weight subgraph. 3D expert weights (MoE) are routed through the
// GatherMatmul-oriented path: dequantized in f16, with constant folding disabled on the chain.
ov::Output<ov::Node> weight_node;
if (is_u4) {
weight_node = make_int4_weights(weights, scales, zp, weights_per_block, use_bias);
weight_node = make_int4_weights(weights, scales, zp, weights_per_block, use_bias, for_gather_matmul);
} else {
weight_node = make_int8_weights(weights, scales, zp, weights_per_block, use_bias);
weight_node = make_int8_weights(weights, scales, zp, weights_per_block, use_bias, for_gather_matmul);
}
auto result = weight_node.get_node_shared_ptr();
@ -702,28 +846,76 @@ std::shared_ptr<ov::Node> requantize_to_buffers(const ggml_tensor * tensor,
ov::Tensor & scales,
ov::Tensor & zp) {
int64_t n_elements = ggml_nelements(tensor);
const int64_t ne0 = tensor->ne[0]; // elements per row
const int64_t n_rows = n_elements / ne0;
const auto * type_traits = ggml_get_type_traits(tensor->type);
const size_t src_row_bytes = ggml_row_size(tensor->type, ne0);
// First dequantize to F32
std::vector<float> weights_f32(n_elements);
ggml_get_type_traits(tensor->type)->to_float(data, weights_f32.data(), n_elements);
// Handle F16 case - just convert and create constant
if (requant_type == ExtraQuantType::F16) {
ggml_get_type_traits(GGML_TYPE_F16)->from_float_ref(weights_f32.data(), weights.data(), n_elements);
auto result = std::make_shared<ov::op::v0::Constant>(weights);
result->set_friendly_name(tensor->name);
return result;
}
// Requantize to target quantized format
bool is_u4 = (requant_type == ExtraQuantType::Q4_0_C || requant_type == ExtraQuantType::Q4_0_128);
if (is_u4) {
quantize_q4_0(weights_f32.data(), weights, scales, zp, n_elements, block_size);
} else if (requant_type == ExtraQuantType::Q8_1_C) {
quantize_q8_1(weights_f32.data(), weights, scales, zp, n_elements, block_size);
// Streaming dequant (opt-in via GGML_OPENVINO_REDUCE_COMPILE_MEM or
// GGML_OPENVINO_MEMORY_OPTIMIZE): instead of
// materializing the full n_elements F32 array (e.g. ~1 GB for token_embd), dequantize
// a chunk of complete rows into a small scratch and quantize/convert it straight into
// the output buffers, capping the transient F32 footprint at CHUNK_ROWS*ne0 floats.
//
// Only valid (and only used) for the Q8_0_C / Q8_1_C / F16 targets whose block size
// divides a row (channel-wise _C uses block_size == ne0) so no target block straddles
// a row boundary, and Q8/F16 have no cross-block packing. The u4 (Q4_0) path packs two
// weights per byte with running zp ORs that assume a single whole-array call, so it is
// never streamed. When the flag is off, behavior is identical to the original
// full-materialization path.
const bool stream_requant = ggml_openvino_reduce_compile_mem_enabled() && !is_u4 &&
!(block_size > 0 && ne0 % block_size != 0);
if (!stream_requant) {
// Full materialization (original behavior): dequantize the whole tensor to F32,
// then convert/quantize in one call.
std::vector<float> weights_f32(n_elements);
type_traits->to_float(data, weights_f32.data(), n_elements);
if (requant_type == ExtraQuantType::F16) {
ggml_get_type_traits(GGML_TYPE_F16)->from_float_ref(weights_f32.data(), weights.data(), n_elements);
auto result = std::make_shared<ov::op::v0::Constant>(weights);
result->set_friendly_name(tensor->name);
return result;
}
if (is_u4) {
quantize_q4_0(weights_f32.data(), weights, scales, zp, n_elements, block_size);
} else if (requant_type == ExtraQuantType::Q8_1_C) {
quantize_q8_1(weights_f32.data(), weights, scales, zp, n_elements, block_size);
} else {
quantize_q8_0(weights_f32.data(), weights, scales, zp, n_elements, block_size);
}
} else {
quantize_q8_0(weights_f32.data(), weights, scales, zp, n_elements, block_size);
// Streaming path for Q8_0_C / Q8_1_C / F16 (covers token_embd, output.weight,
// and per-layer Q6_K/Q5_K requant — the large transient cases).
const int64_t CHUNK_ROWS = std::min<int64_t>(n_rows, 256);
std::vector<float> scratch(CHUNK_ROWS * ne0);
// F16 destination: 2 bytes/element, advanced per chunk by r0*ne0 elements.
auto * f16_base = static_cast<uint8_t *>(weights.data());
for (int64_t r0 = 0; r0 < n_rows; r0 += CHUNK_ROWS) {
const int64_t rows = std::min(CHUNK_ROWS, n_rows - r0);
const int64_t elems = rows * ne0;
const auto * src = static_cast<const uint8_t *>(data) + r0 * src_row_bytes;
type_traits->to_float(src, scratch.data(), elems);
if (requant_type == ExtraQuantType::F16) {
ggml_get_type_traits(GGML_TYPE_F16)
->from_float_ref(scratch.data(), f16_base + (r0 * ne0) * sizeof(uint16_t), elems);
} else {
const int64_t block_offset = (r0 * ne0) / block_size;
if (requant_type == ExtraQuantType::Q8_1_C) {
quantize_q8_1(scratch.data(), weights, scales, zp, elems, block_size, block_offset);
} else {
quantize_q8_0(scratch.data(), weights, scales, zp, elems, block_size, block_offset);
}
}
}
if (requant_type == ExtraQuantType::F16) {
auto result = std::make_shared<ov::op::v0::Constant>(weights);
result->set_friendly_name(tensor->name);
return result;
}
}
// Create the OpenVINO weight subgraph
@ -745,8 +937,11 @@ OvWeight process_weight_tensor(const ggml_tensor * tensor, const void * data, vo
OvWeight result;
// Get 2D shape for weights [rows, cols]
ov::Shape node_shape = {static_cast<size_t>(tensor->ne[1]), static_cast<size_t>(tensor->ne[0])};
// Get shape for weights: [rows, cols], or [n_expert, rows, cols] for 3D MoE expert weights.
ov::Shape node_shape = (tensor->ne[2] > 1) ?
ov::Shape{static_cast<size_t>(tensor->ne[2]), static_cast<size_t>(tensor->ne[1]),
static_cast<size_t>(tensor->ne[0])} :
ov::Shape{static_cast<size_t>(tensor->ne[1]), static_cast<size_t>(tensor->ne[0])};
// Handle F16/F32/BF16 weights
if (tensor->type == GGML_TYPE_F32 || tensor->type == GGML_TYPE_F16 || tensor->type == GGML_TYPE_BF16) {
@ -788,6 +983,35 @@ OvWeight process_weight_tensor(const ggml_tensor * tensor, const void * data, vo
OPENVINO_THROW("Unsupported quantized type: ", ggml_type_name(tensor->type));
}
// 3D MoE expert weights (for_gather_matmul) always use the exact f16 zero-point path (see
// extract_quantized_weights) -- must be kept in sync with the "use_bias || for_gather_matmul"
// check in ggml_openvino_get_extracted_layout, which sizes/offsets the zp slot accordingly.
// Requantized tensors (layout.is_requant) are handled by requantize_to_buffers instead, whose
// zp sizing/type is unaffected by for_gather_matmul, so they are excluded here.
const bool for_gather_matmul = tensor->ne[2] > 1;
const bool zp_is_f16 = !layout.is_requant && (use_bias || for_gather_matmul);
const bool is_3d_mxfp4_moe = tensor->type == GGML_TYPE_MXFP4 && (tensor->ne[2] > 1 || tensor->ne[3] > 1);
if (is_3d_mxfp4_moe) {
ov::Shape packed_shape = {static_cast<size_t>(tensor->ne[3]),
static_cast<size_t>(tensor->ne[2]),
static_cast<size_t>(tensor->ne[1]),
static_cast<size_t>(tensor->ne[0] / MXFP4_BLOCK_SIZE),
MXFP4_BLOCK_BYTES};
const size_t tensor_bytes = ggml_nbytes(tensor);
if (output_base_ptr) {
auto * buf_base = static_cast<uint8_t *>(output_base_ptr);
memcpy(buf_base + layout.weights_offset, data, tensor_bytes);
result.weights = ov::Tensor(ov::element::u8, packed_shape, buf_base + layout.weights_offset);
} else {
result.weights = ov::Tensor(ov::element::u8, packed_shape);
memcpy(result.weights.data(), data, tensor_bytes);
}
result.weight_node = make_mxfp4_moe_packed_weights(result.weights).get_node_shared_ptr();
result.weight_node->set_friendly_name(tensor->name);
return result;
}
if (use_bias) {
OPENVINO_ASSERT(!layout.is_requant,
"use_bias is only used for test-backend-ops, which should not have requantization");
@ -812,24 +1036,44 @@ OvWeight process_weight_tensor(const ggml_tensor * tensor, const void * data, vo
// Quantized path (normal extraction or quantized requant)
// Create weight/scale/zp tensors - shared between both paths
// For symmetric quantization, use signed types (i4/i8) and no ZP tensor
ov::element::Type weight_type = layout.is_symmetric ? (layout.is_u4 ? ov::element::i4 : ov::element::i8) :
(layout.is_u4 ? ov::element::u4 : ov::element::u8);
ov::Shape scale_shape = {node_shape[0], node_shape[1] / layout.weights_per_block};
ov::element::Type weight_type = tensor->type == GGML_TYPE_MXFP4 ?
ov::element::f4e2m1 :
(layout.is_symmetric ? (layout.is_u4 ? ov::element::i4 : ov::element::i8) :
(layout.is_u4 ? ov::element::u4 : ov::element::u8));
ov::Shape scale_shape = node_shape;
scale_shape.back() /= layout.weights_per_block;
if (tensor->type == GGML_TYPE_MXFP4) {
if (tensor->ne[2] == 1 && tensor->ne[3] == 1) {
node_shape = {static_cast<size_t>(tensor->ne[1]), static_cast<size_t>(tensor->ne[0])};
} else {
node_shape.clear();
for (int i = GGML_MAX_DIMS - 1; i >= 0; --i) {
node_shape.push_back(static_cast<size_t>(tensor->ne[i]));
}
}
scale_shape = node_shape;
scale_shape.back() /= layout.weights_per_block;
}
if (output_base_ptr) {
uint8_t * buf_base = static_cast<uint8_t *>(output_base_ptr);
result.weights = ov::Tensor(weight_type, node_shape, buf_base + layout.weights_offset);
result.scales = ov::Tensor(ov::element::f16, scale_shape, buf_base + layout.scales_offset);
const ov::element::Type scale_type = tensor->type == GGML_TYPE_MXFP4 ? ov::element::f8e8m0 : ov::element::f16;
result.scales = ov::Tensor(scale_type, scale_shape, buf_base + layout.scales_offset);
if (!layout.is_symmetric) {
ov::element::Type zp_type = layout.is_u4 ? ov::element::u4 : ov::element::u8;
ov::element::Type zp_type =
zp_is_f16 ? ov::element::f16 : (layout.is_u4 ? ov::element::u4 : ov::element::u8);
result.zp = ov::Tensor(zp_type, scale_shape, buf_base + layout.zp_offset);
}
// else: result.zp remains default-constructed (empty) for symmetric
} else {
result.weights = ov::Tensor(weight_type, node_shape);
result.scales = ov::Tensor(ov::element::f16, scale_shape);
const ov::element::Type scale_type = tensor->type == GGML_TYPE_MXFP4 ? ov::element::f8e8m0 : ov::element::f16;
result.scales = ov::Tensor(scale_type, scale_shape);
if (!layout.is_symmetric) {
if (use_bias) {
if (zp_is_f16) {
result.zp = ov::Tensor(ov::element::f16, scale_shape);
} else {
ov::element::Type zp_type = layout.is_u4 ? ov::element::u4 : ov::element::u8;
@ -939,16 +1183,21 @@ void quantize_q8_0(const float * x,
ov::Tensor & scales_arr,
ov::Tensor & zp_arr,
int64_t k,
int64_t qk) {
int64_t qk,
int64_t block_offset) {
assert(k % qk == 0);
const int nb = k / qk;
auto * weights = static_cast<uint8_t *>(weights_arr.data());
auto * scales = scales_arr.data<ov::element_type_traits<ov::element::f16>::value_type>();
// block_offset lets a caller quantize a chunk of blocks into the right place in the
// output buffers (used for streaming requant). x points at this chunk's first block;
// outputs are advanced by block_offset blocks. Q8 has one scale/zp per block (no
// nibble packing), so any block boundary is safe.
auto * weights = static_cast<uint8_t *>(weights_arr.data()) + block_offset * qk;
auto * scales = scales_arr.data<ov::element_type_traits<ov::element::f16>::value_type>() + block_offset;
bool is_symmetric = (weights_arr.get_element_type() == ov::element::i8); // Signed i8 path
if (!is_symmetric) {
auto * zp = static_cast<uint8_t *>(zp_arr.data());
auto * zp = static_cast<uint8_t *>(zp_arr.data()) + block_offset;
for (int i = 0; i < nb; i++) {
float amax = 0.0f;
for (int j = 0; j < qk; j++) {
@ -990,13 +1239,15 @@ void quantize_q8_1(const float * x,
ov::Tensor & scales_arr,
ov::Tensor & zp_arr,
int64_t k,
int64_t qk) {
int64_t qk,
int64_t block_offset) {
assert(k % qk == 0);
const int nb = k / qk;
auto * weights = static_cast<uint8_t *>(weights_arr.data());
auto * scales = scales_arr.data<ov::element_type_traits<ov::element::f16>::value_type>();
auto * zp = static_cast<uint8_t *>(zp_arr.data());
// See quantize_q8_0: block_offset places this chunk's output at the right block.
auto * weights = static_cast<uint8_t *>(weights_arr.data()) + block_offset * qk;
auto * scales = scales_arr.data<ov::element_type_traits<ov::element::f16>::value_type>() + block_offset;
auto * zp = static_cast<uint8_t *>(zp_arr.data()) + block_offset;
for (int i = 0; i < nb; i++) {
float min = std::numeric_limits<float>::max();
float max = std::numeric_limits<float>::lowest();

39
ggml/src/ggml-openvino/ggml-quants.h
View File

@ -4,6 +4,7 @@
#include <cstdint>
#include <openvino/op/constant.hpp>
#include <openvino/core/node_output.hpp>
#include <openvino/runtime/tensor.hpp>
void unpack_32_4(const uint8_t * data, uint8_t * dst);
@ -49,19 +50,38 @@ void extract_q6_k_data(const ggml_tensor * tensor,
ov::Tensor & scales_arr,
ov::Tensor & zp_arr);
void extract_mxfp4_data(const ggml_tensor * tensor, ov::Tensor & weights_arr, ov::Tensor & scales_arr);
static constexpr size_t GGML_QUANTIZATION_GROUP_SIZE = 32;
// If for_gather_matmul is true, the weight tensor may be N-D (e.g. 3D MoE expert weights
// [n_expert, rows, cols]). The dequantization chain (Convert->[Subtract]->Multiply) is built as
// usual but left in f16 (no final Convert to f32) -- ov::pass::MarkDequantization (registered in
// translate_session.cpp) marks the chain so it survives model-build-time ConstantFolding -- see
// make_int8_weights.cpp/make_int4_weights.cpp. mul_mat_id.cpp constructs ov::op::internal::GatherMatmul
// directly from the resulting f16 dequant chain.
//
// When use_bias is true (explicitly, or implicitly because for_gather_matmul is true), the zp
// tensor is expected to hold an exact f16 bias value (rather than a rounded integer zero point);
// it is converted in place into an exact zero_point = -bias/scale and consumed via Subtract, not
// Add, so the chain still matches OpenVINO's Convert->Subtract->Multiply decompression pattern.
ov::Output<ov::Node> make_int8_weights(ov::Tensor & weight,
ov::Tensor & scales,
ov::Tensor & zp,
size_t group_size = GGML_QUANTIZATION_GROUP_SIZE,
bool use_bias = false);
bool use_bias = false,
bool for_gather_matmul = false);
ov::Output<ov::Node> make_int4_weights(ov::Tensor & weight,
ov::Tensor & scales,
ov::Tensor & zp,
size_t group_size = GGML_QUANTIZATION_GROUP_SIZE,
bool use_bias = false);
bool use_bias = false,
bool for_gather_matmul = false);
ov::Output<ov::Node> make_mxfp4_weights(ov::Tensor & weight, ov::Tensor & scales);
ov::Output<ov::Node> make_mxfp4_moe_packed_weights(ov::Tensor & weight);
// Extract quantized weights from tensor and create weight subgraph
// If weights/scales/zp are provided (non-empty), uses them as output buffers
@ -73,7 +93,9 @@ std::shared_ptr<ov::Node> extract_quantized_weights(
ov::Tensor & weights,
ov::Tensor & scales,
ov::Tensor & zp,
bool use_bias = false); // Use fp bias instead of quantized zero_point (for test-backend-ops)
bool use_bias = false); // Use an exact f16 zero point (vs. a rounded integer one); always
// used for for_gather_matmul (3D MoE expert) weights regardless of
// this flag, and also settable explicitly for test-backend-ops.
// Requantize weights from tensor to target format, writing to provided buffers
// For F16 target, only weights buffer is used (scales/zp ignored)
@ -126,7 +148,10 @@ OvWeight process_weight_tensor(
const ggml_tensor * tensor,
const void * data, // Source data pointer (may differ from tensor->data)
void * output_base_ptr = nullptr, // Base pointer for output buffers (or nullptr for internal allocation)
bool use_bias = false); // Use fp bias instead of quantized zero_point, only used in test-backend-ops
bool use_bias = false); // Use an exact f16 zero point (vs. a rounded integer one);
// always used for for_gather_matmul (3D MoE expert) weights
// regardless of this flag, and also settable explicitly for
// test-backend-ops.
void quantize_q4_0(const float * x,
ov::Tensor & weights_arr,
@ -139,13 +164,15 @@ void quantize_q8_1(const float * x,
ov::Tensor & scales_arr,
ov::Tensor & zp_arr,
int64_t k,
int64_t qk);
int64_t qk,
int64_t block_offset = 0);
void quantize_q8_0(const float * x,
ov::Tensor & weights_arr,
ov::Tensor & scales_arr,
ov::Tensor & zp_arr,
int64_t k,
int64_t qk);
int64_t qk,
int64_t block_offset = 0);
namespace ov {
namespace op {

272
ggml/src/ggml-openvino/model-cache.cpp Normal file
View File

@ -0,0 +1,272 @@
#include "model-cache.h"
#include "ggml-backend-impl.h"
#include "ggml-backend.h"
#include "ggml-impl.h"
#include "ggml-openvino-extra.h"
#include <cerrno>
#include <cstdio>
#include <cstring>
#include <fstream>
#include <openvino/core/version.hpp>
#include <string>
#include <sys/stat.h>
#include <sys/types.h>
#include <vector>
#if defined(_WIN32)
# include <direct.h>
#endif
namespace {
// 64-bit FNV-1a, the mixing primitive for all fingerprints here.
inline uint64_t fnv1a(uint64_t h, const void * data, size_t n) {
const uint8_t * p = static_cast<const uint8_t *>(data);
for (size_t i = 0; i < n; ++i) {
h ^= p[i];
h *= 0x100000001b3ull;
}
return h;
}
inline uint64_t fnv1a_u64(uint64_t h, uint64_t v) {
return fnv1a(h, &v, sizeof(v));
}
constexpr uint64_t FNV_OFFSET = 0xcbf29ce484222325ull;
// Bytes sampled from each end of a weight tensor for the sampled hash. The whole
// model is never hashed (that would cost seconds every run); instead we sample a
// bounded window from the head and tail of each weight's bytes. The manifest
// re-verify (same sample) guards the residual collision risk.
constexpr size_t WEIGHT_SAMPLE_BYTES = 4096;
// Is this src a model weight, mirroring create_weight_nodes()'s selection:
// non-view tensor whose buffer is USAGE_WEIGHTS or whose type is quantized.
bool is_weight_src(const ggml_tensor * src) {
if (src == nullptr || src->view_src != nullptr || src->buffer == nullptr) {
return false;
}
return src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS || ggml_is_quantized(src->type);
}
// Per-weight sampled fingerprint: identity (name/shape/type) + a bounded byte
// sample. Returns FNV offset basis if data is unavailable (kept deterministic).
uint64_t weight_fingerprint(const ggml_tensor * t) {
uint64_t h = FNV_OFFSET;
h = fnv1a(h, t->name, strlen(t->name));
for (int i = 0; i < GGML_MAX_DIMS; ++i) {
h = fnv1a_u64(h, static_cast<uint64_t>(t->ne[i]));
}
h = fnv1a_u64(h, static_cast<uint64_t>(t->type));
const size_t nbytes = ggml_nbytes(t);
h = fnv1a_u64(h, nbytes);
if (t->data != nullptr && nbytes > 0) {
const size_t head = nbytes < WEIGHT_SAMPLE_BYTES ? nbytes : WEIGHT_SAMPLE_BYTES;
h = fnv1a(h, t->data, head);
if (nbytes > WEIGHT_SAMPLE_BYTES) {
const size_t tail = nbytes < 2 * WEIGHT_SAMPLE_BYTES ? nbytes - WEIGHT_SAMPLE_BYTES : WEIGHT_SAMPLE_BYTES;
h = fnv1a(h, static_cast<const uint8_t *>(t->data) + (nbytes - tail), tail);
}
}
return h;
}
// Walk the cgraph and invoke fn(weight_tensor) for each distinct weight, in node
// order. De-duplicates by tensor pointer so a weight used by several nodes is
// fingerprinted once, deterministically.
template <typename F>
void for_each_weight(const ggml_cgraph * cgraph, F && fn) {
std::vector<const ggml_tensor *> seen;
for (int i = 0; i < cgraph->n_nodes; ++i) {
const ggml_tensor * node = cgraph->nodes[i];
for (int s = 0; s < GGML_MAX_SRC; ++s) {
const ggml_tensor * src = node->src[s];
if (!is_weight_src(src)) {
continue;
}
bool dup = false;
for (const auto * p : seen) {
if (p == src) {
dup = true;
break;
}
}
if (dup) {
continue;
}
seen.push_back(src);
fn(src);
}
}
}
std::string ov_version_string() {
const ov::Version v = ov::get_openvino_version();
return std::string(v.buildNumber ? v.buildNumber : "unknown");
}
std::string hex64(uint64_t v) {
char buf[17];
snprintf(buf, sizeof(buf), "%016llx", static_cast<unsigned long long>(v));
return std::string(buf);
}
// Portable mkdir for a single path component. Returns true if the directory
// exists after the call (created now or already present).
bool make_dir(const std::string & path) {
#if defined(_WIN32)
int rc = _mkdir(path.c_str());
#else
int rc = ::mkdir(path.c_str(), 0755);
#endif
if (rc == 0 || errno == EEXIST) {
return true;
}
return false;
}
// Create `path` and any missing parents (like `mkdir -p`). Best-effort:
// returns true only if the full directory exists afterwards.
bool make_dirs(const std::string & path) {
if (path.empty()) {
return false;
}
std::string acc;
for (size_t i = 0; i < path.size(); ++i) {
const char c = path[i];
acc.push_back(c);
const bool sep = (c == '/'
#if defined(_WIN32)
|| c == '\\'
#endif
);
// Create each intermediate component (skip a leading "/" root).
if (sep && acc.size() > 1) {
std::string component = acc.substr(0, acc.size() - 1);
if (!make_dir(component)) {
return false;
}
}
}
return make_dir(path);
}
} // namespace
std::string ggml_openvino_model_cache_dir() {
const char * dir = ggml_openvino_getenv_str("GGML_OPENVINO_COMPILED_MODEL_CACHE_DIR");
if (!dir || strlen(dir) == 0) {
return std::string();
}
std::string path(dir);
// Create the cache directory (and parents) on first use so callers don't
// have to pre-create it; a missing dir would otherwise silently disable the
// cache (manifest/blob writes fail with no directory to write into).
if (!make_dirs(path)) {
GGML_LOG_WARN("ggml-openvino: could not create model cache dir '%s' (errno=%d); caching disabled\n",
path.c_str(), errno);
return std::string();
}
return path;
}
uint64_t ggml_openvino_model_fingerprint(const ggml_cgraph * cgraph,
const std::string & device,
bool fa,
const int32_t * rope_params,
int rope_len,
uint64_t extra_cfg) {
uint64_t h = FNV_OFFSET;
// Topology: node count + each node's op and name (cheap, and distinguishes
// graphs that share weights but differ structurally).
h = fnv1a_u64(h, static_cast<uint64_t>(cgraph->n_nodes));
for (int i = 0; i < cgraph->n_nodes; ++i) {
const ggml_tensor * node = cgraph->nodes[i];
h = fnv1a_u64(h, static_cast<uint64_t>(node->op));
h = fnv1a(h, node->name, strlen(node->name));
}
// Weights: the model identity.
for_each_weight(cgraph, [&](const ggml_tensor * t) { h = fnv1a_u64(h, weight_fingerprint(t)); });
// Config that changes the produced blob.
h = fnv1a(h, device.data(), device.size());
h = fnv1a_u64(h, fa ? 1u : 0u);
if (rope_params && rope_len > 0) {
h = fnv1a(h, rope_params, sizeof(int32_t) * static_cast<size_t>(rope_len));
}
h = fnv1a_u64(h, extra_cfg);
const std::string ver = ov_version_string();
h = fnv1a(h, ver.data(), ver.size());
return h;
}
std::string ggml_openvino_model_cache_blob_path(const std::string & dir, uint64_t fingerprint) {
return dir + "/" + hex64(fingerprint) + ".blob";
}
std::string ggml_openvino_model_cache_manifest_path(const std::string & dir, uint64_t fingerprint) {
return dir + "/" + hex64(fingerprint) + ".manifest";
}
bool ggml_openvino_model_cache_write_manifest(const std::string & path,
const ggml_cgraph * cgraph,
uint64_t fingerprint) {
std::ofstream f(path, std::ios::trunc);
if (!f.is_open()) {
return false;
}
f << "fingerprint " << hex64(fingerprint) << "\n";
f << "ov_version " << ov_version_string() << "\n";
for_each_weight(cgraph, [&](const ggml_tensor * t) {
f << t->name << " " << t->ne[0] << " " << t->ne[1] << " " << t->ne[2] << " " << t->ne[3] << " "
<< static_cast<int>(t->type) << " " << hex64(weight_fingerprint(t)) << "\n";
});
return f.good();
}
bool ggml_openvino_model_cache_verify_manifest(const std::string & path,
const ggml_cgraph * cgraph,
uint64_t fingerprint) {
std::ifstream f(path);
if (!f.is_open()) {
return false;
}
std::string tag, val;
// header: fingerprint
if (!(f >> tag >> val) || tag != "fingerprint" || val != hex64(fingerprint)) {
return false;
}
// header: ov_version
if (!(f >> tag >> val) || tag != "ov_version" || val != ov_version_string()) {
return false;
}
// Build the expected per-weight lines from the live cgraph, then require an
// exact match (same set, same order) against the manifest.
std::vector<std::string> expected;
for_each_weight(cgraph, [&](const ggml_tensor * t) {
expected.push_back(std::string(t->name) + " " + std::to_string(t->ne[0]) + " " + std::to_string(t->ne[1]) +
" " + std::to_string(t->ne[2]) + " " + std::to_string(t->ne[3]) + " " +
std::to_string(static_cast<int>(t->type)) + " " + hex64(weight_fingerprint(t)));
});
size_t idx = 0;
std::string line;
std::getline(f, line); // consume rest of ov_version line
while (std::getline(f, line)) {
if (line.empty()) {
continue;
}
if (idx >= expected.size() || line != expected[idx]) {
return false;
}
++idx;
}
return idx == expected.size();
}

56
ggml/src/ggml-openvino/model-cache.h Normal file
View File

@ -0,0 +1,56 @@
#pragma once
// Frontend-level compiled-model cache (GGML_OPENVINO_COMPILED_MODEL_CACHE_DIR).
//
// The OpenVINO plugin's own ov::cache_dir caches the compiled blob keyed by the
// *OV model*, but producing that model still runs the full frontend every time:
// weight requantization (incl. the large token_embd F32 transient) and the
// ggml->OV graph conversion. This cache keys off a fingerprint computed directly
// from the ggml cgraph, so a hit skips requant + convert + compile entirely and
// instead imports a previously exported CompiledModel blob.
//
// Opt-in and independent from GGML_OPENVINO_CACHE_DIR. Default off.
#include "ggml.h"
#include <cstdint>
#include <string>
// Returns the compiled-model cache directory from GGML_OPENVINO_COMPILED_MODEL_CACHE_DIR,
// or empty if unset/disabled. When empty, callers must not use the cache.
std::string ggml_openvino_model_cache_dir();
// Compute a stable 64-bit fingerprint identifying the model+config that a cgraph
// would compile to. Combines graph topology, a sampled hash of every weight
// tensor (name/shape/dtype + bounded byte sample), and the config that changes
// the produced blob (device, flash-attention, rope params, the compile-memory
// flags, stateful, and the OpenVINO version). `device` is the resolved device
// string; `fa` is the flash-attention flag; `rope_params`/`rope_len` cover the
// model's rope configuration; `extra_cfg` folds in any other blob-affecting bits.
uint64_t ggml_openvino_model_fingerprint(const ggml_cgraph * cgraph,
const std::string & device,
bool fa,
const int32_t * rope_params,
int rope_len,
uint64_t extra_cfg);
// Path to the compiled-blob file for a fingerprint (<dir>/<hex>.blob).
std::string ggml_openvino_model_cache_blob_path(const std::string & dir, uint64_t fingerprint);
// Path to the sidecar manifest (<dir>/<hex>.manifest) holding the per-weight
// fingerprints, used to re-verify a hit before trusting the blob.
std::string ggml_openvino_model_cache_manifest_path(const std::string & dir, uint64_t fingerprint);
// Write/read the manifest. The manifest is a newline-separated list of
// "name ne0 ne1 ne2 ne3 type sample_hash" lines plus a header line with the
// fingerprint and OV version. Returns false on I/O error.
bool ggml_openvino_model_cache_write_manifest(const std::string & path,
const ggml_cgraph * cgraph,
uint64_t fingerprint);
// Verify that the cgraph's weights still match the stored manifest (guards the
// sampled-hash collision risk: a blob is only trusted if every weight's
// name/shape/type/sample-hash matches what was cached). Returns true on match.
bool ggml_openvino_model_cache_verify_manifest(const std::string & path,
const ggml_cgraph * cgraph,
uint64_t fingerprint);

25
ggml/src/ggml-openvino/openvino/decoder.h
View File

@ -6,12 +6,25 @@
#include <openvino/core/partial_shape.hpp>
#include <openvino/core/shape.hpp>
#include <openvino/frontend/decoder.hpp>
#include <set>
#include <string>
namespace ov {
namespace frontend {
namespace ggml {
struct ModelInputInfo {
element::Type type;
PartialShape shape;
};
struct ModelExtraInputInfo {
element::Type type;
Shape shape;
int64_t value;
bool is_parameter;
};
class GgmlDecoder : public DecoderBase {
public:
virtual ov::Any get_attribute(const std::string & name) const = 0;
@ -75,6 +88,10 @@ public:
virtual std::vector<std::string> get_output_names(int node_idx) const = 0;
virtual std::string get_inplace_op_src(int node_idx) const = 0;
virtual bool is_view_like_alias_of(int node_idx, const std::string & view_src_name) const = 0;
virtual const std::string & get_op_type() const = 0;
virtual const std::string & get_op_type(int node_idx) const = 0;
@ -87,15 +104,17 @@ public:
virtual int get_op_case(int node_idx) const = 0;
virtual const std::map<std::string, std::shared_ptr<ov::Node>> & get_model_inputs() const = 0;
virtual const std::map<std::string, std::shared_ptr<ov::Node>> & get_model_extra_inputs() const = 0;
virtual const std::map<std::string, ModelInputInfo> & get_model_inputs() const = 0;
virtual const std::map<std::string, ModelExtraInputInfo> & get_model_extra_inputs() const = 0;
virtual const std::map<std::string, std::shared_ptr<ov::Node>> & get_model_weights() const = 0;
virtual std::vector<std::string> get_model_output_names() const = 0;
virtual std::set<std::string> get_model_output_names() const = 0;
virtual int32_t * get_rope_params() const = 0;
virtual bool has_mixed_rope_params() const = 0;
virtual int get_ssm_state_size() const = 0;
virtual std::map<std::string, std::string> get_kv_param_res_names() const = 0;
virtual bool is_static() const = 0;

2
ggml/src/ggml-openvino/openvino/node_context.h
View File

@ -153,6 +153,8 @@ public:
bool is_stateful() const { return m_decoder->is_stateful(); }
int get_ssm_state_size() const { return m_decoder->get_ssm_state_size(); }
private:
std::shared_ptr<GgmlDecoder> m_decoder;
std::shared_ptr<TensorMap> & m_tensor_map;

45
ggml/src/ggml-openvino/openvino/op/add.cpp Normal file
View File

@ -0,0 +1,45 @@
#include "../node_context.h"
#include "../op_table.h"
#include "../utils.h"
#include <memory>
#include <openvino/op/add.hpp>
#include <openvino/op/constant.hpp>
#include <openvino/op/reduce_sum.hpp>
#include <openvino/op/unsqueeze.hpp>
namespace ov {
namespace frontend {
namespace ggml {
namespace op {
OutputVector translate_add(const NodeContext & context) {
num_inputs_check(context, 2, 2);
if (context.get_op_case() == 1) {
// MoE expert-plane sum (see is_moe_expert_sum_add): input 1 is a VIEW plane of the
// shared base tensor `experts` = [n_embd, n_expert_used, n_tokens, 1] (ggml order) ->
// [1, n_tokens, n_expert_used, n_embd] (OV order). The whole ADD chain is equivalent to
// reducing the expert axis (OV axis 2) of that base, so bypass the chain and the
// per-plane Slices entirely.
size_t view_size = context.get_view_input_size(1);
auto base_name = context.get_view_input_src_name(1, view_size - 1);
auto base = context.get_input(base_name);
auto reduced = std::make_shared<ov::op::v1::ReduceSum>(
base, ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {2}), false);
auto res =
std::make_shared<ov::op::v0::Unsqueeze>(reduced, ov::op::v0::Constant::create(ov::element::i64, {1}, {1}));
return rename_outputs_with_suffix({res}, context.get_name());
}
auto input_0 = process_view_input_new(context, 0);
auto input_1 = process_view_input_new(context, 1);
auto res = std::make_shared<ov::op::v1::Add>(input_0, input_1);
return rename_outputs_with_suffix({res}, context.get_name());
}
} // namespace op
} // namespace ggml
} // namespace frontend
} // namespace ov

157
ggml/src/ggml-openvino/openvino/op/cpy.cpp
View File

@ -2,10 +2,19 @@
#include "../op_table.h"
#include "../utils.h"
#include <climits>
#include <memory>
#include <vector>
#include <openvino/op/add.hpp>
#include <openvino/op/concat.hpp>
#include <openvino/op/constant.hpp>
#include <openvino/op/convert.hpp>
#include <openvino/op/gather.hpp>
#include <openvino/op/multiply.hpp>
#include <openvino/op/negative.hpp>
#include <openvino/op/reshape.hpp>
#include <openvino/op/shape_of.hpp>
#include <openvino/op/slice.hpp>
namespace ov {
namespace frontend {
@ -13,18 +22,158 @@ namespace ggml {
namespace op {
OutputVector translate_cpy(const NodeContext & context) {
auto input = process_view_input_new(context, 0);
auto op_case = context.get_op_case();
auto input_shape = context.get_input_shape(0);
auto output_shape = context.get_output_shape();
auto output_shape = context.get_input_shape(1);
if (op_case == 4) {
auto src = process_view_input_new(context, 0);
auto base = context.get_input(1);
int64_t n_elems = 1;
for (const auto & dim : context.get_output_shape().to_shape()) {
n_elems *= static_cast<int64_t>(dim);
}
const auto output_stride = context.get_output_stride();
const size_t elem_size = output_stride.empty() ? context.get_output_type().size() : output_stride.back();
FRONT_END_OP_CONVERSION_CHECK(elem_size > 0, "CPY conv state view update has invalid element size");
const int64_t begin_val = static_cast<int64_t>(context.get_output_op_offset() / elem_size);
const int64_t end_val = begin_val + n_elems;
auto flat_shape = ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{1, 1, 1, -1});
src = std::make_shared<ov::op::v1::Reshape>(src, flat_shape, false);
if (src.get_element_type() != context.get_output_type()) {
src = std::make_shared<ov::op::v0::Convert>(src, context.get_output_type());
}
auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0});
auto begin = ov::op::v0::Constant::create(ov::element::i64, {1}, {begin_val});
auto end = ov::op::v0::Constant::create(ov::element::i64, {1}, {end_val});
auto int_max = ov::op::v0::Constant::create(ov::element::i64, {1}, {INT_MAX});
auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
auto axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {3});
auto head_part = std::make_shared<ov::op::v8::Slice>(base, zero, begin, one, axis);
auto tail_part = std::make_shared<ov::op::v8::Slice>(base, end, int_max, one, axis);
auto res = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{head_part, src, tail_part}, 3);
return rename_outputs_with_suffix({res}, context.get_name());
}
// Recurrent state cache writeback into a slot block of the cache. Where the block starts and
// where the copied data starts in the source are runtime inputs, so the cached model works for
// any kv head, active sequence count and token count. The result is the full updated cache.
// op_case 1: gated-delta-net state, op_case 2: conv state, op_case 3: defrag remainder.
const std::string slot_begin_name = "rs_slot_begin_" + context.get_name();
const bool slice_assign =
context.has_input(slot_begin_name) && !context.is_stateful() && (op_case >= 1 && op_case <= 3);
if (slice_assign) {
const int64_t slot_axis = 2;
auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0});
auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
auto int_max = ov::op::v0::Constant::create(ov::element::i64, {1}, {INT_MAX});
auto axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {slot_axis});
auto feature = ov::op::v0::Constant::create(ov::element::i64, {4},
std::vector<int64_t>{1, 1, -1, output_shape[3].get_length()});
ov::Output<ov::Node> src;
ov::Output<ov::Node> begin = context.get_input(slot_begin_name);
auto base = context.get_input(1);
if (op_case == 1) {
// GDN packs [attn | state snapshots]; the state part runs from src_begin to the end.
auto src_begin = context.get_input("rs_src_begin_" + context.get_name());
auto state_part = std::make_shared<ov::op::v8::Slice>(context.get_input(0), src_begin, int_max, one, axis);
src = std::make_shared<ov::op::v1::Reshape>(state_part, feature, false);
} else if (op_case == 2) {
// conv_input is [previous conv state | new tokens]; copy the conv_kernel_size - 1 wide
// window starting at src_begin, which is the snapshot this writeback corresponds to.
auto window_size = (int64_t) input_shape[3].get_length();
auto src_begin = context.get_input("rs_src_begin_" + context.get_name());
auto src_end = std::make_shared<ov::op::v1::Add>(
src_begin, ov::op::v0::Constant::create(ov::element::i64, {1}, {window_size}));
auto window = std::make_shared<ov::op::v8::Slice>(context.get_input(0), src_begin, src_end, one,
ov::op::v0::Constant::create(ov::element::i64, {1}, {3}));
const auto base_shape = base.get_partial_shape();
FRONT_END_OP_CONVERSION_CHECK(base_shape.rank().is_static() && base_shape.rank().get_length() == 4,
"CPY conv state cache update requires rank-4 base cache");
FRONT_END_OP_CONVERSION_CHECK(base_shape[3].is_static(),
"CPY conv state cache update requires static feature size");
FRONT_END_OP_CONVERSION_CHECK(input_shape.rank().is_static() && input_shape.rank().get_length() == 4 &&
input_shape[2].is_static() && input_shape[3].is_static(),
"CPY conv state cache update requires static source feature view");
const int64_t full_feature_size = base_shape[3].get_length();
const int64_t update_feature_size = input_shape[2].get_length() * input_shape[3].get_length();
const auto output_stride = context.get_output_stride();
const size_t elem_size = output_stride.empty() ? context.get_output_type().size() : output_stride.back();
FRONT_END_OP_CONVERSION_CHECK(elem_size > 0,
"CPY conv state cache update has invalid element size");
const int64_t feature_begin = static_cast<int64_t>(context.get_output_op_offset() / elem_size) %
full_feature_size;
const int64_t feature_end = feature_begin + update_feature_size;
FRONT_END_OP_CONVERSION_CHECK(feature_begin >= 0 && feature_end <= full_feature_size,
"CPY conv state cache update feature range is out of bounds");
auto partial_feature = ov::op::v0::Constant::create(
ov::element::i64, {4}, std::vector<int64_t>{1, 1, -1, update_feature_size});
src = std::make_shared<ov::op::v1::Reshape>(window, partial_feature, false);
if (src.get_element_type() != context.get_output_type()) {
src = std::make_shared<ov::op::v0::Convert>(src, context.get_output_type());
}
auto src_len = std::make_shared<ov::op::v8::Gather>(
std::make_shared<ov::op::v3::ShapeOf>(src, ov::element::i64), axis,
ov::op::v0::Constant::create(ov::element::i64, {}, {0}));
auto slot_end = std::make_shared<ov::op::v1::Add>(begin, src_len);
auto active_slots = std::make_shared<ov::op::v8::Slice>(base, begin, slot_end, one, axis);
auto feature_axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {3});
auto feature_begin_node = ov::op::v0::Constant::create(ov::element::i64, {1}, {feature_begin});
auto feature_end_node = ov::op::v0::Constant::create(ov::element::i64, {1}, {feature_end});
auto feature_head = std::make_shared<ov::op::v8::Slice>(active_slots, zero, feature_begin_node, one,
feature_axis);
auto feature_tail = std::make_shared<ov::op::v8::Slice>(active_slots, feature_end_node, int_max, one,
feature_axis);
src = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{feature_head, src, feature_tail}, 3);
} else {
// op_case 3: gathered remainder rows already have the cache slot layout [1, 1, extra, feature]
src = context.get_input(0);
}
if (src.get_element_type() != context.get_output_type()) {
src = std::make_shared<ov::op::v0::Convert>(src, context.get_output_type());
}
auto src_len =
std::make_shared<ov::op::v8::Gather>(std::make_shared<ov::op::v3::ShapeOf>(src, ov::element::i64), axis,
ov::op::v0::Constant::create(ov::element::i64, {}, {0}));
auto end = std::make_shared<ov::op::v1::Add>(begin, src_len);
auto head_part = std::make_shared<ov::op::v8::Slice>(base, zero, begin, one, axis);
auto tail_part = std::make_shared<ov::op::v8::Slice>(base, end, int_max, one, axis);
auto res = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{head_part, src, tail_part}, slot_axis);
return rename_outputs_with_suffix({res}, context.get_name());
}
auto input = process_view_input_new(context, 0);
// Non-cast CPY may need a reshape (e.g. [3,192,1,1] -> [576,1,1,1])
if (input_shape != output_shape) {
auto new_shape = ov::op::v0::Constant::create(
ov::element::i64, {static_cast<size_t>(output_shape.rank().get_length())}, output_shape.to_shape());
input = std::make_shared<ov::op::v1::Reshape>(input, new_shape, false);
}
auto res = std::make_shared<ov::op::v0::Convert>(input, context.get_output_type());
ov::Output<Node> res;
if (context.get_input_type(0) != context.get_output_type()) {
res = std::make_shared<ov::op::v0::Convert>(input, context.get_output_type());
} else {
res = input;
}
if (res.get_node_shared_ptr() == context.get_input(0).get_node_shared_ptr()) {
return {res};
}
return rename_outputs_with_suffix({res}, context.get_name());
}

29
ggml/src/ggml-openvino/openvino/op/cumsum.cpp Normal file
View File

@ -0,0 +1,29 @@
#include "../node_context.h"
#include "../op_table.h"
#include "../utils.h"
#include <openvino/op/constant.hpp>
#include <openvino/op/cum_sum.hpp>
namespace ov {
namespace frontend {
namespace ggml {
namespace op {
// GGML cumsum computes prefix sum along dim 0 (the innermost/fastest dimension).
// In OV layout the dims are reversed: ggml [ne0, ne1, ne2, ne3] → OV [ne3, ne2, ne1, ne0],
// so ggml dim 0 maps to OV axis 3 (last axis).
OutputVector translate_cumsum(const NodeContext & context) {
num_inputs_check(context, 1, 1);
auto x = context.get_input(0);
auto axis = ov::op::v0::Constant::create(ov::element::i64, {}, {3});
auto res = std::make_shared<ov::op::v0::CumSum>(x, axis);
return rename_outputs_with_suffix({res}, context.get_name());
}
} // namespace op
} // namespace ggml
} // namespace frontend
} // namespace ov

58
ggml/src/ggml-openvino/openvino/op/diag.cpp Normal file
View File

@ -0,0 +1,58 @@
#include "../node_context.h"
#include "../op_table.h"
#include "../utils.h"
#include <openvino/op/constant.hpp>
#include <openvino/op/equal.hpp>
#include <openvino/op/multiply.hpp>
#include <openvino/op/range.hpp>
#include <openvino/op/reshape.hpp>
#include <openvino/op/select.hpp>
namespace ov {
namespace frontend {
namespace ggml {
namespace op {
// GGML DIAG takes a 1D vector (ne0, 1, ne2, ne3) and produces a diagonal matrix
// of shape (ne0, ne0, ne2, ne3).
// In OV layout (ggml [ne0, ne1, ne2, ne3] → OV [ne3, ne2, ne1, ne0]):
// input: [ne3, ne2, 1, ne0]
// output: [ne3, ne2, ne0, ne0]
// The diagonal: output[..., i, j] = input[..., 0, j] if i == j, else 0.
OutputVector translate_diag(const NodeContext & context) {
num_inputs_check(context, 1, 1);
auto x = context.get_input(0); // OV shape: [ne3, ne2, 1, ne0]
auto out_shape = context.get_output_shape().to_shape();
int64_t n = static_cast<int64_t>(out_shape[3]); // ne0
// Build index range [0, 1, ..., n-1]
auto start = ov::op::v0::Constant::create(ov::element::i64, {}, {int64_t(0)});
auto stop = ov::op::v0::Constant::create(ov::element::i64, {}, {n});
auto step = ov::op::v0::Constant::create(ov::element::i64, {}, {int64_t(1)});
auto range = std::make_shared<ov::op::v4::Range>(start, stop, step, ov::element::i64);
// col_idx shape [1, 1, 1, n]
auto col_shape = ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{1, 1, 1, n});
auto col_idx = std::make_shared<ov::op::v1::Reshape>(range, col_shape, false);
// row_idx shape [1, 1, n, 1]
auto row_shape = ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{1, 1, n, 1});
auto row_idx = std::make_shared<ov::op::v1::Reshape>(range, row_shape, false);
// mask: true where col == row (diagonal)
auto mask = std::make_shared<ov::op::v1::Equal>(col_idx, row_idx);
// Broadcast input from [ne3, ne2, 1, ne0] to [ne3, ne2, ne0, ne0] via select
auto zero = ov::op::v0::Constant::create(ov::element::f32, {}, {0.0f});
auto res = std::make_shared<ov::op::v1::Select>(mask, x, zero);
return rename_outputs_with_suffix({res}, context.get_name());
}
} // namespace op
} // namespace ggml
} // namespace frontend
} // namespace ov

34
ggml/src/ggml-openvino/openvino/op/fill.cpp Normal file
View File

@ -0,0 +1,34 @@
#include "../node_context.h"
#include "../op_table.h"
#include "../utils.h"
#include <openvino/op/broadcast.hpp>
#include <openvino/op/constant.hpp>
namespace ov {
namespace frontend {
namespace ggml {
namespace op {
// GGML FILL sets all elements of a tensor to a constant value.
// The constant is stored as a float in op_params[0].
OutputVector translate_fill(const NodeContext & context) {
num_inputs_check(context, 1, 1);
float c;
memcpy(&c, context.get_output_op_params(), sizeof(float));
auto shape = context.get_input_shape(0).to_shape();
auto val = ov::op::v0::Constant::create(ov::element::f32, {}, {c});
auto target_shape = ov::op::v0::Constant::create(ov::element::i64, {shape.size()},
std::vector<int64_t>(shape.begin(), shape.end()));
auto res = std::make_shared<ov::op::v3::Broadcast>(val, target_shape);
return rename_outputs_with_suffix({res}, context.get_name());
}
} // namespace op
} // namespace ggml
} // namespace frontend
} // namespace ov

100
ggml/src/ggml-openvino/openvino/op/gated_delta_net.cpp
View File

@ -19,6 +19,7 @@
#include <openvino/op/reshape.hpp>
#include <openvino/op/squeeze.hpp>
#include <openvino/op/subtract.hpp>
#include <openvino/op/tile.hpp>
#include <openvino/op/transpose.hpp>
#include <openvino/op/unsqueeze.hpp>
#include <vector>
@ -31,57 +32,76 @@ namespace op {
static OutputVector translate_gated_delta_net_ref(const NodeContext & context);
OutputVector translate_gated_delta_net(const NodeContext & context) {
// auto v_shape = context.get_input_shape(2).to_shape(); // [B, T, H_v, S_v]
// auto q_shape = context.get_input_shape(0).to_shape(); // [B, T, H_k, S_k]
auto v_shape = context.get_input_shape(2).to_shape(); // [B, T, H_v, S_v]
auto q_shape = context.get_input_shape(0).to_shape(); // [B, T, H_k, S_k]
// // Fused GatedDeltaNet op only supports scalar gate (kda=0).
// // Fall back to reference implementation for per-key-dimension gating.
// // if (kda) {
// // return translate_gated_delta_net_ref(context);
// // }
// auto q = context.get_input(0);
// auto k = context.get_input(1);
// auto v = context.get_input(2);
// auto g = context.get_input(3);
// auto beta = context.get_input(4);
// auto state = context.get_input(5);
// Fused GatedDeltaNet op only supports scalar gate (kda=0).
// Fall back to reference implementation for per-key-dimension gating.
// if (kda) {
// return translate_gated_delta_net_ref(context);
// }
// const int64_t B = v_shape[0];
// const int64_t T = v_shape[1];
// const int64_t H_v = v_shape[2];
// const int64_t S_v = v_shape[3];
const int64_t H_v = v_shape[2];
const int64_t S_v = v_shape[3];
const int64_t H_k = q_shape[2];
// const int64_t S_k = q_shape[3];
// // ggml state layout (OV notation): [B, H_v, value_dim, key_dim]
// // GatedDeltaNet op expects: [B, H_v, key_dim, value_dim]
// auto state_reshape_shape =
// ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{B, H_v, S_v, S_k});
// state = std::make_shared<ov::op::v1::Reshape>(state, state_reshape_shape, false);
// auto state_perm = ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{0, 1, 3, 2});
// state = std::make_shared<ov::op::v1::Transpose>(state, state_perm);
auto q = context.get_input(0);
auto k = context.get_input(1);
auto v = process_view_input(context, 2, H_v * S_v);
auto g = context.get_input(3);
auto beta = context.get_input(4);
auto state = context.get_input(5);
// g = std::make_shared<ov::op::v0::Squeeze>(g, ov::op::v0::Constant::create(ov::element::i64, {1}, {3}));
// beta = std::make_shared<ov::op::v0::Squeeze>(beta, ov::op::v0::Constant::create(ov::element::i64, {1}, {3}));
// ggml maps GQA heads in tiled order, while OV GDN maps repeated heads in grouped order.
if (H_v != H_k) {
const int64_t repeat = H_v / H_k;
auto repeats = ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{1, 1, repeat, 1});
q = std::make_shared<ov::op::v0::Tile>(q, repeats);
k = std::make_shared<ov::op::v0::Tile>(k, repeats);
}
// auto gdn = std::make_shared<ov::op::internal::GatedDeltaNet>(q, k, v, state, g, beta);
if (context.get_view_input_size(2)) {
// Same as l2_norm case 1
v = std::make_shared<ov::op::v0::Squeeze>(v, ov::op::v0::Constant::create(ov::element::i64, {1}, {0}));
auto v_shape = context.get_input_shape(2).to_shape();
std::vector<int64_t> reshape_pattern = {0, 0, (int64_t) v_shape[2], (int64_t) v_shape[3]};
v = std::make_shared<ov::op::v1::Reshape>(
v, ov::op::v0::Constant::create(ov::element::i64, {4}, reshape_pattern), true);
}
// auto attn_4d = gdn->output(0);
// auto state_4d = gdn->output(1); // [B, H_v, key_dim, value_dim]
// // Transpose output state back to ggml layout [B, H_v, value_dim, key_dim]
// auto state_transposed = std::make_shared<ov::op::v1::Transpose>(state_4d, state_perm);
// auto flat_shape_1d = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1});
// auto attn = std::make_shared<ov::op::v1::Reshape>(attn_4d, flat_shape_1d, false);
// auto new_state = std::make_shared<ov::op::v1::Reshape>(state_transposed, flat_shape_1d, false);
// auto packed = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{attn, new_state}, 0);
// auto out_shape =
// ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{1, 1, T * B + S_v * B, S_v * H_v});
// auto res = std::make_shared<ov::op::v1::Reshape>(packed, out_shape, false);
// ggml state layout (OV notation): [B, H_v, value_dim, key_dim]
// GatedDeltaNet op expects: [B, H_v, key_dim, value_dim]
auto state_perm = ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{0, 1, 3, 2});
state = std::make_shared<ov::op::v1::Transpose>(state, state_perm);
// return rename_outputs_with_suffix({res}, context.get_name());
g = std::make_shared<ov::op::v0::Squeeze>(g, ov::op::v0::Constant::create(ov::element::i64, {1}, {3}));
beta = std::make_shared<ov::op::v0::Squeeze>(beta, ov::op::v0::Constant::create(ov::element::i64, {1}, {3}));
// The OV version in CI does not have the GatedDeltaNet op, so use reference implementation for now.
return translate_gated_delta_net_ref(context);
// std::cout << "GatedDeltaNet input shapes: q=" << q.get_partial_shape() << ", k=" << k.get_partial_shape()
// << ", v=" << v.get_partial_shape() << ", g=" << g.get_partial_shape()
// << ", beta=" << beta.get_partial_shape() << ", state=" << state.get_partial_shape() << std::endl;
auto gdn = std::make_shared<ov::op::internal::GatedDeltaNet>(q, k, v, state, g, beta);
auto attn_4d = gdn->output(0);
auto state_4d = gdn->output(1); // [B, H_v, key_dim, value_dim]
// std::cout << "GatedDeltaNet output shapes: attn=" << gdn->output(0).get_partial_shape()
// << ", new_state=" << gdn->output(1).get_partial_shape() << std::endl;
// Transpose output state back to ggml layout [B, H_v, value_dim, key_dim]
auto state_transposed = std::make_shared<ov::op::v1::Transpose>(state_4d, state_perm);
auto flat_shape_1d = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1});
auto attn = std::make_shared<ov::op::v1::Reshape>(attn_4d, flat_shape_1d, false);
auto new_state = std::make_shared<ov::op::v1::Reshape>(state_transposed, flat_shape_1d, false);
auto packed = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{attn, new_state}, 0);
auto out_shape = ov::op::v0::Constant::create(ov::element::i64, {4},
std::vector<int64_t>{1, 1, -1 /*T * B + S_v * B*/, S_v * H_v});
auto res = std::make_shared<ov::op::v1::Reshape>(packed, out_shape, false);
return rename_outputs_with_suffix({res}, context.get_name());
}
static OutputVector translate_gated_delta_net_ref(const NodeContext & context) {

43
ggml/src/ggml-openvino/openvino/op/gather_matmul.hpp Normal file
View File

@ -0,0 +1,43 @@
// Copyright (C) 2018-2026 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
// Local mirror of OpenVINO's internal ov::op::internal::GatherMatmul op.
//
// The op class body (validate_and_infer_types / clone_with_new_inputs) is
// provided by the linked libopenvino.so; only the declaration is needed here so
// the backend can construct the node directly (same approach as GatedDeltaNet).
// The class layout must stay in sync with
// openvino/src/common/transformations/include/ov_ops/gather_matmul.hpp
//
// \note GatherMatmul op class is under development and subject to change.
#pragma once
#include "openvino/op/op.hpp"
namespace ov::op::internal {
class OPENVINO_API GatherMatmul : public ov::op::Op {
public:
OPENVINO_OP("GatherMatmul")
GatherMatmul() = default;
GatherMatmul(const ov::Output<Node>& A,
const ov::Output<Node>& B,
const ov::Output<Node>& indices,
const ov::Output<Node>& bias);
GatherMatmul(const ov::Output<Node>& A, const ov::Output<Node>& B, const ov::Output<Node>& indices);
std::shared_ptr<Node> clone_with_new_inputs(const ov::OutputVector& new_args) const override;
void validate_and_infer_types() override;
private:
// the weights matrix B is expected to have the transposed form [group, N, K]
static constexpr bool transp_a = false;
static constexpr bool transp_b = true;
};
} // namespace ov::op::internal

84
ggml/src/ggml-openvino/openvino/op/get_rows.cpp
View File

@ -2,11 +2,16 @@
#include "../op_table.h"
#include "../utils.h"
#include <climits>
#include <openvino/core/node.hpp>
#include <openvino/core/node_output.hpp>
#include <openvino/op/broadcast.hpp>
#include <openvino/op/concat.hpp>
#include <openvino/op/constant.hpp>
#include <openvino/op/convert.hpp>
#include <openvino/op/gather.hpp>
#include <openvino/op/shape_of.hpp>
#include <openvino/op/slice.hpp>
#include <openvino/op/squeeze.hpp>
#include <openvino/op/unsqueeze.hpp>
@ -20,7 +25,27 @@ OutputVector translate_get_rows(const NodeContext & context) {
Output<Node> res;
auto data = process_view_input_new(context, 0);
auto indices = process_view_input_new(context, 1);
auto op_case = context.get_op_case();
ov::Output<ov::Node> indices;
if ((op_case == 1 || op_case == 2) && context.has_input("s_copy_active_slot_len")) {
// Recurrent state reorder (inp->s_copy): slice the active (op_case 1) or extra (op_case 2)
// segment from the s_copy index list at runtime, instead of baking the static view offset,
// so the cached IR works for any number of active sequences.
auto s_copy = context.get_input(1);
auto len = context.get_input("s_copy_active_slot_len");
auto step = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
auto axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {3});
if (op_case == 1) {
auto begin = ov::op::v0::Constant::create(ov::element::i64, {1}, {0});
indices = std::make_shared<ov::op::v8::Slice>(s_copy, begin, len, step, axis);
} else {
auto end = ov::op::v0::Constant::create(ov::element::i64, {1}, {INT_MAX});
indices = std::make_shared<ov::op::v8::Slice>(s_copy, len, end, step, axis);
}
} else {
indices = process_view_input_new(context, 1);
}
// data[1,b,x,y] ind[1,1,b,x'] test-backend-ops case
// data[x,y] ind[1,1,1,x'] normal case
@ -37,7 +62,62 @@ OutputVector translate_get_rows(const NodeContext & context) {
auto axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {1});
data =
std::make_shared<ov::op::v0::Squeeze>(data, ov::op::v0::Constant::create(ov::element::i64, {1}, {0}));
res = std::make_shared<ov::op::v8::Gather>(data, indices, axis, 1);
// data: [batch, rows, ...], indices: [batch, n] - this is a batched gather
// (batch_dims=1) along the rows axis. The data and indices batch dims are
// logically equal (both == n_tokens) but reach this node through independent
// reshapes, so the GPU plugin's gather shape inference cannot prove
// data.shape[0] == indices.shape[0] and rejects the node. We must tie both
// batch dims to the SAME value, and crucially that value must stay DYNAMIC.
const auto data_ps = data.get_partial_shape();
const auto idx_ps = indices.get_partial_shape();
const bool data_batch_static = data_ps.rank().is_static() && data_ps[0].is_static();
const bool idx_batch_dynamic = idx_ps.rank().is_dynamic() || idx_ps[0].is_dynamic();
if (data_batch_static && idx_batch_dynamic) {
// MoE per-expert-scale path: `data` is a statically-tiled REPEAT
// (ggml_repeat_4d(scale, 1, n_expert, n_tokens, 1)) whose batch dim is a
// compile-time-constant n_tokens, and every batch slice is IDENTICAL (it was
// tiled from a single [1, n_expert, 1] scale). `indices` (selected_experts)
// carries the genuinely dynamic token dim. Broadcasting indices up to the
// static data batch (the naive fix) would freeze the token dim to the
// captured prefill length, and that static value then flows through the
// gather into the residual stream, making every following decoder layer
// static -> triggers the GPU in-place-concat KV-cache corruption (only
// layer 0 stays dynamic). A static->dynamic Broadcast cannot expand, so
// instead collapse the redundant data batch to 1 and broadcast 1->dynamic to
// match the indices batch. Mathematically identical (the slices are equal),
// and the whole graph stays dynamic.
auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0});
auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
auto axis0 = ov::op::v0::Constant::create(ov::element::i64, {1}, {0});
auto data_b1 = std::make_shared<ov::op::v8::Slice>(data, zero, one, one, axis0); // [1, rows, ...]
auto idx_shape = std::make_shared<ov::op::v3::ShapeOf>(indices, ov::element::i64);
auto idx_batch = get_dimensions(idx_shape, {0}); // [batch] (dynamic)
auto data_b1_shape = std::make_shared<ov::op::v3::ShapeOf>(data_b1, ov::element::i64);
const auto rank = data_ps.rank().get_length();
std::vector<int> rest_axes;
for (int a = 1; a < rank; ++a) {
rest_axes.push_back(a);
}
auto data_rest = get_dimensions(data_b1_shape, rest_axes); // [rows, ...]
auto data_target = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{idx_batch, data_rest}, 0);
data =
std::make_shared<ov::op::v3::Broadcast>(data_b1, data_target, ov::op::BroadcastType::BIDIRECTIONAL);
res = std::make_shared<ov::op::v8::Gather>(data, indices, axis, 1);
} else {
// General case: tie the indices batch to the data batch (the data batch is
// already dynamic, e.g. the routing-weights gather whose data comes from the
// activations). Broadcast indices to [data_batch, indices_n].
auto data_shape = std::make_shared<ov::op::v3::ShapeOf>(data, ov::element::i64);
auto data_batch = get_dimensions(data_shape, {0}); // [batch]
auto idx_shape = std::make_shared<ov::op::v3::ShapeOf>(indices, ov::element::i64);
auto idx_n = get_dimensions(idx_shape, {1}); // [n]
auto idx_target = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{data_batch, idx_n}, 0);
indices = std::make_shared<ov::op::v3::Broadcast>(indices, idx_target,
ov::op::BroadcastType::BIDIRECTIONAL);
res = std::make_shared<ov::op::v8::Gather>(data, indices, axis, 1);
}
}
} else if (context.is_stateful() && data.get_partial_shape().rank() == 3) {
auto axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {1});

17
ggml/src/ggml-openvino/openvino/op/l2_norm.cpp
View File

@ -8,7 +8,9 @@
#include <openvino/op/maximum.hpp>
#include <openvino/op/multiply.hpp>
#include <openvino/op/reduce_sum.hpp>
#include <openvino/op/reshape.hpp>
#include <openvino/op/sqrt.hpp>
#include <openvino/op/squeeze.hpp>
namespace ov {
namespace frontend {
@ -20,6 +22,21 @@ OutputVector translate_l2_norm(const NodeContext & context) {
auto input_node = process_view_input_new(context, 0);
if (context.get_op_case() == 1) {
// 92: [ 128, 16, 1, 2] VIEW q_conv-1
// [ 6144, 1, 2, 1] 0: UNARY conv_output_silu-1
// 93: [ 128, 16, 1, 2] L2_NORM q_conv_predelta-1
// [ 128, 16, 1, 2] 0: VIEW q_conv-1
auto output_shape = context.get_output_shape().to_shape();
input_node = process_view_input(context, 0, output_shape[2] * output_shape[3]);
input_node =
std::make_shared<ov::op::v0::Squeeze>(input_node, ov::op::v0::Constant::create(ov::element::i64, {1}, {0}));
std::vector<int64_t> reshape_pattern = {0, 0, (int64_t) output_shape[2], (int64_t) output_shape[3]};
input_node = std::make_shared<ov::op::v1::Reshape>(
input_node, ov::op::v0::Constant::create(ov::element::i64, {4}, reshape_pattern), true);
}
auto squared = std::make_shared<ov::op::v1::Multiply>(input_node, input_node);
auto sum_squared = std::make_shared<ov::op::v1::ReduceSum>(

153
ggml/src/ggml-openvino/openvino/op/mul_mat_id.cpp
View File

@ -1,6 +1,8 @@
#include "../node_context.h"
#include "../op_table.h"
#include "../utils.h"
#include "gather_matmul.hpp"
#include "ggml-openvino/ggml-openvino-extra.h"
#include <cstdint>
#include <cstring>
@ -18,6 +20,7 @@
#include <openvino/op/reshape.hpp>
#include <openvino/op/shape_of.hpp>
#include <openvino/op/slice.hpp>
#include <openvino/op/transpose.hpp>
#include <openvino/op/unsqueeze.hpp>
#include <vector>
@ -37,6 +40,70 @@ ov::Output<ov::Node> slice_axis(const ov::Output<ov::Node> & input, int64_t axis
const_i64({axis}));
}
ov::Output<ov::Node> static_shape_dims_or_shapeof(const ov::Output<ov::Node> & input,
const std::vector<int> & dims) {
const auto partial_shape = input.get_partial_shape();
if (partial_shape.is_static()) {
std::vector<int64_t> values;
values.reserve(dims.size());
for (const int64_t dim : dims) {
values.push_back(partial_shape[dim].get_length());
}
return const_i64(values);
}
auto shape = std::make_shared<ov::op::v3::ShapeOf>(input, ov::element::i64);
return get_dimensions(shape, dims);
}
ov::Output<ov::Node> translate_mul_mat_id_gather_matmul_fallback(const NodeContext & context,
ov::Output<ov::Node> expert_weights,
ov::Output<ov::Node> activations,
ov::Output<ov::Node> ids) {
auto gather_axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {0});
ov::Output<ov::Node> selected_weights = std::make_shared<ov::op::v8::Gather>(expert_weights, ids, gather_axis);
const auto output_type = context.get_output_type();
if (selected_weights.get_element_type() != ov::element::f32) {
selected_weights = std::make_shared<ov::op::v0::Convert>(selected_weights, ov::element::f32);
}
if (activations.get_element_type() != ov::element::f32) {
activations = std::make_shared<ov::op::v0::Convert>(activations, ov::element::f32);
}
auto activations_shape = std::make_shared<ov::op::v3::ShapeOf>(activations, ov::element::i64);
auto ids_shape = std::make_shared<ov::op::v3::ShapeOf>(ids, ov::element::i64);
ov::Output<ov::Node> acts_target_dims = std::make_shared<ov::op::v0::Concat>(
ov::OutputVector{
get_dimensions(activations_shape, {0}),
get_dimensions(ids_shape, {1}),
get_dimensions(activations_shape, {2}),
},
0);
ov::Output<ov::Node> acts_broadcasted =
std::make_shared<ov::op::v3::Broadcast>(activations, acts_target_dims, ov::op::BroadcastType::BIDIRECTIONAL);
auto activations_expanded = std::make_shared<ov::op::v0::Unsqueeze>(acts_broadcasted, const_i64({2}));
ov::Output<ov::Node> result =
std::make_shared<ov::op::v0::MatMul>(activations_expanded, selected_weights, false, true);
auto output_shape = context.get_output_shape();
FRONT_END_OP_CONVERSION_CHECK(output_shape.rank().is_static() && output_shape.rank().get_length() == 4,
"Unexpected MUL_MAT_ID output rank");
FRONT_END_OP_CONVERSION_CHECK(output_shape[3].is_static(), "Expected static row dimension for MUL_MAT_ID output");
auto batch_dim = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
auto row_dim = ov::op::v0::Constant::create(ov::element::i64, {1}, {output_shape[3].get_length()});
auto result_target_dims = std::make_shared<ov::op::v0::Concat>(
ov::OutputVector{batch_dim, get_dimensions(ids_shape, {0, 1}), row_dim}, 0);
result = std::make_shared<ov::op::v1::Reshape>(result, result_target_dims, false);
if (result.get_element_type() != output_type) {
result = std::make_shared<ov::op::v0::Convert>(result, output_type);
}
return result;
}
ov::Output<ov::Node> translate_mul_mat_id_mxfp4_packed(const NodeContext & context,
ov::Output<ov::Node> expert_weights,
ov::Output<ov::Node> activations,
@ -144,22 +211,33 @@ OutputVector translate_mul_mat_id(const NodeContext & context) {
context.get_name());
}
// General (non-packed) path: dense F32/F16/BF16 weights, or the f16 dequantization chain for
// quantized MoE experts (see extract_quantized_weights / make_int4_weights / make_int8_weights in
// ggml-quants.cpp). Routed through ov::op::internal::GatherMatmul instead of a naive
// Gather+Broadcast+MatMul, so the selected expert's full weight matrix is never materialized per
// token. The CPU plugin's ConvertGatherMatmulToGatherMatmulCompressed pass (run during
// compile_model) fuses the dequantization chain feeding GatherMatmul's B input into a
// GatherMatmulCompressed node automatically, as long as MarkDequantization has marked the chain --
// see translate_session.cpp's apply_transformations for the MarkDequantization registration.
//
// OpenVINO sees GGML tensors in reversed dimension order:
// weights: [1, n_expert, m, k]
// activations: [1, n_tokens, n_used_or_1, k]
// ids: [1, 1, n_tokens, n_used]
// Rebuild the logical ranks explicitly from the 4D inputs instead of relying
// on fixed squeeze axes: real graphs can arrive through VIEW/RESHAPE chains
// where singleton axes are still represented differently at this point.
auto expert_weights_shape_4d = std::make_shared<ov::op::v3::ShapeOf>(expert_weights, ov::element::i64);
auto activations_shape_4d = std::make_shared<ov::op::v3::ShapeOf>(activations, ov::element::i64);
auto ids_shape_4d = std::make_shared<ov::op::v3::ShapeOf>(ids, ov::element::i64);
// expert_weights is either [1, n_expert, m, k] (4D, e.g. non-quantized weights without a
// pre-built extra) or already [n_expert, m, k] (3D, weights routed through
// process_weight_tensor) -- GatherMatmul's B input expects the latter.
auto expert_weights_rank = expert_weights.get_partial_shape().rank();
FRONT_END_OP_CONVERSION_CHECK(expert_weights_rank.is_static(),
"Expected static rank for MUL_MAT_ID expert weights");
const bool use_gpu_fallback = ggml_openvino_get_device_name() == "GPU";
if (expert_weights_rank.get_length() == 4) {
auto expert_weights_shape_3d = static_shape_dims_or_shapeof(expert_weights, {1, 2, 3});
expert_weights = std::make_shared<ov::op::v1::Reshape>(expert_weights, expert_weights_shape_3d, false);
}
auto expert_weights_shape_3d = get_dimensions(expert_weights_shape_4d, {1, 2, 3});
auto activations_shape_3d = get_dimensions(activations_shape_4d, {1, 2, 3});
auto ids_shape_2d = get_dimensions(ids_shape_4d, {2, 3});
auto activations_shape_3d = static_shape_dims_or_shapeof(activations, {1, 2, 3});
auto ids_shape_2d = static_shape_dims_or_shapeof(ids, {2, 3});
expert_weights = std::make_shared<ov::op::v1::Reshape>(expert_weights, expert_weights_shape_3d, false);
activations = std::make_shared<ov::op::v1::Reshape>(activations, activations_shape_3d, false);
ids = std::make_shared<ov::op::v1::Reshape>(ids, ids_shape_2d, false);
@ -167,51 +245,30 @@ OutputVector translate_mul_mat_id(const NodeContext & context) {
ids = std::make_shared<ov::op::v0::Convert>(ids, ov::element::i32);
}
auto gather_axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {0});
ov::Output<ov::Node> selected_weights = std::make_shared<ov::op::v8::Gather>(expert_weights, ids, gather_axis);
const auto output_type = context.get_output_type();
if (selected_weights.get_element_type() != ov::element::f32) {
selected_weights = std::make_shared<ov::op::v0::Convert>(selected_weights, ov::element::f32);
}
if (activations.get_element_type() != ov::element::f32) {
activations = std::make_shared<ov::op::v0::Convert>(activations, ov::element::f32);
}
auto activations_shape = std::make_shared<ov::op::v3::ShapeOf>(activations, ov::element::i64);
auto ids_shape = std::make_shared<ov::op::v3::ShapeOf>(ids, ov::element::i64);
ov::Output<ov::Node> acts_target_dims = std::make_shared<ov::op::v0::Concat>(
ov::OutputVector{
get_dimensions(activations_shape, {0}),
get_dimensions(ids_shape, {1}),
get_dimensions(activations_shape, {2}),
},
0);
ov::Output<ov::Node> acts_broadcasted =
std::make_shared<ov::op::v3::Broadcast>(activations, acts_target_dims, ov::op::BroadcastType::BIDIRECTIONAL);
if (use_gpu_fallback || !expert_weights.get_partial_shape().is_static() || !activations.get_partial_shape().is_static() ||
!ids.get_partial_shape().is_static()) {
return rename_outputs_with_suffix({translate_mul_mat_id_gather_matmul_fallback(context, expert_weights, activations, ids)},
context.get_name());
}
auto unsqueeze_axes = ov::op::v0::Constant::create(ov::element::i64, {1}, {2});
auto activations_expanded = std::make_shared<ov::op::v0::Unsqueeze>(acts_broadcasted, unsqueeze_axes);
// GatherMatmul's A input is [n_used_or_1, n_tokens, k]; activations_3d is
// [n_tokens, n_used_or_1, k].
auto activations_transpose_order = const_i64({1, 0, 2});
ov::Output<ov::Node> activations_for_gather =
std::make_shared<ov::op::v1::Transpose>(activations, activations_transpose_order);
auto batch_dim = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
auto output_shape = context.get_output_shape();
FRONT_END_OP_CONVERSION_CHECK(output_shape.rank().is_static() && output_shape.rank().get_length() == 4,
"Unexpected MUL_MAT_ID output rank");
FRONT_END_OP_CONVERSION_CHECK(output_shape[3].is_static(), "Expected static row dimension for MUL_MAT_ID output");
const auto row_dim_value = output_shape[3].get_length();
auto row_dim = ov::op::v0::Constant::create(ov::element::i64, {1}, {row_dim_value});
ov::Output<ov::Node> result = std::make_shared<ov::op::internal::GatherMatmul>(activations_for_gather, expert_weights, ids);
ov::Output<ov::Node> result =
std::make_shared<ov::op::v0::MatMul>(activations_expanded, selected_weights, false, true);
auto result_target_dims = std::make_shared<ov::op::v0::Concat>(
ov::OutputVector{
batch_dim,
get_dimensions(ids_shape, {0, 1}),
row_dim,
},
0);
result = std::make_shared<ov::op::v1::Reshape>(result, result_target_dims, false);
// result is [n_used, n_tokens, m]; GGML expects [1, n_tokens, n_used, m].
auto result_transpose_order = const_i64({1, 0, 2});
result = std::make_shared<ov::op::v1::Transpose>(result, result_transpose_order);
auto unsqueeze_axes = ov::op::v0::Constant::create(ov::element::i64, {1}, {0});
result = std::make_shared<ov::op::v0::Unsqueeze>(result, unsqueeze_axes);
if (result.get_element_type() != output_type) {
result = std::make_shared<ov::op::v0::Convert>(result, output_type);

46
ggml/src/ggml-openvino/openvino/op/repeat.cpp
View File

@ -23,47 +23,21 @@ OutputVector translate_repeat(const NodeContext & context) {
auto input = process_view_input_new(context, 0);
const auto input_shape = context.get_input_shape(0);
const auto output_shape = context.get_output_shape();
const auto input_shape = context.get_input_shape(0).to_shape();
const auto output_shape = context.get_output_shape().to_shape();
if (input_shape.rank().is_static() && output_shape.rank().is_static() &&
input_shape.rank() == output_shape.rank()) {
const auto rank = static_cast<size_t>(input_shape.rank().get_length());
std::vector<int64_t> repeats(rank, 1);
bool all_static = true;
std::vector<int64_t> repeats(4, 1);
for (size_t axis = 0; axis < 4; ++axis) {
const int64_t input_dim = input_shape[axis];
const int64_t output_dim = output_shape[axis];
for (size_t axis = 0; axis < rank; ++axis) {
if (!input_shape[axis].is_static() || !output_shape[axis].is_static()) {
all_static = false;
break;
}
FRONT_END_OP_CONVERSION_CHECK(input_dim > 0 && output_dim > 0 && output_dim % input_dim == 0,
"REPEAT input shape ", input_shape, " cannot tile to match ", output_shape);
const int64_t input_dim = input_shape[axis].get_length();
const int64_t output_dim = output_shape[axis].get_length();
FRONT_END_OP_CONVERSION_CHECK(input_dim > 0 && output_dim > 0 && output_dim % input_dim == 0,
"REPEAT input shape ", input_shape, " cannot tile to match ", output_shape);
repeats[axis] = output_dim / input_dim;
}
if (all_static) {
auto repeats_node = ov::op::v0::Constant::create(ov::element::i64, {repeats.size()}, repeats);
ov::Output<ov::Node> res = std::make_shared<ov::op::v0::Tile>(input, repeats_node);
return rename_outputs_with_suffix({res}, context.get_name());
}
repeats[axis] = output_dim / input_dim;
}
// Dynamic fallback: tile by the ratio of output to input shape.
auto input_shape_node = std::make_shared<ov::op::v3::ShapeOf>(input, ov::element::i64);
std::shared_ptr<ov::Node> target_shape_node;
if (output_shape.rank().is_static() && output_shape.is_static()) {
target_shape_node =
ov::op::v0::Constant::create(ov::element::i64, {output_shape.to_shape().size()}, output_shape.to_shape());
} else {
target_shape_node = std::make_shared<ov::op::v3::ShapeOf>(context.get_input(1), ov::element::i64);
}
auto repeats_node = std::make_shared<ov::op::v1::Divide>(target_shape_node, input_shape_node);
auto repeats_node = ov::op::v0::Constant::create(ov::element::i64, {repeats.size()}, repeats);
ov::Output<ov::Node> res = std::make_shared<ov::op::v0::Tile>(input, repeats_node);
return rename_outputs_with_suffix({res}, context.get_name());
}

35
ggml/src/ggml-openvino/openvino/op/reshape.cpp
View File

@ -25,13 +25,12 @@ OutputVector translate_reshape(const NodeContext & context) {
}
int op_case = context.get_op_case();
FRONT_END_CHECK_IMPLEMENTED(
op_case == 1 || op_case == 2 || op_case == 3 || op_case == 4 || op_case == 5 || op_case == 6,
"Unsupported RESHAPE case");
auto output_shape = context.get_output_shape().to_shape();
std::shared_ptr<ov::Node> new_shape_node;
if (op_case == 1) {
if (op_case == 0) {
new_shape_node = ov::op::v0::Constant::create(ov::element::i64, {4}, context.get_output_shape().to_shape());
} else if (op_case == 1) {
if (context.is_stateful()) {
new_shape_node = ov::op::v0::Constant::create(
ov::element::i64, {3}, std::vector<int64_t>{-1, (int64_t) output_shape[2], (int64_t) output_shape[3]});
@ -76,9 +75,33 @@ OutputVector translate_reshape(const NodeContext & context) {
// ov::op::v0::Constant::create(ov::element::i64, {1}, {(int64_t) context.get_output_shape().to_shape()[3]});
// auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
// new_shape_node = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{one, one, token_len, emb_size}, 0);
} else if (op_case == 6) {
new_shape_node = ov::op::v0::Constant::create(ov::element::i64, {4}, context.get_output_shape().to_shape());
// 14: [ 6144, 1, 2, 1] RESHAPE linear_attn_qkv_mixed-0
// [ 6144, 2, 1, 1] 0: MUL_MAT node_13
// reshape to [1, n_slot_active_len, -1, 6144]
if (context.has_input("s_copy_active_slot_len")) {
auto n_slot_active_len = context.get_input("s_copy_active_slot_len");
auto emb_size = ov::op::v0::Constant::create(ov::element::i64, {1},
{(int64_t) context.get_output_shape().to_shape()[3]});
auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
auto neg_one = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1});
new_shape_node =
std::make_shared<ov::op::v0::Concat>(ov::OutputVector{one, n_slot_active_len, neg_one, emb_size}, 0);
} else {
new_shape_node = ov::op::v0::Constant::create(ov::element::i64, {4}, context.get_output_shape().to_shape());
}
} else if (op_case == 7) {
// 57: [ 2048, 2, 1, 1] RESHAPE linear_attn_out-0 (reshaped)
// [ 2048, 1, 2, 1] 0: MUL_MAT linear_attn_out-0
std::vector<int64_t> shape_vec = {1, 1, -1, (int64_t) context.get_output_shape().to_shape()[3]};
new_shape_node = ov::op::v0::Constant::create(ov::element::i64, {4}, shape_vec);
} else if (op_case == 8) {
// 106: [ 128, 128, 16, 2] RESHAPE state_predelta-1
// [ 262144, 2, 1, 1] 0: GET_ROWS node_86
auto output_shape = context.get_output_shape().to_shape();
std::vector<int64_t> shape_vec = {-1, (int64_t) output_shape[1], (int64_t) output_shape[2],
(int64_t) output_shape[3]};
new_shape_node = ov::op::v0::Constant::create(ov::element::i64, {4}, shape_vec);
}
auto res = std::make_shared<ov::op::v1::Reshape>(context.get_input(0), new_shape_node, false);
return rename_outputs_with_suffix({res}, context.get_name());

41
ggml/src/ggml-openvino/openvino/op/rms_norm.cpp
View File

@ -7,8 +7,11 @@
#include <openvino/op/constant.hpp>
#include <openvino/op/divide.hpp>
#include <openvino/op/multiply.hpp>
#include <openvino/op/negative.hpp>
#include <openvino/op/power.hpp>
#include <openvino/op/reduce_mean.hpp>
#include <openvino/op/reshape.hpp>
#include <openvino/op/slice.hpp>
#include <openvino/op/sqrt.hpp>
namespace ov {
@ -19,9 +22,41 @@ namespace op {
OutputVector translate_rms_norm(const NodeContext & context) {
num_inputs_check(context, 1, 1);
auto input_node = process_view_input_new(context, 0);
auto square = std::make_shared<ov::op::v1::Power>(
input_node, ov::op::v0::Constant::create(ov::element::f32, ov::Shape{1}, {2.0f}));
auto op_case = context.get_op_case();
ov::Output<ov::Node> input_node;
if (op_case == 1) {
input_node = process_view_input_new(context, 0);
} else if (op_case == 2) {
auto ssm_state_size = context.get_ssm_state_size();
// The GDN op packs [attn | new_state] along the row axis; the state occupies the last
// ssm_state_size * n_seqs rows. Slice it off (scaling by the active sequence count) to keep
// just the attention output.
ov::Output<ov::Node> state_end;
if (context.has_input("s_copy_active_slot_len")) {
auto len = context.get_input("s_copy_active_slot_len");
auto state_rows = std::make_shared<ov::op::v1::Multiply>(
ov::op::v0::Constant::create(ov::element::i64, {1}, {ssm_state_size}), len);
state_end = std::make_shared<ov::op::v0::Negative>(state_rows);
} else {
state_end = ov::op::v0::Constant::create(ov::element::i64, {1}, {-ssm_state_size});
}
auto gdn_attn_output = std::make_shared<ov::op::v8::Slice>(
context.get_input(0), ov::op::v0::Constant::create(ov::element::i64, {1}, {0}), state_end,
ov::op::v0::Constant::create(ov::element::i64, {1}, {1}),
ov::op::v0::Constant::create(ov::element::i64, {1}, {2}));
auto input_shape = context.get_input_shape(0).to_shape();
input_node = std::make_shared<ov::op::v1::Reshape>(
gdn_attn_output,
ov::op::v0::Constant::create(
ov::element::i64, {4}, std::vector<int64_t>{1, -1, (int64_t) input_shape[2], (int64_t) input_shape[3]}),
false);
} else {
input_node = process_view_input_new(context, 0);
}
auto square = std::make_shared<ov::op::v1::Multiply>(input_node, input_node);
auto mean = std::make_shared<ov::op::v1::ReduceMean>(
square, ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {-1}), true);

110
ggml/src/ggml-openvino/openvino/op/rope.cpp
View File

@ -22,6 +22,7 @@
#include <openvino/op/subtract.hpp>
#include <openvino/op/transpose.hpp>
#include <openvino/op/unsqueeze.hpp>
#include <openvino/op/variadic_split.hpp>
#include <vector>
namespace ov {
@ -40,6 +41,9 @@ OutputVector translate_rope(const NodeContext & context) {
auto output_shape = context.get_output_shape().to_shape();
int32_t * op_params = context.get_output_op_params();
const int mode = op_case;
const int64_t head_dim = static_cast<int64_t>(output_shape[3]);
const int64_t configured_n_dims = static_cast<int64_t>(op_params[1]);
const int64_t n_dims = configured_n_dims == 0 ? head_dim : configured_n_dims;
constexpr int TYPE_NORMAL = 0;
constexpr int TYPE_NEOX = 1;
@ -80,6 +84,9 @@ OutputVector translate_rope(const NodeContext & context) {
data_node = std::make_shared<ov::op::v0::Convert>(data_node, ov::element::f32);
}
FRONT_END_OP_CONVERSION_CHECK(n_dims > 0 && n_dims <= head_dim && (n_dims % 2 == 0),
"ROPE expects even n_dims in [1, head_dim]");
// TODO(openvino-gpu-rope-fusion): TEMPORARY WORKAROUND - do NOT revert until the
// OpenVINO GPU plugin is updated.
//
@ -94,13 +101,18 @@ OutputVector translate_rope(const NodeContext & context) {
// be restored to the captured even/odd translation. Until then, keep both paths:
// the active Flux rewrite here and the previous translation preserved below.
if (mode == TYPE_NORMAL) {
auto axis_last = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1});
auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0});
auto step_one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
// Emit the Flux-style interleaved-RoPE pattern so the GPU plugin's
// RoPEFusionFlux matcher folds this subgraph into ov::op::internal::RoPE:
// x_paired = Reshape(x, [1, S, n_heads, head_size/2, 2])
// x_paired = Reshape(x_rot, [1, S, n_heads, n_dims/2, 2])
// x0, x1 = Split(x_paired, axis=-1, num_splits=2)
// x1_neg = x1 * -1
// x_rotated = Reshape(Concat([x1_neg, x0], axis=-1), [1, S, n_heads, head_size])
// y = x * t_cos + x_rotated * t_sin
// x_rotated = Reshape(Concat([x1_neg, x0], axis=-1), [1, S, n_heads, n_dims])
// y_rot = x_rot * t_cos + x_rotated * t_sin
// y = Concat([y_rot, x_tail], axis=-1) if n_dims < head_dim
// Mathematically equivalent to the even/odd Slice form below.
//
// RoPEFusionFlux requires rank_equals(4) on x, t_cos and t_sin. The cos/sin
@ -114,15 +126,16 @@ OutputVector translate_rope(const NodeContext & context) {
std::vector<int64_t>{1, -1, (int64_t) output_shape[2], (int64_t) output_shape[3]});
data_node = std::make_shared<ov::op::v1::Reshape>(data_node, r4_shape, false);
}
const int64_t head_size = static_cast<int64_t>(output_shape[3]);
const int64_t n_heads = static_cast<int64_t>(output_shape[2]);
const int64_t half = head_size / 2;
const int64_t half = n_dims / 2;
auto rot_end = ov::op::v0::Constant::create(ov::element::i64, {1}, {n_dims});
auto rot_data = std::make_shared<ov::op::v8::Slice>(data_node, zero, rot_end, step_one, axis_last);
auto neg_one_f = ov::op::v0::Constant::create(data_node->get_element_type(), ov::Shape{}, {-1.0f});
auto paired_shape =
ov::op::v0::Constant::create(ov::element::i64, {5}, std::vector<int64_t>{1, -1, n_heads, half, 2});
auto x_paired = std::make_shared<ov::op::v1::Reshape>(data_node, paired_shape, false);
auto paired_shape = ov::op::v0::Constant::create(
ov::element::i64, {5}, std::vector<int64_t>{1, -1, n_heads, half, 2});
auto x_paired = std::make_shared<ov::op::v1::Reshape>(rot_data, paired_shape, false);
auto split_axis = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {-1});
auto data_split = std::make_shared<ov::op::v1::Split>(x_paired, split_axis, 2);
@ -133,28 +146,38 @@ OutputVector translate_rope(const NodeContext & context) {
auto x_rotated_paired = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{x1_neg, x0}, -1);
auto flat_shape =
ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{1, -1, n_heads, head_size});
auto x_rotated = std::make_shared<ov::op::v1::Reshape>(x_rotated_paired, flat_shape, false);
ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{1, -1, n_heads, n_dims});
auto x_rotated =
std::make_shared<ov::op::v1::Reshape>(x_rotated_paired, flat_shape, false);
// Expand cos/sin from [..., head_size/2] to [..., head_size] by repeating each
// Expand cos/sin from [..., n_dims/2] to [..., n_dims] by repeating each
// entry twice. Use special_zero on the final Reshape so the seq dim passes
// through dynamically. Final rank is 4 to satisfy the matcher's predicate.
auto expand_cos_sin = [&](Output<Node> cs) {
auto cs_unsq =
std::make_shared<ov::op::v0::Unsqueeze>(cs, ov::op::v0::Constant::create(ov::element::i64, {1}, {-1}));
auto bcast_target =
ov::op::v0::Constant::create(ov::element::i64, {5}, std::vector<int64_t>{1, 1, 1, half, 2});
auto bcast =
std::make_shared<ov::op::v3::Broadcast>(cs_unsq, bcast_target, ov::op::BroadcastType::BIDIRECTIONAL);
auto flat = ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{0, 0, 0, head_size});
auto cs_unsq = std::make_shared<ov::op::v0::Unsqueeze>(
cs, ov::op::v0::Constant::create(ov::element::i64, {1}, {-1}));
auto bcast_target = ov::op::v0::Constant::create(
ov::element::i64, {5}, std::vector<int64_t>{1, 1, 1, half, 2});
auto bcast = std::make_shared<ov::op::v3::Broadcast>(
cs_unsq, bcast_target, ov::op::BroadcastType::BIDIRECTIONAL);
auto flat = ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{0, 0, 0, n_dims});
return std::make_shared<ov::op::v1::Reshape>(bcast, flat, true);
};
Output<Node> cos_full = expand_cos_sin(cos_theta_node);
Output<Node> sin_full = expand_cos_sin(sin_theta_node);
auto y1 = std::make_shared<ov::op::v1::Multiply>(data_node, cos_full);
auto y1 = std::make_shared<ov::op::v1::Multiply>(rot_data, cos_full);
auto y2 = std::make_shared<ov::op::v1::Multiply>(x_rotated, sin_full);
res = std::make_shared<ov::op::v1::Add>(y1, y2);
auto rotated = std::make_shared<ov::op::v1::Add>(y1, y2);
if (n_dims < head_dim) {
auto tail_start = ov::op::v0::Constant::create(ov::element::i64, {1}, {n_dims});
auto tail_end = ov::op::v0::Constant::create(ov::element::i64, {1}, {head_dim});
auto tail = std::make_shared<ov::op::v8::Slice>(data_node, tail_start, tail_end, step_one, axis_last);
res = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{rotated, tail}, -1);
} else {
res = rotated;
}
}
// PRESERVED PREVIOUS TRANSLATION - Re-enable this branch (and remove the Flux branch above) once
// the GPU plugin's RoPE fusion is updated to recognize the even/odd Slice form;
@ -196,8 +219,27 @@ OutputVector translate_rope(const NodeContext & context) {
// ov::element::i64, {4}, std::vector<int64_t>{1, -1, (int64_t) output_shape[2], (int64_t) output_shape[3]});
// res = std::make_shared<ov::op::v1::Reshape>(stack, data_shape, false);
else if (mode == TYPE_NEOX) {
auto data_split = std::make_shared<ov::op::v1::Split>(
data_node, ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {-1}), 2);
// In stateful mode the data arrives rank-3 ([S, n_heads, head_size]) while the
// cos/sin tables are rank-4 ([1, S, 1, n_dims/2]). The resulting mixed-rank
// broadcast in the Multiply below is miscomputed by the OpenVINO GPU plugin,
// corrupting the rotated Q/K. Lift the data to rank-4 ([1, S, n_heads, head_size])
// first so the RoPE Multiplies are equal-rank, matching the TYPE_NORMAL branch.
// Stateful RoPE already produced rank-4 output, so downstream attention is unaffected.
if (context.is_stateful()) {
auto r4_shape = ov::op::v0::Constant::create(
ov::element::i64, {4},
std::vector<int64_t>{1, -1, (int64_t) output_shape[2], (int64_t) output_shape[3]});
data_node = std::make_shared<ov::op::v1::Reshape>(data_node, r4_shape, false);
}
auto axis_last = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {-1});
std::vector<int64_t> split_lengths = {n_dims / 2, n_dims / 2};
if (n_dims < head_dim) {
split_lengths.push_back(head_dim - n_dims);
}
auto data_split = std::make_shared<ov::op::v1::VariadicSplit>(
data_node, axis_last,
ov::op::v0::Constant::create(ov::element::i64, {split_lengths.size()}, split_lengths));
Output<Node> slice_data_node_0 = data_split->outputs()[0];
Output<Node> slice_data_node_1 = data_split->outputs()[1];
@ -209,16 +251,27 @@ OutputVector translate_rope(const NodeContext & context) {
std::make_shared<ov::op::v1::Multiply>(slice_data_node_0, sin_theta_node),
std::make_shared<ov::op::v1::Multiply>(slice_data_node_1, cos_theta_node));
res = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{first_half_node, second_half_node}, -1);
if (n_dims < head_dim) {
Output<Node> tail = data_split->outputs()[2];
res = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{first_half_node, second_half_node, tail}, -1);
} else {
res = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{first_half_node, second_half_node}, -1);
}
} else if (mode == TYPE_IMROPE) {
int64_t n_dims = data_node->get_output_partial_shape(0)[3].get_length();
auto cos_sin_shape = std::make_shared<ov::op::v0::Constant>(ov::element::i64, ov::Shape{4},
std::vector<int64_t>{1, -1, 1, (n_dims >> 1)});
auto cos_reshaped = std::make_shared<ov::op::v1::Reshape>(cos_theta_node, cos_sin_shape, true);
auto sin_reshaped = std::make_shared<ov::op::v1::Reshape>(sin_theta_node, cos_sin_shape, true);
auto split_axis = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {3});
auto split_a = std::make_shared<ov::op::v1::Split>(data_node, split_axis, 2);
std::vector<int64_t> split_lengths = {n_dims / 2, n_dims / 2};
if (n_dims < head_dim) {
split_lengths.push_back(head_dim - n_dims);
}
auto split_a = std::make_shared<ov::op::v1::VariadicSplit>(
data_node, split_axis,
ov::op::v0::Constant::create(ov::element::i64, {split_lengths.size()}, split_lengths));
auto x0 = split_a->output(0);
auto x1 = split_a->output(1);
auto mul_a = std::make_shared<ov::op::v1::Multiply>(x0, cos_reshaped);
@ -229,7 +282,12 @@ OutputVector translate_rope(const NodeContext & context) {
auto mul_d = std::make_shared<ov::op::v1::Multiply>(x1, cos_reshaped);
auto add = std::make_shared<ov::op::v1::Add>(mul_c, mul_d);
res = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{sub, add}, 3);
if (n_dims < head_dim) {
auto tail = split_a->output(2);
res = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{sub, add, tail}, 3);
} else {
res = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{sub, add}, 3);
}
}
if (res.get_element_type() != output_type) {

45
ggml/src/ggml-openvino/openvino/op/scale.cpp
View File

@ -2,9 +2,24 @@
#include "../op_table.h"
#include "../utils.h"
#include <openvino/core/except.hpp>
#include <openvino/op/add.hpp>
#include <openvino/op/concat.hpp>
#include <openvino/op/constant.hpp>
#include <openvino/op/convert.hpp>
#include <openvino/op/equal.hpp>
#include <openvino/op/gather.hpp>
#include <openvino/op/greater_eq.hpp>
#include <openvino/op/if.hpp>
#include <openvino/op/less.hpp>
#include <openvino/op/logical_or.hpp>
#include <openvino/op/multiply.hpp>
#include <openvino/op/range.hpp>
#include <openvino/op/reshape.hpp>
#include <openvino/op/shape_of.hpp>
#include <openvino/op/slice.hpp>
#include <openvino/op/squeeze.hpp>
#include <openvino/op/unsqueeze.hpp>
#include <vector>
namespace ov {
@ -21,6 +36,36 @@ OutputVector translate_scale(const NodeContext & context) {
memcpy(&bias, (float *) context.get_output_op_params() + 1, sizeof(float));
auto scale_node = std::make_shared<ov::op::v0::Constant>(ov::element::f32, ov::Shape{}, std::vector<float>{scale});
if (context.get_op_case() == 1 && context.has_input("cache_rs_reset_len")) {
auto cache_rs_reset_idx = context.get_input("cache_rs_reset_idx");
auto cache_rs_reset_len = context.get_input("cache_rs_reset_len");
auto cache_rs = context.get_input(0);
auto cache_shape = std::make_shared<ov::op::v3::ShapeOf>(cache_rs, ov::element::i64);
auto n_slots_1d = std::make_shared<ov::op::v8::Gather>(
cache_shape, ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {2}),
ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {0}));
auto n_slots = std::make_shared<ov::op::v0::Squeeze>(n_slots_1d);
auto iota = std::make_shared<ov::op::v4::Range>(
ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {0}), n_slots,
ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {1}), ov::element::i64);
auto idx_plus_len = std::make_shared<ov::op::v1::Add>(cache_rs_reset_idx, cache_rs_reset_len);
auto less_than_idx = std::make_shared<ov::op::v1::Less>(iota, cache_rs_reset_idx);
auto greater_equal_idx_plus_len = std::make_shared<ov::op::v1::GreaterEqual>(iota, idx_plus_len);
auto keep_mask = std::make_shared<ov::op::v1::LogicalOr>(less_than_idx, greater_equal_idx_plus_len);
auto keep_mask_f32 = std::make_shared<ov::op::v0::Convert>(keep_mask, ov::element::f32);
auto keep_mask_reshape = std::make_shared<ov::op::v0::Unsqueeze>(
keep_mask_f32, ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {1}));
auto cleared_cache_rs = std::make_shared<ov::op::v1::Multiply>(cache_rs, keep_mask_reshape);
return rename_outputs_with_suffix({cleared_cache_rs}, context.get_name());
}
auto scaled = std::make_shared<ov::op::v1::Multiply>(context.get_input(0), scale_node);
std::shared_ptr<ov::Node> res;

76
ggml/src/ggml-openvino/openvino/op/set.cpp Normal file
View File

@ -0,0 +1,76 @@
#include "../node_context.h"
#include "../op_table.h"
#include "../utils.h"
#include <cstdint>
#include <openvino/frontend/exception.hpp>
#include <openvino/op/add.hpp>
#include <openvino/op/constant.hpp>
#include <openvino/op/convert.hpp>
#include <openvino/op/range.hpp>
#include <openvino/op/reduce_prod.hpp>
#include <openvino/op/reshape.hpp>
#include <openvino/op/scatter_update.hpp>
#include <openvino/op/shape_of.hpp>
namespace ov {
namespace frontend {
namespace ggml {
namespace op {
// GGML SET writes src1 into a view of src0 and returns the updated tensor.
OutputVector translate_set(const NodeContext & context) {
num_inputs_check(context, 2, 2);
auto dst = process_view_input_new(context, 0);
auto src = process_view_input_new(context, 1);
src = std::make_shared<ov::op::v0::Convert>(src, context.get_output_type());
const auto dst_stride = context.get_input_stride(0);
FRONT_END_OP_CONVERSION_CHECK(dst_stride.size() >= 4, "SET requires 4D destination strides");
const auto * op_params = reinterpret_cast<const uint32_t *>(context.get_output_op_params());
const size_t offset = static_cast<size_t>(op_params[3]);
const size_t elem_size = dst_stride.back();
FRONT_END_OP_CONVERSION_CHECK(elem_size != 0 && offset % elem_size == 0,
"SET offset must be aligned to destination element size");
const int64_t offset_elems = static_cast<int64_t>(offset / elem_size);
auto dst_flat = std::make_shared<ov::op::v1::Reshape>(
dst,
ov::op::v0::Constant::create(ov::element::i64, {1}, {-1}),
false);
auto src_flat = std::make_shared<ov::op::v1::Reshape>(
src,
ov::op::v0::Constant::create(ov::element::i64, {1}, {-1}),
false);
auto src_shape = std::make_shared<ov::op::v3::ShapeOf>(src_flat, ov::element::i64);
auto src_len = std::make_shared<ov::op::v1::ReduceProd>(
src_shape,
ov::op::v0::Constant::create(ov::element::i64, {1}, {0}),
false);
auto start = ov::op::v0::Constant::create(ov::element::i64, {}, {offset_elems});
auto stop = std::make_shared<ov::op::v1::Add>(start, src_len);
auto step = ov::op::v0::Constant::create(ov::element::i64, {}, {1});
auto indices = std::make_shared<ov::op::v4::Range>(start, stop, step, ov::element::i64);
auto axis = ov::op::v0::Constant::create(ov::element::i64, {}, {0});
auto updated_flat = std::make_shared<ov::op::v3::ScatterUpdate>(dst_flat, indices, src_flat, axis);
auto dst_shape = std::make_shared<ov::op::v3::ShapeOf>(dst, ov::element::i64);
auto res = std::make_shared<ov::op::v1::Reshape>(updated_flat, dst_shape, false);
return rename_outputs_with_suffix({res}, context.get_name());
}
} // namespace op
} // namespace ggml
} // namespace frontend
} // namespace ov

41
ggml/src/ggml-openvino/openvino/op/set_rows.cpp
View File

@ -8,11 +8,13 @@
#include <openvino/core/node.hpp>
#include <openvino/core/node_output.hpp>
#include <openvino/frontend/exception.hpp>
#include <openvino/op/broadcast.hpp>
#include <openvino/op/concat.hpp>
#include <openvino/op/constant.hpp>
#include <openvino/op/convert.hpp>
#include <openvino/op/gather.hpp>
#include <openvino/op/reshape.hpp>
#include <openvino/op/scatter_elements_update.hpp>
#include <openvino/op/scatter_update.hpp>
#include <openvino/op/shape_of.hpp>
#include <openvino/op/slice.hpp>
@ -29,20 +31,17 @@ OutputVector translate_set_rows(const NodeContext & context) {
num_inputs_check(context, 3, 3);
auto data = process_view_input_new(context, 0);
auto indices = context.get_input(1);
auto dst = context.get_input(2);
auto indices = process_view_input_new(context, 1);
auto dst = process_view_input_new(context, 2);
data = std::make_shared<ov::op::v0::Convert>(data, context.get_output_type());
auto row_size = context.get_input_shape(2)[3].get_length();
const auto indices_shape = context.get_input_shape(1);
const bool multidim_indices = indices_shape.rank().is_static() &&
indices_shape.rank().get_length() == 4 &&
((indices_shape[1].is_static() && indices_shape[1].get_length() > 1) ||
(indices_shape[2].is_static() && indices_shape[2].get_length() > 1));
auto ind_squeezed =
std::make_shared<ov::op::v0::Squeeze>(indices, ov::op::v0::Constant::create(ov::element::i64, {3}, {0, 1, 2}));
auto data_reshaped = std::make_shared<ov::op::v1::Reshape>(
data,
ov::op::v0::Constant::create(ov::element::i64, {4},
{(int64_t) 1, (int64_t) 1, (int64_t) -1, (int64_t) row_size}),
false);
auto axes = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {2});
Output<Node> res;
@ -53,11 +52,31 @@ OutputVector translate_set_rows(const NodeContext & context) {
data = std::make_shared<ov::op::v1::Reshape>(
data, ov::op::v0::Constant::create(ov::element::i64, {4}, {(int64_t) 1, (int64_t) -1, dim2, dim3}), false);
res = std::make_shared<ov::op::v0::Concat>(OutputVector{dst, data}, concat_axis);
} else if (multidim_indices) {
auto updates_shape = std::make_shared<ov::op::v3::ShapeOf>(data, ov::element::i64);
auto indices_rank3 = std::make_shared<ov::op::v0::Squeeze>(
indices, ov::op::v0::Constant::create(ov::element::i64, {1}, {0}));
auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
auto indices_rank4_shape = std::make_shared<ov::op::v0::Concat>(OutputVector{get_dimensions(updates_shape, {0, 1, 2}), one}, 0);
auto indices_rank4 = std::make_shared<ov::op::v1::Reshape>(indices_rank3, indices_rank4_shape, false);
auto broadcasted_indices = std::make_shared<ov::op::v3::Broadcast>(indices_rank4, updates_shape);
res = std::make_shared<ov::op::v3::ScatterElementsUpdate>(dst, broadcasted_indices, data, axes);
} else {
auto row_size = context.get_input_shape(2)[3].get_length();
auto ind_squeezed = std::make_shared<ov::op::v0::Squeeze>(
indices, ov::op::v0::Constant::create(ov::element::i64, {3}, {0, 1, 2}));
auto data_reshaped = std::make_shared<ov::op::v1::Reshape>(
data,
ov::op::v0::Constant::create(ov::element::i64, {4},
{(int64_t) 1, (int64_t) 1, (int64_t) -1, (int64_t) row_size}),
false);
res = std::make_shared<ov::op::v3::ScatterUpdate>(dst, ind_squeezed, data_reshaped, axes);
}
if (auto dst_reshape = std::dynamic_pointer_cast<ov::op::v1::Reshape>(dst.get_node_shared_ptr())) {
auto dst_reshape = std::dynamic_pointer_cast<ov::op::v1::Reshape>(dst.get_node_shared_ptr());
if (!multidim_indices && dst_reshape) {
// Fix the case of multiple sequences, reshape back to original shape [1, n_seq, ctx_per_seq, emb]
// ctx_per_seq is not fixed due to llama-bench compatibility
auto dst_shape_partial = dst_reshape->get_input_partial_shape(0);

108
ggml/src/ggml-openvino/openvino/op/solve_tri.cpp Normal file
View File

@ -0,0 +1,108 @@
#include "../node_context.h"
#include "../op_table.h"
#include "../utils.h"
#include <openvino/op/broadcast.hpp>
#include <openvino/op/constant.hpp>
#include <openvino/op/divide.hpp>
#include <openvino/op/gather.hpp>
#include <openvino/op/loop.hpp>
#include <openvino/op/matmul.hpp>
#include <openvino/op/scatter_update.hpp>
#include <openvino/op/shape_of.hpp>
#include <openvino/op/subtract.hpp>
namespace ov {
namespace frontend {
namespace ggml {
namespace op {
// GGML SOLVE_TRI: solve Ax = B for lower-triangular A via forward substitution.
// Currently only lower, right, non-unitriangular variant is implemented.
//
// ggml layout: A [n, n, B1, B2], B [k, n, B1, B2] → X [k, n, B1, B2]
// OV layout: A [B2, B1, n, n], B [B2, B1, n, k] → X [B2, B1, n, k]
//
// Forward substitution row i:
// x[i] = (b[i] - sum_{t<i} A[i,t]*x[t]) / A[i,i]
//
// Implemented as an OV Loop op iterating n times with a carried X accumulator.
// Key insight: A is lower-triangular and X starts as zeros, so the full matmul
// A_row_i @ X_partial = sum_{t<i} A[i,t]*x[t] exactly (upper triangle of A
// is zero; unfilled rows of X are zero).
OutputVector translate_solve_tri(const NodeContext & context) {
num_inputs_check(context, 2, 2);
auto A = context.get_input(0); // [B2, B1, n, n]
auto B = context.get_input(1); // [B2, B1, n, k]
auto A_shape = context.get_input_shape(0).to_shape();
int64_t n = static_cast<int64_t>(A_shape[2]);
// Initial X: zeros with shape of B
auto B_shape_node = std::make_shared<ov::op::v3::ShapeOf>(B, ov::element::i64);
auto zero_f32 = ov::op::v0::Constant::create(ov::element::f32, {}, {0.0f});
auto X_init = std::make_shared<ov::op::v3::Broadcast>(zero_f32, B_shape_node);
// --- Loop body parameters ---
// body_iter: iteration counter injected by the Loop op (i64, shape {1})
auto body_iter = std::make_shared<ov::op::v0::Parameter>(ov::element::i64, ov::Shape{1});
auto body_X = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
auto body_A = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
auto body_B_p = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
auto c_axis2 = ov::op::v0::Constant::create(ov::element::i64, {1}, {int64_t(2)});
auto c_axis3 = ov::op::v0::Constant::create(ov::element::i64, {1}, {int64_t(3)});
auto c_axis2_scalar = ov::op::v0::Constant::create(ov::element::i64, {}, {int64_t(2)});
// b_i = B[..., i, :] [B2, B1, 1, k]
auto b_i = std::make_shared<ov::op::v8::Gather>(body_B_p, body_iter, c_axis2);
// A_row_i = A[..., i, :] [B2, B1, 1, n]
auto A_row_i = std::make_shared<ov::op::v8::Gather>(body_A, body_iter, c_axis2);
// sum_i = A_row_i @ X [B2, B1, 1, k]
// (lower-tri zeros + unfilled-X zeros make this equal to the partial sum)
auto sum_i = std::make_shared<ov::op::v0::MatMul>(A_row_i, body_X, false, false);
// diag_i = A[..., i, i] [B2, B1, 1, 1]
auto diag_i = std::make_shared<ov::op::v8::Gather>(A_row_i, body_iter, c_axis3);
// x_i = (b_i - sum_i) / diag_i [B2, B1, 1, k]
auto x_i = std::make_shared<ov::op::v1::Divide>(
std::make_shared<ov::op::v1::Subtract>(b_i, sum_i), diag_i);
// X_updated: scatter x_i into body_X at row i along axis 2
auto X_updated = std::make_shared<ov::op::v3::ScatterUpdate>(body_X, body_iter, x_i, c_axis2_scalar);
auto body_cond = ov::op::v0::Constant::create(ov::element::boolean, ov::Shape{1}, {true});
auto body = std::make_shared<ov::Model>(
ov::OutputVector{body_cond, X_updated},
ov::ParameterVector{body_iter, body_X, body_A, body_B_p});
// --- Assemble Loop ---
auto trip_count = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, std::vector<int64_t>{n});
auto exec_cond = ov::op::v0::Constant::create(ov::element::boolean, ov::Shape{1}, {true});
auto loop = std::make_shared<ov::op::v5::Loop>(trip_count, exec_cond);
loop->set_function(body);
// iter_counter_body_param_idx=0 (body_iter), exec_condition_body_result_idx=0 (body_cond)
loop->set_special_body_ports(ov::op::v5::Loop::SpecialBodyPorts{0, 0});
// Carried state: X feeds back from X_updated each iteration
loop->set_merged_input(body_X, X_init, X_updated);
// Invariant inputs passed through unchanged
loop->set_invariant_input(body_A, A);
loop->set_invariant_input(body_B_p, B);
// Final output: value of X_updated after the last iteration
auto X_final = loop->get_iter_value(X_updated, -1);
return rename_outputs_with_suffix({X_final}, context.get_name());
}
} // namespace op
} // namespace ggml
} // namespace frontend
} // namespace ov

Some files were not shown because too many files have changed in this diff Show More