/*
 * UNARY TRANSFORMER ENGINE v2 - Configurable dimensions
 * 
 * Full Qwen2/Qwen3 forward pass in C with AVX-512 + OpenMP.
 * Supports any model size via runtime config.
 *
 * (c) 2026 OpenTransformers Ltd / Scott Bisset
 */

#include <immintrin.h>
#include <omp.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdio.h>
#include <time.h>

#define MAX_SEQ     4096
#define RMS_EPS     1e-6f

/* ============================================================
 * Config - set at init time
 * ============================================================ */
typedef struct {
    int hidden;
    int inter;
    int n_heads;
    int n_kv_heads;
    int head_dim;
    int n_layers;
    int vocab;
    float rope_theta;
    int has_attn_bias;    /* 1 for Qwen2 (1.5B), 0 for Qwen3 (4B) */
    int tie_embeddings;   /* 1 if lm_head shares embed weights */
} Config;

/* ============================================================
 * Unary linear layer
 * ============================================================ */
typedef struct {
    uint64_t *sign_bits;
    uint64_t *mag_planes;
    float    *scales;
    float    *bias;
    int       out_dim;
    int       in_dim;
    int       n_planes;
} UnaryLinear;

/* FP16 linear (for lm_head when not tied) */
typedef struct {
    uint16_t *weight;
    int       out_dim;
    int       in_dim;
} FP16Linear;

/* ============================================================
 * Transformer layer
 * ============================================================ */
typedef struct {
    UnaryLinear q_proj, k_proj, v_proj, o_proj;
    UnaryLinear gate_proj, up_proj, down_proj;
    float *input_norm;
    float *post_norm;
    float *q_bias, *k_bias, *v_bias;
    float *q_norm, *k_norm;   /* QK-Norm (Qwen3) */
} Layer;

/* ============================================================
 * Full model
 * ============================================================ */
typedef struct {
    Config cfg;
    
    uint16_t *embed;        /* FP16 embeddings */
    Layer    *layers;       /* Dynamic array */
    float    *final_norm;
    FP16Linear lm_head;     /* Only used if !tie_embeddings */
    
    /* KV cache */
    float *k_cache;
    float *v_cache;
    
    /* Scratch buffers */
    float *hidden;
    float *hidden2;
    float *q;
    float *k;
    float *v;
    float *attn_out;
    float *gate;
    float *up;
    float *down_in;
    float *logits;
    float *attn_scores;
    
    int n_planes;
} Model;

/* ============================================================
 * AVX-512 Unary matvec: y = W @ x
 * ============================================================ */
static void unary_matvec(
    const UnaryLinear *layer, const float *x, float *y
) {
    int out_dim = layer->out_dim;
    int in_dim = layer->in_dim;
    int n_planes = layer->n_planes;
    int chunks = (in_dim + 63) / 64;
    int in_padded = (in_dim + 15) & ~15;

    #pragma omp parallel for schedule(dynamic, 64)
    for (int i = 0; i < out_dim; i++) {
        const uint64_t *row_sign = layer->sign_bits + (size_t)i * chunks;
        float total = 0.0f;

        /* Aligned local copy of input for this thread */
        float x_local[in_padded] __attribute__((aligned(64)));
        memcpy(x_local, x, in_dim * sizeof(float));
        if (in_padded > in_dim)
            memset(x_local + in_dim, 0, (in_padded - in_dim) * sizeof(float));

        for (int p = 0; p < n_planes; p++) {
            const uint64_t *plane_row = layer->mag_planes +
                ((size_t)p * out_dim + i) * chunks;
            __m512 acc = _mm512_setzero_ps();

            for (int c = 0; c < chunks; c++) {
                uint64_t mbits = plane_row[c];
                uint64_t sbits = row_sign[c];
                uint64_t pos_bits = mbits & ~sbits;
                uint64_t neg_bits = mbits & sbits;

                for (int g = 0; g < 4 && (c * 64 + g * 16) < in_padded; g++) {
                    int offset = c * 64 + g * 16;
                    __m512 xv = _mm512_load_ps(x_local + offset);
                    __mmask16 pmask = (__mmask16)((pos_bits >> (g * 16)) & 0xFFFF);
                    __mmask16 nmask = (__mmask16)((neg_bits >> (g * 16)) & 0xFFFF);
                    acc = _mm512_mask_add_ps(acc, pmask, acc, xv);
                    acc = _mm512_mask_sub_ps(acc, nmask, acc, xv);
                }
            }
            total += _mm512_reduce_add_ps(acc);
        }
        y[i] = total * layer->scales[i];
        if (layer->bias) y[i] += layer->bias[i];
    }
}

/* FP16 matvec for lm_head */
static void fp16_matvec(const FP16Linear *layer, const float *x, float *y) {
    int out_dim = layer->out_dim;
    int in_dim = layer->in_dim;
    const uint16_t *w = layer->weight;

    #pragma omp parallel for schedule(dynamic, 256)
    for (int i = 0; i < out_dim; i++) {
        __m512 acc = _mm512_setzero_ps();
        int j;
        for (j = 0; j + 16 <= in_dim; j += 16) {
            __m256i h = _mm256_loadu_si256((__m256i*)(w + (size_t)i * in_dim + j));
            __m512 wv = _mm512_cvtph_ps(h);
            __m512 xv = _mm512_loadu_ps(x + j);
            acc = _mm512_fmadd_ps(wv, xv, acc);
        }
        float sum = _mm512_reduce_add_ps(acc);
        for (; j < in_dim; j++) {
            __m128i hv = _mm_set1_epi16(w[(size_t)i * in_dim + j]);
            __m128 fv = _mm_cvtph_ps(hv);
            float wf;
            _mm_store_ss(&wf, fv);
            sum += wf * x[j];
        }
        y[i] = sum;
    }
}

/* ============================================================
 * Basic ops - all AVX-512 vectorized
 * ============================================================ */

static void rmsnorm(const float *x, const float *weight, float *y, int dim) {
    __m512 sum_sq = _mm512_setzero_ps();
    int i;
    for (i = 0; i + 16 <= dim; i += 16) {
        __m512 xv = _mm512_loadu_ps(x + i);
        sum_sq = _mm512_fmadd_ps(xv, xv, sum_sq);
    }
    float ss = _mm512_reduce_add_ps(sum_sq);
    for (; i < dim; i++) ss += x[i] * x[i];
    float rms = 1.0f / sqrtf(ss / dim + RMS_EPS);

    for (i = 0; i + 16 <= dim; i += 16) {
        __m512 xv = _mm512_loadu_ps(x + i);
        __m512 wv = _mm512_loadu_ps(weight + i);
        __m512 rv = _mm512_set1_ps(rms);
        _mm512_storeu_ps(y + i, _mm512_mul_ps(_mm512_mul_ps(xv, rv), wv));
    }
    for (; i < dim; i++) y[i] = x[i] * rms * weight[i];
}

static void silu_inplace(float *x, int n) {
    int i;
    for (i = 0; i + 16 <= n; i += 16) {
        __m512 xv = _mm512_loadu_ps(x + i);
        __m512 neg = _mm512_sub_ps(_mm512_setzero_ps(), xv);
        /* exp(-x) approximation not great with AVX, use scalar */
        float tmp[16];
        _mm512_storeu_ps(tmp, xv);
        for (int j = 0; j < 16; j++)
            tmp[j] = tmp[j] / (1.0f + expf(-tmp[j]));
        _mm512_storeu_ps(x + i, _mm512_loadu_ps(tmp));
    }
    for (; i < n; i++)
        x[i] = x[i] / (1.0f + expf(-x[i]));
}

static void elemwise_mul(const float *a, const float *b, float *c, int n) {
    int i;
    for (i = 0; i + 16 <= n; i += 16) {
        __m512 av = _mm512_loadu_ps(a + i);
        __m512 bv = _mm512_loadu_ps(b + i);
        _mm512_storeu_ps(c + i, _mm512_mul_ps(av, bv));
    }
    for (; i < n; i++) c[i] = a[i] * b[i];
}

static void vec_add(float *y, const float *x, int n) {
    int i;
    for (i = 0; i + 16 <= n; i += 16) {
        __m512 yv = _mm512_loadu_ps(y + i);
        __m512 xv = _mm512_loadu_ps(x + i);
        _mm512_storeu_ps(y + i, _mm512_add_ps(yv, xv));
    }
    for (; i < n; i++) y[i] += x[i];
}

static void apply_rope(float *vec, int pos, int dim, float theta) {
    for (int i = 0; i < dim; i += 2) {
        float freq = 1.0f / powf(theta, (float)i / dim);
        float angle = pos * freq;
        float cos_a = cosf(angle);
        float sin_a = sinf(angle);
        float v0 = vec[i];
        float v1 = vec[i + 1];
        vec[i]     = v0 * cos_a - v1 * sin_a;
        vec[i + 1] = v0 * sin_a + v1 * cos_a;
    }
}

static void softmax(float *x, int n) {
    float max_val = x[0];
    for (int i = 1; i < n; i++) if (x[i] > max_val) max_val = x[i];
    float sum = 0.0f;
    for (int i = 0; i < n; i++) { x[i] = expf(x[i] - max_val); sum += x[i]; }
    float inv = 1.0f / sum;
    for (int i = 0; i < n; i++) x[i] *= inv;
}

/* ============================================================
 * Embedding lookup (FP16 -> FP32)
 * ============================================================ */
static void embed_token(const Model *m, int token_id, float *out) {
    int hidden = m->cfg.hidden;
    const uint16_t *row = m->embed + (size_t)token_id * hidden;
    int i;
    for (i = 0; i + 16 <= hidden; i += 16) {
        __m256i h = _mm256_loadu_si256((__m256i*)(row + i));
        __m512 fv = _mm512_cvtph_ps(h);
        _mm512_storeu_ps(out + i, fv);
    }
    for (; i < hidden; i++) {
        __m128i hv = _mm_set1_epi16(row[i]);
        __m128 fv = _mm_cvtph_ps(hv);
        _mm_store_ss(out + i, fv);
    }
}

/* KV cache helpers */
static float* kv_ptr(float *cache, const Config *c, int layer, int pos, int kv_head) {
    return cache + ((size_t)layer * MAX_SEQ * c->n_kv_heads +
                    (size_t)pos * c->n_kv_heads + kv_head) * c->head_dim;
}

/* ============================================================
 * ATTENTION
 * ============================================================ */
static void attention(Model *m, int layer_idx, int pos) {
    Config *c = &m->cfg;
    Layer *layer = &m->layers[layer_idx];
    int heads_per_kv = c->n_heads / c->n_kv_heads;
    
    unary_matvec(&layer->q_proj, m->hidden2, m->q);
    unary_matvec(&layer->k_proj, m->hidden2, m->k);
    unary_matvec(&layer->v_proj, m->hidden2, m->v);
    
    if (c->has_attn_bias) {
        if (layer->q_bias) vec_add(m->q, layer->q_bias, c->n_heads * c->head_dim);
        if (layer->k_bias) vec_add(m->k, layer->k_bias, c->n_kv_heads * c->head_dim);
        if (layer->v_bias) vec_add(m->v, layer->v_bias, c->n_kv_heads * c->head_dim);
    }
    
    /* QK-Norm (Qwen3): RMSNorm each head's Q and K before RoPE */
    if (layer->q_norm) {
        for (int h = 0; h < c->n_heads; h++)
            rmsnorm(m->q + h * c->head_dim, layer->q_norm, m->q + h * c->head_dim, c->head_dim);
    }
    if (layer->k_norm) {
        for (int h = 0; h < c->n_kv_heads; h++)
            rmsnorm(m->k + h * c->head_dim, layer->k_norm, m->k + h * c->head_dim, c->head_dim);
    }
    
    for (int h = 0; h < c->n_heads; h++)
        apply_rope(m->q + h * c->head_dim, pos, c->head_dim, c->rope_theta);
    for (int h = 0; h < c->n_kv_heads; h++)
        apply_rope(m->k + h * c->head_dim, pos, c->head_dim, c->rope_theta);
    
    for (int h = 0; h < c->n_kv_heads; h++) {
        memcpy(kv_ptr(m->k_cache, c, layer_idx, pos, h),
               m->k + h * c->head_dim, c->head_dim * sizeof(float));
        memcpy(kv_ptr(m->v_cache, c, layer_idx, pos, h),
               m->v + h * c->head_dim, c->head_dim * sizeof(float));
    }
    
    float scale = 1.0f / sqrtf((float)c->head_dim);
    memset(m->attn_out, 0, c->n_heads * c->head_dim * sizeof(float));
    
    for (int h = 0; h < c->n_heads; h++) {
        int kv_h = h / heads_per_kv;
        float *q_head = m->q + h * c->head_dim;
        float *out_head = m->attn_out + h * c->head_dim;
        
        for (int t = 0; t <= pos; t++) {
            float *k_cached = kv_ptr(m->k_cache, c, layer_idx, t, kv_h);
            __m512 acc = _mm512_setzero_ps();
            int d;
            for (d = 0; d + 16 <= c->head_dim; d += 16) {
                __m512 qv = _mm512_loadu_ps(q_head + d);
                __m512 kv = _mm512_loadu_ps(k_cached + d);
                acc = _mm512_fmadd_ps(qv, kv, acc);
            }
            float dot = _mm512_reduce_add_ps(acc);
            for (; d < c->head_dim; d++) dot += q_head[d] * k_cached[d];
            m->attn_scores[t] = dot * scale;
        }
        
        softmax(m->attn_scores, pos + 1);
        
        for (int t = 0; t <= pos; t++) {
            float w = m->attn_scores[t];
            if (w < 1e-8f) continue;
            float *v_cached = kv_ptr(m->v_cache, c, layer_idx, t, kv_h);
            __m512 wv = _mm512_set1_ps(w);
            int d;
            for (d = 0; d + 16 <= c->head_dim; d += 16) {
                __m512 ov = _mm512_loadu_ps(out_head + d);
                __m512 vv = _mm512_loadu_ps(v_cached + d);
                _mm512_storeu_ps(out_head + d, _mm512_fmadd_ps(wv, vv, ov));
            }
            for (; d < c->head_dim; d++) out_head[d] += w * v_cached[d];
        }
    }
    
    unary_matvec(&layer->o_proj, m->attn_out, m->hidden2);
}

/* ============================================================
 * MLP - SwiGLU
 * ============================================================ */
static void mlp(Model *m, int layer_idx) {
    Layer *layer = &m->layers[layer_idx];
    int inter = m->cfg.inter;
    
    unary_matvec(&layer->gate_proj, m->hidden2, m->gate);
    unary_matvec(&layer->up_proj, m->hidden2, m->up);
    
    silu_inplace(m->gate, inter);
    elemwise_mul(m->gate, m->up, m->down_in, inter);
    
    unary_matvec(&layer->down_proj, m->down_in, m->hidden2);
}

/* ============================================================
 * FORWARD ONE TOKEN
 * ============================================================ */
float* forward_token(Model *m, int token_id, int pos) {
    Config *c = &m->cfg;
    
    embed_token(m, token_id, m->hidden);
    
    for (int l = 0; l < c->n_layers; l++) {
        rmsnorm(m->hidden, m->layers[l].input_norm, m->hidden2, c->hidden);
        attention(m, l, pos);
        vec_add(m->hidden, m->hidden2, c->hidden);
        rmsnorm(m->hidden, m->layers[l].post_norm, m->hidden2, c->hidden);
        mlp(m, l);
        vec_add(m->hidden, m->hidden2, c->hidden);
    }
    
    rmsnorm(m->hidden, m->final_norm, m->hidden2, c->hidden);
    
    /* LM head - either tied embeddings or separate FP16 */
    if (c->tie_embeddings) {
        /* Use embed weights as lm_head (FP16 matvec) */
        FP16Linear tied;
        tied.weight = m->embed;
        tied.out_dim = c->vocab;
        tied.in_dim = c->hidden;
        fp16_matvec(&tied, m->hidden2, m->logits);
    } else {
        fp16_matvec(&m->lm_head, m->hidden2, m->logits);
    }
    
    return m->logits;
}

/* ============================================================
 * TOP-P SAMPLING
 * ============================================================ */
static int sample_top_p(float *logits, int vocab, float temperature, float top_p) {
    if (temperature > 0) {
        float inv_t = 1.0f / temperature;
        for (int i = 0; i < vocab; i++) logits[i] *= inv_t;
    }
    softmax(logits, vocab);
    
    float *probs = (float *)malloc(vocab * sizeof(float));
    int *indices = (int *)malloc(vocab * sizeof(int));
    memcpy(probs, logits, vocab * sizeof(float));
    for (int i = 0; i < vocab; i++) indices[i] = i;
    
    int n_keep = 0;
    float cum = 0.0f;
    while (cum < top_p && n_keep < vocab) {
        int best = n_keep;
        for (int i = n_keep + 1; i < vocab; i++)
            if (probs[i] > probs[best]) best = i;
        float tmp_p = probs[n_keep]; probs[n_keep] = probs[best]; probs[best] = tmp_p;
        int tmp_i = indices[n_keep]; indices[n_keep] = indices[best]; indices[best] = tmp_i;
        cum += probs[n_keep];
        n_keep++;
        if (n_keep >= 40) break;
    }
    
    float sum = 0.0f;
    for (int i = 0; i < n_keep; i++) sum += probs[i];
    float r = (float)rand() / RAND_MAX * sum;
    float acc = 0.0f;
    int chosen = indices[0];
    for (int i = 0; i < n_keep; i++) {
        acc += probs[i];
        if (acc >= r) { chosen = indices[i]; break; }
    }
    
    free(probs);
    free(indices);
    return chosen;
}

/* ============================================================
 * GENERATE
 * ============================================================ */
int generate(
    Model *m,
    const int *prompt_ids, int prompt_len,
    int *out_tokens, int max_new_tokens,
    float temperature, float top_p,
    int eos_token
) {
    srand(time(NULL));
    
    for (int i = 0; i < prompt_len; i++) {
        forward_token(m, prompt_ids[i], i);
    }
    
    int pos = prompt_len;
    int generated = 0;
    
    for (int t = 0; t < max_new_tokens; t++) {
        float *logits = m->logits;
        
        int next_token;
        if (temperature <= 0) {
            next_token = 0;
            for (int i = 1; i < m->cfg.vocab; i++)
                if (logits[i] > logits[next_token]) next_token = i;
        } else {
            next_token = sample_top_p(logits, m->cfg.vocab, temperature, top_p);
        }
        
        out_tokens[t] = next_token;
        generated++;
        
        if (next_token == eos_token) break;
        
        forward_token(m, next_token, pos);
        pos++;
    }
    
    return generated;
}

/* ============================================================
 * MODEL ALLOCATION with config
 * ============================================================ */
Model* model_alloc(
    int n_planes,
    int hidden, int inter, int n_heads, int n_kv_heads,
    int head_dim, int n_layers, int vocab,
    float rope_theta, int has_attn_bias, int tie_embeddings
) {
    Model *m = (Model *)calloc(1, sizeof(Model));
    m->n_planes = n_planes;
    
    Config *c = &m->cfg;
    c->hidden = hidden;
    c->inter = inter;
    c->n_heads = n_heads;
    c->n_kv_heads = n_kv_heads;
    c->head_dim = head_dim;
    c->n_layers = n_layers;
    c->vocab = vocab;
    c->rope_theta = rope_theta;
    c->has_attn_bias = has_attn_bias;
    c->tie_embeddings = tie_embeddings;
    
    m->layers = (Layer *)calloc(n_layers, sizeof(Layer));
    
    size_t kv_size = (size_t)n_layers * MAX_SEQ * n_kv_heads * head_dim;
    m->k_cache = (float *)calloc(kv_size, sizeof(float));
    m->v_cache = (float *)calloc(kv_size, sizeof(float));
    
    m->hidden      = (float *)aligned_alloc(64, hidden * sizeof(float));
    m->hidden2     = (float *)aligned_alloc(64, hidden * sizeof(float));
    m->q           = (float *)aligned_alloc(64, n_heads * head_dim * sizeof(float));
    m->k           = (float *)aligned_alloc(64, n_kv_heads * head_dim * sizeof(float));
    m->v           = (float *)aligned_alloc(64, n_kv_heads * head_dim * sizeof(float));
    m->attn_out    = (float *)aligned_alloc(64, n_heads * head_dim * sizeof(float));
    m->gate        = (float *)aligned_alloc(64, inter * sizeof(float));
    m->up          = (float *)aligned_alloc(64, inter * sizeof(float));
    m->down_in     = (float *)aligned_alloc(64, inter * sizeof(float));
    m->logits      = (float *)aligned_alloc(64, vocab * sizeof(float));
    m->attn_scores = (float *)aligned_alloc(64, MAX_SEQ * sizeof(float));
    m->final_norm  = (float *)aligned_alloc(64, hidden * sizeof(float));
    
    size_t kv_mb = kv_size * 2 * sizeof(float) / (1024*1024);
    printf("Model config: hidden=%d inter=%d heads=%d kv_heads=%d layers=%d vocab=%d\n",
           hidden, inter, n_heads, n_kv_heads, n_layers, vocab);
    printf("KV cache: %zu MB, tied_embed=%d, attn_bias=%d\n",
           kv_mb, tie_embeddings, has_attn_bias);
    
    return m;
}

/* Weight setters */
void model_set_embed(Model *m, uint16_t *data) { m->embed = data; }
void model_set_final_norm(Model *m, float *data) { memcpy(m->final_norm, data, m->cfg.hidden * sizeof(float)); }
void model_set_lm_head(Model *m, uint16_t *data, int out_dim, int in_dim) {
    m->lm_head.weight = data;
    m->lm_head.out_dim = out_dim;
    m->lm_head.in_dim = in_dim;
}

void layer_set_norms(Model *m, int l, float *input_norm, float *post_norm) {
    m->layers[l].input_norm = input_norm;
    m->layers[l].post_norm = post_norm;
}

void layer_set_bias(Model *m, int l, float *q_bias, float *k_bias, float *v_bias) {
    m->layers[l].q_bias = q_bias;
    m->layers[l].k_bias = k_bias;
    m->layers[l].v_bias = v_bias;
}

void layer_set_qk_norm(Model *m, int l, float *q_norm, float *k_norm) {
    m->layers[l].q_norm = q_norm;
    m->layers[l].k_norm = k_norm;
}

void layer_set_unary(
    UnaryLinear *ul,
    uint64_t *sign, uint64_t *planes, float *scales,
    int out_dim, int in_dim, int n_planes
) {
    ul->sign_bits = sign;
    ul->mag_planes = planes;
    ul->scales = scales;
    ul->out_dim = out_dim;
    ul->in_dim = in_dim;
    ul->n_planes = n_planes;
    ul->bias = NULL;
}

void layer_set_linears(
    Model *m, int l,
    uint64_t *q_sign, uint64_t *q_planes, float *q_scales, int q_out, int q_in,
    uint64_t *k_sign, uint64_t *k_planes, float *k_scales, int k_out, int k_in,
    uint64_t *v_sign, uint64_t *v_planes, float *v_scales, int v_out, int v_in,
    uint64_t *o_sign, uint64_t *o_planes, float *o_scales, int o_out, int o_in,
    uint64_t *g_sign, uint64_t *g_planes, float *g_scales, int g_out, int g_in,
    uint64_t *u_sign, uint64_t *u_planes, float *u_scales, int u_out, int u_in,
    uint64_t *d_sign, uint64_t *d_planes, float *d_scales, int d_out, int d_in,
    int n_planes
) {
    layer_set_unary(&m->layers[l].q_proj, q_sign, q_planes, q_scales, q_out, q_in, n_planes);
    layer_set_unary(&m->layers[l].k_proj, k_sign, k_planes, k_scales, k_out, k_in, n_planes);
    layer_set_unary(&m->layers[l].v_proj, v_sign, v_planes, v_scales, v_out, v_in, n_planes);
    layer_set_unary(&m->layers[l].o_proj, o_sign, o_planes, o_scales, o_out, o_in, n_planes);
    layer_set_unary(&m->layers[l].gate_proj, g_sign, g_planes, g_scales, g_out, g_in, n_planes);
    layer_set_unary(&m->layers[l].up_proj, u_sign, u_planes, u_scales, u_out, u_in, n_planes);
    layer_set_unary(&m->layers[l].down_proj, d_sign, d_planes, d_scales, d_out, d_in, n_planes);
}

void model_reset_cache(Model *m) {
    size_t kv_size = (size_t)m->cfg.n_layers * MAX_SEQ * m->cfg.n_kv_heads * m->cfg.head_dim;
    memset(m->k_cache, 0, kv_size * sizeof(float));
    memset(m->v_cache, 0, kv_size * sizeof(float));
}

void model_free(Model *m) {
    free(m->k_cache); free(m->v_cache);
    free(m->hidden); free(m->hidden2);
    free(m->q); free(m->k); free(m->v);
    free(m->attn_out); free(m->gate); free(m->up); free(m->down_in);
    free(m->logits); free(m->attn_scores); free(m->final_norm);
    free(m->layers);
    free(m);
}