server/minicpm_server.c

/* minicpm_server.c - MiniCPM5-2B (Llama arch) mixed-precision CPU inference for LAL
 *
 * Architecture: 42 layers, 2048 hidden, 16Q/2KV heads (GQA 8:1), 128 head_dim,
 *               6144 MLP (SwiGLU), 130560 vocab, RoPE theta 5e6, RMSNorm eps 1e-6.
 *
 * Mixed precision (from GPQ8 file, per-tensor qtype):
 *   qtype=0 F32     : norms
 *   qtype=1 Q8 row  : embed_tokens + lm_head (int8 + per-row fp32 scale)
 *   qtype=3 Q8_0    : layer 0 and 41 (8-bit, sensitive edge layers), 34B/32elem
 *   qtype=5 Q4_K    : main body q/k/v/o/down + gate/up outside ternary range, 144B/256elem
 *   qtype=6 TERNARY : gate/up of layers 8..31, 3-value {-s,0,+s}, 18B/64elem
 *                     (2B fp16 absmean scale + 8B bit0 stream + 8B bit1 stream;
 *                      kernel: t = b0 - b1, y = x_scale * sum_b s_b * (S0-S1))
 *
 * Windows (MinGW-w64) port: CreateFileMapping mmap, QueryPerformanceCounter,
 * no dlopen/weak-symbol LAL bridge, no pthread (OpenMP only).
 *
 * Build: gcc -O3 -march=native -fopenmp -I. -o prebuilt/minicpm_server.exe
 *        tools/server/minicpm_server.c -lm
 * Run:   prebuilt/minicpm_server.exe --weights C:/models/minicpm_mixed.bin
 *          --tokenizer C:/models/MiniCPM5-2B-cpu --prompt "..." --n 64 --threads 4
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdint.h>
#include <time.h>
#include <immintrin.h>
#include <omp.h>

#ifdef _WIN32
#include <windows.h>
#include <malloc.h>
#else
#include <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#endif

/* === Architecture constants (MiniCPM5-2B, LlamaForCausalLM) === */
#define N_EMBD       2048
#define N_LAYER      42
#define N_HEAD       16
#define N_KV_HEAD    2
#define HEAD_DIM     128
#define N_Q_PER_KV   (N_HEAD / N_KV_HEAD)  /* 8 */
#define MLP_DIM      6144
#define VOCAB_SIZE   130560
#define N_CTX        2048
#define ROPE_THETA   5000000.0f
#define RMS_EPS      1e-6f
#define KV_DIM       (N_KV_HEAD * HEAD_DIM)  /* 256 */
#define Q_DIM        (N_HEAD * HEAD_DIM)     /* 2048 */

/* === SIMD wrappers (AVX2) === */
typedef __m256 v8f;
static inline float v8f_hsum(v8f v) {
    __m128 hi = _mm256_extractf128_ps(v, 1);
    __m128 lo = _mm256_castps256_ps128(v);
    __m128 s = _mm_add_ps(lo, hi);
    s = _mm_hadd_ps(s, s); s = _mm_hadd_ps(s, s);
    return _mm_cvtss_f32(s);
}

/* === Reusable LAL SDK headers === */
#define XQ_MAX 6144  /* max in_dim across all matmuls = MLP_DIM */
#include "runtime/lal_q8_kernel.h"
#include "runtime/lal_q4k_kernel.h"
#include "runtime/lal_sampling.h"
#include "runtime/lal_dequant.h"
#include "runtime/lal_tokenizer.h"
#include "runtime/lal_simd_optim.h"

/* === Aligned alloc (portable) === */
static void *xalloc(size_t size, size_t align) {
#ifdef _WIN32
    return _aligned_malloc(size ? size : 1, align);
#else
    void *p = NULL;
    if (posix_memalign(&p, align, size ? size : 1)) return NULL;
    return p;
#endif
}
#define memalign(a, s) xalloc((s), (a))

/* === Timer === */
static double now_sec(void) {
#ifdef _WIN32
    static LARGE_INTEGER freq = {0};
    LARGE_INTEGER c;
    if (!freq.QuadPart) QueryPerformanceFrequency(&freq);
    QueryPerformanceCounter(&c);
    return (double)c.QuadPart / (double)freq.QuadPart;
#else
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return ts.tv_sec + ts.tv_nsec * 1e-9;
#endif
}

/* === Windows read-only mmap === */
#ifdef _WIN32
static void *g_mmap_base;
static size_t g_mmap_size;
static void *win_mmap_ro(const char *path, size_t *out_size) {
    HANDLE f = CreateFileA(path, GENERIC_READ, FILE_SHARE_READ, NULL,
                           OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (f == INVALID_HANDLE_VALUE) { fprintf(stderr, "[!] cannot open %s\n", path); exit(1); }
    LARGE_INTEGER sz;
    GetFileSizeEx(f, &sz);
    HANDLE m = CreateFileMappingA(f, NULL, PAGE_READONLY, 0, 0, NULL);
    if (!m) { fprintf(stderr, "[!] CreateFileMapping failed\n"); exit(1); }
    void *p = MapViewOfFile(m, FILE_MAP_READ, 0, 0, 0);
    CloseHandle(m); CloseHandle(f);
    if (!p) { fprintf(stderr, "[!] MapViewOfFile failed\n"); exit(1); }
    *out_size = (size_t)sz.QuadPart;
    return p;
}
#else
static void *g_mmap_base;
static size_t g_mmap_size;
static void *win_mmap_ro(const char *path, size_t *out_size) {
    int fd = open(path, O_RDONLY);
    if (fd < 0) { fprintf(stderr, "[!] cannot open %s\n", path); exit(1); }
    struct stat st;
    fstat(fd, &st);
    void *p = mmap(NULL, st.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
    close(fd);
    if (p == MAP_FAILED) { fprintf(stderr, "[!] mmap failed\n"); exit(1); }
    madvise(p, st.st_size, MADV_SEQUENTIAL);
    *out_size = st.st_size;
    return p;
}
#endif

/* === Global state === */
static int g_n_threads = 1;
static float g_temperature = 0.8f;
static int g_top_k = 40;
static float g_rep_penalty = 1.1f;
static int g_recent[256], g_n_recent = 0;

/* === TERNARY AVX2 kernel (qtype=6) ===
 * Block: 18 bytes / 64 elems = [2B fp16 scale][8B bit0][8B bit1]
 * t_i = bit0_i - bit1_i  in {-1,0,+1};  w_i ~= scale * t_i
 * dot = x_scale * sum_b scale_b * (S0_b - S1_b),  S = sum(bit * xq)  (int32)
 */
static uint8_t g_bit_lut[256][8];
static void bit_lut_init(void) {
    for (int b = 0; b < 256; b++)
        for (int j = 0; j < 8; j++)
            g_bit_lut[b][j] = (uint8_t)((b >> j) & 1);
}

/* binary dot over 32 elems: bits = 4 bytes (little-endian bitstream), xq = 32 int8 */
static inline int32_t bdot32(const uint8_t *bits, const int8_t *xq) {
    uint8_t exp32[32] __attribute__((aligned(32)));
    memcpy(exp32 + 0,  g_bit_lut[bits[0]], 8);
    memcpy(exp32 + 8,  g_bit_lut[bits[1]], 8);
    memcpy(exp32 + 16, g_bit_lut[bits[2]], 8);
    memcpy(exp32 + 24, g_bit_lut[bits[3]], 8);
    __m256i a = _mm256_loadu_si256((const __m256i *)exp32);   /* unsigned 0/1 */
    __m256i w = _mm256_loadu_si256((const __m256i *)xq);     /* signed int8  */
    __m256i r16 = _mm256_maddubs_epi16(a, w);                /* 16 x int16   */
    __m256i r32 = _mm256_madd_epi16(_mm256_set1_epi16(1), r16); /* 8 x int32 */
    __m128i lo = _mm256_castsi256_si128(r32);
    __m128i hi = _mm256_extractf128_si256(r32, 1);
    __m128i s = _mm_add_epi32(lo, hi);
    s = _mm_add_epi32(s, _mm_shuffle_epi32(s, _MM_SHUFFLE(1, 0, 3, 2)));
    s = _mm_add_epi32(s, _mm_shuffle_epi32(s, _MM_SHUFFLE(2, 3, 0, 1)));
    return _mm_cvtsi128_si32(s);
}

static inline float tern_dot_row(const uint8_t *row, const int8_t *xq,
                                 float x_scale, int in_dim) {
    int nb = in_dim / 64;
    float acc = 0;
    for (int b = 0; b < nb; b++) {
        const uint8_t *blk = row + (size_t)b * 18;
        uint16_t s16;
        memcpy(&s16, blk, 2);
        __m128i sh = _mm_set1_epi16((short)s16);
        float s = _mm_cvtss_f32(_mm_cvtph_ps(sh));
        const int8_t *xb = xq + (size_t)b * 64;
        int32_t i0 = bdot32(blk + 2, xb) + bdot32(blk + 6, xb + 32);
        int32_t i1 = bdot32(blk + 10, xb) + bdot32(blk + 14, xb + 32);
        acc += s * (float)(i0 - i1);
    }
    return acc * x_scale;
}

static void parallel_matmul_ternary(float *y, const uint8_t *W,
                                     const int8_t *xq, float x_scale,
                                     int in_dim, int out_dim) {
    int nb = in_dim / 64;
    int row_stride = nb * 18;
    if (g_n_threads <= 1 || out_dim < 1024) {
        for (int r = 0; r < out_dim; r++)
            y[r] = tern_dot_row(W + (size_t)r * row_stride, xq, x_scale, in_dim);
        return;
    }
    #pragma omp parallel num_threads(g_n_threads)
    {
        int tid = omp_get_thread_num();
        int n = omp_get_num_threads();
        int chunk = (out_dim + n - 1) / n;
        int start = tid * chunk;
        int end = start + chunk;
        if (end > out_dim) end = out_dim;
        for (int r = start; r < end; r++)
            y[r] = tern_dot_row(W + (size_t)r * row_stride, xq, x_scale, in_dim);
    }
}

/* === Parallel Q8_0 matmul (packed 34B blocks, kernel quantizes x internally) === */
static void parallel_matmul_q8_0_w(float *y, const uint8_t *q8_0_W,
                                    const float *x, int in_dim, int out_dim) {
    if (g_n_threads <= 1 || out_dim < 1024) {
        lal_matmul_q8_0(y, q8_0_W, x, NULL, in_dim, out_dim);
        return;
    }
    int row_stride = (in_dim / 32) * 34;
    #pragma omp parallel num_threads(g_n_threads)
    {
        int tid = omp_get_thread_num();
        int n = omp_get_num_threads();
        int chunk = (out_dim + n - 1) / n;
        int start = tid * chunk;
        int end = start + chunk;
        if (end > out_dim) end = out_dim;
        if (start < out_dim)
            lal_matmul_q8_0(y + start, q8_0_W + (size_t)start * row_stride,
                            x, NULL, in_dim, end - start);
    }
}

/* === Parallel Q4_K matmul (prepared) === */
static void parallel_matmul_q4_k_p(float *y, const uint8_t *q4k_W, const float *x,
                                    int in_dim, int out_dim,
                                    const int8_t *xq, const int16_t *bsums,
                                    const int8_t *xq_arr, float x_scale) {
    if (g_n_threads <= 1 || out_dim < 1024) {
        lal_matmul_q4_k_prepared(y, q4k_W, x, NULL, in_dim, out_dim,
                                  xq, bsums, xq_arr, x_scale);
        return;
    }
    int row_stride = (in_dim / 256) * 144;
    #pragma omp parallel num_threads(g_n_threads)
    {
        int tid = omp_get_thread_num();
        int n = omp_get_num_threads();
        int chunk = (out_dim + n - 1) / n;
        int start = tid * chunk;
        int end = start + chunk;
        if (end > out_dim) end = out_dim;
        if (start < out_dim)
            lal_matmul_q4_k_prepared(y + start, q4k_W + (size_t)start * row_stride,
                                      x, NULL, in_dim, end - start,
                                      xq, bsums, xq_arr, x_scale);
    }
}

static int8_t *g_wte_q; static float *g_wte_s;   /* Q8row embedding */
static float *g_norm_f_w;
static int8_t *g_lm_head_q;        /* Q8row lm_head */
static float *g_lm_head_s;
static float *g_x, *g_ln, *g_q, *g_k, *g_v, *g_attn_out, *g_proj, *g_mlp_out;
static float *g_logits;
static int8_t *g_xq_cache;
static float **kv_k, **kv_v;
static int g_rope_need = 0;

/* === RMSNorm / RoPE / GQA wrappers === */
static void qwen_rms_norm(float *out, const float *x, const float *w, int n) {
    lal_rms_norm_simd(out, x, w, n, RMS_EPS);
}
static float g_rope_cos[N_CTX][HEAD_DIM / 2];
static float g_rope_sin[N_CTX][HEAD_DIM / 2];
static void rope_init(void) {
    for (int p = 0; p < N_CTX; p++)
        for (int d = 0; d < HEAD_DIM / 2; d++) {
            float theta = (float)p / powf(ROPE_THETA, (float)(2 * d) / HEAD_DIM);
            g_rope_cos[p][d] = cosf(theta);
            g_rope_sin[p][d] = sinf(theta);
        }
}
static void rope_apply(float *q, float *k, int pos) {
    lal_rope_apply_simd(q, k, N_HEAD, N_KV_HEAD, HEAD_DIM, pos,
                         g_rope_cos[pos], g_rope_sin[pos]);
}
static void gqa_attn(float *out, const float *Q, const float *Kn, const float *Vn,
                     int layer, int pos) {
    lal_gqa_attn_simd(out, Q, Kn, Vn, kv_k[layer], kv_v[layer], pos,
                       N_HEAD, N_KV_HEAD, HEAD_DIM, N_Q_PER_KV, KV_DIM, N_CTX);
}

/* === GPQ8 tensor file === */
typedef struct {
    char key[128];
    int ndim, shape[4];
    int qtype;       /* 0=F32, 1=Q8row, 3=Q8_0, 5=Q4_K, 6=TERNARY */
    uint64_t data_len;
    void *data;
    int n_scale;
    float *scale;
} GPQ8Tensor;

static GPQ8Tensor *g_gp_tensors;
static int g_gp_n;

static GPQ8Tensor *gp_find(const char *key) {
    for (int i = 0; i < g_gp_n; i++)
        if (strcmp(g_gp_tensors[i].key, key) == 0) return &g_gp_tensors[i];
    fprintf(stderr, "[!] tensor not found: %s\n", key);
    return NULL;
}

static void load_gpq8(const char *path) {
    printf("[*] mmap-loading %s ...\n", path); fflush(stdout);
    g_mmap_base = win_mmap_ro(path, &g_mmap_size);
    /* warm pages: touch 1 byte per 64KB (fast, avoids inference stalls) */
    {
        volatile char sink = 0;
        const char *base = (const char *)g_mmap_base;
        for (size_t off = 0; off < g_mmap_size; off += 65536) sink += base[off];
        printf("[*] mmap warmed (%.1f MB)\n", g_mmap_size / 1048576.0); fflush(stdout);
    }
    const unsigned char *p = (const unsigned char *)g_mmap_base;
    if (memcmp(p, "GPQ8", 4) != 0) { fprintf(stderr, "[!] bad magic\n"); exit(1); }
    p += 4;
    g_gp_n = *(const int *)p; p += 4;
    printf("[*] %d tensors (%.2f GB mmap'd)\n", g_gp_n, g_mmap_size / 1073741824.0);
    fflush(stdout);
    g_gp_tensors = calloc(g_gp_n, sizeof(GPQ8Tensor));
    for (int i = 0; i < g_gp_n; i++) {
        GPQ8Tensor *t = &g_gp_tensors[i];
        int klen = *(const int *)p; p += 4;
        memcpy(t->key, p, klen); t->key[klen] = 0; p += klen;
        t->ndim = *(const int *)p; p += 4;
        for (int d = 0; d < t->ndim; d++) { t->shape[d] = *(const int *)p; p += 4; }
        t->qtype = *p; p += 1;
        t->data_len = *(const uint64_t *)p; p += 8;
        t->data = (void *)p; p += t->data_len;
        t->n_scale = *(const int *)p; p += 4;
        t->scale = (t->n_scale > 0) ? (float *)p : NULL;
        if (t->n_scale > 0) p += (size_t)t->n_scale * 4;
    }
    printf("[*] all tensors mapped\n"); fflush(stdout);
}

static float *get_f32(const char *key) {
    GPQ8Tensor *t = gp_find(key);
    if (!t || t->qtype != 0) { fprintf(stderr, "[!] %s not F32\n", key); exit(1); }
    return (float *)t->data;
}
static const uint8_t *get_packed(const char *key, int want_qtype) {
    GPQ8Tensor *t = gp_find(key);
    if (!t || t->qtype != want_qtype) {
        fprintf(stderr, "[!] %s qtype=%d expected %d\n", key, t ? t->qtype : -1, want_qtype);
        exit(1);
    }
    return (const uint8_t *)t->data;
}

/* === Layer struct: per-matrix qtype + packed pointer === */
typedef struct {
    float *norm1_w, *norm2_w;
    int tq, tk, tv, to, tgate, tup, tdown;   /* per-matrix qtype (3/5/6) */
    const uint8_t *wq, *wk, *wv, *wo, *wgate, *wup, *wdown;
} Layer;
static Layer g_layers[N_LAYER];

/* === Mixed-precision matmul dispatch === */
static void matmul_single(float *y, int qtype, const uint8_t *W,
                           const float *x, int in_dim, int out_dim) {
    if (qtype == 3) {
        parallel_matmul_q8_0_w(y, W, x, in_dim, out_dim);
        return;
    }
    /* prepare x once (shared by Q4_K and TERNARY) */
    static int8_t xq[XQ_MAX] __attribute__((aligned(32)));
    static int16_t bsums[XQ_MAX / 32] __attribute__((aligned(32)));
    static int8_t xq_arr[XQ_MAX] __attribute__((aligned(32)));
    float xs = lal_q4k_prepare_x(x, in_dim, xq, bsums, xq_arr);
    if (qtype == 5) {
        parallel_matmul_q4_k_p(y, W, x, in_dim, out_dim, xq, bsums, xq_arr, xs);
    } else if (qtype == 6) {
        parallel_matmul_ternary(y, W, xq, xs, in_dim, out_dim);
    } else {
        fprintf(stderr, "[!] bad qtype %d\n", qtype); exit(1);
    }
}

/* QKV: share one prepare across q/k/v when any is Q4_K/TERNARY */
static void matmul_qkv(Layer *L, const float *x) {
    static int8_t xq[XQ_MAX] __attribute__((aligned(32)));
    static int16_t bsums[XQ_MAX / 32] __attribute__((aligned(32)));
    static int8_t xq_arr[XQ_MAX] __attribute__((aligned(32)));
    float xs = 0;
    if (L->tq == 5 || L->tq == 6 || L->tk == 5 || L->tv == 5)
        xs = lal_q4k_prepare_x(x, N_EMBD, xq, bsums, xq_arr);
    if (L->tq == 5)
        parallel_matmul_q4_k_p(g_q, L->wq, x, N_EMBD, Q_DIM, xq, bsums, xq_arr, xs);
    else if (L->tq == 6)
        parallel_matmul_ternary(g_q, L->wq, xq, xs, N_EMBD, Q_DIM);
    else
        parallel_matmul_q8_0_w(g_q, L->wq, x, N_EMBD, Q_DIM);
    if (L->tk == 5)
        parallel_matmul_q4_k_p(g_k, L->wk, x, N_EMBD, KV_DIM, xq, bsums, xq_arr, xs);
    else if (L->tk == 6)
        parallel_matmul_ternary(g_k, L->wk, xq, xs, N_EMBD, KV_DIM);
    else
        parallel_matmul_q8_0_w(g_k, L->wk, x, N_EMBD, KV_DIM);
    if (L->tv == 5)
        parallel_matmul_q4_k_p(g_v, L->wv, x, N_EMBD, KV_DIM, xq, bsums, xq_arr, xs);
    else if (L->tv == 6)
        parallel_matmul_ternary(g_v, L->wv, xq, xs, N_EMBD, KV_DIM);
    else
        parallel_matmul_q8_0_w(g_v, L->wv, x, N_EMBD, KV_DIM);
}

/* MLP: gate/up (mixed) + SiLU + down (prepared from act) */
static void mlp_forward(Layer *L, const float *x, float *out) {
    static float gate_buf[MLP_DIM], up_buf[MLP_DIM], act_buf[MLP_DIM];
    static int8_t xq[XQ_MAX] __attribute__((aligned(32)));
    static int16_t bsums[XQ_MAX / 32] __attribute__((aligned(32)));
    static int8_t xq_arr[XQ_MAX] __attribute__((aligned(32)));
    float xs = lal_q4k_prepare_x(x, N_EMBD, xq, bsums, xq_arr);
    if (L->tgate == 5)
        parallel_matmul_q4_k_p(gate_buf, L->wgate, x, N_EMBD, MLP_DIM, xq, bsums, xq_arr, xs);
    else if (L->tgate == 6)
        parallel_matmul_ternary(gate_buf, L->wgate, xq, xs, N_EMBD, MLP_DIM);
    else
        parallel_matmul_q8_0_w(gate_buf, L->wgate, x, N_EMBD, MLP_DIM);
    if (L->tup == 5)
        parallel_matmul_q4_k_p(up_buf, L->wup, x, N_EMBD, MLP_DIM, xq, bsums, xq_arr, xs);
    else if (L->tup == 6)
        parallel_matmul_ternary(up_buf, L->wup, xq, xs, N_EMBD, MLP_DIM);
    else
        parallel_matmul_q8_0_w(up_buf, L->wup, x, N_EMBD, MLP_DIM);
    if (L->tdown == 5) {
        /* fused SiLU(gate)*up + quantize act + bsums + rearrange for Q4_K down */
        static int8_t d_xq[XQ_MAX] __attribute__((aligned(32)));
        static int16_t d_bsums[XQ_MAX / 32] __attribute__((aligned(32)));
        static int8_t d_xq_arr[XQ_MAX] __attribute__((aligned(32)));
        float ds = lal_silu_mul_prepare_simd(act_buf, gate_buf, up_buf, MLP_DIM,
                                              d_xq, d_bsums, d_xq_arr);
        parallel_matmul_q4_k_p(out, L->wdown, act_buf, MLP_DIM, N_EMBD,
                                d_xq, d_bsums, d_xq_arr, ds);
    } else if (L->tdown == 6) {
        lal_silu_mul_simd(act_buf, gate_buf, up_buf, MLP_DIM);
        static int8_t d_xq[XQ_MAX] __attribute__((aligned(32)));
        static int16_t d_bsums[XQ_MAX / 32] __attribute__((aligned(32)));
        static int8_t d_xq_arr[XQ_MAX] __attribute__((aligned(32)));
        float ds = lal_q4k_prepare_x(act_buf, MLP_DIM, d_xq, d_bsums, d_xq_arr);
        parallel_matmul_ternary(out, L->wdown, d_xq, ds, MLP_DIM, N_EMBD);
    } else {
        lal_silu_mul_simd(act_buf, gate_buf, up_buf, MLP_DIM);
        parallel_matmul_q8_0_w(out, L->wdown, act_buf, MLP_DIM, N_EMBD);
    }
}

static int forward(int tok, int pos) {
    int dbg = getenv("LAL_DEBUG") ? atoi(getenv("LAL_DEBUG")) : 0;
    if (tok < 0 || tok >= VOCAB_SIZE) tok = 0;
    /* embedding lookup (Q8row -> f32) */
    lal_dequant_row_f32(g_wte_q + (size_t)tok * N_EMBD, g_x, g_wte_s[tok], N_EMBD);
    if (dbg) {
        float s = 0; for (int i = 0; i < N_EMBD; i++) s += g_x[i] * g_x[i];
        fprintf(stderr, "[dbg] emb |x|=%.4f x0=%f\n", sqrtf(s), g_x[0]);
    }
    for (int l = 0; l < N_LAYER; l++) {
        Layer *L = &g_layers[l];
        qwen_rms_norm(g_ln, g_x, L->norm1_w, N_EMBD);
        matmul_qkv(L, g_ln);
        if (dbg && l < 2) {
            float s = 0; for (int i = 0; i < Q_DIM; i++) s += g_q[i] * g_q[i];
            fprintf(stderr, "[dbg] L%d q|.|=%.4f q0=%f\n", l, sqrtf(s), g_q[0]);
        }
        rope_apply(g_q, g_k, pos);
        gqa_attn(g_attn_out, g_q, g_k, g_v, l, pos);
        matmul_single(g_proj, L->to, L->wo, g_attn_out, Q_DIM, N_EMBD);
        lal_residual_add_simd(g_x, g_proj, N_EMBD);
        if (dbg && l < 2) {
            float s = 0; for (int i = 0; i < N_EMBD; i++) s += g_x[i] * g_x[i];
            fprintf(stderr, "[dbg] L%d post-attn |x|=%.4f\n", l, sqrtf(s));
        }
        qwen_rms_norm(g_ln, g_x, L->norm2_w, N_EMBD);
        mlp_forward(L, g_ln, g_mlp_out);
        if (dbg && l < 2) {
            float s = 0; for (int i = 0; i < N_EMBD; i++) s += g_mlp_out[i] * g_mlp_out[i];
            fprintf(stderr, "[dbg] L%d mlp|.|=%.4f m0=%f\n", l, sqrtf(s), g_mlp_out[0]);
        }
        lal_residual_add_simd(g_x, g_mlp_out, N_EMBD);
    }
    qwen_rms_norm(g_ln, g_x, g_norm_f_w, N_EMBD);
    if (dbg) {
        float s = 0; for (int i = 0; i < N_EMBD; i++) s += g_ln[i] * g_ln[i];
        fprintf(stderr, "[dbg] final |ln|=%.4f ln0=%f\n", sqrtf(s), g_ln[0]);
    }
    /* int8 LM head with abs-xq trick, parallel over vocab */
    float scale_x = lal_quantize_x_int8(g_ln, g_xq_cache, N_EMBD);
    static uint8_t abs_xq[N_EMBD] __attribute__((aligned(32)));
    lal_compute_abs_xq(g_xq_cache, abs_xq, N_EMBD);
    #pragma omp parallel num_threads(g_n_threads)
    {
        int tid = omp_get_thread_num();
        int n = omp_get_num_threads();
        int v_per = (VOCAB_SIZE + n - 1) / n;
        int v_start = tid * v_per;
        int v_end = v_start + v_per;
        if (v_end > VOCAB_SIZE) v_end = VOCAB_SIZE;
        if (v_start < VOCAB_SIZE)
            lal_lm_head_int8_range_abs(g_logits, g_xq_cache, abs_xq, scale_x,
                                        g_lm_head_q, g_lm_head_s,
                                        v_start, v_end, N_EMBD);
    }
    /* sample */
    if (dbg) {
        int nan = 0;
        for (int v = 0; v < VOCAB_SIZE; v++) {
            float x = g_logits[v];
            if (!(x == x) || x > 1e30f || x < -1e30f) nan++;
        }
        fprintf(stderr, "[dbg] logits nan/inf=%d L[0..7]=%.3f %.3f %.3f %.3f %.3f %.3f %.3f %.3f\n",
                nan, g_logits[0], g_logits[1], g_logits[2], g_logits[3],
                g_logits[4], g_logits[5], g_logits[6], g_logits[7]);
    }
    int next = lal_sample_token(g_logits, VOCAB_SIZE, g_temperature, g_top_k,
                                 g_rep_penalty, g_recent, g_n_recent);
    if (g_n_recent < 256) g_recent[g_n_recent++] = next;
    else { memmove(g_recent, g_recent + 1, 255 * sizeof(int)); g_recent[255] = next; }
    return next;
}

/* === Tokenizer: vocab hash + greedy longest-match over byte-level BPE space === */
typedef struct { char key[512]; int id; } TEntry;
#define TOK_HASH_BITS 18
#define TOK_HASH_SIZE (1 << TOK_HASH_BITS)
static TEntry g_htab[TOK_HASH_SIZE];
static char **g_vocab_str;
static int g_vocab_str_n;

static unsigned tok_hash(const char *s, int len) {
    unsigned h = 2166136261u;
    for (int i = 0; i < len; i++) { h ^= (unsigned char)s[i]; h *= 16777619u; }
    return h & (TOK_HASH_SIZE - 1);
}
static void tins(const char *key, int id) {
    unsigned h = tok_hash(key, (int)strlen(key));
    while (g_htab[h].key[0]) h = (h + 1) & (TOK_HASH_SIZE - 1);
    strncpy(g_htab[h].key, key, 511);
    g_htab[h].key[511] = 0;
    g_htab[h].id = id;
}
static int tok_find(const char *key) {
    unsigned h = tok_hash(key, (int)strlen(key));
    while (g_htab[h].key[0]) {
        if (strcmp(g_htab[h].key, key) == 0) return g_htab[h].id;
        h = (h + 1) & (TOK_HASH_SIZE - 1);
    }
    return -1;
}

static void load_tokenizer(const char *dir) {
    /* vocab.tsv: "id<TAB>token" lines, generated from tokenizer.json
     * (model.vocab + added_tokens merged; token escaped: \\ \t \n) */
    char path[1024];
    snprintf(path, sizeof(path), "%s/vocab.tsv", dir);
    FILE *f = fopen(path, "rb");
    if (!f) { fprintf(stderr, "[!] cannot open %s\n", path); exit(1); }
    int mx = VOCAB_SIZE + 200;
    g_vocab_str = calloc(mx, sizeof(char *));
    char line[1024];
    int cnt = 0;
    while (fgets(line, sizeof line, f)) {
        char *tab = strchr(line, '\t');
        if (!tab) continue;
        *tab = 0;
        int id = atoi(line);
        char *tok = tab + 1;
        int len = (int)strlen(tok);
        while (len && (tok[len - 1] == '\n' || tok[len - 1] == '\r')) tok[--len] = 0;
        char real[512];
        int r = 0;
        for (int i = 0; i < len && r < 511; i++) {
            if (tok[i] == '\\' && i + 1 < len) {
                i++;
                if (tok[i] == 't') real[r++] = '\t';
                else if (tok[i] == 'n') real[r++] = '\n';
                else if (tok[i] == 'r') real[r++] = '\r';
                else real[r++] = tok[i];
            } else real[r++] = tok[i];
        }
        real[r] = 0;
        if (id >= 0 && id < mx) {
            g_vocab_str[id] = strdup(real);
            tins(real, id);
            if (id + 1 > g_vocab_str_n) g_vocab_str_n = id + 1;
            cnt++;
        }
    }
    fclose(f);
    printf("[*] tokenizer: %d tokens (max id %d)\n", cnt, g_vocab_str_n - 1);
}

/* convert raw utf-8 text to byte-level BPE "magic" string */
static int text_to_magic(const char *text, char *out, int cap) {
    int n = 0;
    for (const char *p = text; *p; p++) {
        unsigned cp = lal_bpe_cp_for_byte((unsigned char)*p);
        char tmp[5];
        int m = lal_utf8_encode(cp, tmp);
        for (int i = 0; i < m && n < cap - 1; i++) out[n++] = tmp[i];
    }
    out[n] = 0;
    return n;
}

/* greedy longest-match encoding over the magic string */
static int *encode_text(const char *text, int *n_out) {
    int cap = (int)strlen(text) * 4 + 8;
    char *magic = malloc(cap);
    int mlen = text_to_magic(text, magic, cap);
    int *ids = malloc((mlen + 16) * sizeof(int));
    int n = 0, pos = 0;
    while (pos < mlen) {
        int max_len = mlen - pos;
        if (max_len > 48) max_len = 48;
        int found = -1;
        char tmp[64];
        for (int len = max_len; len >= 1; len--) {
            memcpy(tmp, magic + pos, len);
            tmp[len] = 0;
            found = tok_find(tmp);
            if (found >= 0) { pos += len; break; }
        }
        if (found < 0) { /* single magic char should always match */
            fprintf(stderr, "[!] encode stuck at %d\n", pos);
            pos++;
            continue;
        }
        ids[n++] = found;
    }
    free(magic);
    *n_out = n;
    return ids;
}

static void decode_token(int id, char *out, int maxlen) {
    if (id < 0 || id >= g_vocab_str_n || !g_vocab_str[id]) { out[0] = 0; return; }
    lal_decode_bpe_token(g_vocab_str[id], out, maxlen);
}

/* === Main === */
int main(int argc, char **argv) {
    srand((unsigned)time(NULL));

    const char *weights = "C:/models/minicpm_mixed.bin";
    const char *tokdir = "C:/models/MiniCPM5-2B-cpu";
    const char *prompt = NULL;
    const char *prompt_file = NULL;
    int n_gen = 48;
    for (int i = 1; i < argc; i++) {
        if (!strcmp(argv[i], "--weights") && i + 1 < argc) weights = argv[++i];
        else if (!strcmp(argv[i], "--tokenizer") && i + 1 < argc) tokdir = argv[++i];
        else if (!strcmp(argv[i], "--prompt") && i + 1 < argc) prompt = argv[++i];
        else if (!strcmp(argv[i], "--prompt-file") && i + 1 < argc) prompt_file = argv[++i];
        else if (!strcmp(argv[i], "--n") && i + 1 < argc) n_gen = atoi(argv[++i]);
        else if (!strcmp(argv[i], "--threads") && i + 1 < argc) g_n_threads = atoi(argv[++i]);
        else if (!strcmp(argv[i], "--temp") && i + 1 < argc) g_temperature = (float)atof(argv[++i]);
        else if (!strcmp(argv[i], "--top-k") && i + 1 < argc) g_top_k = atoi(argv[++i]);
        else if (!strcmp(argv[i], "--rep-penalty") && i + 1 < argc) g_rep_penalty = (float)atof(argv[++i]);
        else { fprintf(stderr, "[?] unknown arg %s\n", argv[i]); }
    }
    static char prompt_buf[4096];
    if (prompt_file) {
        FILE *pf = fopen(prompt_file, "rb");
        if (!pf) { fprintf(stderr, "[!] cannot open %s\n", prompt_file); return 1; }
        size_t pn = fread(prompt_buf, 1, sizeof(prompt_buf) - 1, pf);
        while (pn && (prompt_buf[pn - 1] == '\n' || prompt_buf[pn - 1] == '\r')) pn--;
        prompt_buf[pn] = 0;
        fclose(pf);
        prompt = prompt_buf;
    }
    if (!prompt) prompt = "Hello";

    printf("=== MiniCPM5-2B (LAL, mixed-precision: Q4_K body + Q8_0 edges + TERNARY gate/up) ===\n");
    printf("[*] %d layers, %d hidden, %dQ/%dKV heads, %d head_dim, %d MLP, %d vocab\n",
           N_LAYER, N_EMBD, N_HEAD, N_KV_HEAD, HEAD_DIM, MLP_DIM, VOCAB_SIZE);
    printf("[*] threads=%d temp=%.2f top_k=%d\n", g_n_threads, g_temperature, g_top_k);

    bit_lut_init();
    load_gpq8(weights);

    /* wire global tensors */
    g_norm_f_w = get_f32("model.norm.weight");
    {
        GPQ8Tensor *t = gp_find("model.embed_tokens.weight");
        if (!t || t->qtype != 1 || t->n_scale != VOCAB_SIZE) {
            fprintf(stderr, "[!] embed_tokens Q8row mismatch\n"); exit(1);
        }
        g_wte_q = (int8_t *)t->data;
        g_wte_s = t->scale;
    }
    {
        GPQ8Tensor *t = gp_find("lm_head.weight");
        if (!t || t->qtype != 1 || t->n_scale != VOCAB_SIZE) {
            fprintf(stderr, "[!] lm_head Q8row mismatch\n"); exit(1);
        }
        g_lm_head_q = (int8_t *)t->data;
        g_lm_head_s = t->scale;
    }

    /* wire layers (per-matrix mixed precision) */
    char key[256];
    int cnt_q8_0 = 0, cnt_q4k = 0, cnt_tern = 0;
    for (int l = 0; l < N_LAYER; l++) {
        Layer *L = &g_layers[l];
        snprintf(key, sizeof(key), "model.layers.%d.input_layernorm.weight", l);
        L->norm1_w = get_f32(key);
        snprintf(key, sizeof(key), "model.layers.%d.post_attention_layernorm.weight", l);
        L->norm2_w = get_f32(key);
        struct { const char *name; int *t; const uint8_t **w; } mats[7] = {
            {"self_attn.q_proj.weight",  &L->tq,    &L->wq},
            {"self_attn.k_proj.weight",  &L->tk,    &L->wk},
            {"self_attn.v_proj.weight",  &L->tv,    &L->wv},
            {"self_attn.o_proj.weight",  &L->to,    &L->wo},
            {"mlp.gate_proj.weight",     &L->tgate, &L->wgate},
            {"mlp.up_proj.weight",       &L->tup,   &L->wup},
            {"mlp.down_proj.weight",     &L->tdown, &L->wdown},
        };
        for (int m = 0; m < 7; m++) {
            snprintf(key, sizeof(key), "model.layers.%d.%s", l, mats[m].name);
            GPQ8Tensor *t = gp_find(key);
            if (!t) exit(1);
            *mats[m].t = t->qtype;
            *mats[m].w = (const uint8_t *)t->data;
            if (t->qtype == 3) cnt_q8_0++;
            else if (t->qtype == 5) cnt_q4k++;
            else if (t->qtype == 6) cnt_tern++;
        }
    }
    printf("[*] layer matrices: %d x Q8_0 (8-bit edges), %d x Q4_K (4-bit), %d x TERNARY (3-value)\n",
           cnt_q8_0, cnt_q4k, cnt_tern);

    /* working buffers */
    g_x = memalign(32, N_EMBD * sizeof(float));
    g_ln = memalign(32, N_EMBD * sizeof(float));
    g_q = memalign(32, Q_DIM * sizeof(float));
    g_k = memalign(32, KV_DIM * sizeof(float));
    g_v = memalign(32, KV_DIM * sizeof(float));
    g_attn_out = memalign(32, Q_DIM * sizeof(float));
    g_proj = memalign(32, N_EMBD * sizeof(float));
    g_mlp_out = memalign(32, N_EMBD * sizeof(float));
    g_logits = memalign(32, VOCAB_SIZE * sizeof(float));
    g_xq_cache = memalign(32, N_EMBD);

    /* KV cache: 42 x 2048 x 256 x 4B x 2 = 176 MB */
    kv_k = malloc(N_LAYER * sizeof(float *));
    kv_v = malloc(N_LAYER * sizeof(float *));
    for (int l = 0; l < N_LAYER; l++) {
        kv_k[l] = memalign(32, (size_t)N_CTX * KV_DIM * sizeof(float));
        kv_v[l] = memalign(32, (size_t)N_CTX * KV_DIM * sizeof(float));
    }
    printf("[*] KV cache: %.0f MB\n", (double)N_LAYER * N_CTX * KV_DIM * 4 * 2 / 1048576);
    rope_init();
    load_tokenizer(tokdir);

    /* special token ids (resolved from vocab) */
    int tok_bos = tok_find("<s>");
    int tok_im_start = tok_find("<|im_start|>");
    int tok_im_end = tok_find("<|im_end|>");
    printf("[*] special tokens: <s>=%d <|im_start|>=%d <|im_end|>=%d\n",
           tok_bos, tok_im_start, tok_im_end);

    /* build chat prompt (MiniCPM5 template, thinking disabled):
     * <s><|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n */
    int *pids = malloc((strlen(prompt) * 4 + 128) * sizeof(int));
    int n_prompt = 0;
    pids[n_prompt++] = tok_bos;
    int nt;
    int *t;
    t = encode_text("<|im_start|>user\n", &nt);
    for (int i = 0; i < nt; i++) pids[n_prompt++] = t[i];
    free(t);
    t = encode_text(prompt, &nt);
    for (int i = 0; i < nt; i++) pids[n_prompt++] = t[i];
    free(t);
    /* debug: print prompt token ids + text */
    {
        int echo = getenv("ECHO_TOKENS") ? atoi(getenv("ECHO_TOKENS")) : 0;
        if (echo) {
            int base = 1; /* skip bos for display mapping */
            fprintf(stderr, "[echo] prompt pieces:\n");
            char buf[256];
            for (int i = 1; i < n_prompt; i++) {
                decode_token(pids[i], buf, (int)sizeof(buf));
                fprintf(stderr, "  %d -> id=%d text=[%s]\n", i, pids[i], buf);
            }
            (void)base;
        }
    }
    t = encode_text("<|im_end|>\n", &nt);
    for (int i = 0; i < nt; i++) pids[n_prompt++] = t[i];
    free(t);
    t = encode_text("<|im_start|>assistant\n<think>\n\n</think>\n\n", &nt);
    for (int i = 0; i < nt; i++) pids[n_prompt++] = t[i];
    free(t);
    printf("[*] prompt: %d tokens (incl. chat template)\n", n_prompt);
    printf("[*] generating %d tokens...\n\n", n_gen);

    double t0 = now_sec();
    int pos = 0, next = -1;
    for (int i = 0; i < n_prompt; i++) {
        next = forward(pids[i], pos);
        pos++;
        if (pos >= N_CTX) break;
    }
    double t_prefill = now_sec() - t0;
    int gen_count = 0;
    char out_buf[65536] = {0};
    int opos = 0;
    double t_gen0 = now_sec();
    for (int g = 0; g < n_gen && pos < N_CTX; g++) {
        char ts[256];
        if (next >= 0 && next < g_vocab_str_n && g_vocab_str[next])
            lal_decode_bpe_token(g_vocab_str[next], ts, (int)sizeof(ts));
        else ts[0] = 0;
        int slen = (int)strlen(ts);
        if (opos + slen < (int)sizeof(out_buf) - 1) {
            memcpy(out_buf + opos, ts, slen);
            opos += slen;
        }
        if (next == tok_im_end || next == 1) break;  /* EOS */
        printf("%s", ts);
        fflush(stdout);
        next = forward(next, pos);
        pos++;
        gen_count++;
    }
    out_buf[opos] = 0;
    double t_gen = now_sec() - t_gen0;
    printf("\n\n[*] prefill: %d tokens in %.2fs (%.0f tok/s)\n",
           n_prompt, t_prefill, n_prompt / (t_prefill + 1e-9));
    printf("[*] decode: %d tokens in %.2fs (%.1f tok/s, %d threads)\n",
           gen_count, t_gen, gen_count / (t_gen + 1e-9), g_n_threads);
    printf("[*] output: %s\n", out_buf);
    free(pids);
    return 0;
}