char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Charlotte Somsupport aliceai_t5_moe98a01afd9

char/aliceai-t5-moe
39.3 KiB928 linesraw
1#include "common.h"
2#include "log.h"
3#include "ggml-backend.h"
4#include "ggml.h"
5#include "gguf.h"
6#include "ggml-cpp.h"
7#include "llama.h"
8#include "llama-cpp.h"
9
10// TODO: replace with #include "llama-ext.h" in the future
11#include "../src/llama-arch.h"
12#include "../src/llama-model-saver.h"
13
14#include <cinttypes>
15#include <cstddef>
16#include <cstdio>
17#include <cstring>
18#include <cstdint>
19#include <random>
20#include <stdexcept>
21#include <string>
22#include <utility>
23#include <vector>
24
25// normalized mean squared error = mse(a, b) / mse(a, 0)
26static double nmse(const std::vector<float> & a, const std::vector<float> & b) {
27    GGML_ASSERT(a.size() == b.size());
28    double mse_a_b = 0.0;
29    double mse_a_0 = 0.0;
30
31    for (size_t i = 0; i < a.size(); i++) {
32        float a_i = a[i];
33        float b_i = b[i];
34
35        mse_a_b += (a_i - b_i) * (a_i - b_i);
36        mse_a_0 += a_i * a_i;
37    }
38
39    return mse_a_b / mse_a_0;
40}
41
42static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) {
43    size_t seed = *(const size_t *) userdata;
44    std::hash<std::string> hasher;
45    seed ^= hasher(tensor->name);
46    std::mt19937 gen(seed);
47    std::normal_distribution<float> dis(0.0f, 1.0e-2f);
48
49    const int64_t ne = ggml_nelements(tensor);
50    if (tensor->type == GGML_TYPE_F32) {
51        std::vector<float> tmp(ne);
52        for (int64_t i = 0; i < ne; i++) {
53            tmp[i] = dis(gen);
54        }
55        ggml_backend_tensor_set(tensor, tmp.data(), 0, ggml_nbytes(tensor));
56    } else if (tensor->type == GGML_TYPE_F16) {
57        std::vector<ggml_fp16_t> tmp(ne);
58        for (int64_t i = 0; i < ne; i++) {
59            tmp[i] = ggml_fp32_to_fp16(dis(gen));
60        }
61        ggml_backend_tensor_set(tensor, tmp.data(), 0, ggml_nbytes(tensor));
62    } else {
63        GGML_ABORT("fatal error");
64    }
65}
66
67static void usage(char ** argv) {
68    printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v N] [-h/--help]\n", argv[0]);
69}
70
71static std::vector<llama_token> get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed){
72    std::mt19937 gen(seed);
73    std::uniform_int_distribution<> dis(0, n_vocab - 1);
74    std::vector<llama_token> ret;
75    ret.reserve(n_tokens);
76    for (uint32_t i = 0; i < n_tokens; i++) {
77        ret.push_back(dis(gen));
78    }
79    return ret;
80}
81
82static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
83    gguf_context_ptr ret(gguf_init_empty());
84    llama_model_saver ms(arch, ret.get());
85    const uint32_t n_ctx = 256;
86
87    uint32_t n_vocab = 128;
88    uint32_t n_embd  = 256;
89    uint32_t n_head  = 2;
90    uint32_t n_ff    = 384;
91    uint32_t n_layer = 2;
92    if (arch == LLM_ARCH_LLAMA4) {
93        n_layer = 4; // hparams.n_no_rope_layer_step is hard-coded to 4
94    } else if (arch == LLM_ARCH_GEMMA4) {
95        n_embd = 128;
96        n_head = 2;
97        n_ff   = 192;
98        n_layer = 5; // need at least 5 for swa_pattern (every 5th is full_attention)
99    } else if (arch == LLM_ARCH_GEMMA3N) {
100        n_embd = 64;
101        n_head = 1;
102        n_ff   = 96;
103        n_layer = 22; // hparams.n_layer_kv_from_start = 20 is hardcoded
104    } else if (arch == LLM_ARCH_DEEPSEEK4) {
105        // head size 64 so that GPU flash attention kernels support the model
106        n_embd  = 512;
107        n_head  = 8;
108        n_ff    = 1024;
109        n_layer = 4;
110    } else if (arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_LAGUNA) {
111        n_embd = 160; // exercise per-head tensor split granularity with head size 80
112    } else if (arch == LLM_ARCH_QWEN3 || arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE) {
113        n_head = 4;
114    } else if (arch == LLM_ARCH_DEEPSEEK2
115            || arch == LLM_ARCH_DEEPSEEK32
116            || arch == LLM_ARCH_GLM_DSA
117            || arch == LLM_ARCH_DOTS3NOTE
118            || arch == LLM_ARCH_KIMI_LINEAR
119            || arch == LLM_ARCH_BAILINGMOE3
120            || arch == LLM_ARCH_KIMI_K3
121            || arch == LLM_ARCH_MISTRAL4
122            || arch == LLM_ARCH_HY_V4) {
123        n_embd = 128;
124        n_head = 1;
125        n_ff   = 192;
126    } else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) {
127        n_layer = 3;
128    } else if (arch == LLM_ARCH_CHAMELEON) {
129        n_vocab = 10240;
130    } else if (arch == LLM_ARCH_QWEN3TTS) {
131        //n_vocab = 4096; // must be >= the hard-coded codec head size (3072)
132        n_vocab = 3072; // TODO: should be 4096, but user code cannot get `n_vocab_out` yet [TAG_LLAMA_N_VOCAB_OUT]
133    }
134
135    uint32_t n_head_kv = n_head;
136    if (arch == LLM_ARCH_QWEN3) {
137        n_head_kv = 1; // MQA coverage
138    } else if (arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE) {
139        n_head_kv = 2; // GQA coverage
140    }
141    const uint32_t n_embd_head = n_embd / n_head;
142
143    ms.add_kv(LLM_KV_GENERAL_ARCHITECTURE,      llm_arch_name(arch));
144    ms.add_kv(LLM_KV_VOCAB_SIZE,                n_vocab);
145    ms.add_kv(LLM_KV_CONTEXT_LENGTH,            n_ctx);
146    ms.add_kv(LLM_KV_EMBEDDING_LENGTH,          n_embd);
147    ms.add_kv(LLM_KV_FEATURES_LENGTH,           n_embd);
148    ms.add_kv(LLM_KV_BLOCK_COUNT,               n_layer);
149    ms.add_kv(LLM_KV_LEADING_DENSE_BLOCK_COUNT, uint32_t(1));
150
151    if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) {
152        std::vector<uint32_t> n_ff_per_layer;
153        n_ff_per_layer.reserve(n_layer);
154        for (uint32_t il = 0; il < n_layer; il++) {
155            n_ff_per_layer.push_back(il <= 1 ? 0 : n_ff);
156        }
157        ms.add_kv(LLM_KV_FEED_FORWARD_LENGTH, n_ff_per_layer);
158    } else {
159        ms.add_kv(LLM_KV_FEED_FORWARD_LENGTH, n_ff);
160    }
161
162    ms.add_kv(LLM_KV_USE_PARALLEL_RESIDUAL,   false);
163    ms.add_kv(LLM_KV_LOGIT_SCALE,             1.0f);
164    ms.add_kv(LLM_KV_TIME_MIX_EXTRA_DIM,      uint32_t(64));
165    ms.add_kv(LLM_KV_TIME_DECAY_EXTRA_DIM,    uint32_t(128));
166    ms.add_kv(LLM_KV_FULL_ATTENTION_INTERVAL, uint32_t(2));
167
168    if (arch == LLM_ARCH_PLAMO2 || arch == LLM_ARCH_JAMBA || arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE ||
169            arch == LLM_ARCH_GRANITE_HYBRID || arch == LLM_ARCH_LFM2 || arch == LLM_ARCH_LFM2MOE || arch == LLM_ARCH_KIMI_LINEAR ||
170            arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3) {
171        GGML_ASSERT(n_layer >= 2);
172        std::vector<uint32_t> n_head_per_layer;
173        n_head_per_layer.reserve(n_layer);
174        for (uint32_t il = 0; il < n_layer; il++) {
175            n_head_per_layer.push_back(il == 1 ? 0 : n_head);
176        }
177        ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head_per_layer);
178        ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_per_layer);
179    } else {
180        ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head);
181        ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(1) : n_head_kv);
182    }
183
184    ms.add_kv(LLM_KV_ATTENTION_MAX_ALIBI_BIAS, 8.0f);
185    if (arch == LLM_ARCH_DEEPSEEK4) {
186        ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH,   n_embd_head);
187        ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, n_embd_head);
188        ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT,   n_embd_head/2);
189    } else if (arch == LLM_ARCH_DEEPSEEK2
190            || arch == LLM_ARCH_DEEPSEEK32
191            || arch == LLM_ARCH_GLM_DSA
192            || arch == LLM_ARCH_DOTS3NOTE
193            || arch == LLM_ARCH_KIMI_LINEAR
194            || arch == LLM_ARCH_BAILINGMOE3
195            || arch == LLM_ARCH_KIMI_K3
196            || arch == LLM_ARCH_MISTRAL4
197            || arch == LLM_ARCH_HY_V4) {
198        ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH,       uint32_t(576));
199        ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH,     uint32_t(512));
200        ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT,       uint32_t(64));
201        ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA,   uint32_t(192));
202        ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, uint32_t(128));
203        if (arch == LLM_ARCH_DOTS3NOTE) {
204            // SWA layers reuse the same MLA geometry as the full layers in this fixture
205            ms.add_kv(LLM_KV_ATTENTION_KV_LORA_RANK_SWA,     uint32_t(512));
206            ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_SWA,       uint32_t(576));
207            ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_SWA,     uint32_t(512));
208            ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA_SWA,   uint32_t(192));
209            ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA_SWA, uint32_t(128));
210            ms.add_kv(LLM_KV_ROPE_FREQ_BASE_SWA,             10000.0f);
211            // indexer on the full-attention layers (inverse of the swa pattern)
212            std::vector<uint32_t> indexer_types;
213            indexer_types.reserve(n_layer);
214            for (uint32_t il = 0; il < n_layer; il++) {
215                indexer_types.push_back(il % 2 ? 0 : 1);
216            }
217            ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, indexer_types);
218        }
219    } else if (arch == LLM_ARCH_MINIMAX_M3) {
220        // partial rotary: n_rot must not exceed the indexer key length (64)
221        ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT,       uint32_t(64));
222    }
223    ms.add_kv(LLM_KV_ATTENTION_CLAMP_KQV,              1.0f);
224    ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_EPS,          1e-5f);
225    ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS,      1e-5f);
226    ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_EPS,          1e-5f);
227    ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_GROUPS,       uint32_t(8));
228    ms.add_kv(LLM_KV_ATTENTION_Q_LORA_RANK,            arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(64) : uint32_t(512));
229    ms.add_kv(LLM_KV_ATTENTION_KV_LORA_RANK,           uint32_t(512));
230    ms.add_kv(LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT, uint32_t(8));
231    ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW,         n_ctx/8);
232
233    if (arch == LLM_ARCH_GEMMA4) {
234        ms.add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER,      n_embd/2);
235        ms.add_kv(LLM_KV_ATTENTION_SHARED_KV_LAYERS,      uint32_t(0));
236        ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_SWA,        n_embd_head);
237        ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_SWA,      n_embd_head);
238        ms.add_kv(LLM_KV_ROPE_FREQ_BASE_SWA,              10000.0f);
239        // SWA pattern: every 5th layer is full attention (matches E2B layer_types)
240        ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(5));
241    } else if (arch == LLM_ARCH_COHERE2MOE || arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_SPARK2_5 ||
242            arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_GRANITE_SWA || arch == LLM_ARCH_DOTS3NOTE ||
243            arch == LLM_ARCH_MAPLE) {
244        std::vector<uint32_t> pattern;
245        pattern.reserve(n_layer);
246        for (uint32_t il = 0; il < n_layer; il++) {
247            pattern.push_back(il % 2);
248        }
249        ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, pattern);
250    } else {
251        ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(2));
252    }
253
254    // MSA requires one indexer head per GQA (KV) head, unlike the DSA archs where the
255    // indexer head count is independent of the main attention head count.
256    if (arch == LLM_ARCH_QWEN4EXP) {
257        ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT,    uint32_t(4));
258        ms.add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK, uint32_t(8));
259        // without this the QSA layers fall back to dense and go uncovered
260        ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector<uint32_t>(n_layer, 4));
261
262        // has_cell_ext() needs ple_n_heads here: the indexer cache serializes no ext without it
263        const uint32_t ple_ngram_size      = 3;
264        const uint32_t ple_heads_per_ngram = 2;
265        const uint32_t ple_n_heads         = (ple_ngram_size - 1)*ple_heads_per_ngram;
266        GGML_ASSERT(n_embd % ple_n_heads == 0);
267        const uint32_t ple_head_dim = n_embd/ple_n_heads;
268
269        std::vector<uint64_t> ple_head_offsets(ple_n_heads);
270        std::vector<uint64_t> ple_head_vocab_sizes(ple_n_heads, n_vocab);
271        for (uint32_t h = 0; h < ple_n_heads; h++) {
272            ple_head_offsets[h] = uint64_t(h)*n_vocab;
273        }
274
275        // the PLE history lives in the recurrent cache, so it must sit on a linear attention layer
276        ms.add_kv(LLM_KV_PLE_LAYERS,                  std::vector<uint32_t>({ 0 }));
277        ms.add_kv(LLM_KV_PLE_NGRAM_SIZE,              ple_ngram_size);
278        ms.add_kv(LLM_KV_PLE_HEADS_PER_NGRAM,         ple_heads_per_ngram);
279        ms.add_kv(LLM_KV_PLE_CONV_KERNEL,             uint32_t(4));
280        ms.add_kv(LLM_KV_PLE_EOS_TOKEN_ID,            uint32_t(0));
281        ms.add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER,  ple_head_dim);
282        ms.add_kv(LLM_KV_PLE_LAYER_MULTIPLIERS,       std::vector<uint64_t>({ 1, 3, 5 }));
283        ms.add_kv(LLM_KV_PLE_HEAD_OFFSETS,            ple_head_offsets);
284        ms.add_kv(LLM_KV_PLE_HEAD_VOCAB_SIZES,        ple_head_vocab_sizes);
285    }
286
287    // minimax-m3 keeps one indexer head per GQA head; the rest use a fixed 64 to match the fused
288    ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT,   arch == LLM_ARCH_MINIMAX_M3 ? n_head : uint32_t(64));
289    // qwen4exp ropes indexer keys with the main rotary width, so its head can't be < n_rot
290    ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH,
291              arch == LLM_ARCH_QWEN4EXP ? n_embd_head : uint32_t(128));
292
293    ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K,        uint32_t(8));
294    ms.add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE,   uint32_t(4));
295    ms.add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, uint32_t(1));
296    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}));
297
298    if (arch == LLM_ARCH_HY_V4) {
299        ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT,     uint32_t(4));
300        ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON,   1.0e-6f);
301        ms.add_kv(LLM_KV_HYPER_CONNECTION_MAGNITUDE, 2.0f);
302        ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP,           10.0f);
303        ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE,       1.0f);
304        ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM,        true);
305        // layer 0 must own an indexer, the odd layers share it
306        std::vector<uint32_t> indexer_types;
307        indexer_types.reserve(n_layer);
308        for (uint32_t il = 0; il < n_layer; il++) {
309            indexer_types.push_back(il % 2 ? 0 : 1);
310        }
311        ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, indexer_types);
312    }
313
314    if (arch == LLM_ARCH_DEEPSEEK4) {
315        ms.add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,         uint32_t(8));
316        ms.add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK,           uint32_t(32));
317        ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS,            std::vector<uint32_t>({0, 0, 4, 128}));
318        ms.add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE,    160000.0f);
319        ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT,               uint32_t(4));
320        ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, uint32_t(2));
321        ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON,             1.0e-6f);
322        ms.add_kv(LLM_KV_HASH_LAYER_COUNT,                      uint32_t(0));
323        ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP,                      10.0f);
324        ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE,                  1.0f);
325        ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM,                   true);
326    }
327
328    if (arch == LLM_ARCH_MAPLE) {
329        ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, 7.0f);
330    }
331
332    if (arch == LLM_ARCH_ALICEAI_T5_MOE) {
333        ms.add_kv(LLM_KV_DECODER_BLOCK_COUNT, n_layer - 1);
334        ms.add_kv(LLM_KV_DECODER_START_TOKEN_ID, uint32_t(1));
335        ms.add_kv(LLM_KV_ENCODER_ATTENTION_HEAD_COUNT_KV, n_head);
336        ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, uint32_t(1));
337        ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f);
338        ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true);
339    }
340
341    ms.add_kv(LLM_KV_TOKENIZER_MODEL,         "no_vocab");
342    // ms.add_kv(LLM_KV_DENSE_2_FEAT_OUT,     n_embd);
343    // ms.add_kv(LLM_KV_DENSE_3_FEAT_IN,      n_embd);
344
345    if (moe) {
346        ms.add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, n_ff);
347        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
348        ms.add_kv(LLM_KV_EXPERT_LATENT_LENGTH,       n_ff);
349        ms.add_kv(LLM_KV_INTERLEAVE_MOE_LAYER_STEP,  uint32_t(2));
350        ms.add_kv(LLM_KV_EXPERT_COUNT,               uint32_t(2));
351        ms.add_kv(LLM_KV_EXPERT_USED_COUNT,          uint32_t(1));
352        ms.add_kv(LLM_KV_EXPERT_SHARED_COUNT,        uint32_t(1));
353        ms.add_kv(LLM_KV_EXPERT_GATING_FUNC,         arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(4) : uint32_t(2)); // sqrtsoftplus : sigmoid
354        ms.add_kv(LLM_KV_EXPERT_GROUP_SCALE,         1.0f);
355        ms.add_kv(LLM_KV_EXPERTS_PER_GROUP,          uint32_t(1));
356    }
357
358    ms.add_kv(LLM_KV_POSNET_EMBEDDING_LENGTH,   n_embd);
359    ms.add_kv(LLM_KV_POSNET_BLOCK_COUNT,        n_layer);
360    ms.add_kv(LLM_KV_CONVNEXT_EMBEDDING_LENGTH, n_embd);
361    ms.add_kv(LLM_KV_CONVNEXT_BLOCK_COUNT,      n_layer);
362    ms.add_kv(LLM_KV_XIELU_ALPHA_N,             1.0f);
363    ms.add_kv(LLM_KV_XIELU_ALPHA_P,             1.0f);
364    ms.add_kv(LLM_KV_XIELU_BETA,                1.0f);
365    ms.add_kv(LLM_KV_XIELU_EPS,                 1.0e-7f);
366    ms.add_kv(LLM_KV_SSM_INNER_SIZE,            arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_QWEN4EXP ? 256 : 2*n_embd);
367    ms.add_kv(LLM_KV_SSM_CONV_KERNEL,           uint32_t(4));
368    ms.add_kv(LLM_KV_SSM_STATE_SIZE,            uint32_t(128));
369    ms.add_kv(LLM_KV_SSM_TIME_STEP_RANK,        n_head);
370    ms.add_kv(LLM_KV_SSM_GROUP_COUNT,           arch == LLM_ARCH_PLAMO2 ? 0 : uint32_t(2));
371    ms.add_kv(LLM_KV_KDA_HEAD_DIM,              uint32_t(128));
372    ms.add_kv(LLM_KV_KDA_SAFE_GATE,              true);
373    ms.add_kv(LLM_KV_KDA_GATE_LOWER_BOUND,       -5.0f);
374    if (arch == LLM_ARCH_BAILINGMOE3) {
375        ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP,   std::vector<float>({0.0f, 4.0f}));
376        ms.add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, std::vector<float>({0.0f, 5.0f}));
377    }
378    ms.add_kv(LLM_KV_WKV_HEAD_SIZE,             n_embd/n_head);
379    ms.add_kv(LLM_KV_SHORTCONV_L_CACHE,         uint32_t(3));
380    ms.add_kv(LLM_KV_RESIDUAL_SCALE,            3.5565588200778455f);
381    ms.add_kv(LLM_KV_ATTN_RES_BLOCK_SIZE,       uint32_t(12));
382    ms.add_kv(LLM_KV_ACTIVATION_SITU_BETA,      4.0f);
383    ms.add_kv(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, 25.0f);
384    ms.add_kv(LLM_KV_KDA_GATE_LOWER_BOUND,      -5.0f);
385
386    for (uint32_t il = 0; il < n_layer; il++) {
387        ggml_tensor t;
388        memset(&t, 0, sizeof(ggml_tensor));
389        t.type = GGML_TYPE_F16;
390        ggml_format_name(&t, "conv%" PRIu32 "d.weight", il);
391        gguf_add_tensor(ms.gguf_ctx, &t);
392        ggml_format_name(&t, "posnet.%" PRIu32 ".conv1.weight", il);
393        gguf_add_tensor(ms.gguf_ctx, &t);
394        ggml_format_name(&t, "posnet.%" PRIu32 ".conv2.weight", il);
395        gguf_add_tensor(ms.gguf_ctx, &t);
396        ggml_format_name(&t, "convnext.%" PRIu32 ".dw.weight", il);
397        gguf_add_tensor(ms.gguf_ctx, &t);
398    }
399    return ret;
400}
401
402static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) {
403    return true;
404}
405
406static std::pair<llama_model_ptr, llama_context_ptr> get_model_and_ctx(
407        struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector<ggml_backend_dev_t> & devs,
408        const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) {
409    GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr));
410    llama_model_params model_params = llama_model_default_params();
411    model_params.progress_callback = silent_model_load_progress;
412    std::vector<ggml_backend_dev_t> devs_copy = devs;
413    devs_copy.push_back(nullptr);
414    model_params.devices = devs_copy.data();
415    model_params.split_mode = split_mode;
416
417    llama_context_params ctx_params = llama_context_default_params();
418    ctx_params.n_ctx = 0;
419    ctx_params.n_threads = 4;
420    ctx_params.n_threads_batch = 4;
421    if (!encode) {
422        ctx_params.n_ubatch = 64;
423    }
424
425    size_t tmp = seed;
426    llama_model_ptr model(gguf_ctx != nullptr ?
427        llama_model_init_from_user(gguf_ctx, set_tensor_data, &tmp, model_params) :
428        llama_model_load_from_file_ptr(file, model_params));
429    if (!model) {
430        throw std::runtime_error("failed to create llama model");
431    }
432    llama_context_ptr lctx(llama_init_from_model(model.get(), ctx_params));
433    if (!lctx) {
434        throw std::runtime_error("failed to create llama context");
435    }
436    return std::make_pair(std::move(model), std::move(lctx));
437}
438
439static std::vector<float> get_logits(
440        llama_model * model, llama_context * lctx, const std::vector<llama_token> & tokens, bool encode = false) {
441    const uint32_t n_vocab  = llama_vocab_n_tokens(llama_model_get_vocab(model));
442    const uint32_t n_ctx    = llama_n_ctx(lctx);
443    const uint32_t n_tokens = tokens.size();
444    llama_batch batch = llama_batch_init(n_ctx, 0, 1);
445    GGML_ASSERT(n_tokens <= n_ctx);
446    for (uint32_t pos = 0; pos < n_tokens; pos++) {
447        common_batch_add(batch, tokens[pos], pos, {0}, true);
448    }
449    batch.n_tokens = n_tokens;
450    if (encode) {
451        if (llama_encode(lctx, batch)) {
452            llama_batch_free(batch);
453            throw std::runtime_error("failed to encode batch");
454        }
455    }
456    if (llama_decode(lctx, batch)) {
457        llama_batch_free(batch);
458        throw std::runtime_error("failed to decode batch");
459    }
460
461    std::vector<float> ret;
462    ret.reserve(n_tokens*n_vocab);
463    for (uint32_t i = 0; i < n_tokens; i++) {
464        const float * logits_ith = llama_get_logits_ith(lctx, i);
465        for (uint32_t j = 0; j < n_vocab; j++) {
466            ret.push_back(logits_ith[j]);
467        }
468    }
469    llama_batch_free(batch);
470    return ret;
471}
472
473static void test_encoder_decoder_batches(llama_model * model, llama_context * ctx, size_t seed) {
474    if (!llama_model_has_encoder(model) || !llama_model_has_decoder(model)) {
475        return;
476    }
477
478    const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
479    auto decoder = get_tokens(8, n_vocab, seed);
480    for (uint32_t n_enc : {16u, 16u, 7u, 32u}) {
481        auto encoder = get_tokens(n_enc, n_vocab, ++seed);
482        llama_memory_clear(llama_get_memory(ctx), true);
483        GGML_ASSERT(llama_encode(ctx, llama_batch_get_one(encoder.data(), encoder.size())) == 0);
484        const auto expected = get_logits(model, ctx, decoder);
485
486        llama_memory_clear(llama_get_memory(ctx), true);
487        std::vector<float> actual;
488        int8_t outputs[] = {1, 1};
489        uint32_t pos = 0;
490        for (uint32_t count : {1u, 1u, 2u, 2u, 1u, 1u}) {
491            auto batch = llama_batch_get_one(decoder.data() + pos, count);
492            batch.logits = outputs;
493            GGML_ASSERT(llama_decode(ctx, batch) == 0);
494            for (uint32_t i = 0; i < count; ++i) {
495                const float * logits = llama_get_logits_ith(ctx, i);
496                actual.insert(actual.end(), logits, logits + n_vocab);
497            }
498            pos += count;
499        }
500        GGML_ASSERT(nmse(expected, actual) < 1e-4);
501    }
502}
503
504static bool moe_mandatory(const llm_arch arch) {
505    switch (arch) {
506        case LLM_ARCH_LLAMA4:
507        case LLM_ARCH_COHERE2MOE:
508        case LLM_ARCH_GROK:
509        case LLM_ARCH_QWEN2MOE:
510        case LLM_ARCH_QWEN3MOE:
511        case LLM_ARCH_QWEN3NEXT:
512        case LLM_ARCH_QWEN3VLMOE:
513        case LLM_ARCH_QWEN35MOE:
514        case LLM_ARCH_QWEN4EXP:
515        case LLM_ARCH_PHIMOE:
516        case LLM_ARCH_DBRX:
517        case LLM_ARCH_OLMOE:
518        case LLM_ARCH_ARCTIC:
519        case LLM_ARCH_DEEPSEEK:
520        case LLM_ARCH_DEEPSEEK2:
521        case LLM_ARCH_DEEPSEEK32:
522        case LLM_ARCH_DOTS3NOTE:
523        case LLM_ARCH_DEEPSEEK4:
524        case LLM_ARCH_GLM4_MOE:
525        case LLM_ARCH_GLM_DSA:
526        case LLM_ARCH_EXAONE_MOE:
527        case LLM_ARCH_BAILINGMOE:
528        case LLM_ARCH_BAILINGMOE2:
529        case LLM_ARCH_BAILINGMOE3:
530        case LLM_ARCH_DOTS1:
531        case LLM_ARCH_AFMOE:
532        case LLM_ARCH_ERNIE4_5:
533        case LLM_ARCH_ERNIE4_5_MOE:
534        case LLM_ARCH_HUNYUAN_MOE:
535        case LLM_ARCH_HY_V3:
536        case LLM_ARCH_HY_V4:
537        case LLM_ARCH_OPENAI_MOE:
538        case LLM_ARCH_LFM2MOE:
539        case LLM_ARCH_SMALLTHINKER:
540        case LLM_ARCH_LLADA_MOE:
541        case LLM_ARCH_GROVEMOE:
542        case LLM_ARCH_MINIMAX_01:
543        case LLM_ARCH_MINIMAX_M2:
544        case LLM_ARCH_MINIMAX_M3:
545        case LLM_ARCH_RND1:
546        case LLM_ARCH_PADDLEOCR:
547        case LLM_ARCH_MIMO2:
548        case LLM_ARCH_KIMI_LINEAR:
549        case LLM_ARCH_KIMI_K3:
550        case LLM_ARCH_STEP35:
551        case LLM_ARCH_MISTRAL4:
552        case LLM_ARCH_MELLUM:
553        case LLM_ARCH_LAGUNA:
554        case LLM_ARCH_MAPLE:
555        case LLM_ARCH_ALICEAI_T5_MOE:
556            return true;
557        default:
558            return false;
559    }
560}
561
562static bool moe_implemented(const llm_arch arch) {
563    if (moe_mandatory(arch)) {
564        return true;
565    }
566    switch (arch) {
567        case LLM_ARCH_LLAMA:
568        case LLM_ARCH_REFACT:
569        case LLM_ARCH_MINICPM:
570        case LLM_ARCH_GRANITE:
571        case LLM_ARCH_GRANITE_MOE:
572        case LLM_ARCH_MISTRAL3:
573        case LLM_ARCH_LLAMA_EMBED:
574            return true;
575        default:
576            return false;
577    }
578}
579
580static bool arch_supported(const llm_arch arch) {
581    if (arch == LLM_ARCH_CLIP || arch == LLM_ARCH_GPTJ || arch == LLM_ARCH_UNKNOWN) {
582        return false; // These models don't have usable implementations.
583    }
584    if (arch == LLM_ARCH_CHAMELEON) {
585        return false; // Only half-implemented and to be removed in the future.
586    }
587    if (arch == LLM_ARCH_WAVTOKENIZER_DEC) {
588        return false; // FIXME CUDA backend crashes.
589    }
590    if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) {
591        return false; // FIXME @ngxson
592    }
593    if (arch == LLM_ARCH_GRANITE_SWITCH) {
594        return false; // FIXME adapter fixture
595    }
596    if (arch == LLM_ARCH_LLAMA_EMBED || arch == LLM_ARCH_GEMMA_EMBEDDING || arch == LLM_ARCH_T5ENCODER) {
597        return false; // FIXME Embedding (?) models produce inconsistent results.
598    }
599    if (arch == LLM_ARCH_RWKV6 || arch == LLM_ARCH_RWKV6QWEN2 || arch == LLM_ARCH_RWKV7 || arch == LLM_ARCH_ARWKV7) {
600        return false; // FIXME RWKV models hang indefinitely.
601    }
602    if (arch == LLM_ARCH_BERT || arch == LLM_ARCH_MODERN_BERT || arch == LLM_ARCH_NOMIC_BERT || arch == LLM_ARCH_NOMIC_BERT_MOE ||
603            arch == LLM_ARCH_NEO_BERT || arch == LLM_ARCH_JINA_BERT_V2 || arch == LLM_ARCH_JINA_BERT_V3 || arch == LLM_ARCH_EUROBERT) {
604        return false; // TODO vocab
605    }
606    if (arch == LLM_ARCH_PLM) {
607        return false; // TODO tensor shapes
608    }
609    if (arch == LLM_ARCH_DEEPSEEK2OCR) {
610        return false;
611    }
612    // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI.
613#ifdef GGML_USE_WEBGPU
614    if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_DOTS3NOTE || arch == LLM_ARCH_QWEN4EXP ||
615            arch == LLM_ARCH_HY_V4) {
616        return false;
617    }
618#endif // GGML_USE_WEBGPU
619
620    // FIXME: jamba produces incorrect output (~0.55 NMSE vs CPU) on the HIP
621    // backend on RDNA3.5 (gfx1151); the SSM kernels need investigation.
622#ifdef GGML_USE_HIP
623    if (arch == LLM_ARCH_JAMBA) {
624        return false;
625    }
626#endif // GGML_USE_HIP
627
628    return true;
629}
630
631static int save_models(const llm_arch target_arch, const size_t seed, const int verbosity, const std::string & dir) {
632    struct user_data_t {
633        struct {
634            ggml_log_callback callback;
635            void * user_data;
636        } log_old;
637
638        int verbosity;
639
640        user_data_t(int verbosity) : verbosity(verbosity) {
641            llama_log_get(&log_old.callback, &log_old.user_data);
642        }
643    };
644    user_data_t ud(verbosity);
645
646    llama_log_set([](ggml_log_level level, const char * text, void * user_data) {
647        const user_data_t * ud = (const user_data_t *) user_data;
648        int verbosity = common_log_get_verbosity(level);
649        if (verbosity <= ud->verbosity) {
650            ud->log_old.callback(level, text, ud->log_old.user_data);
651        }
652    }, &ud);
653
654    for (const llm_arch & arch : llm_arch_all()) {
655        if (arch == LLM_ARCH_UNKNOWN) {
656            continue;
657        }
658        if (target_arch != LLM_ARCH_UNKNOWN && arch != target_arch) {
659            continue;
660        }
661        if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) {
662            continue; // FIXME: ISWA KV cache initialization needs more fixture params
663        }
664        if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) {
665            continue;
666        }
667        for (bool moe : {false, true}) {
668            if (moe && !moe_implemented(arch)) {
669                continue;
670            }
671            if (!moe && moe_mandatory(arch)) {
672                continue;
673            }
674            if (!llama_model_saver_supports_arch(arch) || !arch_supported(arch)) {
675                LOG_INF("%s: %s model (%s) is unsupported, skipping\n", __func__, llm_arch_name(arch), moe ? "MoE" : "dense");
676                continue;
677            }
678            gguf_context_ptr gguf_ctx = get_gguf_ctx(arch, moe);
679            auto model_and_ctx = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {});
680            const std::string path = dir + "/" + llm_arch_name(arch) + (moe ? "-moe.gguf" : "-dense.gguf");
681            LOG_INF("%s: Saving %s model (%s) to %s...\n", __func__, llm_arch_name(arch), moe ? "MoE" : "dense", path.c_str());
682            llama_model_save_to_file(model_and_ctx.first.get(), path.c_str());
683        }
684    }
685    llama_log_set(ud.log_old.callback, ud.log_old.user_data);
686    return 0;
687}
688
689static int test_backends(const llm_arch target_arch, const size_t seed, const int verbosity) {
690    struct user_data_t {
691        struct {
692            ggml_log_callback callback;
693            void * user_data;
694        } log_old;
695
696        int verbosity;
697
698        user_data_t(int verbosity) : verbosity(verbosity) {
699            llama_log_get(&log_old.callback, &log_old.user_data);
700        }
701    };
702    user_data_t ud(verbosity);
703
704    llama_log_set([](ggml_log_level level, const char * text, void * user_data) {
705        const user_data_t * ud = (const user_data_t *) user_data;
706        int verbosity = common_log_get_verbosity(level);
707        if (verbosity <= ud->verbosity) {
708            ud->log_old.callback(level, text, ud->log_old.user_data);
709        }
710    }, &ud);
711
712    const std::vector<llama_token> tokens = get_tokens(128, 128, seed);
713
714    struct device_config {
715        std::vector<ggml_backend_dev_t> devs;
716        std::string                     label;
717        llama_split_mode                split_mode;
718
719        device_config(std::vector<ggml_backend_dev_t> devs, std::string name, llama_split_mode split_mode)
720            : devs(std::move(devs)), label(std::move(name)), split_mode(split_mode) {}
721    };
722
723    std::vector<device_config> dev_configs;
724    size_t max_device_label_length = 4;
725    {
726        std::vector<ggml_backend_dev_t> devices_meta;
727        {
728            const size_t device_count = ggml_backend_dev_count();
729            for (size_t i = 0; i < device_count; i++) {
730                ggml_backend_dev_t dev = ggml_backend_dev_get(i);
731                dev_configs.emplace_back(std::vector<ggml_backend_dev_t>{dev}, ggml_backend_dev_description(dev), LLAMA_SPLIT_MODE_LAYER);
732                max_device_label_length = std::max(max_device_label_length, dev_configs.back().label.length());
733
734                // cpu-based devices cannot be used in tensor split mode
735                if (ggml_backend_dev_buffer_type(dev) != ggml_backend_cpu_buffer_type()) {
736                    devices_meta.push_back(dev);
737                }
738            }
739        }
740
741        dev_configs.emplace_back(devices_meta, "Meta", LLAMA_SPLIT_MODE_TENSOR);
742    }
743
744    size_t max_arch_name_length = 0;
745    for (const llm_arch & arch : llm_arch_all()) {
746        max_arch_name_length = std::max(max_arch_name_length, strlen(llm_arch_name(arch)));
747    }
748
749    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";
750    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|";
751    const std::string template_row_res = "%15s %10s|%20s|\n";
752
753    bool all_ok = true;
754    common_log_flush(common_log_main());
755    printf(template_header.c_str(), "Model arch.", "Device", "Config", "NMSE vs. CPU", "Roundtrip");
756    printf("|");
757    for (size_t i = 0; i < max_arch_name_length; i++) {
758        printf("-");
759    }
760    printf("|");
761    for (size_t i = 0; i < max_device_label_length; i++) {
762        printf("-");
763    }
764    printf("|------|---------------|---------|\n");
765    for (const llm_arch & arch : llm_arch_all()) {
766        if (arch == LLM_ARCH_UNKNOWN) {
767            continue;
768        }
769        if (target_arch != LLM_ARCH_UNKNOWN && arch != target_arch) {
770            continue;
771        }
772        if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) {
773            continue; // FIXME: ISWA KV cache initialization needs more fixture params
774        }
775        if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) {
776            continue;
777        }
778
779        const bool encode = arch == LLM_ARCH_T5 || arch == LLM_ARCH_ALICEAI_T5_MOE || arch == LLM_ARCH_DREAM || arch == LLM_ARCH_LLADA || arch == LLM_ARCH_LLADA_MOE || arch == LLM_ARCH_RND1;
780        for (bool moe : {false, true}) {
781            if (moe && !moe_implemented(arch)) {
782                continue;
783            }
784            if (!moe && moe_mandatory(arch)) {
785                continue;
786            }
787            const std::string config_name = moe ? "MoE" : "Dense";
788            gguf_context_ptr gguf_ctx = get_gguf_ctx(arch, moe);
789            if (arch == LLM_ARCH_BAILINGMOE3) {
790                GGML_ASSERT(gguf_remove_key(gguf_ctx.get(), "bailingmoe3.kda.safe_gate") >= 0);
791            }
792            std::pair<llama_model_ptr, llama_context_ptr> model_and_ctx_cpu;
793            std::vector<float> logits_cpu;
794            for (device_config & dc : dev_configs) {
795                // print test config first; should anything fail during model loading or inference, at least we know which test case caused it
796                printf(template_row_cfg.c_str(),
797                    llm_arch_name(arch), dc.label.c_str(), config_name.c_str());
798                fflush(stdout);
799
800                std::pair<llama_model_ptr, llama_context_ptr> model_and_ctx_dev;
801                std::vector<float> logits_dev;
802                std::string status_nmse      = "\033[1;33mSKIP\033[0m";
803                std::string status_roundtrip = "\033[1;33mSKIP\033[0m";
804                char nmse_str[12] = {0};
805
806                bool skip = !arch_supported(arch) || (dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && dc.devs.empty());
807                if (!skip) {
808                    if (logits_cpu.empty()) {
809                        model_and_ctx_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, encode);
810                        logits_cpu = get_logits(model_and_ctx_cpu.first.get(), model_and_ctx_cpu.second.get(), tokens, encode);
811                    }
812                    if (dc.split_mode != LLAMA_SPLIT_MODE_TENSOR || llm_arch_supports_sm_tensor(arch)) {
813                        model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, dc.devs, dc.split_mode, encode);
814                        logits_dev = get_logits(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), tokens, encode);
815                        test_encoder_decoder_batches(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), seed);
816                        const double nmse_val = nmse(logits_cpu, logits_dev);
817                        snprintf(nmse_str, sizeof(nmse_str), "(%.2e)", nmse_val);
818                        status_nmse = "\033[1;32mOK\033[0m";
819                        if (nmse_val > 1e-4) {
820                            all_ok = false;
821                            status_nmse = "\033[1;31mFAIL\033[0m";
822                        }
823                    }
824
825                    FILE * file = tmpfile(); // Can be null on Windows without administrator privileges.
826                    // FIXME: when adding a tensor to a gguf_context a copy is made, this changes the pointer which the meta backend
827                    //     in turn uses to map the tensors to their simple equivalents - this is fundamentally incompatible
828                    if (file != nullptr && llama_model_saver_supports_arch(arch) && dc.split_mode != LLAMA_SPLIT_MODE_TENSOR) {
829                        GGML_ASSERT(model_and_ctx_dev.first && model_and_ctx_dev.second);
830                        llama_model_saver ms = llama_model_saver(model_and_ctx_dev.first.get());
831                        ms.add_kv_from_model();
832                        ms.add_tensors_from_model();
833                        ms.save(file);
834                        rewind(file);
835
836                        auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, dc.devs, dc.split_mode, encode);
837                        const std::vector<float> logits_roundtrip = get_logits(
838                            model_and_ctx_roundtrip.first.get(), model_and_ctx_roundtrip.second.get(), tokens, encode);
839                        status_roundtrip = "\033[1;32mOK\033[0m";
840                        GGML_ASSERT(logits_roundtrip.size() == logits_dev.size());
841                        for (size_t i = 0; i < logits_roundtrip.size(); i++) {
842                            if (logits_roundtrip[i] != logits_dev[i]) {
843                                all_ok = false;
844                                status_roundtrip = "\033[1;31mFAIL\033[0m";
845                                break;
846                            }
847                        }
848                    }
849                }
850
851                // log the results for this test case
852                printf(template_row_res.c_str(),
853                    status_nmse.c_str(), nmse_str, status_roundtrip.c_str());
854            }
855        }
856    }
857    llama_log_set(ud.log_old.callback, ud.log_old.user_data);
858    return all_ok ? 0 : 1;
859}
860
861int main(int argc, char ** argv) {
862    // init the logger at max verbosity. filter with a custom callback respecting the user-configure verbosity
863    common_log_set_verbosity_thold(LOG_LEVEL_DEBUG);
864    common_init();
865
866    std::random_device rd;
867
868    llm_arch arch = LLM_ARCH_UNKNOWN;
869    size_t seed = rd();
870    std::string out;
871
872    int verbosity = LOG_LEVEL_ERROR;
873
874    for (int i = 1; i < argc; i++) {
875        if (strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--help") == 0) {
876            usage(argv);
877            return 0;
878        }
879        if (strcmp(argv[i], "-a") == 0 || strcmp(argv[i], "--arch") == 0) {
880            if (i + 1 < argc) {
881                const std::string arch_name = argv[++i];
882                arch = llm_arch_from_string(arch_name);
883                if (arch == LLM_ARCH_UNKNOWN) {
884                    LOG_ERR("%s: unkown LLM architecture: %s\n", __func__, arch_name.c_str());
885                    return 1;
886                }
887            } else {
888                usage(argv);
889                return 1;
890            }
891        }
892        if (strcmp(argv[i], "-s") == 0 || strcmp(argv[i], "--seed") == 0) {
893            if (i + 1 < argc) {
894                seed = std::stoull(argv[++i]);
895            } else {
896                usage(argv);
897                return 1;
898            }
899        }
900        if (strcmp(argv[i], "-v") == 0) {
901            if (i + 1 < argc) {
902                verbosity = std::stoull(argv[++i]);
903            } else {
904                usage(argv);
905                return 1;
906            }
907        }
908        if (strcmp(argv[i], "-o") == 0 || strcmp(argv[i], "--out") == 0) {
909            if (i + 1 < argc) {
910                out = argv[++i];
911            } else {
912                usage(argv);
913                return 1;
914            }
915        }
916    }
917    printf("%s: using seed %zu\n", __func__, seed);
918
919    try {
920        if (!out.empty()) {
921            return save_models(arch, seed, verbosity, out);
922        }
923        return test_backends(arch, seed, verbosity);
924    } catch (const std::exception & err) {
925        fprintf(stderr, "encountered runtime error: %s\n", err.what());
926        return -1;
927    }
928}