#include "common.h"
#include "log.h"
#include "ggml-backend.h"
#include "ggml.h"
#include "gguf.h"
#include "ggml-cpp.h"
#include "llama.h"
#include "llama-cpp.h"
// TODO: replace with #include "llama-ext.h" in the future
#include "../src/llama-arch.h"
#include "../src/llama-model-saver.h"
#include <cinttypes>
#include <cstddef>
#include <cstdio>
#include <cstring>
#include <cstdint>
#include <random>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>
// normalized mean squared error = mse(a, b) / mse(a, 0)
static double nmse(const std::vector<float> & a, const std::vector<float> & b) {
GGML_ASSERT(a.size() == b.size());
double mse_a_b = 0.0;
double mse_a_0 = 0.0;
for (size_t i = 0; i < a.size(); i++) {
float a_i = a[i];
float b_i = b[i];
mse_a_b += (a_i - b_i) * (a_i - b_i);
mse_a_0 += a_i * a_i;
}
return mse_a_b / mse_a_0;
}
static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) {
size_t seed = *(const size_t *) userdata;
std::hash<std::string> hasher;
seed ^= hasher(tensor->name);
std::mt19937 gen(seed);
std::normal_distribution<float> dis(0.0f, 1.0e-2f);
const int64_t ne = ggml_nelements(tensor);
if (tensor->type == GGML_TYPE_F32) {
std::vector<float> tmp(ne);
for (int64_t i = 0; i < ne; i++) {
tmp[i] = dis(gen);
}
ggml_backend_tensor_set(tensor, tmp.data(), 0, ggml_nbytes(tensor));
} else if (tensor->type == GGML_TYPE_F16) {
std::vector<ggml_fp16_t> tmp(ne);
for (int64_t i = 0; i < ne; i++) {
tmp[i] = ggml_fp32_to_fp16(dis(gen));
}
ggml_backend_tensor_set(tensor, tmp.data(), 0, ggml_nbytes(tensor));
} else {
GGML_ABORT("fatal error");
}
}
static void usage(char ** argv) {
printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-v/--verbose]\n", argv[0]);
}
static std::vector<llama_token> get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed){
std::mt19937 gen(seed);
std::uniform_int_distribution<> dis(0, n_vocab - 1);
std::vector<llama_token> ret;
ret.reserve(n_tokens);
for (uint32_t i = 0; i < n_tokens; i++) {
ret.push_back(dis(gen));
}
return ret;
}
static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
gguf_context_ptr ret(gguf_init_empty());
llama_model_saver ms(arch, ret.get());
const uint32_t n_ctx = 128;
uint32_t n_vocab = 128;
uint32_t n_embd = 256;
uint32_t n_head = 2;
uint32_t n_ff = 384;
uint32_t n_layer = 2;
if (arch == LLM_ARCH_LLAMA4) {
n_layer = 4; // hparams.n_no_rope_layer_step is hard-coded to 4
} else if (arch == LLM_ARCH_GEMMA4) {
n_embd = 128;
n_head = 2;
n_ff = 192;
n_layer = 5; // need at least 5 for swa_pattern (every 5th is full_attention)
} else if (arch == LLM_ARCH_GEMMA3N) {
n_embd = 64;
n_head = 1;
n_ff = 96;
n_layer = 22; // hparams.n_layer_kv_from_start = 20 is hardcoded
} else if (arch == LLM_ARCH_DEEPSEEK4) {
// head size 64 so that GPU flash attention kernels support the model
n_embd = 512;
n_head = 8;
n_ff = 1024;
n_layer = 4;
} else if (arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_LAGUNA) {
n_embd = 160; // exercise per-head tensor split granularity with head size 80
} else if (arch == LLM_ARCH_QWEN3 || arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE) {
n_head = 4;
} else if (arch == LLM_ARCH_DEEPSEEK2
|| arch == LLM_ARCH_DEEPSEEK32
|| arch == LLM_ARCH_GLM_DSA
|| arch == LLM_ARCH_DOTS3NOTE
|| arch == LLM_ARCH_KIMI_LINEAR
|| arch == LLM_ARCH_BAILINGMOE3
|| arch == LLM_ARCH_KIMI_K3
|| arch == LLM_ARCH_MISTRAL4) {
n_embd = 128;
n_head = 1;
n_ff = 192;
} else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) {
n_layer = 3;
} else if (arch == LLM_ARCH_CHAMELEON) {
n_vocab = 10240;
} else if (arch == LLM_ARCH_QWEN3TTS) {
n_vocab = 4096; // must be >= the hard-coded codec head size (3072)
}
uint32_t n_head_kv = n_head;
if (arch == LLM_ARCH_QWEN3) {
n_head_kv = 1; // MQA coverage
} else if (arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE) {
n_head_kv = 2; // GQA coverage
}
const uint32_t n_embd_head = n_embd / n_head;
ms.add_kv(LLM_KV_GENERAL_ARCHITECTURE, llm_arch_name(arch));
ms.add_kv(LLM_KV_VOCAB_SIZE, n_vocab);
ms.add_kv(LLM_KV_CONTEXT_LENGTH, n_ctx);
ms.add_kv(LLM_KV_EMBEDDING_LENGTH, n_embd);
ms.add_kv(LLM_KV_FEATURES_LENGTH, n_embd);
ms.add_kv(LLM_KV_BLOCK_COUNT, n_layer);
ms.add_kv(LLM_KV_LEADING_DENSE_BLOCK_COUNT, uint32_t(1));
if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) {
std::vector<uint32_t> n_ff_per_layer;
n_ff_per_layer.reserve(n_layer);
for (uint32_t il = 0; il < n_layer; il++) {
n_ff_per_layer.push_back(il <= 1 ? 0 : n_ff);
}
ms.add_kv(LLM_KV_FEED_FORWARD_LENGTH, n_ff_per_layer);
} else {
ms.add_kv(LLM_KV_FEED_FORWARD_LENGTH, n_ff);
}
ms.add_kv(LLM_KV_USE_PARALLEL_RESIDUAL, false);
ms.add_kv(LLM_KV_LOGIT_SCALE, 1.0f);
ms.add_kv(LLM_KV_TIME_MIX_EXTRA_DIM, uint32_t(64));
ms.add_kv(LLM_KV_TIME_DECAY_EXTRA_DIM, uint32_t(128));
ms.add_kv(LLM_KV_FULL_ATTENTION_INTERVAL, uint32_t(2));
if (arch == LLM_ARCH_PLAMO2 || arch == LLM_ARCH_JAMBA || arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE ||
arch == LLM_ARCH_GRANITE_HYBRID || arch == LLM_ARCH_LFM2 || arch == LLM_ARCH_LFM2MOE || arch == LLM_ARCH_KIMI_LINEAR ||
arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3) {
GGML_ASSERT(n_layer >= 2);
std::vector<uint32_t> n_head_per_layer;
n_head_per_layer.reserve(n_layer);
for (uint32_t il = 0; il < n_layer; il++) {
n_head_per_layer.push_back(il == 1 ? 0 : n_head);
}
ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head_per_layer);
ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_per_layer);
} else {
ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head);
ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(1) : n_head_kv);
}
ms.add_kv(LLM_KV_ATTENTION_MAX_ALIBI_BIAS, 8.0f);
if (arch == LLM_ARCH_DEEPSEEK4) {
ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, n_embd_head);
ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, n_embd_head);
ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, n_embd_head/2);
} else if (arch == LLM_ARCH_DEEPSEEK2
|| arch == LLM_ARCH_DEEPSEEK32
|| arch == LLM_ARCH_GLM_DSA
|| arch == LLM_ARCH_DOTS3NOTE
|| arch == LLM_ARCH_KIMI_LINEAR
|| arch == LLM_ARCH_BAILINGMOE3
|| arch == LLM_ARCH_KIMI_K3
|| arch == LLM_ARCH_MISTRAL4) {
ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, uint32_t(576));
ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, uint32_t(512));
ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64));
ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA, uint32_t(192));
ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, uint32_t(128));
if (arch == LLM_ARCH_DOTS3NOTE) {
// SWA layers reuse the same MLA geometry as the full layers in this fixture
ms.add_kv(LLM_KV_ATTENTION_KV_LORA_RANK_SWA, uint32_t(512));
ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_SWA, uint32_t(576));
ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_SWA, uint32_t(512));
ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA_SWA, uint32_t(192));
ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA_SWA, uint32_t(128));
ms.add_kv(LLM_KV_ROPE_FREQ_BASE_SWA, 10000.0f);
// indexer on the full-attention layers (inverse of the swa pattern)
std::vector<uint32_t> indexer_types;
indexer_types.reserve(n_layer);
for (uint32_t il = 0; il < n_layer; il++) {
indexer_types.push_back(il % 2 ? 0 : 1);
}
ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, indexer_types);
}
} else if (arch == LLM_ARCH_MINIMAX_M3) {
// partial rotary: n_rot must not exceed the indexer key length (64)
ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64));
}
ms.add_kv(LLM_KV_ATTENTION_CLAMP_KQV, 1.0f);
ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_EPS, 1e-5f);
ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, 1e-5f);
ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_EPS, 1e-5f);
ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_GROUPS, uint32_t(8));
ms.add_kv(LLM_KV_ATTENTION_Q_LORA_RANK, arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(64) : uint32_t(512));
ms.add_kv(LLM_KV_ATTENTION_KV_LORA_RANK, uint32_t(512));
ms.add_kv(LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT, uint32_t(8));
ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW, n_ctx/8);
if (arch == LLM_ARCH_GEMMA4) {
ms.add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, n_embd/2);
ms.add_kv(LLM_KV_ATTENTION_SHARED_KV_LAYERS, uint32_t(0));
ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_SWA, n_embd_head);
ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_SWA, n_embd_head);
ms.add_kv(LLM_KV_ROPE_FREQ_BASE_SWA, 10000.0f);
// SWA pattern: every 5th layer is full attention (matches E2B layer_types)
ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(5));
} else if (arch == LLM_ARCH_COHERE2MOE || arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_STEP35 ||
arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_GRANITE_SWA || arch == LLM_ARCH_DOTS3NOTE) {
std::vector<uint32_t> pattern;
pattern.reserve(n_layer);
for (uint32_t il = 0; il < n_layer; il++) {
pattern.push_back(il % 2);
}
ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, pattern);
} else {
ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(2));
}
// MSA requires one indexer head per GQA (KV) head, unlike the DSA archs where the
// indexer head count is independent of the main attention head count.
ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, arch == LLM_ARCH_MINIMAX_M3 || arch == LLM_ARCH_DEEPSEEK4 ? n_head : uint32_t(1));
ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, uint32_t(64));
ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(8));
ms.add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, uint32_t(4));
ms.add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, uint32_t(1));
ms.add_kv(LLM_KV_ROPE_DIMENSION_SECTIONS, std::vector<uint32_t>({n_embd_head/4, n_embd_head/4, n_embd_head/4, n_embd_head/4}));
if (arch == LLM_ARCH_DEEPSEEK4) {
ms.add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, uint32_t(8));
ms.add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, uint32_t(32));
ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector<uint32_t>({0, 0, 4, 128}));
ms.add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, 160000.0f);
ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4));
ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, uint32_t(2));
ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, 1.0e-6f);
ms.add_kv(LLM_KV_HASH_LAYER_COUNT, uint32_t(0));
ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, 10.0f);
ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f);
ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true);
}
ms.add_kv(LLM_KV_TOKENIZER_MODEL, "no_vocab");
// ms.add_kv(LLM_KV_DENSE_2_FEAT_OUT, n_embd);
// ms.add_kv(LLM_KV_DENSE_3_FEAT_IN, n_embd);
if (moe) {
ms.add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, n_ff);
ms.add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, n_ff / 2); // distinct from n_ff so a saver key-clobber surfaces on reload
ms.add_kv(LLM_KV_EXPERT_LATENT_LENGTH, n_ff);
ms.add_kv(LLM_KV_INTERLEAVE_MOE_LAYER_STEP, uint32_t(2));
ms.add_kv(LLM_KV_EXPERT_COUNT, uint32_t(2));
ms.add_kv(LLM_KV_EXPERT_USED_COUNT, uint32_t(1));
ms.add_kv(LLM_KV_EXPERT_SHARED_COUNT, uint32_t(1));
ms.add_kv(LLM_KV_EXPERT_GATING_FUNC, arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(4) : uint32_t(2)); // sqrtsoftplus : sigmoid
ms.add_kv(LLM_KV_EXPERT_GROUP_SCALE, 1.0f);
ms.add_kv(LLM_KV_EXPERTS_PER_GROUP, uint32_t(1));
}
ms.add_kv(LLM_KV_POSNET_EMBEDDING_LENGTH, n_embd);
ms.add_kv(LLM_KV_POSNET_BLOCK_COUNT, n_layer);
ms.add_kv(LLM_KV_CONVNEXT_EMBEDDING_LENGTH, n_embd);
ms.add_kv(LLM_KV_CONVNEXT_BLOCK_COUNT, n_layer);
ms.add_kv(LLM_KV_XIELU_ALPHA_N, 1.0f);
ms.add_kv(LLM_KV_XIELU_ALPHA_P, 1.0f);
ms.add_kv(LLM_KV_XIELU_BETA, 1.0f);
ms.add_kv(LLM_KV_XIELU_EPS, 1.0e-7f);
ms.add_kv(LLM_KV_SSM_INNER_SIZE, arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE ? 256 : 2*n_embd);
ms.add_kv(LLM_KV_SSM_CONV_KERNEL, uint32_t(4));
ms.add_kv(LLM_KV_SSM_STATE_SIZE, uint32_t(128));
ms.add_kv(LLM_KV_SSM_TIME_STEP_RANK, n_head);
ms.add_kv(LLM_KV_SSM_GROUP_COUNT, arch == LLM_ARCH_PLAMO2 ? 0 : uint32_t(2));
ms.add_kv(LLM_KV_KDA_HEAD_DIM, uint32_t(128));
ms.add_kv(LLM_KV_KDA_SAFE_GATE, true);
ms.add_kv(LLM_KV_KDA_GATE_LOWER_BOUND, -5.0f);
if (arch == LLM_ARCH_BAILINGMOE3) {
ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, std::vector<float>({0.0f, 4.0f}));
ms.add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, std::vector<float>({0.0f, 5.0f}));
}
ms.add_kv(LLM_KV_WKV_HEAD_SIZE, n_embd/n_head);
ms.add_kv(LLM_KV_SHORTCONV_L_CACHE, uint32_t(3));
ms.add_kv(LLM_KV_RESIDUAL_SCALE, 3.5565588200778455f);
ms.add_kv(LLM_KV_ATTN_RES_BLOCK_SIZE, uint32_t(12));
ms.add_kv(LLM_KV_ACTIVATION_SITU_BETA, 4.0f);
ms.add_kv(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, 25.0f);
ms.add_kv(LLM_KV_KDA_GATE_LOWER_BOUND, -5.0f);
for (uint32_t il = 0; il < n_layer; il++) {
ggml_tensor t;
memset(&t, 0, sizeof(ggml_tensor));
t.type = GGML_TYPE_F16;
ggml_format_name(&t, "conv%" PRIu32 "d.weight", il);
gguf_add_tensor(ms.gguf_ctx, &t);
ggml_format_name(&t, "posnet.%" PRIu32 ".conv1.weight", il);
gguf_add_tensor(ms.gguf_ctx, &t);
ggml_format_name(&t, "posnet.%" PRIu32 ".conv2.weight", il);
gguf_add_tensor(ms.gguf_ctx, &t);
ggml_format_name(&t, "convnext.%" PRIu32 ".dw.weight", il);
gguf_add_tensor(ms.gguf_ctx, &t);
}
return ret;
}
static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) {
return true;
}
static std::pair<llama_model_ptr, llama_context_ptr> get_model_and_ctx(
struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector<ggml_backend_dev_t> & devs,
const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) {
GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr));
llama_model_params model_params = llama_model_default_params();
model_params.progress_callback = silent_model_load_progress;
std::vector<ggml_backend_dev_t> devs_copy = devs;
devs_copy.push_back(nullptr);
model_params.devices = devs_copy.data();
model_params.split_mode = split_mode;
llama_context_params ctx_params = llama_context_default_params();
ctx_params.n_ctx = 0;
ctx_params.n_threads = 4;
ctx_params.n_threads_batch = 4;
if (!encode) {
ctx_params.n_ubatch = 64;
}
size_t tmp = seed;
llama_model_ptr model(gguf_ctx != nullptr ?
llama_model_init_from_user(gguf_ctx, set_tensor_data, &tmp, model_params) :
llama_model_load_from_file_ptr(file, model_params));
if (!model) {
throw std::runtime_error("failed to create llama model");
}
llama_context_ptr lctx(llama_init_from_model(model.get(), ctx_params));
if (!lctx) {
throw std::runtime_error("failed to create llama context");
}
return std::make_pair(std::move(model), std::move(lctx));
}
static std::vector<float> get_logits(
llama_model * model, llama_context * lctx, const std::vector<llama_token> & tokens, bool encode = false) {
const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
const uint32_t n_ctx = llama_n_ctx(lctx);
const uint32_t n_tokens = tokens.size();
llama_batch batch = llama_batch_init(n_ctx, 0, 1);
GGML_ASSERT(n_tokens <= n_ctx);
for (uint32_t pos = 0; pos < n_tokens; pos++) {
common_batch_add(batch, tokens[pos], pos, {0}, true);
}
batch.n_tokens = n_tokens;
if (encode) {
if (llama_encode(lctx, batch)) {
llama_batch_free(batch);
throw std::runtime_error("failed to encode batch");
}
}
if (llama_decode(lctx, batch)) {
llama_batch_free(batch);
throw std::runtime_error("failed to decode batch");
}
std::vector<float> ret;
ret.reserve(n_tokens*n_vocab);
for (uint32_t i = 0; i < n_tokens; i++) {
const float * logits_ith = llama_get_logits_ith(lctx, i);
for (uint32_t j = 0; j < n_vocab; j++) {
ret.push_back(logits_ith[j]);
}
}
llama_batch_free(batch);
return ret;
}
static bool moe_mandatory(const llm_arch arch) {
switch (arch) {
case LLM_ARCH_LLAMA4:
case LLM_ARCH_COHERE2MOE:
case LLM_ARCH_GROK:
case LLM_ARCH_QWEN2MOE:
case LLM_ARCH_QWEN3MOE:
case LLM_ARCH_QWEN3NEXT:
case LLM_ARCH_QWEN3VLMOE:
case LLM_ARCH_QWEN35MOE:
case LLM_ARCH_PHIMOE:
case LLM_ARCH_DBRX:
case LLM_ARCH_OLMOE:
case LLM_ARCH_ARCTIC:
case LLM_ARCH_DEEPSEEK:
case LLM_ARCH_DEEPSEEK2:
case LLM_ARCH_DEEPSEEK32:
case LLM_ARCH_DOTS3NOTE:
case LLM_ARCH_DEEPSEEK4:
case LLM_ARCH_GLM4_MOE:
case LLM_ARCH_GLM_DSA:
case LLM_ARCH_EXAONE_MOE:
case LLM_ARCH_BAILINGMOE:
case LLM_ARCH_BAILINGMOE2:
case LLM_ARCH_BAILINGMOE3:
case LLM_ARCH_DOTS1:
case LLM_ARCH_AFMOE:
case LLM_ARCH_ERNIE4_5:
case LLM_ARCH_ERNIE4_5_MOE:
case LLM_ARCH_HUNYUAN_MOE:
case LLM_ARCH_HY_V3:
case LLM_ARCH_OPENAI_MOE:
case LLM_ARCH_LFM2MOE:
case LLM_ARCH_SMALLTHINKER:
case LLM_ARCH_LLADA_MOE:
case LLM_ARCH_GROVEMOE:
case LLM_ARCH_MINIMAX_01:
case LLM_ARCH_MINIMAX_M2:
case LLM_ARCH_MINIMAX_M3:
case LLM_ARCH_RND1:
case LLM_ARCH_PADDLEOCR:
case LLM_ARCH_MIMO2:
case LLM_ARCH_KIMI_LINEAR:
case LLM_ARCH_KIMI_K3:
case LLM_ARCH_STEP35:
case LLM_ARCH_MISTRAL4:
case LLM_ARCH_MELLUM:
case LLM_ARCH_LAGUNA:
return true;
default:
return false;
}
}
static bool moe_implemented(const llm_arch arch) {
if (moe_mandatory(arch)) {
return true;
}
switch (arch) {
case LLM_ARCH_LLAMA:
case LLM_ARCH_REFACT:
case LLM_ARCH_MINICPM:
case LLM_ARCH_GRANITE:
case LLM_ARCH_GRANITE_MOE:
case LLM_ARCH_MISTRAL3:
case LLM_ARCH_LLAMA_EMBED:
return true;
default:
return false;
}
}
static bool arch_supported(const llm_arch arch) {
if (arch == LLM_ARCH_CLIP || arch == LLM_ARCH_GPTJ || arch == LLM_ARCH_UNKNOWN) {
return false; // These models don't have usable implementations.
}
if (arch == LLM_ARCH_CHAMELEON) {
return false; // Only half-implemented and to be removed in the future.
}
if (arch == LLM_ARCH_WAVTOKENIZER_DEC) {
return false; // FIXME CUDA backend crashes.
}
if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) {
return false; // FIXME @ngxson
}
if (arch == LLM_ARCH_GRANITE_SWITCH) {
return false; // FIXME adapter fixture
}
if (arch == LLM_ARCH_LLAMA_EMBED || arch == LLM_ARCH_GEMMA_EMBEDDING || arch == LLM_ARCH_T5ENCODER) {
return false; // FIXME Embedding (?) models produce inconsistent results.
}
if (arch == LLM_ARCH_RWKV6 || arch == LLM_ARCH_RWKV6QWEN2 || arch == LLM_ARCH_RWKV7 || arch == LLM_ARCH_ARWKV7) {
return false; // FIXME RWKV models hang indefinitely.
}
if (arch == LLM_ARCH_BERT || arch == LLM_ARCH_MODERN_BERT || arch == LLM_ARCH_NOMIC_BERT || arch == LLM_ARCH_NOMIC_BERT_MOE ||
arch == LLM_ARCH_NEO_BERT || arch == LLM_ARCH_JINA_BERT_V2 || arch == LLM_ARCH_JINA_BERT_V3 || arch == LLM_ARCH_EUROBERT) {
return false; // TODO vocab
}
if (arch == LLM_ARCH_PLM) {
return false; // TODO tensor shapes
}
if (arch == LLM_ARCH_DEEPSEEK2OCR) {
return false;
}
// FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI.
#ifdef GGML_USE_WEBGPU
if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_DOTS3NOTE) {
return false;
}
#endif // GGML_USE_WEBGPU
// FIXME: jamba produces incorrect output (~0.55 NMSE vs CPU) on the HIP
// backend on RDNA3.5 (gfx1151); the SSM kernels need investigation.
#ifdef GGML_USE_HIP
if (arch == LLM_ARCH_JAMBA) {
return false;
}
#endif // GGML_USE_HIP
return true;
}
static int save_models(const llm_arch target_arch, const size_t seed, const ggml_log_level log_level, const std::string & dir) {
struct user_data_t {
struct {
ggml_log_callback callback;
void * user_data;
} original_logger;
ggml_log_level min_level; // prints below this log level go to debug log
};
user_data_t ud;
llama_log_get(&ud.original_logger.callback, &ud.original_logger.user_data);
ud.min_level = log_level;
llama_log_set([](ggml_log_level level, const char * text, void * user_data) {
const user_data_t * ud = (const user_data_t *) user_data;
const ggml_log_level level_eff = level >= ud->min_level ? level : GGML_LOG_LEVEL_DEBUG;
ud->original_logger.callback(level_eff, text, ud->original_logger.user_data);
}, &ud);
for (const llm_arch & arch : llm_arch_all()) {
if (arch == LLM_ARCH_UNKNOWN) {
continue;
}
if (target_arch != LLM_ARCH_UNKNOWN && arch != target_arch) {
continue;
}
if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) {
continue; // FIXME: ISWA KV cache initialization needs more fixture params
}
if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) {
continue;
}
for (bool moe : {false, true}) {
if (moe && !moe_implemented(arch)) {
continue;
}
if (!moe && moe_mandatory(arch)) {
continue;
}
if (!llama_model_saver_supports_arch(arch) || !arch_supported(arch)) {
LOG_INF("%s: %s model (%s) is unsupported, skipping\n", __func__, llm_arch_name(arch), moe ? "MoE" : "dense");
continue;
}
gguf_context_ptr gguf_ctx = get_gguf_ctx(arch, moe);
auto model_and_ctx = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {});
const std::string path = dir + "/" + llm_arch_name(arch) + (moe ? "-moe.gguf" : "-dense.gguf");
LOG_INF("%s: Saving %s model (%s) to %s...\n", __func__, llm_arch_name(arch), moe ? "MoE" : "dense", path.c_str());
llama_model_save_to_file(model_and_ctx.first.get(), path.c_str());
}
}
llama_log_set(ud.original_logger.callback, ud.original_logger.user_data);
return 0;
}
static int test_backends(const llm_arch target_arch, const size_t seed, const ggml_log_level log_level) {
struct user_data_t {
struct {
ggml_log_callback callback;
void * user_data;
} original_logger;
ggml_log_level min_level; // prints below this log level go to debug log
};
user_data_t ud;
llama_log_get(&ud.original_logger.callback, &ud.original_logger.user_data);
ud.min_level = log_level;
llama_log_set([](ggml_log_level level, const char * text, void * user_data) {
const user_data_t * ud = (const user_data_t *) user_data;
const ggml_log_level level_eff = level >= ud->min_level ? level : GGML_LOG_LEVEL_DEBUG;
ud->original_logger.callback(level_eff, text, ud->original_logger.user_data);
}, &ud);
const std::vector<llama_token> tokens = get_tokens(128, 128, seed);
struct device_config {
std::vector<ggml_backend_dev_t> devs;
std::string label;
llama_split_mode split_mode;
device_config(std::vector<ggml_backend_dev_t> devs, std::string name, llama_split_mode split_mode)
: devs(std::move(devs)), label(std::move(name)), split_mode(split_mode) {}
};
std::vector<device_config> dev_configs;
size_t max_device_label_length = 4;
{
std::vector<ggml_backend_dev_t> devices_meta;
{
const size_t device_count = ggml_backend_dev_count();
for (size_t i = 0; i < device_count; i++) {
ggml_backend_dev_t dev = ggml_backend_dev_get(i);
dev_configs.emplace_back(std::vector<ggml_backend_dev_t>{dev}, ggml_backend_dev_description(dev), LLAMA_SPLIT_MODE_LAYER);
max_device_label_length = std::max(max_device_label_length, dev_configs.back().label.length());
// cpu-based devices cannot be used in tensor split mode
if (ggml_backend_dev_buffer_type(dev) != ggml_backend_cpu_buffer_type()) {
devices_meta.push_back(dev);
}
}
}
dev_configs.emplace_back(devices_meta, "Meta", LLAMA_SPLIT_MODE_TENSOR);
}
size_t max_arch_name_length = 0;
for (const llm_arch & arch : llm_arch_all()) {
max_arch_name_length = std::max(max_arch_name_length, strlen(llm_arch_name(arch)));
}
const std::string template_header = std::string("|%" + std::to_string(max_arch_name_length) + "s|%") + std::to_string(max_device_label_length) + "s|%6s|%15s|%9s|\n";
const std::string template_row_cfg = std::string("|%" + std::to_string(max_arch_name_length) + "s|%") + std::to_string(max_device_label_length) + "s|%6s|";
const std::string template_row_res = "%15s %10s|%20s|\n";
bool all_ok = true;
common_log_flush(common_log_main());
printf(template_header.c_str(), "Model arch.", "Device", "Config", "NMSE vs. CPU", "Roundtrip");
printf("|");
for (size_t i = 0; i < max_arch_name_length; i++) {
printf("-");
}
printf("|");
for (size_t i = 0; i < max_device_label_length; i++) {
printf("-");
}
printf("|------|---------------|---------|\n");
for (const llm_arch & arch : llm_arch_all()) {
if (arch == LLM_ARCH_UNKNOWN) {
continue;
}
if (target_arch != LLM_ARCH_UNKNOWN && arch != target_arch) {
continue;
}
if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) {
continue; // FIXME: ISWA KV cache initialization needs more fixture params
}
if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) {
continue;
}
const bool encode = arch == LLM_ARCH_T5 || arch == LLM_ARCH_DREAM || arch == LLM_ARCH_LLADA || arch == LLM_ARCH_LLADA_MOE || arch == LLM_ARCH_RND1;
for (bool moe : {false, true}) {
if (moe && !moe_implemented(arch)) {
continue;
}
if (!moe && moe_mandatory(arch)) {
continue;
}
const std::string config_name = moe ? "MoE" : "Dense";
gguf_context_ptr gguf_ctx = get_gguf_ctx(arch, moe);
if (arch == LLM_ARCH_BAILINGMOE3) {
GGML_ASSERT(gguf_remove_key(gguf_ctx.get(), "bailingmoe3.kda.safe_gate") >= 0);
}
std::pair<llama_model_ptr, llama_context_ptr> model_and_ctx_cpu;
std::vector<float> logits_cpu;
for (device_config & dc : dev_configs) {
// print test config first; should anything fail during model loading or inference, at least we know which test case caused it
printf(template_row_cfg.c_str(),
llm_arch_name(arch), dc.label.c_str(), config_name.c_str());
fflush(stdout);
std::pair<llama_model_ptr, llama_context_ptr> model_and_ctx_dev;
std::vector<float> logits_dev;
std::string status_nmse = "\033[1;33mSKIP\033[0m";
std::string status_roundtrip = "\033[1;33mSKIP\033[0m";
char nmse_str[12] = {0};
bool skip = !arch_supported(arch) || (dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && dc.devs.empty());
if (!skip) {
if (logits_cpu.empty()) {
model_and_ctx_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, encode);
logits_cpu = get_logits(model_and_ctx_cpu.first.get(), model_and_ctx_cpu.second.get(), tokens, encode);
}
if (dc.split_mode != LLAMA_SPLIT_MODE_TENSOR || llm_arch_supports_sm_tensor(arch)) {
model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, dc.devs, dc.split_mode, encode);
logits_dev = get_logits(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), tokens, encode);
const double nmse_val = nmse(logits_cpu, logits_dev);
snprintf(nmse_str, sizeof(nmse_str), "(%.2e)", nmse_val);
status_nmse = "\033[1;32mOK\033[0m";
if (nmse_val > 1e-4) {
all_ok = false;
status_nmse = "\033[1;31mFAIL\033[0m";
}
}
FILE * file = tmpfile(); // Can be null on Windows without administrator privileges.
// FIXME: when adding a tensor to a gguf_context a copy is made, this changes the pointer which the meta backend
// in turn uses to map the tensors to their simple equivalents - this is fundamentally incompatible
if (file != nullptr && llama_model_saver_supports_arch(arch) && dc.split_mode != LLAMA_SPLIT_MODE_TENSOR) {
GGML_ASSERT(model_and_ctx_dev.first && model_and_ctx_dev.second);
llama_model_saver ms = llama_model_saver(model_and_ctx_dev.first.get());
ms.add_kv_from_model();
ms.add_tensors_from_model();
ms.save(file);
rewind(file);
auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, dc.devs, dc.split_mode, encode);
const std::vector<float> logits_roundtrip = get_logits(
model_and_ctx_roundtrip.first.get(), model_and_ctx_roundtrip.second.get(), tokens, encode);
status_roundtrip = "\033[1;32mOK\033[0m";
GGML_ASSERT(logits_roundtrip.size() == logits_dev.size());
for (size_t i = 0; i < logits_roundtrip.size(); i++) {
if (logits_roundtrip[i] != logits_dev[i]) {
all_ok = false;
status_roundtrip = "\033[1;31mFAIL\033[0m";
break;
}
}
}
}
// log the results for this test case
printf(template_row_res.c_str(),
status_nmse.c_str(), nmse_str, status_roundtrip.c_str());
}
}
}
llama_log_set(ud.original_logger.callback, ud.original_logger.user_data);
return all_ok ? 0 : 1;
}
int main(int argc, char ** argv) {
// FIXME these tests are disabled in the CI for macOS-latest-cmake-arm64 because they are segfaulting
common_init();
std::random_device rd;
llm_arch arch = LLM_ARCH_UNKNOWN;
size_t seed = rd();
ggml_log_level log_level = GGML_LOG_LEVEL_ERROR;
std::string out;
for (int i = 1; i < argc; i++) {
if (strcmp(argv[i], "-a") == 0 || strcmp(argv[i], "--arch") == 0) {
if (i + 1 < argc) {
const std::string arch_name = argv[++i];
arch = llm_arch_from_string(arch_name);
if (arch == LLM_ARCH_UNKNOWN) {
LOG_ERR("%s: unkown LLM architecture: %s\n", __func__, arch_name.c_str());
return 1;
}
} else {
usage(argv);
return 1;
}
}
if (strcmp(argv[i], "-s") == 0 || strcmp(argv[i], "--seed") == 0) {
if (i + 1 < argc) {
seed = std::stoull(argv[++i]);
} else {
usage(argv);
return 1;
}
}
if (strcmp(argv[i], "-v") == 0 || strcmp(argv[i], "--verbose") == 0) {
log_level = GGML_LOG_LEVEL_INFO;
continue;
}
if (strcmp(argv[i], "-o") == 0 || strcmp(argv[i], "--out") == 0) {
if (i + 1 < argc) {
out = argv[++i];
} else {
usage(argv);
return 1;
}
}
}
printf("%s: using seed %zu\n", __func__, seed);
try {
if (!out.empty()) {
return save_models(arch, seed, log_level, out);
}
return test_backends(arch, seed, log_level);
} catch (const std::exception & err) {
fprintf(stderr, "encountered runtime error: %s\n", err.what());
return -1;
}
}