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

talk-llama : fix build (#0)

This commit is contained in:
Georgi Gerganov 2026-08-14 19:28:27 +03:00
parent 733f281987
commit 1fe009caed
29 changed files with 2010 additions and 414 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)

10
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" },
@ -145,6 +147,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_MELLUM, "mellum" },
{ LLM_ARCH_NANBEIGE, "nanbeige" },
{ LLM_ARCH_QWEN3TTS, "qwen3tts" },
{ LLM_ARCH_POCKETTTS, "pockettts" },
{ LLM_ARCH_UNKNOWN, "(unknown)" },
};
@ -220,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" },
@ -993,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;

8
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,
@ -150,6 +152,7 @@ enum llm_arch {
LLM_ARCH_DFLASH,
LLM_ARCH_NANBEIGE,
LLM_ARCH_QWEN3TTS,
LLM_ARCH_POCKETTTS,
LLM_ARCH_UNKNOWN,
};
@ -225,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

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);

28
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;
}

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);

41
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:
@ -114,6 +116,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
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:
@ -174,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:
@ -234,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:
@ -1114,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);
@ -1265,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);
@ -1912,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);
@ -2228,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,
@ -2238,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;
@ -2319,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(); };
}
@ -2442,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,
@ -2591,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:
@ -2610,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

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

469
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) {
@ -3395,6 +3632,8 @@ 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, 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) {
@ -3614,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(
@ -3715,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);
@ -3756,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;
@ -3783,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(
@ -4022,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) {
@ -4039,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) {

5
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,6 +35,9 @@ 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(
float dry_multiplier,
float dry_base,

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;

50
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);
@ -1499,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");
}

39
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
@ -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);

96
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;
@ -697,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;
@ -1028,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;
@ -1461,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;
};
@ -1596,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);
}