Buckets:
| // tests/test_engine.c - the GPU engine (Metal or CUDA, whichever is linked) against the CPU reference forward | |
| // (nslm/mova_ref.h), through engine_api.h only. | |
| // | |
| // A small MoVA model folder with random weights in every encoding the engine reads (SEED4P4 routed experts, Q4 value | |
| // experts, Q8 projections / embedding / head, BF16 norms and routers) is written to out/test/engine_model. Then: | |
| // prompt scoring (eng_score: the GEMM and prefill-attention kernels) against the reference logits, | |
| // token-by-token decode (eng_prefill + eng_step: the matvec and split-key attention kernels) against them, | |
| // the engine's router choices (eng_mova_routes) against the reference's, | |
| // prompt-lookup decode (ENG_MODE_PL) committing exactly the tokens of plain greedy decode, | |
| // batched decode of several sequence slots (eng_step_batch) against decoding each alone. | |
| // Logits agree to the engines' f32 accumulation (relative error well under a BF16 step of the logit scale) wherever | |
| // the router choices agree; a choice can flip on a near tie, which the test counts and bounds. | |
| static int fails = 0; | |
| static const char* CONFIG = | |
| "{\n \"architectures\": [\"K2HorizonForCausalLM\"],\n \"attention_gate_func\": \"softplus\",\n" | |
| " \"decoder_sparse_step\": 1,\n \"head_dim\": 128,\n \"hidden_act\": \"silu\",\n \"hidden_size\": 512,\n" | |
| " \"intermediate_size\": 1024,\n \"layernorm_num_groups\": 2,\n \"mlp_only_layers\": [0],\n" | |
| " \"model_type\": \"k2_horizon\",\n \"moe_gate_bias\": true,\n \"moe_intermediate_size\": 256,\n" | |
| " \"mova_num_experts\": 8,\n \"mova_num_experts_per_tok\": 2,\n \"norm_topk_prob\": true,\n" | |
| " \"num_attention_heads\": 8,\n \"num_experts\": 16,\n \"num_experts_per_tok\": 4,\n \"num_hidden_layers\": 5,\n" | |
| " \"num_key_value_heads\": 2,\n \"num_shared_experts\": 1,\n \"query_key_norm\": false,\n \"rms_norm_eps\": 1e-06,\n" | |
| " \"rope_head_dim\": 128,\n \"rope_parameters\": {\"rope_theta\": 10000000.0, \"rope_type\": \"default\"},\n" | |
| " \"router_scaling_factor\": 2.5,\n \"router_score_func\": \"sigmoid\",\n \"tie_word_embeddings\": false,\n" | |
| " \"use_sliding_window\": false,\n \"vocab_size\": 2048\n}\n"; | |
| static uint32_t rng_u32(uint64_t* s) { | |
| *s = *s * 6364136223846793005ull + 1442695040888963407ull; | |
| return (uint32_t) (*s >> 33); | |
| } | |
| static double rng_unit(uint64_t* s) { return (rng_u32(s) & 0xFFFFFF) / 16777216.0; } // [0, 1) | |
| static uint16_t f2bf(float f) { return nslm_f2bf(f); } | |
| typedef struct { | |
| const NsSpec* t; | |
| const MovaTensor* mt; | |
| } Gen; | |
| // Deterministic contents per (tensor, stream): weights of moderate size so activations stay O(1) through the layers. | |
| static int fill(void* ctx, int ti, int s, uint8_t* dst, uint64_t len) { | |
| const Gen* g = (const Gen*) ctx; | |
| const NsSpec* t = &g->t[ti]; | |
| const MovaTensor* mt = &g->mt[ti]; | |
| uint64_t st = 0x9E3779B97F4A7C15ull * (uint64_t) (ti + 1) + 0x632BE59BD9B4E019ull * (uint64_t) (s + 1); | |
| if (t->enc == NS_BF16) { | |
| uint16_t* o = (uint16_t*) dst; | |
| for (uint64_t i = 0; i < len / 2; ++i) { | |
| const double u = rng_unit(&st) * 2 - 1; | |
| const float v = mt->kind == MOVA_K_NORM ? (float) (1.0 + 0.1 * u) | |
| : mt->kind == MOVA_K_ROUTER_BIAS ? (float) (0.02 * u) | |
| : mt->kind == MOVA_K_HEADS ? (float) (0.08 * u) | |
| : strstr(mt->name, ".mla.") ? (float) (0.04 * u) // MLA projections: BF16, Q8's magnitude | |
| : (float) (1.0 * u); // routers: well-separated scores | |
| o[i] = f2bf(v); | |
| } | |
| return 0; | |
| } | |
| if (t->enc == NS_Q8 || t->enc == NS_Q4) { | |
| if (s == 0) { for (uint64_t i = 0; i < len; ++i) dst[i] = (uint8_t) rng_u32(&st); return 0; } | |
| const float sc = t->enc == NS_Q8 ? 0.0003f : 0.005f, mid = t->enc == NS_Q8 ? 128.0f : 8.0f; | |
| uint16_t* o = (uint16_t*) dst; | |
| uint64_t s1 = 0x1234567ull * (uint64_t) (ti + 1); // scales and biases from one stream, so they pair up | |
| for (uint64_t i = 0; i < len / 2; ++i) { | |
| const float sv = sc * (float) (0.75 + 0.5 * rng_unit(&s1)); | |
| o[i] = f2bf(s == 1 ? sv : -mid * sv * (float) (0.9 + 0.2 * rng_unit(&st))); | |
| } | |
| return 0; | |
| } | |
| // SEED4P4: seeds in [1, 65535], any coefficients, exponents in [bias, bias + 2] | |
| if (s == 0) { uint16_t* o = (uint16_t*) dst; for (uint64_t i = 0; i < len / 2; ++i) o[i] = (uint16_t) (1 + rng_u32(&st) % 65535); } | |
| else if (s == 1) { uint16_t* o = (uint16_t*) dst; for (uint64_t i = 0; i < len / 2; ++i) o[i] = (uint16_t) rng_u32(&st); } | |
| else if (s == 2) { int32_t* o = (int32_t*) dst; for (uint64_t i = 0; i < len / 4; ++i) o[i] = -10; } | |
| else for (uint64_t i = 0; i < len; ++i) dst[i] = (uint8_t) ((rng_u32(&st) % 3) | ((rng_u32(&st) % 3) << 4)); | |
| return 0; | |
| } | |
| // skip: leave out the tensors whose names contain it (NULL: none) | |
| static int write_model(const char* dir, const char* config, const char* skip, char* err, int errlen) { | |
| mkdir("out", 0755); | |
| mkdir("out/test", 0755); | |
| mkdir(dir, 0755); | |
| char path[1024]; | |
| snprintf(path, sizeof path, "%s/config.json", dir); | |
| FILE* f = fopen(path, "wb"); | |
| if (!f) { snprintf(err, (size_t) errlen, "cannot write %.400s", path); return -1; } | |
| fputs(config, f); | |
| fclose(f); | |
| MovaCfg c; | |
| if (mova_cfg_parse(&c, config, err, errlen)) return -1; | |
| MovaTensor* mt = NULL; | |
| const int n = mova_tensors(&c, &mt); | |
| NsSpec* sp = (NsSpec*) calloc((size_t) n, sizeof(NsSpec)); | |
| int m = 0; | |
| for (int i = 0; i < n; ++i) { | |
| if (skip && strstr(mt[i].name, skip)) continue; | |
| mt[m] = mt[i]; | |
| const int k = mt[m].kind; | |
| sp[m].name = mt[m].name; | |
| sp[m].enc = k == MOVA_K_EXPERTS ? NS_SEED4P4 : k == MOVA_K_VEXPERTS ? NS_Q4 | |
| : (k == MOVA_K_NORM || k == MOVA_K_ROUTER || k == MOVA_K_ROUTER_BIAS || k == MOVA_K_HEADS) ? NS_BF16 | |
| : strstr(mt[m].name, ".mla.") ? NS_BF16 : NS_Q8; // MLA tensors: BF16, as exported (no *.weight name) | |
| sp[m].slices = mt[m].slices; | |
| sp[m].rows = mt[m].rows; | |
| sp[m].cols = mt[m].cols; | |
| ++m; | |
| } | |
| Gen g = {sp, mt}; | |
| const int rc = ns_write(dir, sp, m, 64ull << 20, NULL, fill, &g, err, errlen); | |
| free(sp); | |
| free(mt); | |
| return rc; | |
| } | |
| // Relative logit error of one row: max |a - b| over the row, divided by max |b|. | |
| static double row_err(const float* a, const float* b, int V) { | |
| double m = 0, s = 0; | |
| for (int i = 0; i < V; ++i) { m = fmax(m, fabs((double) a[i] - b[i])); s = fmax(s, fabs((double) b[i])); } | |
| return s > 0 ? m / s : m; | |
| } | |
| // The router choices of one row (every sparse layer), as sets: selection order among near-equal scores does not matter. | |
| static int same_choice(const int32_t* a, const int32_t* b, int n, int k) { | |
| for (int i = 0; i < n; ++i) { | |
| int x[16], y[16]; | |
| memcpy(x, a + i * k, sizeof(int32_t) * (size_t) k); | |
| memcpy(y, b + i * k, sizeof(int32_t) * (size_t) k); | |
| for (int p = 1; p < k; ++p) for (int q = p; q > 0 && x[q] < x[q - 1]; --q) { const int t = x[q]; x[q] = x[q - 1]; x[q - 1] = t; } | |
| for (int p = 1; p < k; ++p) for (int q = p; q > 0 && y[q] < y[q - 1]; --q) { const int t = y[q]; y[q] = y[q - 1]; y[q - 1] = t; } | |
| if (memcmp(x, y, sizeof(int) * (size_t) k)) return 0; | |
| } | |
| return 1; | |
| } | |
| static int argmax(const float* l, int V) { | |
| int b = 0; | |
| for (int i = 1; i < V; ++i) if (l[i] > l[b]) b = i; | |
| return b; | |
| } | |
| static void set_env(const char* k, const char* v) { // v = NULL: unset | |
| _putenv_s(k, v ? v : ""); | |
| if (v) setenv(k, v, 1); else unsetenv(k); | |
| } | |
| // A second engine opened with one setting changed (env k = v, and the KV format) runs the same scoring and decoding as | |
| // e; returns the number of differing logits and tokens and the largest relative logit difference. Twice: the second | |
| // round starts from caches the first one left. | |
| static int compare_engines(Eng* e, const EngOpts* o0, int kv_format, const char* k, const char* v, uint64_t* st, int V, | |
| double* rel, const char* what) { | |
| EngOpts o = *o0; | |
| o.kv_format = kv_format; | |
| char err[512] = ""; | |
| if (k) set_env(k, v); | |
| Eng* e2 = eng_open(&o, err, sizeof err); | |
| if (k) set_env(k, NULL); | |
| if (!e2) { printf("FAIL: eng_open (%s): %s\n", what, err); ++fails; return -1; } | |
| printf("%s: %s\n", what, eng_describe(e2)); | |
| const int NP = 300, cnt = 40, seq = 0; | |
| int32_t* pid = (int32_t*) malloc(sizeof(int32_t) * NP); | |
| for (int i = 0; i < NP; ++i) pid[i] = (int32_t) (rng_u32(st) % (uint32_t) V); | |
| float* la = (float*) malloc(sizeof(float) * (size_t) cnt * V); | |
| float* lb = (float*) malloc(sizeof(float) * (size_t) cnt * V); | |
| int32_t ga[32], gb[32]; | |
| int diff = 0; | |
| *rel = 0; | |
| for (int round = 0; round < 2; ++round) { | |
| CHECK(eng_score(e, 0, pid, NP - cnt, cnt, la) == 0 && eng_score(e2, 0, pid, NP - cnt, cnt, lb) == 0, "scoring"); | |
| for (size_t i = 0; i < (size_t) cnt * V; ++i) diff += la[i] != lb[i]; | |
| for (int t = 0; t < cnt; ++t) *rel += row_err(lb + (size_t) t * V, la + (size_t) t * V, V) / (2.0 * cnt); // mean over both rounds | |
| CHECK(eng_prefill(e, 0, pid, NP / 2) == 0 && eng_generate(e, &seq, 1, 32, ENG_MODE_AR, ga, NULL) == 0, "decode"); | |
| CHECK(eng_prefill(e2, 0, pid, NP / 2) == 0 && eng_generate(e2, &seq, 1, 32, ENG_MODE_AR, gb, NULL) == 0, "decode"); | |
| diff += memcmp(ga, gb, sizeof ga) != 0; | |
| } | |
| free(pid); free(la); free(lb); | |
| eng_close(e2); | |
| return diff; | |
| } | |
| // 9. batched decode (eng_step_batch): NB sequences of mixed lengths in their own slots (more than one forward's MV_MAXT | |
| // rows), filled by chunked prefill (eng_prefill_begin / eng_prefill_next, interleaved), then stepped together with | |
| // teacher-forced tokens. Each row against the same sequence run alone (same chunks, same steps) in a one-slot engine: | |
| // every row must agree bit for bit (a dense matvec adds each row's products in the same order for any row count), and | |
| // prompt caches must not depend on the prefill chunks either (a server's chunks start wherever a slot's cache reuse | |
| // ends). Route flips are still counted and bounded, and argmax checked, so a failure says how far apart the rows are. | |
| static void test_batch(const char* dir, const char* what, int expand_min, int kv_format) { | |
| enum { NB = 11, STEPS = 4, CAP = 256, CH = 50, NSP = 4, TK = 4, TKV = 2 }; | |
| EngOpts o; | |
| memset(&o, 0, sizeof o); | |
| o.model_dir = dir; | |
| o.resource_dir = getenv("NSLM_RES") ? getenv("NSLM_RES") : "out/res"; | |
| o.max_seqs = 1; | |
| o.kv_tokens = CAP; | |
| o.mla_expand_min = expand_min; | |
| o.kv_format = kv_format; | |
| char err[512] = ""; | |
| Eng* e1 = eng_open(&o, err, sizeof err); | |
| o.max_seqs = NB + 1; | |
| o.kv_tokens = (int64_t) CAP * (NB + 1); | |
| Eng* eb = e1 ? eng_open(&o, err, sizeof err) : NULL; | |
| if (!e1 || !eb) { ++fails; printf("FAIL: %s batch: eng_open: %s\n", what, err); eng_close(e1); return; } | |
| const int V = eng_vocab(eb); | |
| uint64_t st = 4242; | |
| int len[NB], seqs[NB]; | |
| int32_t ids[NB][CAP], tf[NB][STEPS]; | |
| for (int i = 0; i < NB; ++i) { | |
| len[i] = 2 + (i * 37) % 190; // 2 .. 189 tokens: chunks of CH rows and short tails (the decode kernels) | |
| seqs[i] = NB - i; // slots 11 .. 1, slot 0 unused: nonzero cache bases | |
| for (int k = 0; k < len[i]; ++k) ids[i][k] = (int32_t) (rng_u32(&st) % (uint32_t) V); | |
| for (int j = 0; j < STEPS; ++j) tf[i][j] = (int32_t) (rng_u32(&st) % (uint32_t) V); | |
| } | |
| // alone: each sequence in the one-slot engine | |
| float* ref = (float*) malloc(sizeof(float) * (size_t) NB * STEPS * V), *lb = (float*) malloc(sizeof(float) * (size_t) NB * V); | |
| int32_t rm[NB][STEPS][NSP * TK], rv[NB][STEPS][NSP * TKV], bm[NB][NSP * TK], bv[NB][NSP * TKV]; | |
| for (int i = 0; i < NB; ++i) { | |
| int reused = -1, left; | |
| CHECK(eng_prefill_begin(e1, 0, ids[i], len[i], &reused) == 0 && reused == 0, "%s batch: prefill_begin alone", what); | |
| while ((left = eng_prefill_next(e1, 0, CH)) > 0) {} | |
| CHECK(left == 0, "%s batch: prefill_next alone", what); | |
| for (int j = 0; j < STEPS; ++j) { | |
| eng_mova_routes(e1, 1, 1); | |
| CHECK(eng_step(e1, 0, ref + ((size_t) i * STEPS + j) * V) == 0, "%s batch: step alone", what); | |
| CHECK(eng_mova_routes_read(e1, 1, rm[i][j], rv[i][j], NULL, NULL) == 0, "%s batch: routes alone", what); | |
| eng_push(e1, 0, tf[i][j]); | |
| } | |
| } | |
| eng_mova_routes(e1, 0, 0); | |
| for (int k = 0; k < 2; ++k) { // sequence 5 (187 tokens) prefilled in chunks of 7, and in one call | |
| const int i = 5; | |
| int left = 0; | |
| if (k == 0) { | |
| CHECK(eng_prefill_begin(e1, 0, ids[i], len[i], NULL) == 0, "%s batch: prefill_begin", what); | |
| while ((left = eng_prefill_next(e1, 0, 7)) > 0) {} | |
| } else left = eng_prefill(e1, 0, ids[i], len[i]); | |
| CHECK(left == 0 && eng_step(e1, 0, lb) == 0, "%s batch: prefill in chunks", what); | |
| const int same = !memcmp(lb, ref + (size_t) i * STEPS * V, sizeof(float) * (size_t) V); | |
| const double re = row_err(lb, ref + (size_t) i * STEPS * V, V); | |
| printf("%s prefill %s: %s (logit error %.2e)\n", what, k ? "in one call" : "in chunks of 7", same ? "bit-equal to chunks of 50" : "differs", re); | |
| CHECK(same, "%s batch: prefill chunking changed the logits (error %.3e)", what, re); | |
| } | |
| // together: chunked prefills interleaved, then batched steps | |
| for (int i = 0; i < NB; ++i) CHECK(eng_prefill_begin(eb, seqs[i], ids[i], len[i], NULL) == 0, "%s batch: prefill_begin", what); | |
| CHECK(eng_step_batch(eb, seqs, NB, lb) != 0, "%s batch: a step over slots still prefilling must fail", what); | |
| for (int busy = 1; busy;) { | |
| busy = 0; | |
| for (int i = 0; i < NB; ++i) { | |
| const int left = eng_prefill_next(eb, seqs[i], CH); | |
| CHECK(left >= 0, "%s batch: prefill_next", what); | |
| busy |= left > 0; | |
| } | |
| } | |
| const int dup[3] = {seqs[0], seqs[1], seqs[0]}; | |
| CHECK(eng_step_batch(eb, dup, 3, lb) != 0, "%s batch: a sequence listed twice must fail", what); | |
| int rows = 0, exact = 0, flips = 0, am_same = 0; | |
| double worst = 0; | |
| for (int j = 0; j < STEPS; ++j) { | |
| eng_mova_routes(eb, 1, NB); | |
| CHECK(eng_step_batch(eb, seqs, NB, lb) == 0, "%s batch: eng_step_batch", what); | |
| CHECK(eng_mova_routes_read(eb, NB, &bm[0][0], &bv[0][0], NULL, NULL) == 0, "%s batch: routes", what); | |
| for (int i = 0; i < NB; ++i) { | |
| const float* r = ref + ((size_t) i * STEPS + j) * V, *b = lb + (size_t) i * V; | |
| ++rows; | |
| exact += !memcmp(r, b, sizeof(float) * (size_t) V); | |
| if (!same_choice(bm[i], rm[i][j], NSP, TK) || !same_choice(bv[i], rv[i][j], NSP, TKV)) { ++flips; continue; } | |
| worst = fmax(worst, row_err(b, r, V)); | |
| am_same += argmax(b, V) == argmax(r, V); | |
| } | |
| for (int i = 0; i < NB; ++i) eng_push(eb, seqs[i], tf[i][j]); | |
| } | |
| eng_mova_routes(eb, 0, 0); | |
| printf("%s batched decode: %d sequences x %d steps, %d of %d rows bit-equal to decoding alone, %d route flips, logit " | |
| "error max %.2e (other rows), argmax equal %d\n", what, NB, STEPS, exact, rows, flips, worst, am_same); | |
| CHECK(flips <= rows / 8, "%s batch: %d rows with different router choices", what, flips); | |
| CHECK(worst < 2e-2, "%s batch: logit error %.3e against decoding alone", what, worst); | |
| CHECK(am_same == rows - flips, "%s batch: argmax differs in %d rows", what, rows - flips - am_same); | |
| CHECK(exact == rows, "%s batch: %d of %d rows differ from decoding alone", what, rows - exact, rows); | |
| CHECK(eng_prefill(eb, NB + 1, ids[0], 4) != 0 && eng_prefill(eb, 1, ids[0], CAP + 1) != 0, "%s batch: slot bounds", what); | |
| { // a slot's cache copied out (eng_kv_read) and into another slot (eng_kv_write): that slot then reuses it and | |
| // decodes as the original sequence did alone | |
| const int i = 5, P = 128, src = seqs[i]; | |
| void* kv = malloc((size_t) eng_kv_bytes(eb) * P); | |
| int reused = -1, left; | |
| void* hi = (char*) kv + (size_t) eng_kv_bytes(eb) * (P / 2); // positions [P / 2, P) (the layout is per layer) | |
| CHECK(eng_kv_bytes(eb) > 0 && eng_kv_read(eb, src, 0, P / 2, kv) == 0 && eng_kv_read(eb, src, P / 2, P, hi) == 0, | |
| "%s batch: eng_kv_read", what); | |
| CHECK(eng_kv_read(eb, src, 0, CAP, kv) != 0, "%s batch: eng_kv_read past the cache", what); | |
| CHECK(eng_kv_write(eb, 0, ids[i], 0, P / 2, kv) == 0 && eng_len(eb, 0) == P / 2, "%s batch: eng_kv_write", what); | |
| CHECK(eng_kv_write(eb, 0, ids[i], P / 2 + 1, P, hi) != 0, "%s batch: eng_kv_write past the cache", what); | |
| CHECK(eng_kv_write(eb, 0, ids[i], P / 2, P, hi) == 0, "%s batch: eng_kv_write", what); | |
| CHECK(eng_prefill_begin(eb, 0, ids[i], len[i], &reused) == 0 && reused == P, "%s batch: reuse of a written cache (%d)", what, reused); | |
| while ((left = eng_prefill_next(eb, 0, CH)) > 0) {} | |
| CHECK(left == 0 && eng_step(eb, 0, lb) == 0, "%s batch: step after eng_kv_write", what); | |
| const double re = row_err(lb, ref + (size_t) i * STEPS * V, V); | |
| const int same = !memcmp(lb, ref + (size_t) i * STEPS * V, sizeof(float) * (size_t) V); | |
| printf("%s cache copied between slots: %s (logit error %.2e)\n", what, same ? "bit-equal" : "differs", re); | |
| CHECK(same, "%s batch: a copied cache decodes differently (error %.3e)", what, re); | |
| free(kv); | |
| } | |
| eng_close(e1); | |
| eng_close(eb); | |
| free(ref); | |
| free(lb); | |
| } | |
| // 8. MLA (a TransMLA conversion): the same random model with latent attention (ranks 64..320, one 128-dim RoPE key), | |
| // against the C reference: prompt scoring (one-pass latent attention) and decode (split-key + reduce), the cache size. | |
| // A hybrid rank list (0: the layer keeps GQA attention and a BF16 K / V cache) runs the same checks. | |
| static int64_t layer_kv_bytes(int rank, int fmt) { // bytes per position of one layer's cache (the test model: 2 KV heads) | |
| if (!rank) return 2 * 2 * 128 * 2; // GQA layer: BF16 K and V, whatever the MLA format | |
| if (fmt == ENG_KV_BF16) return (128 + rank) * 2; | |
| const int lead = fmt == ENG_KV_FP8 || rank < KVQ_FP4_LEAD ? rank : KVQ_FP4_LEAD; | |
| return 128 + 4 + kvq_row_bytes(lead, rank) + kvq_row_scales(lead, rank); | |
| } | |
| static void test_mla_ranks(const int ranks[5], const char* dir) { | |
| int hybrid = 0; | |
| for (int l = 0; l < 5; ++l) hybrid |= ranks[l] == 0; | |
| char* config = (char*) malloc(strlen(CONFIG) + 128); | |
| sprintf(config, "{\"mla_ranks\": [%d, %d, %d, %d, %d], \"mla_rope_dim\": 128,%s", ranks[0], ranks[1], ranks[2], ranks[3], | |
| ranks[4], CONFIG + 1); | |
| char err[512] = ""; | |
| if (write_model(dir, config, NULL, err, sizeof err)) { ++fails; printf("FAIL: MLA model folder: %s\n", err); free(config); return; } | |
| // A folder without one layer's per-head query map: the reference (and below, the engine) must refuse it, naming | |
| // the tensor | |
| const char* bad = "out/test/engine_model_mla_missing"; | |
| const int bad_ok = !hybrid && write_model(bad, config, "layers.2.self_attn.mla.q_lat", err, sizeof err) == 0; | |
| CHECK(bad_ok || hybrid, "MLA folder without q_lat: %s", err); | |
| if (bad_ok) { | |
| MovaRef* rb = mova_ref_open(bad, 1, err, sizeof err); | |
| CHECK(!rb && strstr(err, "layers.2.self_attn.mla.q_lat"), "MLA reference without q_lat: %s", rb ? "opened" : err); | |
| if (rb) mova_ref_close(rb); | |
| } | |
| free(config); | |
| // The per-head maps are stacked 3-D BF16 tensors [n_head][rows][cols], as the conversion exports them: the engine | |
| // maps them from the folder (as from the real converted model's model-mla-delta.safetensors), and the checkpoint | |
| // reader (the packer's, and the engine's fallback) reads them whole | |
| { | |
| MovaCkpt* ck = mova_ckpt_open(dir, err, sizeof err); | |
| NsModel nm; | |
| CHECK(ck && ns_open(&nm, dir, err, sizeof err) == 0, "MLA folder as a checkpoint: %s", err); | |
| if (ck) { | |
| int l1 = 0; // the first MLA layer | |
| while (!ranks[l1]) ++l1; | |
| char nm1[64]; | |
| snprintf(nm1, sizeof nm1, "model.layers.%d.self_attn.mla.q_lat", l1); | |
| const uint16_t* w3 = mova_ckpt_bf16_3d(ck, nm1, 8, ranks[l1], 128, err, sizeof err); | |
| const NsTensor* t = ns_find(&nm, nm1); | |
| CHECK(w3 && t && t->slices == 8 && !memcmp(w3, t->s[0].p, (size_t) 8 * ranks[1] * 128 * 2), "3-D read of %s: %s", nm1, err); | |
| CHECK(!mova_ckpt_bf16(ck, nm1, ranks[l1], 128, err, sizeof err), "a 3-D tensor read as 2-D"); | |
| ns_close(&nm); | |
| mova_ckpt_close(ck); | |
| } | |
| } | |
| MovaRef* ref = mova_ref_open(dir, 8, err, sizeof err); | |
| if (!ref) { ++fails; printf("FAIL: MLA reference: %s\n", err); return; } | |
| EngOpts o; | |
| memset(&o, 0, sizeof o); | |
| o.model_dir = dir; | |
| o.resource_dir = getenv("NSLM_RES") ? getenv("NSLM_RES") : "out/res"; | |
| o.max_seqs = 1; | |
| o.kv_tokens = 1024; | |
| Eng* e = eng_open(&o, err, sizeof err); | |
| if (!e) { ++fails; printf("FAIL: MLA eng_open: %s\n", err); mova_ref_close(ref); return; } | |
| printf("%s\n", eng_describe(e)); | |
| if (bad_ok) { | |
| EngOpts ob = o; | |
| ob.model_dir = bad; | |
| ob.kv_tokens = 64; | |
| Eng* eb = eng_open(&ob, err, sizeof err); | |
| CHECK(!eb && strstr(err, "layers.2.self_attn.mla.q_lat"), "MLA engine without q_lat: %s", eb ? "opened" : err); | |
| if (eb) eng_close(eb); | |
| } | |
| int64_t want_kv = 0; | |
| for (int l = 0; l < 5; ++l) want_kv += (int64_t) o.kv_tokens * layer_kv_bytes(ranks[l], ENG_KV_BF16); | |
| EngMem mem; | |
| eng_mem(e, &mem); | |
| CHECK(mem.kv == want_kv, "MLA KV cache %lld bytes, want %lld", (long long) mem.kv, (long long) want_kv); | |
| o.kv_format = ENG_KV_Q8; | |
| Eng* e8 = eng_open(&o, err, sizeof err); | |
| CHECK(!e8 && strstr(err, "MLA"), "MLA with the 8-bit cache must be refused (not ignored): %s", e8 ? "opened" : err); | |
| if (e8) eng_close(e8); | |
| o.kv_format = ENG_KV_BF16; | |
| const int V = eng_vocab(e), N = 48, NS = 4, TK = 4, TKV = 2; | |
| uint64_t st = 11; | |
| int32_t ids[64]; | |
| for (int i = 0; i < N; ++i) ids[i] = (int32_t) (rng_u32(&st) % (uint32_t) V); | |
| float* rl = (float*) malloc(sizeof(float) * (size_t) N * V), *el = (float*) malloc(sizeof(float) * (size_t) N * V); | |
| int32_t* rm = (int32_t*) malloc(sizeof(int32_t) * (size_t) N * NS * TK), *rv = (int32_t*) malloc(sizeof(int32_t) * (size_t) N * NS * TKV); | |
| int32_t* em = (int32_t*) malloc(sizeof(int32_t) * (size_t) N * NS * TK), *ev = (int32_t*) malloc(sizeof(int32_t) * (size_t) N * NS * TKV); | |
| CHECK(mova_ref_forward(ref, ids, N - 1, 0, rl, rm, rv) == 0, "MLA reference forward"); | |
| // Prompt rows: the CUDA engine decompresses the keys per head (kc_mla_prefill), and with NSLM_MLA_DEC_KEYS=0 attends | |
| // absorbed; both against the reference (absorbed, in double). (Metal's eng_score expands the latent per head; its | |
| // absorbed prompt rows are checked by the decode below.) | |
| const int cuda = strstr(eng_describe(e), "CUDA") != NULL; | |
| for (int path = 0; path < 1 + cuda; ++path) { | |
| Eng* es = e; | |
| if (path) { | |
| EngOpts oa = o; | |
| oa.kv_format = ENG_KV_BF16; | |
| set_env("NSLM_MLA_DEC_KEYS", "0"); | |
| es = eng_open(&oa, err, sizeof err); | |
| set_env("NSLM_MLA_DEC_KEYS", NULL); | |
| if (!es) { ++fails; printf("FAIL: MLA eng_open (absorbed prompts): %s\n", err); break; } | |
| } | |
| const char* pn = !cuda ? "" : path ? " (absorbed prompts)" : " (prompt keys decompressed)"; | |
| eng_mova_routes(es, 1, N); | |
| CHECK(eng_score(es, 0, ids, 1, N - 1, el) == 0, "MLA eng_score%s", pn); | |
| CHECK(eng_mova_routes_read(es, N - 1, em, ev, NULL, NULL) == 0, "MLA routes%s", pn); | |
| eng_mova_routes(es, 0, 0); | |
| int flips = 0, clean = 0, am_same = 0; | |
| double worst = 0, mean = 0; | |
| for (int t = 0; t < N - 1; ++t) { | |
| const int same = same_choice(em + t * NS * TK, rm + t * NS * TK, NS, TK) && same_choice(ev + t * NS * TKV, rv + t * NS * TKV, NS, TKV); | |
| flips += !same; | |
| const double re = row_err(el + (size_t) t * V, rl + (size_t) t * V, V); | |
| am_same += argmax(el + (size_t) t * V, V) == argmax(rl + (size_t) t * V, V); | |
| mean += re; | |
| if (!flips) { ++clean; worst = fmax(worst, re); } | |
| } | |
| mean /= N - 1; | |
| printf("MLA scoring%s: %d rows, %d before the first router flip (%d flipped); logit error max %.2e (clean rows), mean " | |
| "%.2e; argmax equal %d / %d\n", pn, N - 1, clean, flips, worst, mean, am_same, N - 1); | |
| // The latent attention computes in f32 against the reference's double (GQA's prefill kernel instead reproduces | |
| // MLX's rounding points, which the reference mirrors), and the decompressed path rounds its keys and values to | |
| // BF16 where the reference rounds q_lat q and the latent output, so 1-ulp differences start at row 0 and tip a | |
| // near-tie route early. Over five token sequences on an RTX 4080 the first flip came at rows 3-47 (absorbed) and | |
| // 2-12 (decompressed), with up to 8 flips, clean-row errors up to 1.3e-2 and argmax equal in 44 or more of 47 | |
| // rows on either path. The kernels themselves are pinned to BF16 ulps by tests/test_mova_kernels.c; this test | |
| // checks the wiring, which these bounds already rule out. | |
| CHECK(flips <= (N - 1) / 4 && clean >= 2, "MLA scoring%s: %d flips, %d clean rows", pn, flips, clean); | |
| CHECK(worst < 0.02 && mean < 0.05, "MLA scoring%s: logit error max %.3e, mean %.3e", pn, worst, mean); | |
| CHECK(am_same >= (N - 1) * 9 / 10, "MLA scoring%s: argmax equal in %d of %d rows", pn, am_same, N - 1); | |
| if (path) eng_close(es); | |
| } | |
| // decode: prefill P0 rows, then one row at a time; prompts of 30 rows with the latent expanded (mla_expand_min 1) | |
| // and in latent space (the default, for prompts this short) | |
| float* lg = (float*) malloc(sizeof(float) * (size_t) V); | |
| EngOpts ox = o; | |
| ox.mla_expand_min = 1; | |
| ox.kv_format = ENG_KV_BF16; | |
| Eng* ex = eng_open(&ox, err, sizeof err); | |
| CHECK(ex != NULL, "MLA eng_open (expanded prompts): %s", err); | |
| for (int pass = 0; pass < 3; ++pass) { | |
| Eng* ed = pass == 2 ? ex : e; | |
| const int P0 = pass ? 30 : 8; | |
| if (!ed) continue; | |
| CHECK(eng_prefill(ed, 0, ids, P0) == 0, "MLA prefill"); | |
| double dworst = 0, dmean = 0; | |
| int dsame = 0; | |
| for (int t = P0 - 1; t < N - 1; ++t) { | |
| CHECK(eng_step(ed, 0, lg) == 0, "MLA step"); | |
| const double re = row_err(lg, rl + (size_t) t * V, V); | |
| dworst = fmax(dworst, re); | |
| dmean += re / (N - P0); | |
| dsame += argmax(lg, V) == argmax(rl + (size_t) t * V, V); | |
| eng_push(ed, 0, ids[t + 1]); | |
| } | |
| printf("MLA decode after a %d-row prompt%s: %d steps, logit error mean %.2e max %.2e, argmax equal %d\n", P0, | |
| pass == 2 ? " (expanded)" : "", N - P0, dmean, dworst, dsame); | |
| CHECK(dmean < 0.05 && dworst < 0.25, "MLA decode: logit error mean %.3e, max %.3e", dmean, dworst); | |
| } | |
| if (ex) eng_close(ex); | |
| // the FP8 / FP4 caches (nslm/kvq.h): their size, and decode after a 30-row prompt against the reference (BF16 cache) | |
| for (int f = ENG_KV_FP8; f <= ENG_KV_FP4; ++f) { | |
| EngOpts oq = o; | |
| oq.kv_format = f; | |
| Eng* eq = eng_open(&oq, err, sizeof err); | |
| const char* fn = f == ENG_KV_FP8 ? "fp8" : "fp4"; | |
| if (!eq && strstr(err, "not supported by the CUDA")) { printf("MLA %s cache: not in this engine\n", fn); continue; } | |
| CHECK(eq != NULL, "MLA eng_open (%s cache): %s", fn, err); | |
| if (!eq) continue; | |
| int64_t want = 0; // the RoPE key FP8; the latent FP8, or (fp4) FP8 for its first KVQ_FP4_LEAD values; GQA BF16 | |
| for (int l = 0; l < 5; ++l) want += oq.kv_tokens * layer_kv_bytes(ranks[l], f); | |
| EngMem mq; | |
| eng_mem(eq, &mq); | |
| CHECK(mq.kv == want, "MLA %s KV cache %lld bytes, want %lld", fn, (long long) mq.kv, (long long) want); | |
| const int P0 = 30; | |
| CHECK(eng_prefill(eq, 0, ids, P0) == 0, "MLA prefill (%s)", fn); | |
| double dworst = 0, dmean = 0; | |
| int dsame = 0; | |
| for (int t = P0 - 1; t < N - 1; ++t) { | |
| CHECK(eng_step(eq, 0, lg) == 0, "MLA step (%s)", fn); | |
| const double re = row_err(lg, rl + (size_t) t * V, V); | |
| dworst = fmax(dworst, re); | |
| dmean += re / (N - P0); | |
| dsame += argmax(lg, V) == argmax(rl + (size_t) t * V, V); | |
| eng_push(eq, 0, ids[t + 1]); | |
| } | |
| printf("MLA decode, %s cache (%lld bytes, BF16 %lld): %d steps, logit error mean %.2e max %.2e, argmax equal %d\n", fn, | |
| (long long) mq.kv, (long long) want_kv, N - P0, dmean, dworst, dsame); | |
| // This random model's latents are uniform (no outliers), so these bounds only guard the wiring (measured: fp8 mean | |
| // 8.3e-2, fp4 2.2e-1); quality is the real model's KL divergence (nslm-mova-score --kv). | |
| CHECK(dmean < (f == ENG_KV_FP8 ? 0.12 : 0.3) && dsame >= (N - P0) * 2 / 3, "MLA decode (%s cache): logit error mean %.3e, " | |
| "argmax equal %d of %d", fn, dmean, dsame, N - P0); | |
| eng_close(eq); | |
| } | |
| // prompt lookup (multi-row verify forwards: split-key latent attention over several rows) must commit the tokens of | |
| // plain greedy decode | |
| { | |
| int32_t rep[40], ar[24], pl[24]; | |
| for (int i = 0; i < 40; ++i) rep[i] = ids[i % 10]; | |
| const int seq = 0; | |
| EngStats sp; | |
| memset(&sp, 0, sizeof sp); | |
| CHECK(eng_prefill(e, 0, rep, 40) == 0 && eng_generate(e, &seq, 1, 24, ENG_MODE_AR, ar, NULL) == 0, "MLA AR generate"); | |
| CHECK(eng_prefill(e, 0, rep, 40) == 0 && eng_generate(e, &seq, 1, 24, ENG_MODE_PL, pl, &sp) == 0, "MLA PL generate"); | |
| // The CUDA engine's one-row and multi-row forwards differ by an ulp here and there (its matvec sums one row in | |
| // another order than several), so a verify forward may break a near-tie the other way: tokens up to the first | |
| // difference must match, and there greedy decode's choice may lead the verify's by at most that noise (the | |
| // logits of one and several rows differ by < 2e-2 on this model). | |
| int k = 0; | |
| while (k < 24 && ar[k] == pl[k]) ++k; | |
| if (k < 24) { | |
| float* lk = (float*) malloc(sizeof(float) * (size_t) V); | |
| CHECK(eng_prefill(e, 0, rep, 40) == 0, "MLA AR replay"); | |
| for (int j = 0; j <= k; ++j) { | |
| CHECK(eng_step(e, 0, lk) == 0, "MLA AR replay step"); | |
| eng_push(e, 0, ar[j]); | |
| } | |
| const double gap = (double) lk[ar[k]] - lk[pl[k]]; | |
| printf("MLA prompt lookup: differs from greedy decode at token %d, a near-tie (logit gap %.4f)\n", k, gap); | |
| CHECK(gap >= 0 && gap < 0.04, "MLA: prompt lookup committed different tokens than greedy decode (token %d, logit gap %.4f)", k, gap); | |
| free(lk); | |
| } | |
| printf("MLA prompt lookup: %lld forwards for 24 tokens, %lld of %lld proposals accepted\n", (long long) sp.forwards, | |
| (long long) sp.accepted, (long long) sp.proposals); | |
| CHECK(sp.accepted > 0, "MLA prompt lookup: no proposal accepted, so no multi-row verify was tested"); | |
| } | |
| { // placement: with the cache mostly in host memory (the latents of every width staged per layer for prompts) the | |
| // engine scores and decodes exactly as with all of it in fast memory (engines with unified memory ignore this) | |
| double rel; | |
| const int diff = compare_engines(e, &o, ENG_KV_BF16, "NSLM_KV_VRAM_MB", "0.2", &st, V, &rel, "MLA KV mostly in host memory"); | |
| printf("MLA KV mostly in host memory: %d differences\n", diff); | |
| CHECK(diff == 0, "MLA: KV placement changed the results"); | |
| } | |
| eng_close(e); | |
| mova_ref_close(ref); | |
| free(rl); free(el); free(rm); free(rv); free(em); free(ev); free(lg); | |
| { // MLA capture (mova_ext.h): per-layer sums of squares of the projections' inputs, the same whether the prompt is | |
| // computed in one call or in chunks; value-expert inputs only on MoVA layers | |
| EngOpts oc = o; | |
| oc.kv_format = ENG_KV_BF16; | |
| Eng* ec = eng_open(&oc, err, sizeof err); | |
| int rc = ec ? eng_mla_capture(ec, 1) : -1; | |
| if (ec && rc) printf("MLA capture: not in this engine\n"); | |
| else if (ec) { | |
| enum { NL = 5, D = 512, KV = 256, Q = 1024, OMAX = 8 * 320 }; | |
| static double a[2][NL][D + KV + Q + OMAX]; | |
| int64_t rows[2] = {0, 0}; | |
| for (int pass = 0; pass < 2; ++pass) { | |
| eng_mla_capture(ec, 0); | |
| eng_mla_capture(ec, 1); // fresh sums (a new capture) | |
| if (pass == 0) CHECK(eng_prefill(ec, 0, ids, 40) == 0, "MLA capture prefill"); | |
| else { | |
| int left; | |
| eng_free(ec, 0); | |
| CHECK(eng_prefill_begin(ec, 0, ids, 40, NULL) == 0, "MLA capture prefill_begin"); | |
| while ((left = eng_prefill_next(ec, 0, 7)) > 0) {} | |
| } | |
| for (int l = 0; l < NL; ++l) { // a GQA layer has no MLA inputs: its read fails | |
| const int rc = eng_mla_capture_read(ec, l, a[pass][l], a[pass][l] + D, a[pass][l] + D + KV, a[pass][l] + D + KV + Q, &rows[pass]); | |
| CHECK(ranks[l] ? rc == 0 : rc != 0, "MLA capture read, layer %d (rank %d): %d", l, ranks[l], rc); | |
| } | |
| } | |
| double worst = 0; | |
| int zero = 0, vbad = 0; | |
| for (int l = 0; l < NL; ++l) { | |
| if (!ranks[l]) continue; | |
| const int n = D + KV + Q + 8 * ranks[l]; | |
| for (int i = 0; i < n; ++i) { | |
| const double x = a[0][l][i], y = a[1][l][i]; | |
| if (i >= D && i < D + KV) { vbad += l == 0 ? x != 0 : !(x > 0); continue; } // layer 0: dense | |
| zero += !(x > 0); | |
| worst = fmax(worst, fabs(x - y) / fmax(fabs(x), 1e-30)); | |
| } | |
| } | |
| printf("MLA capture: %lld rows; sums in one call vs chunks of 7: worst relative difference %.2e\n", (long long) rows[0], worst); | |
| CHECK(rows[0] == 39 && rows[1] == 39, "MLA capture rows %lld, %lld", (long long) rows[0], (long long) rows[1]); | |
| CHECK(!zero && !vbad && worst < 1e-4, "MLA capture: %d empty sums, %d bad value-expert sums, difference %.3e", zero, vbad, worst); | |
| // mode 2 adds each input's X^T X (GPTQ): symmetric, its diagonal the sums above, the same in chunks of 7 | |
| static float hx[2][8 * 320 * 320]; | |
| double wd = 0, wc = 0; | |
| int asym = 0, nox = 0; | |
| eng_mla_capture(ec, 0); | |
| const int xtx = eng_mla_capture(ec, 2) == 0; | |
| if (!xtx) printf("MLA capture X^T X: not in this engine\n"); | |
| for (int l = 0; l < NL && xtx; ++l) | |
| for (int site = 0; site < 4 && ranks[l]; ++site) { | |
| const int dim = site == 0 ? D : site == 1 ? KV : site == 2 ? 128 : ranks[l], nb = site < 2 ? 1 : 8; | |
| const size_t n = (size_t) nb * dim * dim; | |
| const double* sums = a[0][l] + (site == 0 ? 0 : site == 1 ? D : site == 2 ? D + KV : D + KV + Q); | |
| if (site == 1 && l == 0) continue; // dense layer: no value experts | |
| for (int pass = 0; pass < 2; ++pass) { | |
| eng_mla_capture(ec, 0); | |
| eng_mla_capture(ec, 2); | |
| eng_free(ec, 0); | |
| if (pass == 0) CHECK(eng_prefill(ec, 0, ids, 40) == 0, "MLA X^T X prefill"); | |
| else { | |
| CHECK(eng_prefill_begin(ec, 0, ids, 40, NULL) == 0, "MLA X^T X prefill_begin"); | |
| while (eng_prefill_next(ec, 0, 7) > 0) {} | |
| } | |
| if (eng_mla_capture_xtx(ec, l, site, hx[pass])) { ++nox; break; } | |
| } | |
| double sc = 0; | |
| for (size_t i = 0; i < n; ++i) sc = fmax(sc, fabs(hx[0][i])); | |
| for (size_t i = 0; i < n; ++i) wc = fmax(wc, fabs(hx[1][i] - hx[0][i]) / sc); | |
| for (int b = 0; b < nb; ++b) { | |
| const float* h = hx[0] + (size_t) b * dim * dim; | |
| for (int i = 0; i < dim; ++i) { | |
| wd = fmax(wd, fabs(h[i * dim + i] - sums[b * dim + i]) / fmax(sums[b * dim + i], 1e-30)); | |
| for (int j = 0; j < i; ++j) asym += h[i * dim + j] != h[j * dim + i]; | |
| } | |
| } | |
| } | |
| if (xtx) printf("MLA capture X^T X: diagonal vs sums %.2e, one call vs chunks of 7 %.2e, %d asymmetric\n", wd, wc, asym); | |
| CHECK(!xtx || (!nox && !asym && wd < 1e-4 && wc < 1e-4), "MLA capture X^T X: %d unread, %d asymmetric, diagonal %.3e, chunks %.3e", | |
| nox, asym, wd, wc); | |
| } | |
| if (ec) eng_close(ec); | |
| } | |
| { // paged KV (Metal on macOS 26.4+): positions are backed by memory as they are written, and freed with the slot; | |
| // the device's allocated size shows it. Engines that allocate the whole cache at open skip this. | |
| EngOpts op = o; | |
| op.kv_format = ENG_KV_BF16; | |
| op.kv_tokens = 1 << 16; | |
| op.max_seqs = 2; | |
| Eng* ep = eng_open(&op, err, sizeof err); | |
| CHECK(ep != NULL, "MLA eng_open (64k-token cache): %s", err); | |
| if (ep) { | |
| EngMem m0, m1, m2, m3; | |
| int32_t* big = (int32_t*) malloc(sizeof(int32_t) * 3000); | |
| for (int i = 0; i < 3000; ++i) big[i] = ids[i % N]; | |
| const int64_t per = eng_kv_bytes(ep); | |
| eng_mem(ep, &m0); | |
| CHECK(eng_prefill(ep, 1, big, 3000) == 0, "MLA prefill of 3000 tokens"); | |
| eng_mem(ep, &m1); | |
| CHECK(eng_prefill(ep, 1, big, 100) == 0, "MLA prefill of 100 tokens in the same slot"); | |
| eng_mem(ep, &m2); | |
| eng_free(ep, 1); | |
| eng_mem(ep, &m3); | |
| const double grew = (double) (m1.gpu_allocated - m0.gpu_allocated), kept = (double) (m2.gpu_allocated - m0.gpu_allocated); | |
| if (grew < 0.5 * 3000 * per) printf("MLA KV pages: not paged in this engine\n"); | |
| else { | |
| printf("MLA KV pages: 3000 positions back %.1f MB (%.1f MB of KV), 100 keep %.1f MB, freed slot %.1f MB\n", | |
| grew / 1e6, 3000.0 * per / 1e6, kept / 1e6, (double) (m3.gpu_allocated - m0.gpu_allocated) / 1e6); | |
| CHECK(grew <= 3000.0 * per + 64e6 + 2e6 * 4 * 5, "MLA KV pages: %.0f bytes backed for 3000 positions", grew); | |
| // (this model's 3000 positions fit one 64 MB heap, so the 100-token prompt keeps it: no partial hand-back here) | |
| CHECK(kept <= grew && m3.gpu_allocated <= m0.gpu_allocated, "MLA KV pages: not given back (%.0f, %.0f)", | |
| kept, (double) (m3.gpu_allocated - m0.gpu_allocated)); | |
| } | |
| free(big); | |
| eng_close(ep); | |
| } | |
| } | |
| { // paged KV unavailable after its page alignment was worked out (no mapping queue: NSLM_KVP_NO_QUEUE fakes that): | |
| // the dense cache's slots must use the dense layout. A cache copied from slot 0 into slot 1 decodes the same. | |
| EngOpts op = o; | |
| op.kv_format = ENG_KV_BF16; | |
| op.kv_tokens = 1024; | |
| op.max_seqs = 2; | |
| set_env("NSLM_KVP_NO_QUEUE", "1"); | |
| Eng* ep = eng_open(&op, err, sizeof err); | |
| set_env("NSLM_KVP_NO_QUEUE", NULL); | |
| CHECK(ep != NULL, "MLA eng_open (no paging queue): %s", err); | |
| if (ep) { | |
| const int P = N - 1; | |
| const size_t kb = (size_t) eng_kv_bytes(ep) * P; | |
| uint8_t* k0 = (uint8_t*) malloc(kb), *k1 = (uint8_t*) malloc(kb); | |
| float* l0 = (float*) malloc(sizeof(float) * (size_t) V), *l1 = (float*) malloc(sizeof(float) * (size_t) V); | |
| CHECK(eng_prefill(ep, 0, ids, N) == 0 && eng_kv_read(ep, 0, 0, P, k0) == 0, "MLA no paging queue: slot 0"); | |
| CHECK(eng_kv_write(ep, 1, ids, 0, P, k0) == 0, "MLA no paging queue: eng_kv_write to slot 1"); | |
| CHECK(eng_kv_read(ep, 1, 0, P, k1) == 0 && !memcmp(k0, k1, kb), "MLA no paging queue: slot 1's cache differs"); | |
| CHECK(eng_step(ep, 0, l0) == 0 && eng_push(ep, 1, ids[P]) == 0 && eng_step(ep, 1, l1) == 0 && | |
| !memcmp(l0, l1, sizeof(float) * (size_t) V), "MLA no paging queue: slot 1 decodes differently"); | |
| free(k0); free(k1); free(l0); free(l1); | |
| eng_close(ep); | |
| } | |
| } | |
| test_batch(dir, "MLA", 0, ENG_KV_BF16); // these prompts (< 256 rows): attention in latent space | |
| test_batch(dir, "MLA expanded", 1, ENG_KV_BF16); // every prompt with the latent expanded per head | |
| test_batch(dir, "MLA blocks of 16", 16, ENG_KV_BF16); // expanded up to the last multiple of 16 (the copied cache: 128) | |
| test_batch(dir, "MLA fp8", 16, ENG_KV_FP8); | |
| test_batch(dir, "MLA fp4", 16, ENG_KV_FP4); | |
| } | |
| int main(void) { | |
| const char* dir = "out/test/engine_model"; | |
| char err[512] = ""; | |
| if (write_model(dir, CONFIG, NULL, err, sizeof err)) { printf("FAIL: model folder: %s\n", err); return 1; } | |
| MovaRef* ref = mova_ref_open(dir, 8, err, sizeof err); | |
| if (!ref) { printf("FAIL: reference: %s\n", err); return 1; } | |
| { // a folder without one layer's stacked experts (as an original checkpoint, which stores them per expert): the | |
| // reference refuses it, naming the tensor (was a NULL dereference in forward) | |
| const char* nx = "out/test/engine_model_no_experts"; | |
| char e2[512] = ""; | |
| const int ok = write_model(nx, CONFIG, "layers.2.mlp.experts.up_proj", e2, sizeof e2) == 0; | |
| CHECK(ok, "folder without experts: %s", e2); | |
| MovaRef* rx = ok ? mova_ref_open(nx, 1, e2, sizeof e2) : NULL; | |
| CHECK(ok && !rx && strstr(e2, "layers.2.mlp.experts.up_proj") && strstr(e2, "packed"), "reference without experts: %s", | |
| rx ? "opened" : e2); | |
| if (rx) mova_ref_close(rx); | |
| } | |
| EngOpts o; | |
| memset(&o, 0, sizeof o); | |
| o.model_dir = dir; | |
| o.resource_dir = getenv("NSLM_RES") ? getenv("NSLM_RES") : "out/res"; | |
| o.max_seqs = 1; | |
| o.kv_tokens = 1024; | |
| Eng* e = eng_open(&o, err, sizeof err); | |
| if (!e && (strstr(err, "no CUDA device") || strstr(err, "no Metal device"))) { printf("SKIP: %s\n", err); return 77; } // a machine without a GPU | |
| if (!e) { printf("FAIL: eng_open: %s\n", err); return 1; } | |
| printf("%s\n", eng_describe(e)); | |
| const int V = eng_vocab(e), N = 48, NS = 4, TK = 4, TKV = 2; | |
| uint64_t st = 7; | |
| int32_t ids[64]; | |
| for (int i = 0; i < N; ++i) ids[i] = (int32_t) (rng_u32(&st) % (uint32_t) V); | |
| // reference: logits for rows 0 .. N-2 (predicting ids[1 ..]) and its routes | |
| float* rl = (float*) malloc(sizeof(float) * (size_t) N * V); | |
| int32_t* rm = (int32_t*) malloc(sizeof(int32_t) * (size_t) N * NS * TK); | |
| int32_t* rv = (int32_t*) malloc(sizeof(int32_t) * (size_t) N * NS * TKV); | |
| // two references: the prefill kernels' attention rounding (scoring forwards of more than 8 rows) and the decode | |
| // kernels' (single rows) | |
| float* rd = (float*) malloc(sizeof(float) * (size_t) N * V); | |
| mova_ref_attn_rounding(ref, 1); | |
| CHECK(mova_ref_forward(ref, ids, N - 1, 0, rl, rm, rv) == 0, "reference forward"); | |
| mova_ref_attn_rounding(ref, 0); | |
| CHECK(mova_ref_forward(ref, ids, N - 1, 0, rd, NULL, NULL) == 0, "reference forward"); | |
| // 1. prompt scoring: one forward of N - 1 rows | |
| float* el = (float*) malloc(sizeof(float) * (size_t) N * V); | |
| int32_t* em = (int32_t*) malloc(sizeof(int32_t) * (size_t) N * NS * TK); | |
| int32_t* ev = (int32_t*) malloc(sizeof(int32_t) * (size_t) N * NS * TKV); | |
| eng_mova_routes(e, 1, N); | |
| CHECK(eng_score(e, 0, ids, 1, N - 1, el) == 0, "eng_score"); | |
| CHECK(eng_mova_routes_read(e, N - 1, em, ev, NULL, NULL) == 0, "routes"); | |
| eng_mova_routes(e, 0, 0); | |
| // A row whose own router choices differ is not comparable; nor is any row after one: a flip changes that row's | |
| // value-expert output, which later rows read through the KV cache. Rows before the first flip must agree to the | |
| // engines' rounding; all rows together, to a BF16 network's chaotic drift (a stray ulp moves later roundings). | |
| int flips = 0, clean = 0, am_same = 0; | |
| double worst = 0, mean = 0; | |
| for (int t = 0; t < N - 1; ++t) { | |
| const int same = same_choice(em + t * NS * TK, rm + t * NS * TK, NS, TK) && same_choice(ev + t * NS * TKV, rv + t * NS * TKV, NS, TKV); | |
| flips += !same; | |
| const double re = row_err(el + (size_t) t * V, rl + (size_t) t * V, V); | |
| am_same += argmax(el + (size_t) t * V, V) == argmax(rl + (size_t) t * V, V); | |
| mean += re; | |
| if (!flips) { ++clean; worst = fmax(worst, re); } | |
| } | |
| mean /= N - 1; | |
| printf("scoring: %d rows, %d before the first router flip (%d rows flipped); logit error: those rows max %.2e, all rows " | |
| "mean %.2e; argmax equal %d / %d\n", N - 1, clean, flips, worst, mean, am_same, N - 1); | |
| CHECK(flips <= (N - 1) / 4, "scoring: %d rows with a different router choice", flips); | |
| CHECK(clean >= (N - 1) / 4, "scoring: only %d rows before the first router flip", clean); // else worst says nothing | |
| CHECK(worst < 0.02, "scoring: logit error %.3e before the first flip", worst); | |
| CHECK(mean < 0.05, "scoring: mean logit error %.3e", mean); | |
| CHECK(am_same >= (N - 1) * 9 / 10, "scoring: argmax equal in %d of %d rows", am_same, N - 1); | |
| // 2. decode: prefill 8 tokens, then step through the rest one row at a time | |
| const int P0 = 8; | |
| CHECK(eng_prefill(e, 0, ids, P0) == 0, "prefill"); | |
| float* lg = (float*) malloc(sizeof(float) * (size_t) V); | |
| double dworst = 0, dmean = 0; | |
| int dsame = 0; | |
| for (int t = P0 - 1; t < N - 1; ++t) { | |
| CHECK(eng_step(e, 0, lg) == 0, "step"); | |
| const double re = row_err(lg, rd + (size_t) t * V, V); | |
| dworst = fmax(dworst, re); | |
| dmean += re / (N - P0); | |
| dsame += argmax(lg, V) == argmax(rd + (size_t) t * V, V); | |
| eng_push(e, 0, ids[t + 1]); | |
| } | |
| printf("decode: %d steps, logit error max %.2e, mean %.2e (all rows), argmax equal %d / %d\n", N - P0, dworst, dmean, | |
| dsame, N - P0); | |
| // the same bounds as scoring's all-row ones (rows after a router flip may differ more: a gross bound for the max) | |
| CHECK(dmean < 0.05, "decode: mean logit error %.3e", dmean); | |
| CHECK(dworst < 0.25, "decode: logit error %.3e", dworst); | |
| CHECK(dsame >= (N - P0) * 9 / 10, "decode: argmax equal in %d of %d steps", dsame, N - P0); | |
| // 3. prompt lookup must commit the tokens of plain greedy decode (a repetitive prompt so drafts are offered) | |
| int32_t rep[64]; | |
| for (int i = 0; i < 40; ++i) rep[i] = ids[i % 10]; | |
| int32_t ar[24], pl[24]; | |
| const int seq = 0; | |
| EngStats sp; | |
| memset(&sp, 0, sizeof sp); | |
| CHECK(eng_prefill(e, 0, rep, 40) == 0 && eng_generate(e, &seq, 1, 24, ENG_MODE_AR, ar, NULL) == 0, "AR generate"); | |
| CHECK(eng_prefill(e, 0, rep, 40) == 0 && eng_generate(e, &seq, 1, 24, ENG_MODE_PL, pl, &sp) == 0, "PL generate"); | |
| CHECK(!memcmp(ar, pl, sizeof ar), "prompt lookup committed different tokens than greedy decode"); | |
| printf("prompt lookup: %lld forwards for 24 tokens, %lld of %lld proposals accepted\n", (long long) sp.forwards, | |
| (long long) sp.accepted, (long long) sp.proposals); | |
| // 4. long prompts: scoring with route capture (forward by forward) and without (an engine may then run the prompt | |
| // layer by layer) must give the same logits bit for bit: same kernels, same chunks, same positions. 600 rows: two | |
| // chunks; 517: a last chunk of 5 rows, which takes the decode kernels (and must not disturb the layer prefetch); | |
| // also with an expert cache too small for one layer's experts, so that every layer's experts are copied in | |
| { | |
| EngOpts os = o; | |
| set_env("NSLM_EXPERT_VRAM_MB", "13"); | |
| Eng* es = eng_open(&os, err, sizeof err); | |
| set_env("NSLM_EXPERT_VRAM_MB", NULL); | |
| CHECK(es != NULL, "eng_open (small expert cache): %s", err); | |
| const int lens[2] = {600, 517}; | |
| for (int k = 0; k < 4; ++k) { | |
| Eng* ek = k < 2 ? e : es; | |
| const int NL = lens[k & 1]; | |
| if (!ek) continue; | |
| int32_t* lid = (int32_t*) malloc(sizeof(int32_t) * NL); | |
| for (int i = 0; i < NL; ++i) lid[i] = (int32_t) (rng_u32(&st) % (uint32_t) V); | |
| const int from = NL - 70, cnt = 70; | |
| float* la = (float*) malloc(sizeof(float) * (size_t) cnt * V); | |
| float* lb = (float*) malloc(sizeof(float) * (size_t) cnt * V); | |
| CHECK(eng_score(ek, 0, lid, from, cnt, la) == 0, "long scoring"); | |
| eng_mova_routes(ek, 1, NL); | |
| CHECK(eng_score(ek, 0, lid, from, cnt, lb) == 0, "long scoring with routes"); | |
| eng_mova_routes(ek, 0, 0); | |
| int diff = 0; | |
| for (size_t i = 0; i < (size_t) cnt * V; ++i) diff += la[i] != lb[i]; | |
| printf("long prompt (%d tokens, %s expert cache): %d of %d logits differ between the two forwards\n", NL, | |
| k < 2 ? "full" : "small", diff, cnt * V); | |
| CHECK(diff == 0, "long prompt (%d tokens): layer-major and chunked forwards differ", NL); | |
| free(lid); free(la); free(lb); | |
| } | |
| if (es) eng_close(es); | |
| } | |
| // 5. cached prefill: a shared prefix is reused and the result equals a full prefill's next step | |
| int reused = 0; | |
| CHECK(eng_prefill(e, 0, ids, 30) == 0, "prefill 30"); | |
| CHECK(eng_prefill_cached(e, 0, ids, 36, &reused) == 0 && reused == 29, "cached prefill reused %d", reused); | |
| CHECK(eng_len(e, 0) == 36, "length after cached prefill"); | |
| // 5b. a server's sequence: a short request, then a long one sharing a prefix (a cached prefill of many rows), then | |
| // decoding; the long request's tokens must equal a fresh prefill's | |
| { | |
| const int NL = 300; | |
| int32_t* lp = (int32_t*) malloc(sizeof(int32_t) * NL); | |
| for (int i = 0; i < NL; ++i) lp[i] = i < 4 ? ids[i] : (int32_t) (rng_u32(&st) % (uint32_t) V); | |
| int32_t ga[16], gb[16]; | |
| int reused2 = 0; | |
| CHECK(eng_prefill(e, 0, ids, 20) == 0 && eng_generate(e, &seq, 1, 8, ENG_MODE_AR, ga, NULL) == 0, "short request"); | |
| CHECK(eng_prefill_cached(e, 0, lp, NL, &reused2) == 0 && reused2 == 4, "long cached prefill (reused %d)", reused2); | |
| CHECK(eng_generate(e, &seq, 1, 16, ENG_MODE_AR, ga, NULL) == 0, "decode after the long cached prefill"); | |
| CHECK(eng_prefill(e, 0, lp, NL) == 0 && eng_generate(e, &seq, 1, 16, ENG_MODE_AR, gb, NULL) == 0, "fresh prefill"); | |
| CHECK(!memcmp(ga, gb, sizeof ga), "a cached prefill decodes differently from a fresh one"); | |
| printf("long cached prefill: %s\n", memcmp(ga, gb, sizeof ga) ? "DIFFERS" : "same tokens as a fresh prefill"); | |
| free(lp); | |
| } | |
| // 6. placement must not change results: an engine allowed to keep only part of the experts (NSLM_EXPERT_VRAM_MB) | |
| // or of the KV cache (NSLM_KV_VRAM_MB) in fast memory scores and decodes exactly as this one does (engines with | |
| // unified memory ignore both) | |
| { | |
| double rel; | |
| int diff = compare_engines(e, &o, ENG_KV_BF16, "NSLM_EXPERT_VRAM_MB", "13", &st, V, &rel, "small expert cache"); | |
| printf("small expert cache: %d differences\n", diff); | |
| CHECK(diff == 0, "a smaller expert cache changed the results"); | |
| diff = compare_engines(e, &o, ENG_KV_BF16, "NSLM_KV_VRAM_MB", "1", &st, V, &rel, "KV mostly in host memory"); | |
| printf("KV mostly in host memory: %d differences\n", diff); | |
| CHECK(diff == 0, "KV placement changed the results"); | |
| // the prefill GEMM with exact (f32) seed weights: a different rounding point, so close, not equal | |
| diff = compare_engines(e, &o, ENG_KV_BF16, "NSLM_SEED_GEMM_F32", "1", &st, V, &rel, "f32 seed weights in the GEMM"); | |
| printf("f32 seed weights in the GEMM: %d values differ, relative logit difference mean %.2e\n", diff, rel); | |
| CHECK(rel < 0.05, "f32 seed weights in the GEMM: mean relative logit difference %.3e", rel); | |
| } | |
| // 7. the 8-bit KV cache: close to BF16 (its rounding is of BF16's size), and placement-invariant itself | |
| { | |
| EngOpts o8 = o; | |
| o8.kv_format = ENG_KV_Q8; | |
| Eng* e8 = eng_open(&o8, err, sizeof err); | |
| if (!e8 && strstr(err, "not supported")) printf("Q8 KV: skipped (%s)\n", err); // a BF16-only engine (Metal) | |
| else CHECK(e8 != NULL, "eng_open (Q8 KV): %s", err); | |
| if (e8) { | |
| double rel; | |
| const int diff = compare_engines(e, &o, ENG_KV_Q8, NULL, NULL, &st, V, &rel, "Q8 KV against BF16"); | |
| printf("Q8 KV against BF16 KV: %d values differ, relative logit difference mean %.2e\n", diff, rel); | |
| // a gross-error bound only: the random model amplifies 8-bit rounding chaotically; the 8-bit cache's quality is | |
| // measured on the real model (nslm-mova-refcheck --kv q8: KL 1e-2 against the CPU reference, 5e-3 for BF16) | |
| CHECK(rel < 0.25, "Q8 KV: mean relative logit difference %.3e", rel); | |
| const int d2 = compare_engines(e8, &o8, ENG_KV_Q8, "NSLM_KV_VRAM_MB", "0.5", &st, V, &rel, "Q8 KV mostly in host memory"); | |
| printf("Q8 KV mostly in host memory: %d differences\n", d2); | |
| CHECK(d2 == 0, "Q8 KV placement changed the results"); | |
| eng_close(e8); | |
| } | |
| } | |
| eng_close(e); | |
| mova_ref_close(ref); | |
| free(rl); free(rd); free(rm); free(rv); free(el); free(em); free(ev); free(lg); | |
| test_batch(dir, "GQA", 0, ENG_KV_BF16); | |
| { | |
| static const int ranks[5] = {64, 96, 64, 320, 96}; // 320: the fp4 cache's FP4 latent dims (past KVQ_FP4_LEAD) | |
| // hybrid: layers 0 (dense) and 4 (MoVA) keep GQA attention. A GQA layer between MLA layers works too (the real | |
| // T-03 model has ten), but in this small random model it turns the MLA layers' float32-vs-double noise into | |
| // near-tie expert flips (39 of 47 rows with GQA in layer 2, against 9 for the all-MLA model) that these bounds | |
| // do not allow; its first decode steps match the reference to 5e-3 per layer. The real model's check is its | |
| // held-out score (nslm-mova-score) against the MLX loader's. | |
| static const int hybrid[5] = {0, 96, 64, 320, 0}; | |
| test_mla_ranks(ranks, "out/test/engine_model_mla"); | |
| test_mla_ranks(hybrid, "out/test/engine_model_mla_hybrid"); | |
| } | |
| printf("test_engine: %s\n", fails ? "FAIL" : "PASS"); | |
| return fails != 0; | |
| } | |
Xet Storage Details
- Size:
- 55.7 kB
- Xet hash:
- f08a1709f97e7e9a38148e3563fea7d29c10d42291ac275f56abe772d06d3171
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.