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)
{ 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" },
{ 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" },
{ LLM_ARCH_MELLUM, "mellum" },
{ LLM_ARCH_NANBEIGE, "nanbeige" },
{ LLM_ARCH_QWEN3TTS, "qwen3tts" },
+ { LLM_ARCH_POCKETTTS, "pockettts" },
{ LLM_ARCH_UNKNOWN, "(unknown)" },
};
{ 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" },
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;
LLM_ARCH_OLMO,
LLM_ARCH_OLMO2,
LLM_ARCH_OLMOE,
+ LLM_ARCH_MUSE_GLIMMER,
LLM_ARCH_OPENELM,
LLM_ARCH_ARCTIC,
LLM_ARCH_DEEPSEEK,
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,
LLM_ARCH_DFLASH,
LLM_ARCH_NANBEIGE,
LLM_ARCH_QWEN3TTS,
+ LLM_ARCH_POCKETTTS,
LLM_ARCH_UNKNOWN,
};
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,
#include "llama-mmap.h"
#include "llama-model.h"
#include "llama-ext.h"
+#include "llama-sampler.h"
#include "llama.h"
#include <cinttypes>
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;
}
}
}
- // 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;
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;
}
}
- 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",
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;
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);
}
}
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;
}
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;
}
}
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;
}
}
- // 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;
//
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 ||
(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;
}
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);
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();
{
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;
/*.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,
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;
}
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
#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"
}
}
}
- 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);
}
}
}
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);
+ for (const auto & entry : samplers) {
+ if (entry.second->iface->backend_reset) {
+ entry.second->iface->backend_reset(entry.second);
+ }
+ }
- struct llama_sampler_data data = {
- /*.logits =*/ logits_seq,
- /*.probs =*/ nullptr,
- /*.sampled =*/ nullptr,
- /*.candidates =*/ nullptr,
- };
+ static const std::vector<uint32_t> dummy_row = { 0 };
- assert(sampler->iface->backend_apply);
- sampler->iface->backend_apply(sampler, ctx0, gf, &data);
+ for (const auto & [seq_id, sampler] : samplers) {
+ const auto it = sampling_rows.find(seq_id);
- 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);
- }
+ // 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) {
- 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.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) {
- 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.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) {
- res->t_candidates[seq_id] = data.candidates;
- outs[1] = data.candidates;
- 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 llama_sampler_accept_ggml after all samplers have been applied.
+ // 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);
}
}
}
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;
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;
}
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;
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;
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);
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()) {
} 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:
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:
{
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;
}
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;
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);
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:
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:
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:
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:
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);
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);
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);
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,
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;
};
}
- if (mtp_on_hybrid_qwen) {
+ if (mtp_on_hybrid_qwen || mtp_on_hybrid_nemotron) {
filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
}
/*.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,
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:
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
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;
struct llama_layer_shortconv shortconv;
struct llama_layer_nextn nextn;
+
+ struct llama_layer_switch_lora switch_lora;
};
struct llama_device {
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
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;
}
/* .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) {
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;
bool support;
};
-// check if all ggml ops used by the sampler are supported by the backend
-static bool llama_sampler_backend_support(
- llama_sampler * smpl,
- ggml_backend_buffer_type_t buft) {
- auto * device = ggml_backend_buft_get_device(buft);
- if (!device) {
- // CPU backend always supported
- return true;
- }
+// .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 =*/ 128*ggml_tensor_overhead() + ggml_graph_overhead(),
- /*.mem_buffer =*/ NULL,
+ /*.mem_size =*/ max_nodes * ggml_tensor_overhead() + ggml_graph_overhead_custom(max_nodes, false),
+ /*.mem_buffer =*/ nullptr,
/*.no_alloc =*/ true,
};
throw std::runtime_error(format("failed to create ggml context"));
}
- ggml_context * ctx = ctx_ptr.get();
-
- const int64_t n = 1024*1024;
+ 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),
- /*.probs = */ nullptr,
- /*.sampled = */ nullptr,
- /*.candidates = */ ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n),
+ /*.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,
};
- ggml_cgraph * gf = ggml_new_graph(ctx);
-
- smpl->iface->backend_apply(smpl, ctx, gf, &data);
+ if (sampler->iface->backend_reset) {
+ sampler->iface->backend_reset(sampler);
+ }
+ sampler->iface->backend_apply(sampler, ctx, gf, &data);
- if (data.logits) {
- ggml_build_forward_expand(gf, data.logits);
+ for (auto * output : { data.logits, data.probs, data.sampled, data.candidates }) {
+ if (output) {
+ ggml_build_forward_expand(gf, output);
+ }
}
- if (data.probs) {
- ggml_build_forward_expand(gf, data.probs);
+ if (sampler->iface->backend_reset) {
+ sampler->iface->backend_reset(sampler);
}
- if (data.sampled) {
- ggml_build_forward_expand(gf, data.sampled);
+ 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;
}
- if (data.candidates) {
- ggml_build_forward_expand(gf, data.candidates);
+ 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,
+ ggml_backend_buffer_type_t buft) {
+ auto * device = ggml_backend_buft_get_device(buft);
+ if (!device) {
+ // CPU backend always supported
+ return true;
}
- for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
- struct ggml_tensor * op = ggml_graph_node(gf, i);
+ auto probe = llama_sampler_backend_probe_graph(smpl, 1024*1024, GGML_DEFAULT_GRAPH_SIZE, true);
+
+ 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",
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");
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;
}
}
}
+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,
/* .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);
// 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;
}
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);
/* .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() {
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) {
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;
}
// 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;
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) {
{
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;
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;
}
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));
// 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
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);
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,
/* .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) {
/* .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 {
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);
/* .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) {
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);
/* .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) {
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);
/* .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) {
/* .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) {
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);
/* .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) {
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);
/* .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) {
/* .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) {
// 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;
/* .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) {
/* .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) {
/* .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(
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,
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);
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,
/* .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(
/* .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) {
/* .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) {
/* .backend_accept = */ nullptr,
/* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
+ /* .backend_reset = */ nullptr,
+ /* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_adaptive_p(
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);
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;
/* .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(
/* .reset = */ nullptr,
/* .clone = */ llama_sampler_infill_clone,
/* .free = */ llama_sampler_infill_free,
- /* .backend_apply = */ nullptr,
+ /* .backend_init = */ nullptr,
/* .backend_accept = */ nullptr,
+ /* .backend_apply = */ nullptr,
/* .backend_set_input = */ nullptr,
- /* .backend_init = */ nullptr,
+ /* .backend_reset = */ nullptr,
+ /* .copy_state = */ nullptr,
};
struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * vocab) {
);
}
+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) {
// has .backend_init() been called?
bool is_init = false;
+ uint32_t n_nodes = 0;
+
struct info {
bool is_backend;
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,
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:
}
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);
}
return ggml_backend_reg_by_name("RPC") != nullptr;
}
+const char * llama_version(void) {
+ return LLAMA_VERSION;
+}
+
void llama_backend_init(void) {
ggml_time_init();
}
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;
};
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);
// 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`
// 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);
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,
//
// 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);
// [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)(
// 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 {
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);
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);
--- /dev/null
+#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");
+}
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);
// 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;
}
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)
}
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
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);
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);
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;
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;
#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;
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;
--- /dev/null
+#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;
+}
#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,
// 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);
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;
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());
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);
// => {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
// (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);
};
+// 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;
};
+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;
};
+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;
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;
};
};
+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;
--- /dev/null
+#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);
+}
#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);
+}
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);
}
}
-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;
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 {
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
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);
}
}
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;
// 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);
--- /dev/null
+#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);
+}
overview / pkg / ggml / sources / whisper.cpp / commitdiff